add findCUDA method

Author: MJ10

pytorch/edgeml_pytorch/graph/rnn.py

@@ -9,7 +9,7 @@
 
 import edgeml_pytorch.utils as utils
 
-if "CUDA_HOME" in os.environ:
+if utils.findCUDA() is not None:
     import fastgrnn_cuda
 
 def onnx_exportable_rnn(input, fargs, cell, output):
@@ -320,7 +320,7 @@ class FastGRNNCUDACell(RNNCell):
     '''
     def __init__(self, input_size, hidden_size, zetaInit=1.0, nuInit=-4.0, name="FastGRNNCUDACell"):
         super(FastGRNNCUDACell, self).__init__(input_size, hidden_size, "sigmoid", "tanh", 1, 1, 2)
-        if not "CUDA_HOME" in os.environ:
+        if utils.findCUDA() is None:
             raise Exception('FastGRNNCUDACell is supported only on GPU devices.')
         self._input_size = input_size
         self._hidden_size = hidden_size

pytorch/edgeml_pytorch/utils.py

@@ -1,12 +1,40 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT license.
-
+import sys
+import os
 import numpy as np
 import torch
 import torch.nn.functional as F
 import torch.optim as optim
 
 
+def findCUDA():
+    '''Finds the CUDA install path.'''
+    # Guess #1
+    IS_WINDOWS = sys.platform == 'win32'
+    cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
+    if cuda_home is None:
+        # Guess #2
+        try:
+            which = 'where' if IS_WINDOWS else 'which'
+            nvcc = subprocess.check_output(
+                [which, 'nvcc']).decode().rstrip('\r\n')
+            cuda_home = os.path.dirname(os.path.dirname(nvcc))
+        except Exception:
+            # Guess #3
+            if IS_WINDOWS:
+                cuda_homes = glob.glob(
+                    'C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v*.*')
+                if len(cuda_homes) == 0:
+                    cuda_home = ''
+                else:
+                    cuda_home = cuda_homes[0]
+            else:
+                cuda_home = '/usr/local/cuda'
+            if not os.path.exists(cuda_home):
+                cuda_home = None
+    return cuda_home
+
 def multiClassHingeLoss(logits, labels):
     '''
     MultiClassHingeLoss to match C++ Version - No pytorch internal version

pytorch/setup.py

@@ -1,8 +1,9 @@
 import setuptools #enables develop
 import os
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+from edgeml_pytorch.utils import findCUDA
 
-if "CUDA_HOME" in os.environ:
+if findCUDA() is not None:
     setup(
         name='fastgrnn_cuda',
         ext_modules=[