[Mlir-commits] [mlir] 50db478 - [mlir][sparse][gpu] refined build setup for cusparse

[Aart Bik]

Author: Aart Bik
Date: 2023-06-07T11:09:22-07:00
New Revision: 50db4789a81b678bf250b1397e9c8fc3ca9e1e81

URL: https://github.com/llvm/llvm-project/commit/50db4789a81b678bf250b1397e9c8fc3ca9e1e81
DIFF: https://github.com/llvm/llvm-project/commit/50db4789a81b678bf250b1397e9c8fc3ca9e1e81.diff

LOG: [mlir][sparse][gpu] refined build setup for cusparse

Reviewed By: K-Wu

Differential Revision: https://reviews.llvm.org/D152387

Added: 
    

Modified: 
    mlir/lib/ExecutionEngine/CMakeLists.txt
    mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/ExecutionEngine/CMakeLists.txt b/mlir/lib/ExecutionEngine/CMakeLists.txt
index 0eccdf58c0015..39cc9c060db78 100644
--- a/mlir/lib/ExecutionEngine/CMakeLists.txt
+++ b/mlir/lib/ExecutionEngine/CMakeLists.txt
@@ -191,9 +191,6 @@ if(LLVM_ENABLE_PIC)
     # We need the libcuda.so library.
     find_library(CUDA_RUNTIME_LIBRARY cuda HINTS ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES} REQUIRED)
 
-
-
-
     add_mlir_library(mlir_cuda_runtime
       SHARED
       CudaRuntimeWrappers.cpp
@@ -212,18 +209,22 @@ if(LLVM_ENABLE_PIC)
     )
 
     if(MLIR_ENABLE_CUDA_CUSPARSE)
-      
       # Find the libcusparse.so library if CUSPARSE build is requested.
       find_library(CUDA_CUSPARSE_LIBRARY cusparse HINTS ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES} REQUIRED)
+
       target_link_libraries(mlir_cuda_runtime
         PRIVATE
         ${CUDA_CUSPARSE_LIBRARY}
       )
+      target_compile_definitions(mlir_cuda_runtime
+        PRIVATE
+        MLIR_ENABLE_CUDA_CUSPARSE=1
+      )
 
       if(MLIR_ENABLE_CUDA_CUSPARSELT)
-        # Find the libcusparseLt.so library in package manager default path if 
-        # CUSPARSELT build is requested. libcusparseLt.so provides sm80+ tensor 
-        # core support for 2:4 sparsity acceleration. 
+        # Find the libcusparseLt.so library in package manager default path if
+        # CUSPARSELT build is requested. libcusparseLt.so provides sm80+ tensor
+        # core support for 2:4 sparsity acceleration.
         find_library(CUDA_CUSPARSELT_LIBRARY cusparseLt HINTS ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES} REQUIRED)
         find_path(CUDA_CUSPARSELT_HEADER cusparseLt.h HINTS ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} REQUIRED)
         target_include_directories(mlir_cuda_runtime
@@ -234,6 +235,10 @@ if(LLVM_ENABLE_PIC)
           PRIVATE
           ${CUDA_CUSPARSELT_LIBRARY}
         )
+        target_compile_definitions(mlir_cuda_runtime
+          PRIVATE
+          MLIR_ENABLE_CUDA_CUSPARSELT=1
+        )
       endif()
     endif()
   endif()

diff  --git a/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp b/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp
index 7fc73f7e7ca20..f654f6a2a5ef7 100644
--- a/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp
+++ b/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp
@@ -20,9 +20,9 @@
 #include "cuda_bf16.h"
 #include "cuda_fp16.h"
 
-#if MLIR_ENABLE_CUDA_CUSPARSE
+#ifdef MLIR_ENABLE_CUDA_CUSPARSE
 #include "cusparse.h"
-#if MLIR_ENABLE_CUDA_CUSPARSELT
+#ifdef MLIR_ENABLE_CUDA_CUSPARSELT
 #include "cusparseLt.h"
 #endif // MLIR_ENABLE_CUDA_CUSPARSELT
 #endif // MLIR_ENABLE_CUDA_CUSPARSE
@@ -232,7 +232,7 @@ extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSetDefaultDevice(int32_t device) {
   defaultDevice = device;
 }
 
-#if MLIR_ENABLE_CUDA_CUSPARSE
+#ifdef MLIR_ENABLE_CUDA_CUSPARSE
 
 ///
 /// Wrapper methods for the cuSparse library.
@@ -454,7 +454,7 @@ mgpuSDDMM(void *h, int32_t ma, int32_t mb, void *a, void *b, void *c,
                                          CUSPARSE_SDDMM_ALG_DEFAULT, buf))
 }
 
-#if MLIR_ENABLE_CUDA_CUSPARSELT
+#ifdef MLIR_ENABLE_CUDA_CUSPARSELT
 
 ///
 /// Wrapper methods for the cuSparseLt library.