From f9961fbc8f55953030b4f0d8acf856596e012bfd Mon Sep 17 00:00:00 2001
From: "Xia, Yu" <xiayu.x@bytedance.com>
Date: Wed, 7 Jul 2021 16:55:27 +0800
Subject: [PATCH] support to specify CUDA/ROCm root in python pkg building

---
 doc/install.md | 2 ++
 setup.py       | 8 +++++++-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/doc/install.md b/doc/install.md
index cd0c944161..a45d8c7d9f 100644
--- a/doc/install.md
+++ b/doc/install.md
@@ -137,6 +137,8 @@ One may set the following environment variables before executing `pip`:
 | --------------------- | ---------------------- | ------------- | -------------------------- |
 | DP_VARIANT            | `cpu`, `cuda`, `rocm`  | `cpu`         | Build CPU variant or GPU variant with CUDA or ROCM support. |
 | DP_FLOAT_PREC         | `high`, `low`          | `high`        | Build high (double) or low (float) precision. |
+| CUDA_TOOLKIT_ROOT_DIR | Path                   | Detected automatically | The path to the CUDA toolkit directory. |
+| ROCM_ROOT             | Path                   | Detected automatically | The path to the ROCM toolkit directory. |
 
 To test the installation, one should firstly jump out of the source directory
 ```
diff --git a/setup.py b/setup.py
index f574e46e06..7f917ba11d 100644
--- a/setup.py
+++ b/setup.py
@@ -41,9 +41,15 @@
     pass
 elif dp_variant == "cuda":
     cmake_args.append("-DUSE_CUDA_TOOLKIT:BOOL=TRUE")
+    cuda_root = os.environ.get("CUDA_TOOLKIT_ROOT_DIR")
+    if cuda_root:
+        cmake_args.append(f"-DCUDA_TOOLKIT_ROOT_DIR:STRING={cuda_root}")
 elif dp_variant == "rocm":
     cmake_args.append("-DUSE_ROCM_TOOLKIT:BOOL=TRUE")
-elif dp_variant != "":
+    rocm_root = os.environ.get("ROCM_ROOT")
+    if rocm_root:
+        cmake_args.append(f"-DROCM_ROOT:STRING={rocm_root}")
+else:
     raise RuntimeError("Unsupported DP_VARIANT option: %s" % dp_variant)
 
 # FLOAT_PREC