| //===- 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 <cstdio> | |
| #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; | |
| /// Helper method that checks environment value for debugging. | |
| static bool isDebugEnabled() { | |
| const char *kDebugEnvironmentVariable = "MLIR_CUDA_DEBUG"; | |
| static bool 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 | |
| static 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, size_t /*assmeblySize*/) { | |
| 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) { | |
| // Tolerate CUDA_ERROR_DEINITIALIZED: this can run from a global destructor | |
| // after the CUDA primary context has already been torn down. | |
| CUresult result = cuModuleUnload(module); | |
| if (result != CUDA_SUCCESS && result != CUDA_ERROR_DEINITIALIZED) | |
| CUDA_REPORT_IF_ERROR(result); | |
| } | |
| 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) | |
| return reinterpret_cast<void *>(ptr); | |
| if (isHostShared) { | |
| CUDA_REPORT_IF_ERROR( | |
| cuMemAllocManaged(&ptr, sizeBytes, CU_MEM_ATTACH_GLOBAL)); | |
| return reinterpret_cast<void *>(ptr); | |
| } | |
| 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 mgpuLaunchKernelCooperative( | |
| CUfunction function, intptr_t gridX, intptr_t gridY, intptr_t gridZ, | |
| intptr_t clusterX, intptr_t clusterY, intptr_t clusterZ, intptr_t blockX, | |
| intptr_t blockY, intptr_t blockZ, int32_t smem, CUstream stream, | |
| void **params, void **extra) { | |
| ScopedContext scopedContext; | |
| if (smem > 0) { | |
| 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 launchAttrs[3]; | |
| int numAttrs = 0; | |
| bool hasCluster = clusterX > 0 && clusterY > 0 && clusterZ > 0; | |
| if (hasCluster) { | |
| launchAttrs[numAttrs].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; | |
| launchAttrs[numAttrs].value.clusterDim.x = clusterX; | |
| launchAttrs[numAttrs].value.clusterDim.y = clusterY; | |
| launchAttrs[numAttrs].value.clusterDim.z = clusterZ; | |
| numAttrs++; | |
| launchAttrs[numAttrs].id = | |
| CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE; | |
| launchAttrs[numAttrs].value.clusterSchedulingPolicyPreference = | |
| CU_CLUSTER_SCHEDULING_POLICY_SPREAD; | |
| numAttrs++; | |
| } | |
| launchAttrs[numAttrs].id = CU_LAUNCH_ATTRIBUTE_COOPERATIVE; | |
| launchAttrs[numAttrs].value.cooperative = 1; | |
| numAttrs++; | |
| config.numAttrs = numAttrs; | |
| config.attrs = launchAttrs; | |
| debug_print("Launching cooperative kernel (cluster=%d), " | |
| "grid=%ld,%ld,%ld, " | |
| "threads: %ld, %ld, %ld, " | |
| "smem: %dkb\n", | |
| hasCluster, gridX, gridY, gridZ, blockX, blockY, blockZ, smem); | |
| CUDA_REPORT_IF_ERROR(cuLaunchKernelEx(&config, function, params, extra)); | |
| } | |
| 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 |