| //===- CudaRuntimeWrappers.cpp - MLIR CUDA API wrapper library ------------===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // Implements C wrappers around the CUDA library for easy linking in ORC jit. |
| // Also adds some debugging helpers that are helpful when writing MLIR code to |
| // run on GPUs. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "mlir/ExecutionEngine/CRunnerUtils.h" |
| |
| #include <stdio.h> |
| |
| #include "cuda.h" |
| #include "cuda_bf16.h" |
| #include "cuda_fp16.h" |
| |
| #ifdef MLIR_ENABLE_CUDA_CUSPARSE |
| #include "cusparse.h" |
| #ifdef MLIR_ENABLE_CUDA_CUSPARSELT |
| #include "cusparseLt.h" |
| #endif // MLIR_ENABLE_CUDA_CUSPARSELT |
| #endif // MLIR_ENABLE_CUDA_CUSPARSE |
| |
| #ifdef _WIN32 |
| #include <malloc.h> |
| #define MLIR_CUDA_WRAPPERS_EXPORT __declspec(dllexport) |
| #else |
| #define MLIR_CUDA_WRAPPERS_EXPORT __attribute__((visibility("default"))) |
| #endif // _WIN32 |
| |
| #define CUDA_REPORT_IF_ERROR(expr) \ |
| [](CUresult result) { \ |
| if (!result) \ |
| return; \ |
| const char *name = nullptr; \ |
| cuGetErrorName(result, &name); \ |
| if (!name) \ |
| name = "<unknown>"; \ |
| fprintf(stderr, "'%s' failed with '%s'\n", #expr, name); \ |
| }(expr) |
| |
| #define CUSPARSE_REPORT_IF_ERROR(expr) \ |
| { \ |
| cusparseStatus_t status = (expr); \ |
| if (status != CUSPARSE_STATUS_SUCCESS) { \ |
| fprintf(stderr, "cuSPARSE '%s' failed with '%s'\n", #expr, \ |
| cusparseGetErrorString(status)); \ |
| } \ |
| } |
| |
| thread_local static int32_t defaultDevice = 0; |
| |
| const char *kDebugEnvironmentVariable = "MLIR_CUDA_DEBUG"; |
| |
| /// Helper method that checks environment value for debugging. |
| bool isDebugEnabled() { |
| static bool isInitialized = false; |
| static bool isEnabled = false; |
| if (!isInitialized) |
| isEnabled = getenv(kDebugEnvironmentVariable) != nullptr; |
| return isEnabled; |
| } |
| |
| #define debug_print(fmt, ...) \ |
| do { \ |
| if (isDebugEnabled()) \ |
| fprintf(stderr, "%s:%d:%s(): " fmt, "CudaRuntimeWrappers.cpp", __LINE__, \ |
| __func__, __VA_ARGS__); \ |
| } while (0) |
| |
| // Returns default CUdevice |
| CUdevice getDefaultCuDevice() { |
| CUdevice device; |
| CUDA_REPORT_IF_ERROR(cuDeviceGet(&device, /*ordinal=*/defaultDevice)); |
| return device; |
| } |
| |
| // Make the primary context of the current default device current for the |
| // duration |
| // of the instance and restore the previous context on destruction. |
| class ScopedContext { |
| public: |
| ScopedContext() { |
| // Static reference to CUDA primary context for device ordinal |
| // defaultDevice. |
| static CUcontext context = [] { |
| CUDA_REPORT_IF_ERROR(cuInit(/*flags=*/0)); |
| CUcontext ctx; |
| // Note: this does not affect the current context. |
| CUDA_REPORT_IF_ERROR( |
| cuDevicePrimaryCtxRetain(&ctx, getDefaultCuDevice())); |
| return ctx; |
| }(); |
| |
| CUDA_REPORT_IF_ERROR(cuCtxPushCurrent(context)); |
| } |
| |
| ~ScopedContext() { CUDA_REPORT_IF_ERROR(cuCtxPopCurrent(nullptr)); } |
| }; |
| |
| #ifdef MLIR_ENABLE_CUDA_CUSPARSE |
| // Note that (1) Nvidia confirms the safety to share handle across multiple |
| // instances, and streams. (2) Clients are responsible to call the @mgpu |
| // environment initialization/destruction in a thread-safe manner, e.g., |
| // at the beginning of the program before multi-threads are created. |
| static cusparseHandle_t cusparse_env = nullptr; |
| |
| #ifdef MLIR_ENABLE_CUDA_CUSPARSELT |
| // cusparseLtHandle_t is not a pointer type, so we need an additional flag to |
| // indicate whether it is initialized. |
| static cusparseLtHandle_t cusparseLt_env; |
| static bool cusparseLt_initiated = false; |
| |
| #endif // MLIR_ENABLE_CUDA_CUSPARSELT |
| #endif // MLIR_ENABLE_CUDA_CUSPARSE |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT CUmodule |
| mgpuModuleLoad(void *data, size_t /*gpuBlobSize*/) { |
| ScopedContext scopedContext; |
| CUmodule module = nullptr; |
| CUDA_REPORT_IF_ERROR(cuModuleLoadData(&module, data)); |
| return module; |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT CUmodule mgpuModuleLoadJIT(void *data, |
| int optLevel) { |
| ScopedContext scopedContext; |
| CUmodule module = nullptr; |
| char jitErrorBuffer[4096] = {0}; |
| CUjit_option jitOptions[] = {CU_JIT_ERROR_LOG_BUFFER, |
| CU_JIT_ERROR_LOG_BUFFER_SIZE_BYTES, |
| CU_JIT_OPTIMIZATION_LEVEL}; |
| void *jitOptionsVals[] = {jitErrorBuffer, |
| reinterpret_cast<void *>(sizeof(jitErrorBuffer)), |
| reinterpret_cast<void *>(optLevel)}; |
| |
| CUresult result = |
| cuModuleLoadDataEx(&module, data, 3, jitOptions, jitOptionsVals); |
| if (result) { |
| fprintf(stderr, "JIT compilation failed with: '%s'\n", jitErrorBuffer); |
| CUDA_REPORT_IF_ERROR(result); |
| } |
| return module; |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuModuleUnload(CUmodule module) { |
| CUDA_REPORT_IF_ERROR(cuModuleUnload(module)); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT CUfunction |
| mgpuModuleGetFunction(CUmodule module, const char *name) { |
| CUfunction function = nullptr; |
| CUDA_REPORT_IF_ERROR(cuModuleGetFunction(&function, module, name)); |
| return function; |
| } |
| |
| // The wrapper uses intptr_t instead of CUDA's unsigned int to match |
| // the type of MLIR's index type. This avoids the need for casts in the |
| // generated MLIR code. |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuLaunchKernel(CUfunction function, intptr_t gridX, intptr_t gridY, |
| intptr_t gridZ, intptr_t blockX, intptr_t blockY, |
| intptr_t blockZ, int32_t smem, CUstream stream, void **params, |
| void **extra, size_t /*paramsCount*/) { |
| ScopedContext scopedContext; |
| if (smem > 0) { |
| // Avoid checking driver as it's more expensive than if statement |
| int32_t maxShmem = 0; |
| CUdevice device = getDefaultCuDevice(); |
| CUDA_REPORT_IF_ERROR(cuDeviceGet(&device, /*ordinal=*/defaultDevice)); |
| CUDA_REPORT_IF_ERROR(cuDeviceGetAttribute( |
| &maxShmem, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, |
| device)); |
| if (maxShmem < smem) { |
| fprintf(stderr, |
| "Requested shared memory (%dkb) is larger than maximum allowed " |
| "shared memory (%dkb) for this device\n", |
| smem, maxShmem); |
| } |
| CUDA_REPORT_IF_ERROR(cuFuncSetAttribute( |
| function, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem)); |
| } |
| debug_print("Launching kernel, grid=%ld,%ld,%ld, " |
| "threads: %ld, %ld, %ld, " |
| "smem: %dkb\n", |
| gridX, gridY, gridZ, blockX, blockY, blockZ, smem); |
| CUDA_REPORT_IF_ERROR(cuLaunchKernel(function, gridX, gridY, gridZ, blockX, |
| blockY, blockZ, smem, stream, params, |
| extra)); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT CUstream mgpuStreamCreate() { |
| ScopedContext scopedContext; |
| CUstream stream = nullptr; |
| CUDA_REPORT_IF_ERROR(cuStreamCreate(&stream, CU_STREAM_NON_BLOCKING)); |
| return stream; |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuStreamDestroy(CUstream stream) { |
| CUDA_REPORT_IF_ERROR(cuStreamDestroy(stream)); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuStreamSynchronize(CUstream stream) { |
| CUDA_REPORT_IF_ERROR(cuStreamSynchronize(stream)); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuStreamWaitEvent(CUstream stream, |
| CUevent event) { |
| CUDA_REPORT_IF_ERROR(cuStreamWaitEvent(stream, event, /*flags=*/0)); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT CUevent mgpuEventCreate() { |
| ScopedContext scopedContext; |
| CUevent event = nullptr; |
| CUDA_REPORT_IF_ERROR(cuEventCreate(&event, CU_EVENT_DISABLE_TIMING)); |
| return event; |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuEventDestroy(CUevent event) { |
| CUDA_REPORT_IF_ERROR(cuEventDestroy(event)); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuEventSynchronize(CUevent event) { |
| CUDA_REPORT_IF_ERROR(cuEventSynchronize(event)); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuEventRecord(CUevent event, |
| CUstream stream) { |
| CUDA_REPORT_IF_ERROR(cuEventRecord(event, stream)); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void * |
| mgpuMemAlloc(uint64_t sizeBytes, CUstream /*stream*/, bool /*isHostShared*/) { |
| ScopedContext scopedContext; |
| CUdeviceptr ptr = 0; |
| if (sizeBytes != 0) |
| CUDA_REPORT_IF_ERROR(cuMemAlloc(&ptr, sizeBytes)); |
| return reinterpret_cast<void *>(ptr); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuMemFree(void *ptr, |
| CUstream /*stream*/) { |
| CUDA_REPORT_IF_ERROR(cuMemFree(reinterpret_cast<CUdeviceptr>(ptr))); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuMemcpy(void *dst, void *src, size_t sizeBytes, CUstream stream) { |
| CUDA_REPORT_IF_ERROR(cuMemcpyAsync(reinterpret_cast<CUdeviceptr>(dst), |
| reinterpret_cast<CUdeviceptr>(src), |
| sizeBytes, stream)); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuMemset32(void *dst, unsigned int value, size_t count, CUstream stream) { |
| CUDA_REPORT_IF_ERROR(cuMemsetD32Async(reinterpret_cast<CUdeviceptr>(dst), |
| value, count, stream)); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuMemset16(void *dst, unsigned short value, size_t count, CUstream stream) { |
| CUDA_REPORT_IF_ERROR(cuMemsetD16Async(reinterpret_cast<CUdeviceptr>(dst), |
| value, count, stream)); |
| } |
| |
| /// |
| /// Helper functions for writing mlir example code |
| /// |
| |
| // Allows to register byte array with the CUDA runtime. Helpful until we have |
| // transfer functions implemented. |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) { |
| ScopedContext scopedContext; |
| CUDA_REPORT_IF_ERROR(cuMemHostRegister(ptr, sizeBytes, /*flags=*/0)); |
| } |
| |
| /// Registers a memref with the CUDA runtime. `descriptor` is a pointer to a |
| /// ranked memref descriptor struct of rank `rank`. Helpful until we have |
| /// transfer functions implemented. |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuMemHostRegisterMemRef(int64_t rank, StridedMemRefType<char, 1> *descriptor, |
| int64_t elementSizeBytes) { |
| // Only densely packed tensors are currently supported. |
| #ifdef _WIN32 |
| int64_t *denseStrides = (int64_t *)_alloca(rank * sizeof(int64_t)); |
| #else |
| int64_t *denseStrides = (int64_t *)alloca(rank * sizeof(int64_t)); |
| #endif // _WIN32 |
| int64_t *sizes = descriptor->sizes; |
| for (int64_t i = rank - 1, runningStride = 1; i >= 0; i--) { |
| denseStrides[i] = runningStride; |
| runningStride *= sizes[i]; |
| } |
| uint64_t sizeBytes = sizes[0] * denseStrides[0] * elementSizeBytes; |
| int64_t *strides = &sizes[rank]; |
| (void)strides; |
| for (unsigned i = 0; i < rank; ++i) |
| assert(strides[i] == denseStrides[i] && |
| "Mismatch in computed dense strides"); |
| |
| auto *ptr = descriptor->data + descriptor->offset * elementSizeBytes; |
| mgpuMemHostRegister(ptr, sizeBytes); |
| } |
| |
| // Allows to unregister byte array with the CUDA runtime. |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuMemHostUnregister(void *ptr) { |
| ScopedContext scopedContext; |
| CUDA_REPORT_IF_ERROR(cuMemHostUnregister(ptr)); |
| } |
| |
| /// Unregisters a memref with the CUDA runtime. `descriptor` is a pointer to a |
| /// ranked memref descriptor struct of rank `rank` |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuMemHostUnregisterMemRef(int64_t rank, |
| StridedMemRefType<char, 1> *descriptor, |
| int64_t elementSizeBytes) { |
| auto *ptr = descriptor->data + descriptor->offset * elementSizeBytes; |
| mgpuMemHostUnregister(ptr); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSetDefaultDevice(int32_t device) { |
| defaultDevice = device; |
| } |
| |
| /// |
| /// Runtime methods using CUDA 12.0+ driver |
| /// |
| |
| #if (CUDA_VERSION >= 12000) |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuLaunchClusterKernel( |
| CUfunction function, intptr_t clusterX, intptr_t clusterY, |
| intptr_t clusterZ, intptr_t gridX, intptr_t gridY, intptr_t gridZ, |
| intptr_t blockX, intptr_t blockY, intptr_t blockZ, int32_t smem, |
| CUstream stream, void **params, void **extra, size_t /*paramsCount*/) { |
| ScopedContext scopedContext; |
| if (smem > 0) { |
| // Avoid checking driver as it's more expensive than if statement |
| int32_t maxShmem = 0; |
| CUdevice device = getDefaultCuDevice(); |
| CUDA_REPORT_IF_ERROR(cuDeviceGet(&device, /*ordinal=*/defaultDevice)); |
| CUDA_REPORT_IF_ERROR(cuDeviceGetAttribute( |
| &maxShmem, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, |
| device)); |
| if (maxShmem < smem) { |
| fprintf(stderr, |
| "Requested shared memory (%dkb) is larger than maximum allowed " |
| "shared memory (%dkb) for this device\n", |
| smem, maxShmem); |
| } |
| CUDA_REPORT_IF_ERROR(cuFuncSetAttribute( |
| function, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem)); |
| } |
| CUlaunchConfig config; |
| config.gridDimX = gridX; |
| config.gridDimY = gridY; |
| config.gridDimZ = gridZ; |
| config.blockDimX = blockX; |
| config.blockDimY = blockY; |
| config.blockDimZ = blockZ; |
| config.sharedMemBytes = smem; |
| config.hStream = stream; |
| CUlaunchAttribute launchAttr[2]; |
| launchAttr[0].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; |
| launchAttr[0].value.clusterDim.x = clusterX; |
| launchAttr[0].value.clusterDim.y = clusterY; |
| launchAttr[0].value.clusterDim.z = clusterZ; |
| launchAttr[1].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE; |
| launchAttr[1].value.clusterSchedulingPolicyPreference = |
| CU_CLUSTER_SCHEDULING_POLICY_SPREAD; |
| config.numAttrs = 2; |
| config.attrs = launchAttr; |
| |
| debug_print("Launching kernel," |
| "cluster: %ld, %ld, %ld, " |
| "grid=%ld,%ld,%ld, " |
| "threads: %ld, %ld, %ld, " |
| "smem: %dkb\n", |
| clusterX, clusterY, clusterZ, gridX, gridY, gridZ, blockX, blockY, |
| blockZ, smem); |
| |
| CUDA_REPORT_IF_ERROR(cuLaunchKernelEx(&config, function, params, extra)); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuTensorMapEncodeTiled( |
| CUtensorMap *tensorMap, // Tensor map object |
| CUtensorMapDataType tensorDataType, // Tensor data type |
| cuuint32_t tensorRank, // Dimensionality of tensor |
| void *globalAddress, // Starting address |
| const cuuint64_t *globalDim, // Tensor size (number of elements) |
| const cuuint64_t *globalStrides, // Stride size (in bytes) |
| const cuuint32_t *boxDim, // Traversal box (number of elments) |
| const cuuint32_t *elementStrides, // Traversal stride |
| CUtensorMapInterleave interleave, // Type of interleaved layout |
| CUtensorMapSwizzle swizzle, // Bank swizzling pattern |
| CUtensorMapL2promotion l2Promotion, // L2 promotion size |
| CUtensorMapFloatOOBfill oobFill // Padding zfill or NaN fill |
| ) { |
| ScopedContext scopedContext; |
| CUDA_REPORT_IF_ERROR(cuTensorMapEncodeTiled( |
| tensorMap, tensorDataType, tensorRank, globalAddress, globalDim, |
| globalStrides, boxDim, elementStrides, interleave, swizzle, l2Promotion, |
| oobFill)); |
| debug_print("Created TMA descriptor\n Addr: %p\n" |
| "data type : %d\n" |
| "rank : %d\n" |
| "globalDim[5]: %zu, %zu, %zu, %zu, %zu\n" |
| "globalStrides[5]: %zu, %zu, %zu, %zu, %zu\n" |
| "boxDim[5]: %u, %u, %u, %u, %u\n" |
| "elementStrides[5]: %u, %u, %u, %u, %u\n" |
| "interleave: %u \n" |
| "swizzle: %u \n" |
| "l2Promotion: %u \n" |
| "oobFill: %u \n", |
| (void *)&tensorMap, tensorDataType, tensorRank, globalDim[0], |
| globalDim[1], globalDim[2], globalDim[3], globalDim[4], |
| globalStrides[0], globalStrides[1], globalStrides[2], |
| globalStrides[3], globalStrides[4], boxDim[0], boxDim[1], |
| boxDim[2], boxDim[3], boxDim[4], elementStrides[0], |
| elementStrides[1], elementStrides[2], elementStrides[3], |
| elementStrides[4], interleave, swizzle, l2Promotion, oobFill); |
| } |
| |
| template <int Rank> |
| void mgpuGetMemRefDataAndShape(void *rawDescriptor, char **addr, |
| uint64_t *globalDim, uint64_t *globalStrides, |
| const CUtensorMapDataType tensorDataType) { |
| auto descriptor = |
| reinterpret_cast<StridedMemRefType<char, Rank> *>(rawDescriptor); |
| *addr = descriptor->data; |
| for (int i = 0; i < Rank; ++i) { |
| globalDim[i] = static_cast<uint64_t>(descriptor->sizes[Rank - i - 1]); |
| } |
| static constexpr int elementSizeInBytes[] = {1, 2, 4, 4, 8, 8, 2, |
| 4, 8, 2, 4, 4, 4}; |
| for (int i = 0; i < Rank - 1; ++i) { |
| globalStrides[i] = static_cast<uint64_t>( |
| descriptor->strides[Rank - i - 2] * elementSizeInBytes[tensorDataType]); |
| } |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void *mgpuTensorMapEncodeTiledMemref( |
| int64_t tensorRank, // Dimensionality of tensor |
| void *rankedDescriptor, // Ranked MemRef descriptor |
| const CUtensorMapDataType tensorDataType, // Stride size (in bytes) |
| CUtensorMapInterleave interleave, // Type of interleaved layout |
| CUtensorMapSwizzle swizzle, // Bank swizzling pattern |
| CUtensorMapL2promotion l2Promotion, // L2 promotion size |
| CUtensorMapFloatOOBfill oobFill, // Padding zfill or NaN fill |
| int64_t *inputBoxDims // Tensor size (number of elements) |
| ) { |
| CUtensorMap tensorMap; |
| |
| uint32_t boxDim[5] = {1, 1, 1, 1, 1}, elementStrides[5] = {1, 1, 1, 1, 1}; |
| uint64_t globalDim[5] = {1, 1, 1, 1, 1}, globalStrides[5] = {0}; |
| uint32_t tensorRank32 = uint32_t(tensorRank); |
| |
| char *globalAddress = nullptr; |
| switch (tensorRank) { |
| case 1: |
| mgpuGetMemRefDataAndShape<1>(rankedDescriptor, &globalAddress, globalDim, |
| globalStrides, tensorDataType); |
| break; |
| case 2: |
| mgpuGetMemRefDataAndShape<2>(rankedDescriptor, &globalAddress, globalDim, |
| globalStrides, tensorDataType); |
| break; |
| case 3: |
| mgpuGetMemRefDataAndShape<3>(rankedDescriptor, &globalAddress, globalDim, |
| globalStrides, tensorDataType); |
| break; |
| case 4: |
| mgpuGetMemRefDataAndShape<4>(rankedDescriptor, &globalAddress, globalDim, |
| globalStrides, tensorDataType); |
| break; |
| case 5: |
| mgpuGetMemRefDataAndShape<5>(rankedDescriptor, &globalAddress, globalDim, |
| globalStrides, tensorDataType); |
| break; |
| default: |
| fprintf( |
| stderr, |
| "'mgpuTensorMapEncodeTiledMemref' failed with 'rank is too high'\n"); |
| return nullptr; |
| } |
| |
| for (int64_t r = 0; r < tensorRank; ++r) { |
| boxDim[r] = static_cast<uint32_t>(inputBoxDims[tensorRank - r - 1]); |
| } |
| |
| ScopedContext scopedContext; |
| mgpuTensorMapEncodeTiled(&tensorMap, tensorDataType, tensorRank32, |
| globalAddress, globalDim, globalStrides, boxDim, |
| elementStrides, interleave, swizzle, l2Promotion, |
| oobFill); |
| // Copy created tensor map to device |
| CUdeviceptr dTensorMap; |
| CUDA_REPORT_IF_ERROR(cuMemAlloc(&dTensorMap, sizeof(CUtensorMap))); |
| CUDA_REPORT_IF_ERROR(cuMemcpy(dTensorMap, |
| reinterpret_cast<CUdeviceptr>(&tensorMap), |
| sizeof(CUtensorMap))); |
| return reinterpret_cast<void *>(dTensorMap); |
| } |
| #endif |
| |
| #ifdef MLIR_ENABLE_CUDA_CUSPARSE |
| |
| /// |
| /// Wrapper methods for the cuSparse library. |
| /// |
| |
| // Some macro magic to get float/double alpha and beta on host. |
| // TODO: add support to passing alpha and beta as arguments |
| #define ALPHABETA(dtp, alpha, beta) \ |
| __nv_bfloat16(alpha##16bf) = 1.0f; \ |
| __nv_bfloat16(beta##16bf) = 1.0f; \ |
| __half(alpha##16f) = 1.0f; \ |
| __half(beta##16f) = 1.0f; \ |
| float(alpha##f) = 1.0f; \ |
| float(beta##f) = 1.0f; \ |
| double(alpha##d) = 1.0; \ |
| double(beta##d) = 1.0; \ |
| const void *(alpha##p) = nullptr; \ |
| const void *(beta##p) = nullptr; \ |
| if (dtp == CUDA_R_16BF || dtp == CUDA_C_16BF) { \ |
| (alpha##p) = reinterpret_cast<void *>(&(alpha##16bf)); \ |
| (beta##p) = reinterpret_cast<void *>(&(beta##16bf)); \ |
| } else if (dtp == CUDA_R_16F || dtp == CUDA_C_16F) { \ |
| (alpha##p) = reinterpret_cast<void *>(&(alpha##16f)); \ |
| (beta##p) = reinterpret_cast<void *>(&(beta##16f)); \ |
| } else if (dtp == CUDA_R_32F || dtp == CUDA_C_32F) { \ |
| (alpha##p) = reinterpret_cast<void *>(&(alpha##f)); \ |
| (beta##p) = reinterpret_cast<void *>(&(beta##f)); \ |
| } else { \ |
| (alpha##p) = reinterpret_cast<void *>(&(alpha##d)); \ |
| (beta##p) = reinterpret_cast<void *>(&(beta##d)); \ |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuCreateSparseEnv() { |
| // ScopedContext is for cuda initialization. |
| ScopedContext scopedContext; |
| assert(!cusparse_env && "client called mgpuCreateSparseEnv() twice"); |
| CUSPARSE_REPORT_IF_ERROR(cusparseCreate(&cusparse_env)); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuDestroySparseEnv() { |
| assert(cusparse_env && "client did not call mgpuCreateSparseEnv()"); |
| CUSPARSE_REPORT_IF_ERROR(cusparseDestroy(cusparse_env)); |
| cusparse_env = nullptr; |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void * |
| mgpuCreateDnVec(intptr_t size, void *values, int32_t dtp, CUstream /*stream*/) { |
| cusparseDnVecDescr_t vec = nullptr; |
| auto dTp = static_cast<cudaDataType_t>(dtp); |
| CUSPARSE_REPORT_IF_ERROR(cusparseCreateDnVec(&vec, size, values, dTp)) |
| return reinterpret_cast<void *>(vec); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuDestroyDnVec(void *v, CUstream /*stream*/) { |
| cusparseDnVecDescr_t vec = reinterpret_cast<cusparseDnVecDescr_t>(v); |
| CUSPARSE_REPORT_IF_ERROR(cusparseDestroyDnVec(vec)) |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void * |
| mgpuCreateDnMat(intptr_t rows, intptr_t cols, void *values, int32_t dtp, |
| CUstream /*stream*/) { |
| cusparseDnMatDescr_t mat = nullptr; |
| auto dTp = static_cast<cudaDataType_t>(dtp); |
| CUSPARSE_REPORT_IF_ERROR(cusparseCreateDnMat(&mat, rows, cols, /*ld=*/cols, |
| values, dTp, CUSPARSE_ORDER_ROW)) |
| return reinterpret_cast<void *>(mat); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuDestroyDnMat(void *m, CUstream /*stream*/) { |
| cusparseDnMatDescr_t mat = reinterpret_cast<cusparseDnMatDescr_t>(m); |
| CUSPARSE_REPORT_IF_ERROR(cusparseDestroyDnMat(mat)) |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void * |
| mgpuCreateCoo(intptr_t rows, intptr_t cols, intptr_t nnz, void *rowIdxs, |
| void *colIdxs, void *values, int32_t itp, int32_t dtp, |
| CUstream /*stream*/) { |
| cusparseSpMatDescr_t mat = nullptr; |
| auto iTp = static_cast<cusparseIndexType_t>(itp); |
| auto dTp = static_cast<cudaDataType_t>(dtp); |
| CUSPARSE_REPORT_IF_ERROR(cusparseCreateCoo(&mat, rows, cols, nnz, rowIdxs, |
| colIdxs, values, iTp, |
| CUSPARSE_INDEX_BASE_ZERO, dTp)) |
| return reinterpret_cast<void *>(mat); |
| } |
| |
| #ifdef CUSPARSE_COO_AOS // deprecated in cuSPARSE 11.2 |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void * |
| mgpuCreateCooAoS(intptr_t rows, intptr_t cols, intptr_t nnz, void *idxs, |
| void *values, int32_t itp, int32_t dtp, CUstream /*stream*/) { |
| cusparseSpMatDescr_t mat = nullptr; |
| auto iTp = static_cast<cusparseIndexType_t>(itp); |
| auto dTp = static_cast<cudaDataType_t>(dtp); |
| CUSPARSE_REPORT_IF_ERROR(cusparseCreateCooAoS( |
| &mat, rows, cols, nnz, idxs, values, iTp, CUSPARSE_INDEX_BASE_ZERO, dTp)) |
| return reinterpret_cast<void *>(mat); |
| } |
| #endif // CUSPARSE_COO_AOS |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void * |
| mgpuCreateCsr(intptr_t rows, intptr_t cols, intptr_t nnz, void *rowPos, |
| void *colIdxs, void *values, int32_t ptp, int32_t itp, |
| int32_t dtp, CUstream /*stream*/) { |
| cusparseSpMatDescr_t mat = nullptr; |
| auto pTp = static_cast<cusparseIndexType_t>(ptp); |
| auto iTp = static_cast<cusparseIndexType_t>(itp); |
| auto dTp = static_cast<cudaDataType_t>(dtp); |
| CUSPARSE_REPORT_IF_ERROR(cusparseCreateCsr(&mat, rows, cols, nnz, rowPos, |
| colIdxs, values, pTp, iTp, |
| CUSPARSE_INDEX_BASE_ZERO, dTp)) |
| return reinterpret_cast<void *>(mat); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void * |
| mgpuCreateCsc(intptr_t rows, intptr_t cols, intptr_t nnz, void *colPos, |
| void *rowIdxs, void *values, int32_t ptp, int32_t itp, |
| int32_t dtp, CUstream /*stream*/) { |
| cusparseSpMatDescr_t mat = nullptr; |
| auto pTp = static_cast<cusparseIndexType_t>(ptp); |
| auto iTp = static_cast<cusparseIndexType_t>(itp); |
| auto dTp = static_cast<cudaDataType_t>(dtp); |
| CUSPARSE_REPORT_IF_ERROR(cusparseCreateCsc(&mat, rows, cols, nnz, colPos, |
| rowIdxs, values, pTp, iTp, |
| CUSPARSE_INDEX_BASE_ZERO, dTp)) |
| return reinterpret_cast<void *>(mat); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void * |
| mgpuCreateBsr(intptr_t brows, intptr_t bcols, intptr_t bnnz, intptr_t rBsz, |
| intptr_t cBsz, void *rowPos, void *colIdxs, void *values, |
| int32_t ptp, int32_t itp, int32_t dtp, CUstream /*stream*/) { |
| cusparseSpMatDescr_t mat = nullptr; |
| #if CUSPARSE_VERSION >= 12100 |
| auto pTp = static_cast<cusparseIndexType_t>(ptp); |
| auto iTp = static_cast<cusparseIndexType_t>(itp); |
| auto dTp = static_cast<cudaDataType_t>(dtp); |
| CUSPARSE_REPORT_IF_ERROR(cusparseCreateBsr( |
| &mat, brows, bcols, bnnz, rBsz, cBsz, rowPos, colIdxs, values, pTp, iTp, |
| CUSPARSE_INDEX_BASE_ZERO, dTp, CUSPARSE_ORDER_ROW)) |
| #endif |
| return reinterpret_cast<void *>(mat); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuDestroySpMat(void *m, CUstream /*stream*/) { |
| cusparseSpMatDescr_t mat = reinterpret_cast<cusparseSpMatDescr_t>(m); |
| CUSPARSE_REPORT_IF_ERROR(cusparseDestroySpMat(mat)) |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t mgpuSpMVBufferSize( |
| int32_t ma, void *a, void *x, void *y, int32_t ctp, CUstream /*stream*/) { |
| assert(cusparse_env && "client did not call mgpuCreateSparseEnv()"); |
| cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma); |
| cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a); |
| cusparseDnVecDescr_t vecX = reinterpret_cast<cusparseDnVecDescr_t>(x); |
| cusparseDnVecDescr_t vecY = reinterpret_cast<cusparseDnVecDescr_t>(y); |
| cudaDataType_t cTp = static_cast<cudaDataType_t>(ctp); |
| ALPHABETA(cTp, alpha, beta) |
| size_t bufferSize = 0; |
| CUSPARSE_REPORT_IF_ERROR(cusparseSpMV_bufferSize( |
| cusparse_env, modeA, alphap, matA, vecX, betap, vecY, cTp, |
| CUSPARSE_SPMV_ALG_DEFAULT, &bufferSize)) |
| return bufferSize; |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSpMV(int32_t ma, void *a, void *x, |
| void *y, int32_t ctp, |
| void *buf, |
| CUstream /*stream*/) { |
| assert(cusparse_env && "client did not call mgpuCreateSparseEnv()"); |
| cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma); |
| cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a); |
| cusparseDnVecDescr_t vecX = reinterpret_cast<cusparseDnVecDescr_t>(x); |
| cusparseDnVecDescr_t vecY = reinterpret_cast<cusparseDnVecDescr_t>(y); |
| cudaDataType_t cTp = static_cast<cudaDataType_t>(ctp); |
| ALPHABETA(cTp, alpha, beta) |
| CUSPARSE_REPORT_IF_ERROR(cusparseSpMV(cusparse_env, modeA, alphap, matA, vecX, |
| betap, vecY, cTp, |
| CUSPARSE_SPMV_ALG_DEFAULT, buf)) |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t |
| mgpuSpMMBufferSize(int32_t ma, int32_t mb, void *a, void *b, void *c, |
| int32_t ctp, CUstream /*stream*/) { |
| assert(cusparse_env && "client did not call mgpuCreateSparseEnv()"); |
| cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma); |
| cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb); |
| cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a); |
| cusparseDnMatDescr_t matB = reinterpret_cast<cusparseDnMatDescr_t>(b); |
| cusparseDnMatDescr_t matC = reinterpret_cast<cusparseDnMatDescr_t>(c); |
| cudaDataType_t cTp = static_cast<cudaDataType_t>(ctp); |
| ALPHABETA(cTp, alpha, beta) |
| size_t bufferSize = 0; |
| CUSPARSE_REPORT_IF_ERROR(cusparseSpMM_bufferSize( |
| cusparse_env, modeA, modeB, alphap, matA, matB, betap, matC, cTp, |
| CUSPARSE_SPMM_ALG_DEFAULT, &bufferSize)) |
| return bufferSize; |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSpMM(int32_t ma, int32_t mb, |
| void *a, void *b, void *c, |
| int32_t ctp, void *buf, |
| CUstream /*stream*/) { |
| assert(cusparse_env && "client did not call mgpuCreateSparseEnv()"); |
| cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma); |
| cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb); |
| cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a); |
| cusparseDnMatDescr_t matB = reinterpret_cast<cusparseDnMatDescr_t>(b); |
| cusparseDnMatDescr_t matC = reinterpret_cast<cusparseDnMatDescr_t>(c); |
| cudaDataType_t cTp = static_cast<cudaDataType_t>(ctp); |
| ALPHABETA(cTp, alpha, beta) |
| CUSPARSE_REPORT_IF_ERROR(cusparseSpMM(cusparse_env, modeA, modeB, alphap, |
| matA, matB, betap, matC, cTp, |
| CUSPARSE_SPMM_ALG_DEFAULT, buf)) |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t |
| mgpuSDDMMBufferSize(int32_t ma, int32_t mb, void *a, void *b, void *c, |
| int32_t ctp, CUstream /*stream*/) { |
| assert(cusparse_env && "client did not call mgpuCreateSparseEnv()"); |
| cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma); |
| cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb); |
| cusparseDnMatDescr_t matA = reinterpret_cast<cusparseDnMatDescr_t>(a); |
| cusparseDnMatDescr_t matB = reinterpret_cast<cusparseDnMatDescr_t>(b); |
| cusparseSpMatDescr_t matC = reinterpret_cast<cusparseSpMatDescr_t>(c); |
| auto cTp = static_cast<cudaDataType_t>(ctp); |
| ALPHABETA(cTp, alpha, beta) |
| size_t bufferSize = 0; |
| CUSPARSE_REPORT_IF_ERROR(cusparseSDDMM_bufferSize( |
| cusparse_env, modeA, modeB, alphap, matA, matB, betap, matC, cTp, |
| CUSPARSE_SDDMM_ALG_DEFAULT, &bufferSize)) |
| return bufferSize; |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSDDMM(int32_t ma, int32_t mb, |
| void *a, void *b, void *c, |
| int32_t ctp, void *buf, |
| CUstream /*stream*/) { |
| assert(cusparse_env && "client did not call mgpuCreateSparseEnv()"); |
| cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma); |
| cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb); |
| cusparseDnMatDescr_t matA = reinterpret_cast<cusparseDnMatDescr_t>(a); |
| cusparseDnMatDescr_t matB = reinterpret_cast<cusparseDnMatDescr_t>(b); |
| cusparseSpMatDescr_t matC = reinterpret_cast<cusparseSpMatDescr_t>(c); |
| auto cTp = static_cast<cudaDataType_t>(ctp); |
| ALPHABETA(cTp, alpha, beta) |
| CUSPARSE_REPORT_IF_ERROR(cusparseSDDMM(cusparse_env, modeA, modeB, alphap, |
| matA, matB, betap, matC, cTp, |
| CUSPARSE_SDDMM_ALG_DEFAULT, buf)) |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void * |
| mgpuSpGEMMCreateDescr(CUstream /*stream*/) { |
| cusparseSpGEMMDescr_t spgemmDesc = nullptr; |
| CUSPARSE_REPORT_IF_ERROR(cusparseSpGEMM_createDescr(&spgemmDesc)) |
| return reinterpret_cast<void *>(spgemmDesc); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuSpGEMMDestroyDescr(void *s, CUstream /*stream*/) { |
| cusparseSpGEMMDescr_t spgemmDesc = reinterpret_cast<cusparseSpGEMMDescr_t>(s); |
| CUSPARSE_REPORT_IF_ERROR(cusparseSpGEMM_destroyDescr(spgemmDesc)) |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t mgpuSpGEMMWorkEstimation( |
| void *s, int32_t ma, int32_t mb, void *a, void *b, void *c, int32_t ctp, |
| intptr_t bs, void *buf, CUstream /*stream*/) { |
| cusparseSpGEMMDescr_t spgemmDesc = reinterpret_cast<cusparseSpGEMMDescr_t>(s); |
| cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma); |
| cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb); |
| cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a); |
| cusparseSpMatDescr_t matB = reinterpret_cast<cusparseSpMatDescr_t>(b); |
| cusparseSpMatDescr_t matC = reinterpret_cast<cusparseSpMatDescr_t>(c); |
| auto cTp = static_cast<cudaDataType_t>(ctp); |
| ALPHABETA(cTp, alpha, beta) |
| size_t newBufferSize = bs; |
| CUSPARSE_REPORT_IF_ERROR(cusparseSpGEMM_workEstimation( |
| cusparse_env, modeA, modeB, alphap, matA, matB, betap, matC, cTp, |
| CUSPARSE_SPGEMM_DEFAULT, spgemmDesc, &newBufferSize, buf)) |
| return newBufferSize; |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t |
| mgpuSpGEMMCompute(void *s, int32_t ma, int32_t mb, void *a, void *b, void *c, |
| int32_t ctp, intptr_t bsz2, void *buf2, CUstream /*stream*/) { |
| cusparseSpGEMMDescr_t spgemmDesc = reinterpret_cast<cusparseSpGEMMDescr_t>(s); |
| cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma); |
| cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb); |
| cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a); |
| cusparseSpMatDescr_t matB = reinterpret_cast<cusparseSpMatDescr_t>(b); |
| cusparseSpMatDescr_t matC = reinterpret_cast<cusparseSpMatDescr_t>(c); |
| auto cTp = static_cast<cudaDataType_t>(ctp); |
| ALPHABETA(cTp, alpha, beta) |
| size_t newBufferSize2 = bsz2; |
| CUSPARSE_REPORT_IF_ERROR(cusparseSpGEMM_compute( |
| cusparse_env, modeA, modeB, alphap, matA, matB, betap, matC, cTp, |
| CUSPARSE_SPGEMM_DEFAULT, spgemmDesc, &newBufferSize2, buf2)) |
| return newBufferSize2; |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuSpGEMMCopy(void *s, int32_t ma, int32_t mb, void *a, void *b, void *c, |
| int32_t ctp, CUstream /*stream*/) { |
| cusparseSpGEMMDescr_t spgemmDesc = reinterpret_cast<cusparseSpGEMMDescr_t>(s); |
| cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma); |
| cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb); |
| cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a); |
| cusparseSpMatDescr_t matB = reinterpret_cast<cusparseSpMatDescr_t>(b); |
| cusparseSpMatDescr_t matC = reinterpret_cast<cusparseSpMatDescr_t>(c); |
| auto cTp = static_cast<cudaDataType_t>(ctp); |
| ALPHABETA(cTp, alpha, beta) |
| CUSPARSE_REPORT_IF_ERROR( |
| cusparseSpGEMM_copy(cusparse_env, modeA, modeB, alphap, matA, matB, betap, |
| matC, cTp, CUSPARSE_SPGEMM_DEFAULT, spgemmDesc)) |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuSpMatGetSize(void *m, void *r, void *c, void *n, CUstream /*stream*/) { |
| cusparseConstSpMatDescr_t matDescr = |
| reinterpret_cast<cusparseConstSpMatDescr_t>(m); |
| int64_t *rows = reinterpret_cast<int64_t *>(r); |
| int64_t *cols = reinterpret_cast<int64_t *>(c); |
| int64_t *nnz = reinterpret_cast<int64_t *>(n); |
| CUSPARSE_REPORT_IF_ERROR(cusparseSpMatGetSize(matDescr, rows, cols, nnz)); |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuSetCsrPointers(void *m, void *p, void *c, void *v, CUstream /*stream*/) { |
| cusparseSpMatDescr_t matDescr = reinterpret_cast<cusparseSpMatDescr_t>(m); |
| CUSPARSE_REPORT_IF_ERROR(cusparseCsrSetPointers(matDescr, p, c, v)); |
| } |
| |
| #ifdef MLIR_ENABLE_CUDA_CUSPARSELT |
| |
| /// |
| /// Wrapper methods for the cuSparseLt library. |
| /// |
| |
| struct cusparseLtSpMatHandleAndData { |
| cusparseLtMatDescriptor_t mat; |
| // TODO: the following three are 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; |
| void *values{nullptr}; |
| }; |
| |
| struct cusparseLtDnMatHandleAndData { |
| cusparseLtMatDescriptor_t mat; |
| void *values{nullptr}; |
| }; |
| |
| static_assert(sizeof(cusparseLtHandle_t) == 11024, |
| "Unexpected cusparseLt handle size"); |
| static_assert(sizeof(cusparseLtSpMatHandleAndData) == 44104, |
| "Unexpected cusparseLt sparse matrix handle size"); |
| static_assert(sizeof(cusparseLtDnMatHandleAndData) == 11032, |
| "Unexpected cusparseLt dense matrix handle size"); |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuCreateSparseLtEnv() { |
| // ScopedContext is for cuda initialization. |
| ScopedContext scopedContext; |
| assert(!cusparseLt_initiated && |
| "client called mgpuCreateSparseLtEnv() twice"); |
| // Note that cuSparseLt still uses cusparseStatus_t. |
| CUSPARSE_REPORT_IF_ERROR(cusparseLtInit(&cusparseLt_env)); |
| cusparseLt_initiated = true; |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuDestroySparseLtEnv() { |
| assert(cusparseLt_initiated && "client did not call mgpuCreateSparseLtEnv()"); |
| CUSPARSE_REPORT_IF_ERROR(cusparseLtDestroy(&cusparseLt_env)); |
| cusparseLt_initiated = false; |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuCreateCuSparseLtDnMat(void *dh, intptr_t rows, intptr_t cols, void *values, |
| int32_t dtp, CUstream /*stream*/) { |
| assert(cusparseLt_initiated && "client did not call mgpuCreateSparseLtEnv()"); |
| auto dnmat_handle = reinterpret_cast<cusparseLtDnMatHandleAndData *>(dh); |
| dnmat_handle->values = values; |
| auto dTp = static_cast<cudaDataType_t>(dtp); |
| // Assume row-major when deciding lda. |
| const uint32_t alignment = 16; |
| CUSPARSE_REPORT_IF_ERROR(cusparseLtDenseDescriptorInit( |
| &cusparseLt_env, &(dnmat_handle->mat), rows, cols, /*lda=*/cols, |
| alignment, dTp, CUSPARSE_ORDER_ROW)) |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuDestroyCuSparseLtDnMat(void *dh, CUstream /*stream*/) { |
| auto dnmat_handle = reinterpret_cast<cusparseLtDnMatHandleAndData *>(dh); |
| CUSPARSE_REPORT_IF_ERROR(cusparseLtMatDescriptorDestroy(&(dnmat_handle->mat))) |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuCusparseLtCreate2To4SpMat(void *sh, intptr_t rows, intptr_t cols, |
| void *values, int32_t dtp, CUstream /*stream*/) { |
| assert(cusparseLt_initiated && "client did not call mgpuCreateSparseLtEnv()"); |
| auto spmat_handle = reinterpret_cast<cusparseLtSpMatHandleAndData *>(sh); |
| spmat_handle->values = values; |
| auto dTp = static_cast<cudaDataType_t>(dtp); |
| // Assume row-major when deciding lda. |
| const uint32_t alignment = 16; |
| CUSPARSE_REPORT_IF_ERROR(cusparseLtStructuredDescriptorInit( |
| &cusparseLt_env, &(spmat_handle->mat), rows, cols, /*ld=*/cols, alignment, |
| dTp, CUSPARSE_ORDER_ROW, CUSPARSELT_SPARSITY_50_PERCENT)) |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuDestroyCuSparseLtSpMat(void *sh, CUstream /*stream*/) { |
| auto spmat_handle = reinterpret_cast<cusparseLtSpMatHandleAndData *>(sh); |
| CUSPARSE_REPORT_IF_ERROR(cusparseLtMatDescriptorDestroy(&(spmat_handle->mat))) |
| } |
| |
| // Several things are being done in this stage, algorithm selection, planning, |
| // and returning workspace and compressed matrices data buffer sizes. |
| // The parameter prune_flag is used to indicate whether pruning and pruning |
| // check will happen 0 means not prune or prune check, 1 means prune, 2 means |
| // prune & prune check |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuCuSparseLtSpMMBufferSize(void *bs, int32_t ma, int32_t mb, void *a, void *b, |
| void *c, int32_t ctp, int32_t prune_flag, |
| CUstream stream) { |
| assert(cusparseLt_initiated && "client did not call mgpuCreateSparseLtEnv()"); |
| // TODO: support more advanced settings, e.g., the input right operand is a |
| // sparse matrix assuming matA is the sparse matrix |
| auto matA = reinterpret_cast<cusparseLtSpMatHandleAndData *>(a); |
| auto matB = reinterpret_cast<cusparseLtDnMatHandleAndData *>(b); |
| auto matC = reinterpret_cast<cusparseLtDnMatHandleAndData *>(c); |
| auto workspace_size = reinterpret_cast<size_t *>(bs); |
| auto compressed_size = &(reinterpret_cast<size_t *>(bs)[1]); |
| auto compressed_buffer_size = &(reinterpret_cast<size_t *>(bs)[2]); |
| auto cTp = static_cast<cusparseComputeType>(ctp); |
| |
| cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma); |
| cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb); |
| CUSPARSE_REPORT_IF_ERROR(cusparseLtMatmulDescriptorInit( |
| &cusparseLt_env, &(matA->matmul), modeA, modeB, &(matA->mat), |
| &(matB->mat), &(matC->mat), &(matC->mat), cTp)) |
| CUSPARSE_REPORT_IF_ERROR(cusparseLtMatmulAlgSelectionInit( |
| &cusparseLt_env, &(matA->alg_sel), &(matA->matmul), |
| CUSPARSELT_MATMUL_ALG_DEFAULT)) |
| int alg = 0; |
| CUSPARSE_REPORT_IF_ERROR(cusparseLtMatmulAlgSetAttribute( |
| &cusparseLt_env, &(matA->alg_sel), CUSPARSELT_MATMUL_ALG_CONFIG_ID, &alg, |
| sizeof(alg))) |
| |
| CUSPARSE_REPORT_IF_ERROR(cusparseLtMatmulPlanInit( |
| &cusparseLt_env, &(matA->plan), &(matA->matmul), &(matA->alg_sel))) |
| |
| // Pruning step (in-place). |
| if (prune_flag > 0) |
| CUSPARSE_REPORT_IF_ERROR(cusparseLtSpMMAPrune( |
| &cusparseLt_env, &(matA->matmul), matA->values, matA->values, |
| CUSPARSELT_PRUNE_SPMMA_STRIP, stream)) |
| |
| // Check structure of A. |
| // Note that this adds a synchronization on the stream. |
| // TODO: Do we want that? |
| if (prune_flag == 2) { |
| int *dvalid = (int *)mgpuMemAlloc(sizeof(int), stream, false); |
| CUSPARSE_REPORT_IF_ERROR(cusparseLtSpMMAPruneCheck( |
| &cusparseLt_env, &(matA->matmul), matA->values, dvalid, stream)) |
| int valid = 0; |
| mgpuMemcpy(&valid, dvalid, sizeof(int), stream); |
| mgpuStreamSynchronize(stream); |
| mgpuMemFree(dvalid, stream); |
| if (valid != 0) |
| fprintf(stderr, "CUPARSE-LT: sparse matrix is not 2:4; computed results " |
| "will be invalid\n"); |
| } |
| |
| CUSPARSE_REPORT_IF_ERROR(cusparseLtMatmulGetWorkspace( |
| &cusparseLt_env, &(matA->plan), workspace_size)) |
| CUSPARSE_REPORT_IF_ERROR(cusparseLtSpMMACompressedSize( |
| &cusparseLt_env, &(matA->plan), compressed_size, compressed_buffer_size)) |
| } |
| |
| extern "C" MLIR_CUDA_WRAPPERS_EXPORT void |
| mgpuCuSparseLtSpMM(void *a, void *b, void *c, void *d_workspace, |
| void *dA_compressed, void *dA_compressedBuffer, |
| CUstream stream) { |
| assert(cusparseLt_initiated && "client did not call mgpuCreateSparseLtEnv()"); |
| auto matA = reinterpret_cast<cusparseLtSpMatHandleAndData *>(a); |
| auto matB = reinterpret_cast<cusparseLtDnMatHandleAndData *>(b); |
| auto matC = reinterpret_cast<cusparseLtDnMatHandleAndData *>(c); |
| |
| ALPHABETA(CUDA_R_32F, alpha, beta) |
| CUSPARSE_REPORT_IF_ERROR( |
| cusparseLtSpMMACompress(&cusparseLt_env, &(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 |
| CUSPARSE_REPORT_IF_ERROR( |
| cusparseLtMatmul(&cusparseLt_env, &(matA->plan), alphap, dA_compressed, |
| matB->values, betap, matC->values, |
| /*dD*/ matC->values, d_workspace, nullptr, 0)) |
| |
| CUSPARSE_REPORT_IF_ERROR(cusparseLtMatDescriptorDestroy(&(matA->mat))) |
| // destroy the plan associated with the sparse matrix |
| CUSPARSE_REPORT_IF_ERROR(cusparseLtMatmulPlanDestroy(&(matA->plan))) |
| } |
| |
| #endif // MLIR_ENABLE_CUDA_CUSPARSELT |
| #endif // MLIR_ENABLE_CUDA_CUSPARSE |
