diff --git a/mlir/lib/ExecutionEngine/CMakeLists.txt b/mlir/lib/ExecutionEngine/CMakeLists.txt --- a/mlir/lib/ExecutionEngine/CMakeLists.txt +++ b/mlir/lib/ExecutionEngine/CMakeLists.txt @@ -204,11 +204,23 @@ PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} ) - target_link_libraries(mlir_cuda_runtime - PRIVATE - ${CUDA_RUNTIME_LIBRARY} - ${CUDA_CUSPARSE_LIBRARY} - ) + + # We need the cusparseLT to provide 2:4 sparsity support. + if (DEFINED CUDA_CUSPARSELT_LIBRARY) + target_link_libraries(mlir_cuda_runtime + PRIVATE + ${CUDA_RUNTIME_LIBRARY} + ${CUDA_CUSPARSE_LIBRARY} + ${CUDA_CUSPARSELT_LIBRARY} + ) + else() + target_link_libraries(mlir_cuda_runtime + PRIVATE + ${CUDA_RUNTIME_LIBRARY} + ${CUDA_CUSPARSE_LIBRARY} + ) + endif() + add_definitions(-DMLIR_CUDA_CUSPARSELT_ENABLED=(defined(CUDA_CUSPARSELT_LIBRARY))) endif() if(MLIR_ENABLE_ROCM_RUNNER) diff --git a/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp b/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp --- a/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp +++ b/mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp @@ -19,6 +19,10 @@ #include "cuda.h" #include "cusparse.h" +#if MLIR_CUDA_CUSPARSELT_ENABLED +#include "cusparseLt.h" +#endif // MLIR_CUDA_CUSPARSELT_ENABLED + #ifdef _WIN32 #define MLIR_CUDA_WRAPPERS_EXPORT __declspec(dllexport) #else @@ -438,3 +442,5 @@ matB, betap, matC, dtp, CUSPARSE_SDDMM_ALG_DEFAULT, buf)) } + +// operations from cusparseLt \ No newline at end of file