1 //===- cuda-runtime-wrappers.cpp - MLIR CUDA runner wrapper library -------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // Implements C wrappers around the CUDA library for easy linking in ORC jit.
10 // Also adds some debugging helpers that are helpful when writing MLIR code to
11 // run on GPUs.
12 //
13 //===----------------------------------------------------------------------===//
14
15 #include <cassert>
16 #include <numeric>
17
18 #include "mlir/ExecutionEngine/CRunnerUtils.h"
19 #include "llvm/ADT/ArrayRef.h"
20 #include "llvm/Support/raw_ostream.h"
21
22 #include "cuda.h"
23
24 #define CUDA_REPORT_IF_ERROR(expr) \
25 [](CUresult result) { \
26 if (!result) \
27 return; \
28 const char *name = nullptr; \
29 cuGetErrorName(result, &name); \
30 if (!name) \
31 name = "<unknown>"; \
32 llvm::errs() << "'" << #expr << "' failed with '" << name << "'\n"; \
33 }(expr)
34
35 // Static initialization of CUDA context for device ordinal 0.
__anon4f5cbb9f0102null36 static auto InitializeCtx = [] {
37 CUDA_REPORT_IF_ERROR(cuInit(/*flags=*/0));
38 CUdevice device;
39 CUDA_REPORT_IF_ERROR(cuDeviceGet(&device, /*ordinal=*/0));
40 CUcontext context;
41 CUDA_REPORT_IF_ERROR(cuCtxCreate(&context, /*flags=*/0, device));
42 return 0;
43 }();
44
mgpuModuleLoad(void * data)45 extern "C" CUmodule mgpuModuleLoad(void *data) {
46 CUmodule module = nullptr;
47 CUDA_REPORT_IF_ERROR(cuModuleLoadData(&module, data));
48 return module;
49 }
50
mgpuModuleUnload(CUmodule module)51 extern "C" void mgpuModuleUnload(CUmodule module) {
52 CUDA_REPORT_IF_ERROR(cuModuleUnload(module));
53 }
54
mgpuModuleGetFunction(CUmodule module,const char * name)55 extern "C" CUfunction mgpuModuleGetFunction(CUmodule module, const char *name) {
56 CUfunction function = nullptr;
57 CUDA_REPORT_IF_ERROR(cuModuleGetFunction(&function, module, name));
58 return function;
59 }
60
61 // The wrapper uses intptr_t instead of CUDA's unsigned int to match
62 // the type of MLIR's index type. This avoids the need for casts in the
63 // generated MLIR code.
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)64 extern "C" void mgpuLaunchKernel(CUfunction function, intptr_t gridX,
65 intptr_t gridY, intptr_t gridZ,
66 intptr_t blockX, intptr_t blockY,
67 intptr_t blockZ, int32_t smem, CUstream stream,
68 void **params, void **extra) {
69 CUDA_REPORT_IF_ERROR(cuLaunchKernel(function, gridX, gridY, gridZ, blockX,
70 blockY, blockZ, smem, stream, params,
71 extra));
72 }
73
mgpuStreamCreate()74 extern "C" CUstream mgpuStreamCreate() {
75 CUstream stream = nullptr;
76 CUDA_REPORT_IF_ERROR(cuStreamCreate(&stream, CU_STREAM_NON_BLOCKING));
77 return stream;
78 }
79
mgpuStreamDestroy(CUstream stream)80 extern "C" void mgpuStreamDestroy(CUstream stream) {
81 CUDA_REPORT_IF_ERROR(cuStreamDestroy(stream));
82 }
83
mgpuStreamSynchronize(CUstream stream)84 extern "C" void mgpuStreamSynchronize(CUstream stream) {
85 CUDA_REPORT_IF_ERROR(cuStreamSynchronize(stream));
86 }
87
mgpuStreamWaitEvent(CUstream stream,CUevent event)88 extern "C" void mgpuStreamWaitEvent(CUstream stream, CUevent event) {
89 CUDA_REPORT_IF_ERROR(cuStreamWaitEvent(stream, event, /*flags=*/0));
90 }
91
mgpuEventCreate()92 extern "C" CUevent mgpuEventCreate() {
93 CUevent event = nullptr;
94 CUDA_REPORT_IF_ERROR(cuEventCreate(&event, CU_EVENT_DISABLE_TIMING));
95 return event;
96 }
97
mgpuEventDestroy(CUevent event)98 extern "C" void mgpuEventDestroy(CUevent event) {
99 CUDA_REPORT_IF_ERROR(cuEventDestroy(event));
100 }
101
mgpuEventSynchronize(CUevent event)102 extern "C" void mgpuEventSynchronize(CUevent event) {
103 CUDA_REPORT_IF_ERROR(cuEventSynchronize(event));
104 }
105
mgpuEventRecord(CUevent event,CUstream stream)106 extern "C" void mgpuEventRecord(CUevent event, CUstream stream) {
107 CUDA_REPORT_IF_ERROR(cuEventRecord(event, stream));
108 }
109
110 /// Helper functions for writing mlir example code
111
112 // Allows to register byte array with the CUDA runtime. Helpful until we have
113 // transfer functions implemented.
mgpuMemHostRegister(void * ptr,uint64_t sizeBytes)114 extern "C" void mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) {
115 CUDA_REPORT_IF_ERROR(cuMemHostRegister(ptr, sizeBytes, /*flags=*/0));
116 }
117
118 // Allows to register a MemRef with the CUDA runtime. Helpful until we have
119 // transfer functions implemented.
120 extern "C" void
mgpuMemHostRegisterMemRef(int64_t rank,StridedMemRefType<char,1> * descriptor,int64_t elementSizeBytes)121 mgpuMemHostRegisterMemRef(int64_t rank, StridedMemRefType<char, 1> *descriptor,
122 int64_t elementSizeBytes) {
123
124 llvm::SmallVector<int64_t, 4> denseStrides(rank);
125 llvm::ArrayRef<int64_t> sizes(descriptor->sizes, rank);
126 llvm::ArrayRef<int64_t> strides(sizes.end(), rank);
127
128 std::partial_sum(sizes.rbegin(), sizes.rend(), denseStrides.rbegin(),
129 std::multiplies<int64_t>());
130 auto sizeBytes = denseStrides.front() * elementSizeBytes;
131
132 // Only densely packed tensors are currently supported.
133 std::rotate(denseStrides.begin(), denseStrides.begin() + 1,
134 denseStrides.end());
135 denseStrides.back() = 1;
136 assert(strides == llvm::makeArrayRef(denseStrides));
137
138 auto ptr = descriptor->data + descriptor->offset * elementSizeBytes;
139 mgpuMemHostRegister(ptr, sizeBytes);
140 }
141