From a2fdbfc13ccf5e2a0ea7bb356c6bc287054a0bf7 Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Tue, 25 May 2021 16:31:49 -0700
Subject: [PATCH] [Cuda][Codegen] Check for cuda include dir in /usr/include.

Currently, on linux platforms, only checks for cuda install directory
in /usr/local/cuda/include.  The `nvidia-cuda-dev` package of Ubuntu
20.04 installs at /usr/include, so it would be good to check that
location as well.
---
 src/target/opt/build_cuda_on.cc | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/target/opt/build_cuda_on.cc b/src/target/opt/build_cuda_on.cc
index 1a0f08920fb6..4a2917daa5ed 100644
--- a/src/target/opt/build_cuda_on.cc
+++ b/src/target/opt/build_cuda_on.cc
@@ -67,6 +67,10 @@ std::string FindCUDAIncludePath() {
   if (stat(cuda_include_path.c_str(), &st) == 0) {
     return cuda_include_path;
   }
+
+  if (stat("/usr/include/cuda.h", &st) == 0) {
+    return "/usr/include";
+  }
 #endif
   LOG(FATAL) << "Cannot find cuda include path."
              << "CUDA_PATH is not set or CUDA is not installed in the default installation path."