From 0b048afa79db763609eff2bc246f2a0d6f503fec Mon Sep 17 00:00:00 2001
From: pensun <peng.sun@amd.com>
Date: Thu, 3 Mar 2016 06:23:17 -0600
Subject: [PATCH] Fix for hipEnvVar test for titan machine

---
 hipamd/tests/src/hipEnvVar.cpp       | 8 ++------
 hipamd/tests/src/hipEnvVarDriver.cpp | 7 +++++++
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/hipamd/tests/src/hipEnvVar.cpp b/hipamd/tests/src/hipEnvVar.cpp
index 2b2b0afcff..c1b6589fe2 100644
--- a/hipamd/tests/src/hipEnvVar.cpp
+++ b/hipamd/tests/src/hipEnvVar.cpp
@@ -44,7 +44,7 @@ int main(int argc, char **argv)
     //std::cout << "The current env HIP_VISIBLE_DEVICES is"<<str << std::endl;
     extern char *optarg;
     extern int optind;
-    int c, err = 0;
+    int c = 0;
     int retDevCnt=0, retDevInfo=0, setEnvVar=0;
     int device=0;
     string env;
@@ -64,20 +64,16 @@ int main(int argc, char **argv)
         case 'h':
             usage();
             return 0;
-            break;
         default :
             //usage();
             return -1;
-            break;
-        case '?':
-            err = 1;
-            break;
         }
 
     if (setEnvVar ) {
         //env = "export HIP_VISIBLE_DEVICES=" + env;
         //cout<<"The received env var is: "<<env<<endl;
         setenv("HIP_VISIBLE_DEVICES",env.c_str(),1);
+        setenv("CUDA_VISIBLE_DEVICES",env.c_str(),1);
         cout<<"set env HIP_VISIBLE_DEVICES = "<< env.c_str()<<endl;
         //verify if the environment variable is set
         char* pPath;
diff --git a/hipamd/tests/src/hipEnvVarDriver.cpp b/hipamd/tests/src/hipEnvVarDriver.cpp
index 21884f00b1..245b23b52e 100644
--- a/hipamd/tests/src/hipEnvVarDriver.cpp
+++ b/hipamd/tests/src/hipEnvVarDriver.cpp
@@ -53,6 +53,7 @@ int getDevicePCIBusNum(int deviceID){
 
 int main() {
     unsetenv("HIP_VISIBLE_DEVICES");
+    unsetenv("CUDA_VISIBLE_DEVICES");
     //collect the device pci bus ID for all devices
     int totalDeviceNum = getDeviceNumber();
     std::cout << "The total number of available devices is " << totalDeviceNum<< std::endl
@@ -68,6 +69,7 @@ int main() {
     //query the returned device pci bus number, check if match the database
     for (int i = 0; i < totalDeviceNum ; i++) {
         setenv("HIP_VISIBLE_DEVICES",(char*)std::to_string(i).c_str(),1);
+        setenv("CUDA_VISIBLE_DEVICES",(char*)std::to_string(i).c_str(),1);
         //cout<<"HIP_VISIBLE_DEVICES is "<<i<<" data in vector is "<<devPCINum[i]<<endl;
         //std::cout <<"Returned pci number is"<< getDevicePCIBusNum(0) << std::endl;
         if (devPCINum[i] != getDevicePCIBusNum(0)) {
@@ -81,21 +83,26 @@ int main() {
 
     //check when set an invalid device number
     setenv("HIP_VISIBLE_DEVICES","1000,0,1",1);
+    setenv("CUDA_VISIBLE_DEVICES","1000,0,1",1);
     assert(getDeviceNumber() == 0);
 
     if(totalDeviceNum > 2){
         setenv("HIP_VISIBLE_DEVICES","0,1,1000,2",1);
+        setenv("CUDA_VISIBLE_DEVICES","0,1,1000,2",1);
         assert(getDeviceNumber() == 2);
 
         setenv("HIP_VISIBLE_DEVICES","0,1,2",1);
+        setenv("CUDA_VISIBLE_DEVICES","0,1,2",1);
         assert(getDeviceNumber() == 3);
         // test if CUDA_VISIBLE_DEVICES will be accepted by the runtime
         unsetenv("HIP_VISIBLE_DEVICES");
+        unsetenv("CUDA_VISIBLE_DEVICES");
         setenv("CUDA_VISIBLE_DEVICES","0,1,2",1);
         assert(getDeviceNumber() == 3);
     }
 
     setenv("HIP_VISIBLE_DEVICES","-100,0,1",1);
+    setenv("CUDA_VISIBLE_DEVICES","-100,0,1",1);
     assert(getDeviceNumber() == 0);
 
     std::cout << "PASSED" << std::endl;