From 578320d9e3dab8f74bd87769edb05b1fe1b1ae86 Mon Sep 17 00:00:00 2001
From: Alex Yu <sxyu@berkeley.edu>
Date: Tue, 19 Oct 2021 17:24:29 -0700
Subject: [PATCH] Make install work on CUDA < 11

---
 setup.py | 26 ++++++++++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index 3d33f3f8..fa4438b8 100644
--- a/setup.py
+++ b/setup.py
@@ -1,5 +1,7 @@
 from setuptools import setup
+import os
 import os.path as osp
+import warnings
 
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
@@ -10,6 +12,26 @@
 
 CUDA_FLAGS = []
 INSTALL_REQUIREMENTS = []
+include_dirs = [osp.join(ROOT_DIR, "svox2", "csrc", "include")]
+
+# From PyTorch3D
+cub_home = os.environ.get("CUB_HOME", None)
+if cub_home is None:
+	prefix = os.environ.get("CONDA_PREFIX", None)
+	if prefix is not None and os.path.isdir(prefix + "/include/cub"):
+		cub_home = prefix + "/include"
+
+if cub_home is None:
+	warnings.warn(
+		"The environment variable `CUB_HOME` was not found."
+		"Installation will fail if your system CUDA toolkit version is less than 11."
+		"NVIDIA CUB can be downloaded "
+		"from `https://github.com/NVIDIA/cub/releases`. You can unpack "
+		"it to a location of your choice and set the environment variable "
+		"`CUB_HOME` to the folder containing the `CMakeListst.txt` file."
+	)
+else:
+	include_dirs.append(os.path.realpath(cub_home).replace("\\ ", " "))
 
 try:
     ext_modules = [
@@ -20,8 +42,8 @@
             'svox2/csrc/misc_kernel.cu',
             'svox2/csrc/loss_kernel.cu',
             'svox2/csrc/optim_kernel.cu',
-        ], include_dirs=[osp.join(ROOT_DIR, "svox2", "csrc", "include"),],
-        optional=True),
+        ], include_dirs=include_dirs,
+        optional=False),
     ]
 except:
     import warnings