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 @@ -1789,6 +1789,41 @@ }]; } + +def GPU_Create2To4SpMatOp : GPU_Op<"create_2to4_spmat", [GPU_AsyncOpInterface]> { + let summary = "Create sparse matrix with 2:4 sparsity operation"; + let description = [{ + The `gpu.create_2to4_spmat` operation initializes a sparse matrix in dense + format with 2:4 sparsity. + The buffers must already be copied from the host to the device prior to + using this operation. The operation returns a handle to the sparse + matrix descriptor. + + If the `async` keyword is present, the op is executed asynchronously (i.e. + it does not block until the execution has finished on the device). In + that case, it returns a !gpu.async.token in addition to the environment. + + Example: + + ```mlir + %spmat, %token = gpu.create_2to4_spmat async [%dep] %env, %rows, %cols, %mem : memref + ``` + }]; + + let arguments = (ins Variadic:$asyncDependencies, + GPU_SparseEnvHandle:$env, + Index:$rows, + Index:$cols, + AnyMemRef:$memref); + let results = (outs Res:$spMat, + Optional:$asyncToken); + + let assemblyFormat = [{ + custom(type($asyncToken), $asyncDependencies) + $env `,` $rows `,` $cols `,` $memref attr-dict `:` type($memref) + }]; +} + def GPU_DestroySpMatOp : GPU_Op<"destroy_sp_mat", [GPU_AsyncOpInterface]> { let summary = "Destroy sparse matrix operation"; let description = [{ @@ -1960,7 +1995,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 +2006,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 +2020,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 +2060,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 +2070,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 @@ -230,6 +230,42 @@ {llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, llvmPointerType, llvmPointerType, llvmPointerType, llvmInt32Type, llvmInt32Type, llvmInt32Type, llvmPointerType /* void *stream */}}; + FunctionCallBuilder createSparseLtEnvCallBuilder = { + "mgpuCreateSparseLtEnv", + llvmPointerType, + {llvmPointerType /* void *stream */}}; + FunctionCallBuilder destroySparseLtEnvCallBuilder = { + "mgpuDestroySparseLtEnv", + llvmVoidType, + {llvmPointerType, llvmPointerType /* void *stream */}}; + FunctionCallBuilder createLtDnMatCallBuilder = { + "mgpuCreateLtDnMat", + llvmPointerType, + {llvmPointerType, llvmIntPtrType, llvmIntPtrType, llvmPointerType, + llvmInt32Type, llvmPointerType /* void *stream */}}; + FunctionCallBuilder destroyCuSparseLtSpMatBuilder = { + "mgpuDestroyCuSparseLtSpMat", + llvmVoidType, + {llvmPointerType, llvmPointerType /* void *stream */}}; + FunctionCallBuilder destroyCuSparseLtDnMatBuilder = { + "mgpuDestroyCuSparseLtDnMat", + llvmVoidType, + {llvmPointerType, llvmPointerType /* void *stream */}}; + FunctionCallBuilder create2To4SpMatCallBuilder = { + "mgpuCusparseLtCreate2To4SpMat", + llvmPointerType, + {llvmPointerType, llvmIntPtrType, llvmIntPtrType, llvmPointerType, + llvmInt32Type, llvmPointerType /* void *stream */}}; + FunctionCallBuilder cuSparseLtSpmmBufferSizeBuilder = { + "mgpuCuSparseLtSpmmBufferSize", + llvmPointerType, + {llvmPointerType, llvmPointerType, llvmPointerType /*void *stream*/}}; + FunctionCallBuilder cuSparseLtSpmmBuilder = { + "mgpuCuSparseLtSpmm", + llvmVoidType, + {llvmPointerType, llvmPointerType, llvmPointerType, llvmPointerType, + llvmInt32Type, llvmPointerType, llvmPointerType, llvmPointerType, + llvmPointerType /*void *stream*/}}; FunctionCallBuilder destroySpMatCallBuilder = { "mgpuDestroySpMat", llvmVoidType, @@ -559,6 +595,20 @@ ConversionPatternRewriter &rewriter) const override; }; +class ConvertCreate2To4SpMatOpToGpuRuntimeCallPattern + : public ConvertOpToGpuRuntimeCallPattern { +public: + ConvertCreate2To4SpMatOpToGpuRuntimeCallPattern( + LLVMTypeConverter &typeConverter) + : ConvertOpToGpuRuntimeCallPattern( + typeConverter) {} + +private: + LogicalResult + matchAndRewrite(gpu::Create2To4SpMatOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + class ConvertDestroySpMatOpToGpuRuntimeCallPattern : public ConvertOpToGpuRuntimeCallPattern { public: @@ -688,6 +738,47 @@ return builder.create(loc, function, arguments); } +static Value getEnvForCuSparseLtDefs(Value op) { + for (Operation *user : op.getUsers()) { + auto spmmOp = dyn_cast(user); + if (!spmmOp) + continue; + return spmmOp.getEnv(); + } + llvm_unreachable("cannot find the spmm op that uses the op result"); +} + +// TODO: We may want a run-time (of the mlir compiler) disablement/warning: +// cusparseLt currently won't work for cuda architecture <8.0 and will trigger a +// runtime (of the CUDA program) error , but it might be great if we could at +// least output a warning when we found the target architecture is <8.0 and the +// user still wants to use cusparseLt. to make sure when lowering gpu sparse +// dialect to llvm calls, the cusparselt calls are disabled for cuda +// architecture <8.0 +static bool is2To4Sparsity(Value spMat) { + if (auto op = spMat.getDefiningOp()) + return true; + if (auto op = spMat.getDefiningOp()) + return false; + if (auto op = spMat.getDefiningOp()) + return false; + // print the spMat defining op + spMat.getDefiningOp()->print(llvm::errs()); + llvm_unreachable("cannot find spmat def"); +} + +static std::string inferSpMMType(Value op) { + for (Operation *user : op.getUsers()) { + auto spmmOp = dyn_cast(user); + // if the other operator is 50% sparsity then we should use cusparseLt + if (!spmmOp) + continue; + if (is2To4Sparsity(spmmOp.getSpmatA())) + return "cusparseLt"; + } + return "cusparse"; +} + // Returns whether all operands are of LLVM type. static LogicalResult areAllLLVMTypes(Operation *op, ValueRange operands, ConversionPatternRewriter &rewriter) { @@ -1234,6 +1325,8 @@ return llvm::cast(op.getValues().getType()).getElementType(); if (auto op = spMat.getDefiningOp()) return llvm::cast(op.getValues().getType()).getElementType(); + if (auto op = spMat.getDefiningOp()) + return op.getMemref().getType().getElementType(); llvm_unreachable("cannot find spmat def"); } @@ -1252,8 +1345,16 @@ return failure(); Location loc = op.getLoc(); auto stream = adaptor.getAsyncDependencies().front(); - auto handle = - createSparseEnvCallBuilder.create(loc, rewriter, {stream}).getResult(); + // use the cusparseLt create call if the dnmat is used with spmat with + // 2:4 sparsity + Value handle; + if (inferSpMMType(op.getEnv()) == "cusparseLt") { + handle = createSparseLtEnvCallBuilder.create(loc, rewriter, {stream}) + .getResult(); + } else { + handle = + createSparseEnvCallBuilder.create(loc, rewriter, {stream}).getResult(); + } rewriter.replaceOp(op, {handle, stream}); return success(); } @@ -1266,7 +1367,15 @@ return failure(); Location loc = op.getLoc(); auto stream = adaptor.getAsyncDependencies().front(); - destroySparseEnvCallBuilder.create(loc, rewriter, {adaptor.getEnv(), stream}); + // use the cusparseLt destroy call if the dnmat is used with spmat with + // 2:4 sparsity + if (inferSpMMType(op.getEnv()) == "cusparseLt") { + destroySparseLtEnvCallBuilder.create(loc, rewriter, + {adaptor.getEnv(), stream}); + } else { + destroySparseEnvCallBuilder.create(loc, rewriter, + {adaptor.getEnv(), stream}); + } rewriter.replaceOp(op, {stream}); return success(); } @@ -1324,11 +1433,27 @@ llvm::cast(op.getMemref().getType()).getElementType(); auto dw = rewriter.create(loc, llvmInt32Type, dType.getIntOrFloatBitWidth()); - auto handle = - createDnMatCallBuilder - .create(loc, rewriter, - {adaptor.getRows(), adaptor.getCols(), pMat, dw, stream}) - .getResult(); + // TODO: For now, we track the use of the handle and lower it to cusparse / + // cusparseLt accordingly. If in a block, both cusparse and cusparseLt are + // used, we require two separate Creation ops to be the correct logic. In + // future, we may add support to using one handle in sparse tensor / GPU + // dialect in both cusparse and cusparseLt. use the cusparseLt create call if + // the dnmat is used with spmat with 2:4 sparsity + Value handle; + if (inferSpMMType(op.getResult(0)) == "cusparseLt") { + auto envHandle = getEnvForCuSparseLtDefs(op.getResult(0)); + handle = createLtDnMatCallBuilder + .create(loc, rewriter, + {envHandle, adaptor.getRows(), adaptor.getCols(), pMat, + dw, stream}) + .getResult(); + } else { + handle = + createDnMatCallBuilder + .create(loc, rewriter, + {adaptor.getRows(), adaptor.getCols(), pMat, dw, stream}) + .getResult(); + } rewriter.replaceOp(op, {handle, stream}); return success(); } @@ -1341,7 +1466,14 @@ return failure(); Location loc = op.getLoc(); auto stream = adaptor.getAsyncDependencies().front(); - destroyDnMatCallBuilder.create(loc, rewriter, {adaptor.getDmat(), stream}); + // use the cusparseLt destroy call if the dnmat is used with spmat with + // 2:4 sparsity + if (inferSpMMType(op.getDmat()) == "cusparseLt") { + destroyCuSparseLtDnMatBuilder.create(loc, rewriter, + {adaptor.getDmat(), stream}); + } else { + destroyDnMatCallBuilder.create(loc, rewriter, {adaptor.getDmat(), stream}); + } rewriter.replaceOp(op, {stream}); return success(); } @@ -1424,6 +1556,33 @@ return success(); } +// TODO: work on this +LogicalResult ConvertCreate2To4SpMatOpToGpuRuntimeCallPattern::matchAndRewrite( + gpu::Create2To4SpMatOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) || + failed(isAsyncWithOneDependency(rewriter, op))) + return failure(); + Location loc = op.getLoc(); + auto stream = adaptor.getAsyncDependencies().front(); + Value pMat = + MemRefDescriptor(adaptor.getMemref()).allocatedPtr(rewriter, loc); + if (!getTypeConverter()->useOpaquePointers()) + pMat = rewriter.create(loc, llvmPointerType, pMat); + Type dType = + llvm::cast(op.getMemref().getType()).getElementType(); + auto dw = rewriter.create(loc, llvmInt32Type, + dType.getIntOrFloatBitWidth()); + auto envHandle = adaptor.getEnv(); + auto handle = create2To4SpMatCallBuilder + .create(loc, rewriter, + {envHandle, adaptor.getRows(), adaptor.getCols(), + pMat, dw, stream}) + .getResult(); + rewriter.replaceOp(op, {handle, stream}); + return success(); +} + LogicalResult ConvertDestroySpMatOpToGpuRuntimeCallPattern::matchAndRewrite( gpu::DestroySpMatOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const { @@ -1432,7 +1591,14 @@ return failure(); Location loc = op.getLoc(); auto stream = adaptor.getAsyncDependencies().front(); - destroySpMatCallBuilder.create(loc, rewriter, {adaptor.getSpmat(), stream}); + // use the cusparseLt destroy call if the spmat is 2:4 sparsity + if (is2To4Sparsity(op.getSpmat())) { + destroyCuSparseLtSpMatBuilder.create(loc, rewriter, + {adaptor.getSpmat(), stream}); + + } else { + destroySpMatCallBuilder.create(loc, rewriter, {adaptor.getSpmat(), stream}); + } rewriter.replaceOp(op, {stream}); return success(); } @@ -1496,12 +1662,20 @@ auto dw = rewriter.create(loc, llvmInt32Type, dType.getIntOrFloatBitWidth()); auto stream = adaptor.getAsyncDependencies().front(); - auto bufferSize = - spMMBufferSizeCallBuilder - .create(loc, rewriter, - {adaptor.getEnv(), modeA, modeB, adaptor.getSpmatA(), - adaptor.getDnmatB(), adaptor.getDnmatC(), dw, stream}) - .getResult(); + Value bufferSize; + if (is2To4Sparsity(op.getSpmatA())) { + bufferSize = cuSparseLtSpmmBufferSizeBuilder + .create(loc, rewriter, + {adaptor.getEnv(), adaptor.getSpmatA(), stream}) + .getResult(); + } else { + bufferSize = + spMMBufferSizeCallBuilder + .create(loc, rewriter, + {adaptor.getEnv(), modeA, modeB, adaptor.getSpmatA(), + adaptor.getDnmatB(), adaptor.getDnmatC(), dw, stream}) + .getResult(); + } rewriter.replaceOp(op, {bufferSize, stream}); return success(); } @@ -1542,14 +1716,30 @@ auto dw = rewriter.create(loc, llvmInt32Type, dType.getIntOrFloatBitWidth()); auto stream = adaptor.getAsyncDependencies().front(); - Value pBuf = - MemRefDescriptor(adaptor.getBuffer()).allocatedPtr(rewriter, loc); - if (!getTypeConverter()->useOpaquePointers()) - pBuf = rewriter.create(loc, llvmPointerType, pBuf); - spMMCallBuilder.create(loc, rewriter, - {adaptor.getEnv(), modeA, modeB, adaptor.getSpmatA(), - adaptor.getDnmatB(), adaptor.getDnmatC(), dw, pBuf, - stream}); + + // lower to cusparseLt if applicable + if (is2To4Sparsity(op.getSpmatA())) { + SmallVector pBufs; + for (Value buffer : adaptor.getBuffers()) { + Value pBuf = MemRefDescriptor(buffer).allocatedPtr(rewriter, loc); + if (!getTypeConverter()->useOpaquePointers()) + pBuf = rewriter.create(loc, llvmPointerType, pBuf); + pBufs.push_back(pBuf); + } + cuSparseLtSpmmBufferSizeBuilder.create( + loc, rewriter, + {adaptor.getEnv(), adaptor.getSpmatA(), adaptor.getDnmatB(), + adaptor.getDnmatC(), dw, pBufs[0], pBufs[1], pBufs[2], stream}); + } else { + Value pBuf = MemRefDescriptor(adaptor.getBuffers().front()) + .allocatedPtr(rewriter, loc); + if (!getTypeConverter()->useOpaquePointers()) + pBuf = rewriter.create(loc, llvmPointerType, pBuf); + spMMCallBuilder.create(loc, rewriter, + {adaptor.getEnv(), modeA, modeB, adaptor.getSpmatA(), + adaptor.getDnmatB(), adaptor.getDnmatC(), dw, pBuf, + stream}); + } rewriter.replaceOp(op, {stream}); return success(); } @@ -1615,6 +1805,7 @@ ConvertDestroyDnMatOpToGpuRuntimeCallPattern, ConvertCreateCooOpToGpuRuntimeCallPattern, ConvertCreateCsrOpToGpuRuntimeCallPattern, + ConvertCreate2To4SpMatOpToGpuRuntimeCallPattern, ConvertDestroySpMatOpToGpuRuntimeCallPattern, ConvertSpMVBufferSizeOpToGpuRuntimeCallPattern, ConvertSpMVOpToGpuRuntimeCallPattern, 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,148 @@ matB, betap, matC, dtp, CUSPARSE_SDDMM_ALG_DEFAULT, buf)) } + +/// +/// Wrapper methods for the cuSparseLt library. +/// +#if MLIR_CUDA_CUSPARSELT_ENABLED +struct cusparseLtSpMatHandleAndData { + cusparseLtMatDescriptor_t mat; + void *values{nullptr}; + // TODO: the following is associated with the SpMM operator rather than the + // sparse matrix. Create workspace buffers and pass them to the SpMM + // execution. + cusparseLtMatmulAlgSelection_t alg_sel; + cusparseLtMatmulPlan_t plan; + cusparseLtMatmulDescriptor_t matmul; +}; +struct cusparseLtDnMatHandleAndData { + cusparseLtMatDescriptor_t mat; + void *values{nullptr}; +}; +struct cusparseLtWorkspaceSizes { + size_t workspace_size; + size_t compressed_size; + size_t compressed_buffer_size; +}; + +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)) +} + +extern "C" MLIR_CUDA_WRAPPERS_EXPORT void * +mgpuCreateCuSparseLtDnMat(void *h, intptr_t rows, intptr_t cols, void *values, + int32_t dw, CUstream /*stream*/) { + cusparseLtMatDescriptor_t mat; + auto handle = reinterpret_cast(h); + cudaDataType_t dtp = dataTp(dw); + // assuming row-major when deciding lda + CUSPARSE_REPORT_IF_ERROR( + cusparseLtDenseDescriptorInit(handle, &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 matAndData = reinterpret_cast(m); + CUSPARSE_REPORT_IF_ERROR(cusparseLtMatDescriptorDestroy(&(mat->mat))) + // destroy the plan associated with the sparse matrix + CUSPARSE_REPORT_IF_ERROR(cusparseLtMatmulPlanDestroy(&(mat->plan))) +} + +extern "C" MLIR_CUDA_WRAPPERS_EXPORT void +mgpuDestroyCuSparseLtDnMat(void *m, CUstream /*stream*/) { + auto matAndData = reinterpret_cast(m); + CUSPARSE_REPORT_IF_ERROR(cusparseLtMatDescriptorDestroy(&(mat->mat))) +} + +extern "C" MLIR_CUDA_WRAPPERS_EXPORT void * +mgpuCusparseLtCreate2To4SpMat(void *h, intptr_t rows, intptr_t cols, + void *values, int32_t dw, CUstream /*stream*/) { + cusparseLtSpMatHandleAndData matWithData; + matWithData.values = values; + auto handle = reinterpret_cast(h); + cudaDataType_t dtp = dataTp_cusparseLt(dw); + // assuming row-major when deciding lda + CUSPARSE_REPORT_IF_ERROR(cusparseLtStructuredDescriptorInit( + handle, &(matWithData.mat), rows, cols, /*ld=*/cols, /*alignment=*/16, + dtp, CUSPARSE_ORDER_ROW, CUSPARSELT_SPARSITY_50_PERCENT)) + + return reinterpret_cast(matWithData); +} + +// 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 void * +mgpuCuSparseLtSpMMBufferSize(void *h, void *a, CUstream /*stream*/) { + // TODO: support more advanced settings, e.g., the input right operand is a + // sparse matrix assuming matA is the sparse matrix + auto handle = reinterpret_cast(h); + auto matA = reinterpret_cast(a); + + CHECK_CUSPARSE(cusparseLtMatmulAlgSelectionInit( + handle, &(matWithData.alg_sel), &matmul, CUSPARSELT_MATMUL_ALG_DEFAULT)) + int alg = 0; + CHECK_CUSPARSE(cusparseLtMatmulAlgSetAttribute( + handle, &(matWithData.alg_sel), CUSPARSELT_MATMUL_ALG_CONFIG_ID, &alg, + sizeof(alg))) + // TODO: add transpose support + CHECK_CUSPARSE(cusparseLtMatmulDescriptorInit( + handle, &(matA.matmul), c, CUSPARSE_OPERATION_NON_TRANSPOSE, &(matA->mat), + &matB, &matC, &matC, compute_type)) + CHECK_CUSPARSE(cusparseLtMatmulPlanInit(handle, &(matWithData.plan), &matmul, + &(matWithData.alg_sel))) + + CHECK_CUSPARSE(cusparseLtMatmulGetWorkspace(handle, &(matA.plan), + &(sizes.workspace_size))) + CHECK_CUSPARSE(cusparseLtSpMMACompressedSize(handle, &(matA.plan), + &(sizes.compressed_size), + &(sizes.compressed_buffer_size))) + // avoid zero-alloc + sizes.workspace_size = (sizes.workspace_size == 0 ? 1 : sizes.workspace_size); + sizes.compressed_size = + (sizes.compressed_size == 0 ? 1 : sizes.compressed_size); + sizes.compressed_buffer_size = + (sizes.compressed_buffer_size == 0 ? 1 : sizes.compressed_buffer_size); + return reinterpret_cast(sizes); +} + +extern "C" MLIR_CUDA_WRAPPERS_EXPORT void +mgpuCuSparseLtSpMM(void *h, void *a, void *b, void *c, int32_t dw, void *buf, + void *dA_compressed, void *dA_compressedBuffer, + CUstream stream) { + auto handle = reinterpret_cast(h); + auto matA = reinterpret_cast(a); + auto matB = reinterpret_cast(b); + auto matC = reinterpret_cast(c); + ALPHABETA(dw, alpha, beta) + + CHECK_CUSPARSE(cusparseLtSpMMACompress(handle, &(matA->plan), &(matA->values), + dA_compressed, dA_compressedBuffer, + stream)) + + // TODO: add support to multi-stream execution + // Perform the matrix multiplication. D = A*B+C using C==D for now + CHECK_CUSPARSE(cusparseLtMatmul(handle, &(matA->plan), &alpha, dA_compressed, + dB, &beta, matC->values, /*dD*/ matC->values, + d_workspace, &stream, 1)) +} + +#endif // MLIR_CUDA_CUSPARSELT_ENABLED \ No newline at end of file diff --git a/mlir/test/CMakeLists.txt b/mlir/test/CMakeLists.txt --- a/mlir/test/CMakeLists.txt +++ b/mlir/test/CMakeLists.txt @@ -28,6 +28,7 @@ option(MLIR_RUN_X86VECTOR_TESTS "Run X86Vector tests.") option(MLIR_RUN_CUDA_TENSOR_CORE_TESTS "Run CUDA Tensor core WMMA tests.") option(MLIR_RUN_CUDA_SM80_TESTS "Run CUDA A100 tests.") + option(MLIR_RUN_CUDA_SPARSE_2TO4_TESTS "run CUDA sparse 2to4 tests") option(MLIR_RUN_ARM_SVE_TESTS "Run Arm SVE tests.") option(MLIR_RUN_ARM_SME_TESTS "Run Arm SME tests.") diff --git a/mlir/test/Conversion/GPUCommon/2To4Sparsity/lit.local.cfg b/mlir/test/Conversion/GPUCommon/2To4Sparsity/lit.local.cfg new file mode 100644 --- /dev/null +++ b/mlir/test/Conversion/GPUCommon/2To4Sparsity/lit.local.cfg @@ -0,0 +1,2 @@ +if not config.enable_cuda_runner or not config.mlir_run_cuda_sm80_tests or not mlir_run_cuda_sparse_2to4_tests: + config.unsupported = True diff --git a/mlir/test/Conversion/GPUCommon/2To4Sparsity/lower-2to4-sparse-to-gpu-runtime-calls.mlir b/mlir/test/Conversion/GPUCommon/2To4Sparsity/lower-2to4-sparse-to-gpu-runtime-calls.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Conversion/GPUCommon/2To4Sparsity/lower-2to4-sparse-to-gpu-runtime-calls.mlir @@ -0,0 +1,35 @@ +// RUN: mlir-opt %s --gpu-to-llvm='use-opaque-pointers=1' | FileCheck %s + +module attributes {gpu.container_module} { + + // CHECK-LABEL: func @matmul + // CHECK: llvm.call @mgpuStreamCreate + // CHECK: llvm.call @mgpuMemAlloc + // CHECK: llvm.call @mgpuMemAlloc + // CHECK: llvm.call @mgpuCreateSparseLtEnv + // CHECK: llvm.call @mgpuDestroyCuSparseLtSpMat + // CHECK: llvm.call @mgpuCreateCuSparseLtDnMat + // CHECK: llvm.call @mgpuCuSparseLtSpMMBufferSize + // CHECK: llvm.call @mgpuCuSparseLtSpMM + // CHECK: llvm.call @mgpuDestroyCuSparseLtSpMat + // CHECK: llvm.call @mgpuDestroyCuSparseLtDnMat + // CHECK: llvm.call @mgpuDestroySparseLtEnv + // CHECK: llvm.call @mgpuStreamSynchronize + // CHECK: llvm.call @mgpuStreamDestroy + func.func @matmul(%arg0: index) { + %token0 = gpu.wait async + %mem1, %token1 = gpu.alloc async [%token0] (%arg0) : memref + %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref + %env, %token3 = gpu.create_sparse_env async [%token2] + %spmat, %token4 = gpu.create_2to4_spmat async [%token3] %env, %arg0, %arg0, %mem1: memref + %dnmat, %token5 = gpu.create_dn_mat async [%token4] %arg0, %arg0, %mem2 : memref + %bufferSzs, %token6 = gpu.spmm_buffer_size async [%token5] %env, %spmat, %dnmat, %dnmat : tuple + %token7 = gpu.spmm async [%token6] %env, %spmat, %dnmat, %dnmat, <%mem2,%mem2,%mem2> : memref + %token8 = gpu.destroy_sp_mat async [%token7] %spmat + %token9 = gpu.destroy_dn_mat async [%token8] %dnmat + %token10 = gpu.destroy_sparse_env async [%token9] %env + gpu.wait [%token10] + return + } + +} 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]] diff --git a/mlir/test/lit.site.cfg.py.in b/mlir/test/lit.site.cfg.py.in --- a/mlir/test/lit.site.cfg.py.in +++ b/mlir/test/lit.site.cfg.py.in @@ -41,6 +41,7 @@ config.mlir_run_riscv_vector_tests = "@MLIR_RUN_RISCV_VECTOR_TESTS@" config.mlir_run_cuda_tensor_core_tests = @MLIR_RUN_CUDA_TENSOR_CORE_TESTS@ config.mlir_run_cuda_sm80_tests = @MLIR_RUN_CUDA_SM80_TESTS@ +config.mlir_run_cuda_sparse_2to4_tests =@MLIR_RUN_CUDA_SPARSE_2TO4_TESTS@ config.mlir_include_integration_tests = @MLIR_INCLUDE_INTEGRATION_TESTS@ config.arm_emulator_executable = "@ARM_EMULATOR_EXECUTABLE@" config.arm_emulator_options = "@ARM_EMULATOR_OPTIONS@"