Fix compilation when using MKL. (#358)

"USE_MKL" needs to be defined for "mkl_alternate.hpp" to pick the right header.
NVIDIA · May 26, 2021 · 583fef8 · 583fef8
1 parent 2eb5d28
commit 583fef8
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/setup.py b/setup.py
@@ -184,6 +184,8 @@ def _argparse(pattern, argv, is_flag=True, is_list=False):
     assert BLAS in BLAS_LIST, f"Blas option {BLAS} not in valid options {BLAS_LIST}"
     if BLAS == "mkl":
         libraries.append("mkl_rt")
+        CC_FLAGS.append("-DUSE_MKL")
+        NVCC_FLAGS.append("-DUSE_MKL")
     else:
         libraries.append(BLAS)
     if not (BLAS_INCLUDE_DIRS is False):