diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td --- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td +++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td @@ -1960,7 +1960,7 @@ Example: ```mlir - %buffersz, %token = gpu.spmm_buffersize async [%dep] %env, %spmatA{TRANSPOSE}, %dnmatB{TRANSPOSE}, %dnmatC + %bufferszs, %token = gpu.spmm_buffersize async [%dep] %env, %spmatA{TRANSPOSE}, %dnmatB{TRANSPOSE}, %dnmatC : i64 ``` }]; @@ -1971,11 +1971,12 @@ GPU_SparseSpMatHandle:$spmatA, GPU_SparseDnMatHandle:$dnmatB, GPU_SparseDnMatHandle:$dnmatC); - let results = (outs Res:$bufferSz, + let results = (outs Res]>>:$bufferSzs, Optional:$asyncToken); let builders = [OpBuilder<(ins - "::mlir::Type":$bufferSz, + "::mlir::Type":$bufferSzs, "::mlir::Type":$asyncToken, "::mlir::ValueRange":$asyncDependencies, "::mlir::Value":$env, @@ -1984,17 +1985,17 @@ "::mlir::Value":$dnmatC), [{ auto modeA = gpu::TransposeMode::NON_TRANSPOSE; auto modeB = gpu::TransposeMode::NON_TRANSPOSE; - return build($_builder, $_state, bufferSz, asyncToken, asyncDependencies, + return build($_builder, $_state, bufferSzs, asyncToken, asyncDependencies, env, modeA, modeB, spmatA, dnmatB, dnmatC);}]> ]; let assemblyFormat = [{ custom(type($asyncToken), $asyncDependencies) - $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $dnmatC attr-dict + $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $dnmatC attr-dict `:` type($bufferSzs) }]; } -def GPU_SpMMOp : GPU_Op<"spmm", [GPU_AsyncOpInterface]> { +def GPU_SpMMOp : GPU_Op<"spmm", [GPU_AsyncOpInterface, AttrSizedOperandSegments]> { let summary = "SpMM operation"; let description = [{ The `gpu.spmm` operation performs the SpMM operation on the given sparse and @@ -2024,7 +2025,7 @@ GPU_SparseSpMatHandle:$spmatA, GPU_SparseDnMatHandle:$dnmatB, GPU_SparseDnMatHandle:$dnmatC, - AnyMemRef:$buffer); + Variadic:$buffers); let results = (outs Optional:$asyncToken); let builders = [OpBuilder<(ins @@ -2034,16 +2035,16 @@ "::mlir::Value":$spmatA, "::mlir::Value":$dnmatB, "::mlir::Value":$dnmatC, - "::mlir::Value":$buffer), [{ + "::mlir::ValueRange":$buffers), [{ auto modeA = gpu::TransposeMode::NON_TRANSPOSE; auto modeB = gpu::TransposeMode::NON_TRANSPOSE; return build($_builder, $_state, asyncToken, asyncDependencies, env, modeA, - modeB, spmatA, dnmatB, dnmatC, buffer);}]> + modeB, spmatA, dnmatB, dnmatC, buffers);}]> ]; let assemblyFormat = [{ custom(type($asyncToken), $asyncDependencies) - $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $dnmatC `,` $buffer attr-dict `:` type($buffer) + $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $dnmatC `,` $buffers attr-dict `:` type($buffers) }]; } diff --git a/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp b/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp --- a/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp +++ b/mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp @@ -1542,8 +1542,8 @@ auto dw = rewriter.create(loc, llvmInt32Type, dType.getIntOrFloatBitWidth()); auto stream = adaptor.getAsyncDependencies().front(); - Value pBuf = - MemRefDescriptor(adaptor.getBuffer()).allocatedPtr(rewriter, loc); + Value pBuf = MemRefDescriptor(adaptor.getBuffers().front()) + .allocatedPtr(rewriter, loc); if (!getTypeConverter()->useOpaquePointers()) pBuf = rewriter.create(loc, llvmPointerType, pBuf); spMMCallBuilder.create(loc, rewriter, 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 @@ -200,15 +200,36 @@ EXCLUDE_FROM_LIBMLIR ) set_property(TARGET mlir_cuda_runtime PROPERTY CXX_STANDARD 14) - target_include_directories(mlir_cuda_runtime - 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. + # As of the pre-1.0 version, we suppose the cusparselt is downloaded as an + # archive and extracted in an exclusive directory CUDA_CUSPARSELT_DIR, rather + # than installed by the package manager. This is the same as Nvidia examples. + if (DEFINED CUDA_CUSPARSELT_DIR) + target_include_directories(mlir_cuda_runtime + PRIVATE + ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} + ${CUDA_CUSPARSELT_DIR}/include + ) + target_link_libraries(mlir_cuda_runtime + PRIVATE + ${CUDA_RUNTIME_LIBRARY} + ${CUDA_CUSPARSE_LIBRARY} + ${CUDA_CUSPARSELT_DIR}/lib64 + ) + else() + target_include_directories(mlir_cuda_runtime + PRIVATE + ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} + ) + target_link_libraries(mlir_cuda_runtime + PRIVATE + ${CUDA_RUNTIME_LIBRARY} + ${CUDA_CUSPARSE_LIBRARY} + ) + endif() + add_definitions(-DMLIR_CUDA_CUSPARSELT_ENABLED=(defined(CUDA_CUSPARSELT_DIR))) 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,144 @@ matB, betap, matC, dtp, CUSPARSE_SDDMM_ALG_DEFAULT, buf)) } + +/// +/// Wrapper methods for the cuSparseLt library. +/// +#if MLIR_CUDA_CUSPARSELT_ENABLED + +extern "C" MLIR_CUDA_WRAPPERS_EXPORT void * +mgpuCreateSparseLtEnv(CUstream /*stream*/) { + cusparseLtHandle_t handle = nullptr; + // note that cuSparseLt still uses cusparseStatus_t + CUSPARSE_REPORT_IF_ERROR(cusparseLtInit(&handle)) + return reinterpret_cast(handle); +} + +extern "C" MLIR_CUDA_WRAPPERS_EXPORT void +mgpuDestroySparseLtEnv(void *h, CUstream /*stream*/) { + cusparseLtHandle_t handle = reinterpret_cast(h); + CUSPARSE_REPORT_IF_ERROR(cusparseLtDestroy(handle)) +} + +struct cusparseLtSpMatHandleAndData { + cusparseLtMatDescriptor_t mat; + void *rowPos; + void *colIdxs; + void *values; +}; +struct cusparseLtDnMatHandleAndData { + cusparseLtMatDescriptor_t mat; + void *values; +}; + +// TODO: pass handle ptr +extern "C" MLIR_CUDA_WRAPPERS_EXPORT void * +mgpuCreateCuSparseLtDnMat(intptr_t rows, intptr_t cols, void *values, + int32_t dw, CUstream /*stream*/) { + cusparseLtMatDescriptor_t mat; + cudaDataType_t dtp = dataTp(dw); + // assuming row-major when deciding lda + CUSPARSE_REPORT_IF_ERROR( + cusparseLtDenseDescriptorInit(handlePtr, &mat, rows, cols, /*lda=*/cols, + /*alignment=*/16, dtp, CUSPARSE_ORDER_ROW)) + cusparseLtDnMatHandleAndData matWithData{ + .mat = mat, + .values = values, + }; + return reinterpret_cast(matWithData); +} + +// This can be used to destroy both dense matrices and sparse matrices in +// cusparseLt +extern "C" MLIR_CUDA_WRAPPERS_EXPORT void +mgpuDestroyCuSparseLtSpMat(void *m, CUstream /*stream*/) { + auto mat = reinterpret_cast(m); + CUSPARSE_REPORT_IF_ERROR(cusparseLtMatDescriptorDestroy(mat)) +} + +// TODO: pass handle ptr +extern "C" MLIR_CUDA_WRAPPERS_EXPORT void *mgpuCusparseLtCreate224SpMat( + intptr_t rows, intptr_t cols, intptr_t nnz, void *rowPos, void *colIdxs, + void *values, int32_t pw, int32_t iw, int32_t dw, CUstream /*stream*/) { + cusparseLtMatDescriptor_t mat; + cusparseIndexType_t ptp = idxTp(pw); + cusparseIndexType_t itp = idxTp(iw); + cudaDataType_t dtp = dataTp(dw); + // assuming row-major when deciding lda + CUSPARSE_REPORT_IF_ERROR(cusparseLtStructuredDescriptorInit( + handlePtr, &mat, rows, cols, /*ld=*/cols, /*alignment=*/16, dtp, + CUSPARSE_ORDER_ROW, CUSPARSELT_SPARSITY_50_PERCENT)) + cusparseLtSpMatHandleAndData matWithData{ + .mat = mat, + .rowPos = rowPos, + .colIdxs = colIdxs, + .values = values, + }; + return reinterpret_cast(matWithData); +} + +struct cusparseLtWorkspaceSizes{ + size_t workspace_size; + size_t compressed_size; + size_t compressed_buffer_size; +}; + +// Several things are being done in this stage, algorithm selection, planning, +// and returning workspace and compressed matrices data buffer sizes. +extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t +mgpuCuSparseLtSpMMBufferSize(void *h, int32_t ma, int32_t mb, void *a, void *b, + void *c, int32_t dw, CUstream /*stream*/) { + // cusparseHandle_t handle = reinterpret_cast(h); + // cusparseOperation_t modeA = static_cast(ma); + // cusparseOperation_t modeB = static_cast(mb); + // cusparseSpMatDescr_t matA = reinterpret_cast(a); + // cusparseDnMatDescr_t matB = reinterpret_cast(b); + // cusparseDnMatDescr_t matC = reinterpret_cast(c); + // cudaDataType_t dtp = dataTp(dw); + // ALPHABETA(dw, alpha, beta) + // size_t bufferSize = 0; + // CUSPARSE_REPORT_IF_ERROR(cusparseSpMM_bufferSize( + // handle, modeA, modeB, alphap, matA, matB, betap, matC, dtp, + // CUSPARSE_SPMM_ALG_DEFAULT, &bufferSize)) + // return bufferSize == 0 ? 1 : bufferSize; // avoid zero-alloc + + cusparseLtMatmulAlgSelection_t alg_sel; + CHECK_CUSPARSE(cusparseLtMatmulAlgSelectionInit( + &handle, &alg_sel, &matmul, CUSPARSELT_MATMUL_ALG_DEFAULT)) + int alg = 0; + CHECK_CUSPARSE(cusparseLtMatmulAlgSetAttribute( + &handle, &alg_sel, CUSPARSELT_MATMUL_ALG_CONFIG_ID, &alg, sizeof(alg))) + size_t cusparseLtWorkspaceSizes sizes; + CHECK_CUSPARSE(cusparseLtMatmulPlanInit(&handle, &plan, &matmul, &alg_sel)) + CHECK_CUSPARSE(cusparseLtMatmulGetWorkspace(&handle, &plan, &(sizes.workspace_size))) + CHECK_CUSPARSE(cusparseLtSpMMACompressedSize(&handle, &plan, &(sizes.compressed_size), + &(sizes.compressed_buffer_size))) + return reinterpret_cast(sizes); +} + +extern "C" MLIR_CUDA_WRAPPERS_EXPORT void +mgpuCuSparseLtSpMM(void *h, int32_t ma, int32_t mb, void *a, void *b, void *c, + int32_t dw, void *buf, CUstream /*stream*/) { + // cusparseHandle_t handle = reinterpret_cast(h); + // cusparseOperation_t modeA = static_cast(ma); + // cusparseOperation_t modeB = static_cast(mb); + // cusparseSpMatDescr_t matA = reinterpret_cast(a); + // cusparseDnMatDescr_t matB = reinterpret_cast(b); + // cusparseDnMatDescr_t matC = reinterpret_cast(c); + // cudaDataType_t dtp = dataTp(dw); + ALPHABETA(dw, alpha, beta) + // CUSPARSE_REPORT_IF_ERROR(cusparseSpMM(handle, modeA, modeB, alphap, matA, + // matB, betap, matC, dtp, + // CUSPARSE_SPMM_ALG_DEFAULT, buf)) + + CHECK_CUSPARSE(cusparseLtSpMMACompress(&handle, &plan, dA, dA_compressed, + dA_compressedBuffer, stream)) + + // Perform the matrix multiplication + CHECK_CUSPARSE(cusparseLtMatmul(&handle, &plan, &alpha, dA_compressed, dB, + &beta, dC, dD, d_workspace, streams, + num_streams)) +} + +#endif // MLIR_CUDA_CUSPARSELT_ENABLED \ No newline at end of file diff --git a/mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir b/mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir --- a/mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir +++ b/mlir/test/Conversion/GPUCommon/lower-sparse-to-gpu-runtime-calls.mlir @@ -53,7 +53,7 @@ %env, %token3 = gpu.create_sparse_env async [%token2] %spmat, %token4 = gpu.create_csr async [%token3] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref, memref, memref %dnmat, %token5 = gpu.create_dn_mat async [%token4] %arg0, %arg0, %mem2 : memref - %bufferSz, %token6 = gpu.spmm_buffer_size async [%token5] %env, %spmat, %dnmat, %dnmat + %bufferSz, %token6 = gpu.spmm_buffer_size async [%token5] %env, %spmat, %dnmat, %dnmat : index %token7 = gpu.spmm async [%token6] %env, %spmat, %dnmat, %dnmat, %mem2 : memref %token8 = gpu.destroy_sp_mat async [%token7] %spmat %token9 = gpu.destroy_dn_mat async [%token8] %dnmat diff --git a/mlir/test/Dialect/GPU/ops.mlir b/mlir/test/Dialect/GPU/ops.mlir --- a/mlir/test/Dialect/GPU/ops.mlir +++ b/mlir/test/Dialect/GPU/ops.mlir @@ -341,7 +341,7 @@ // CHECK: gpu.create_dn_mat async %dnmat, %token9 = gpu.create_dn_mat async [%token8] %arg0, %arg0, %mem2 : memref // CHECK: gpu.spmm_buffer_size async - %bufferSz2, %token10 = gpu.spmm_buffer_size async [%token9] %env, %spmat, %dnmat, %dnmat + %bufferSz2, %token10 = gpu.spmm_buffer_size async [%token9] %env, %spmat, %dnmat, %dnmat : index // CHECK: gpu.spmm async %token11 = gpu.spmm async [%token10] %env, %spmat, %dnmat, %dnmat, %mem2 : memref // CHECK: gpu.sddmm_buffer_size async diff --git a/mlir/test/Dialect/SparseTensor/GPU/gpu_matmul_lib.mlir b/mlir/test/Dialect/SparseTensor/GPU/gpu_matmul_lib.mlir --- a/mlir/test/Dialect/SparseTensor/GPU/gpu_matmul_lib.mlir +++ b/mlir/test/Dialect/SparseTensor/GPU/gpu_matmul_lib.mlir @@ -49,7 +49,7 @@ // CHECK: %[[VAL_44:.*]], %[[VAL_45:.*]] = gpu.create_csr async {{\[}}%[[VAL_43]]] %[[VAL_6]], %[[VAL_7]], %[[VAL_5]], %[[VAL_14]], %[[VAL_19]], %[[VAL_24]] : memref, memref, memref // CHECK: %[[VAL_46:.*]], %[[VAL_47:.*]] = gpu.create_dn_mat async {{\[}}%[[VAL_45]]] %[[VAL_7]], %[[VAL_8]], %[[VAL_31]] : memref // CHECK: %[[VAL_48:.*]], %[[VAL_49:.*]] = gpu.create_dn_mat async {{\[}}%[[VAL_47]]] %[[VAL_6]], %[[VAL_8]], %[[VAL_38]] : memref -// CHECK: %[[VAL_50:.*]], %[[VAL_51:.*]] = gpu.spmm_buffer_size async {{\[}}%[[VAL_49]]] %[[VAL_42]], %[[VAL_44]], %[[VAL_46]], %[[VAL_48]] +// CHECK: %[[VAL_50:.*]], %[[VAL_51:.*]] = gpu.spmm_buffer_size async {{\[}}%[[VAL_49]]] %[[VAL_42]], %[[VAL_44]], %[[VAL_46]], %[[VAL_48]] : index // CHECK: %[[VAL_52:.*]], %[[VAL_53:.*]] = gpu.alloc async {{\[}}%[[VAL_51]]] (%[[VAL_50]]) : memref // CHECK: %[[VAL_54:.*]] = gpu.spmm async {{\[}}%[[VAL_53]]] %[[VAL_42]], %[[VAL_44]], %[[VAL_46]], %[[VAL_48]], %[[VAL_52]] : memref // CHECK: %[[VAL_55:.*]] = gpu.destroy_sp_mat async {{\[}}%[[VAL_54]]] %[[VAL_44]]