211 lines · cpp
1//===- RocmRuntimeWrappers.cpp - MLIR ROCM runtime 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-exception6//7//===----------------------------------------------------------------------===//8//9// Implements C wrappers around the ROCM library for easy linking in ORC jit.10// Also adds some debugging helpers that are helpful when writing MLIR code to11// 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 21#include "hip/hip_runtime.h"22 23#define HIP_REPORT_IF_ERROR(expr) \24 [](hipError_t result) { \25 if (!result) \26 return; \27 const char *name = hipGetErrorName(result); \28 if (!name) \29 name = "<unknown>"; \30 fprintf(stderr, "'%s' failed with '%s'\n", #expr, name); \31 }(expr)32 33thread_local static int32_t defaultDevice = 0;34 35extern "C" hipModule_t mgpuModuleLoad(void *data, size_t /*gpuBlobSize*/) {36 hipModule_t module = nullptr;37 HIP_REPORT_IF_ERROR(hipModuleLoadData(&module, data));38 return module;39}40 41extern "C" hipModule_t mgpuModuleLoadJIT(void *data, int optLevel) {42 assert(false && "This function is not available in HIP.");43 return nullptr;44}45 46extern "C" void mgpuModuleUnload(hipModule_t module) {47 HIP_REPORT_IF_ERROR(hipModuleUnload(module));48}49 50extern "C" hipFunction_t mgpuModuleGetFunction(hipModule_t module,51 const char *name) {52 hipFunction_t function = nullptr;53 HIP_REPORT_IF_ERROR(hipModuleGetFunction(&function, module, name));54 return function;55}56 57// The wrapper uses intptr_t instead of ROCM's unsigned int to match58// the type of MLIR's index type. This avoids the need for casts in the59// generated MLIR code.60extern "C" void mgpuLaunchKernel(hipFunction_t function, intptr_t gridX,61 intptr_t gridY, intptr_t gridZ,62 intptr_t blockX, intptr_t blockY,63 intptr_t blockZ, int32_t smem,64 hipStream_t stream, void **params,65 void **extra, size_t /*paramsCount*/) {66 HIP_REPORT_IF_ERROR(hipModuleLaunchKernel(function, gridX, gridY, gridZ,67 blockX, blockY, blockZ, smem,68 stream, params, extra));69}70 71extern "C" hipStream_t mgpuStreamCreate() {72 hipStream_t stream = nullptr;73 HIP_REPORT_IF_ERROR(hipStreamCreate(&stream));74 return stream;75}76 77extern "C" void mgpuStreamDestroy(hipStream_t stream) {78 HIP_REPORT_IF_ERROR(hipStreamDestroy(stream));79}80 81extern "C" void mgpuStreamSynchronize(hipStream_t stream) {82 return HIP_REPORT_IF_ERROR(hipStreamSynchronize(stream));83}84 85extern "C" void mgpuStreamWaitEvent(hipStream_t stream, hipEvent_t event) {86 HIP_REPORT_IF_ERROR(hipStreamWaitEvent(stream, event, /*flags=*/0));87}88 89extern "C" hipEvent_t mgpuEventCreate() {90 hipEvent_t event = nullptr;91 HIP_REPORT_IF_ERROR(hipEventCreateWithFlags(&event, hipEventDisableTiming));92 return event;93}94 95extern "C" void mgpuEventDestroy(hipEvent_t event) {96 HIP_REPORT_IF_ERROR(hipEventDestroy(event));97}98 99extern "C" void mgpuEventSynchronize(hipEvent_t event) {100 HIP_REPORT_IF_ERROR(hipEventSynchronize(event));101}102 103extern "C" void mgpuEventRecord(hipEvent_t event, hipStream_t stream) {104 HIP_REPORT_IF_ERROR(hipEventRecord(event, stream));105}106 107extern "C" void *mgpuMemAlloc(uint64_t sizeBytes, hipStream_t /*stream*/,108 bool /*isHostShared*/) {109 void *ptr;110 HIP_REPORT_IF_ERROR(hipMalloc(&ptr, sizeBytes));111 return ptr;112}113 114extern "C" void mgpuMemFree(void *ptr, hipStream_t /*stream*/) {115 HIP_REPORT_IF_ERROR(hipFree(ptr));116}117 118extern "C" void mgpuMemcpy(void *dst, void *src, size_t sizeBytes,119 hipStream_t stream) {120 HIP_REPORT_IF_ERROR(121 hipMemcpyAsync(dst, src, sizeBytes, hipMemcpyDefault, stream));122}123 124extern "C" void mgpuMemset32(void *dst, int value, size_t count,125 hipStream_t stream) {126 HIP_REPORT_IF_ERROR(hipMemsetD32Async(reinterpret_cast<hipDeviceptr_t>(dst),127 value, count, stream));128}129 130extern "C" void mgpuMemset16(void *dst, int short value, size_t count,131 hipStream_t stream) {132 HIP_REPORT_IF_ERROR(hipMemsetD16Async(reinterpret_cast<hipDeviceptr_t>(dst),133 value, count, stream));134}135 136/// Helper functions for writing mlir example code137 138// Allows to register byte array with the ROCM runtime. Helpful until we have139// transfer functions implemented.140extern "C" void mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) {141 HIP_REPORT_IF_ERROR(hipHostRegister(ptr, sizeBytes, /*flags=*/0));142}143 144// Allows to register a MemRef with the ROCm runtime. Helpful until we have145// transfer functions implemented.146extern "C" void147mgpuMemHostRegisterMemRef(int64_t rank, StridedMemRefType<char, 1> *descriptor,148 int64_t elementSizeBytes) {149 150 llvm::SmallVector<int64_t, 4> denseStrides(rank);151 llvm::ArrayRef<int64_t> sizes(descriptor->sizes, rank);152 llvm::ArrayRef<int64_t> strides(sizes.end(), rank);153 154 std::partial_sum(sizes.rbegin(), sizes.rend(), denseStrides.rbegin(),155 std::multiplies<int64_t>());156 auto sizeBytes = denseStrides.front() * elementSizeBytes;157 158 // Only densely packed tensors are currently supported.159 std::rotate(denseStrides.begin(), denseStrides.begin() + 1,160 denseStrides.end());161 denseStrides.back() = 1;162 assert(strides == llvm::ArrayRef(denseStrides));163 164 auto ptr = descriptor->data + descriptor->offset * elementSizeBytes;165 mgpuMemHostRegister(ptr, sizeBytes);166}167 168// Allows to unregister byte array with the ROCM runtime. Helpful until we have169// transfer functions implemented.170extern "C" void mgpuMemHostUnregister(void *ptr) {171 HIP_REPORT_IF_ERROR(hipHostUnregister(ptr));172}173 174// Allows to unregister a MemRef with the ROCm runtime. Helpful until we have175// transfer functions implemented.176extern "C" void177mgpuMemHostUnregisterMemRef(int64_t rank,178 StridedMemRefType<char, 1> *descriptor,179 int64_t elementSizeBytes) {180 auto ptr = descriptor->data + descriptor->offset * elementSizeBytes;181 mgpuMemHostUnregister(ptr);182}183 184template <typename T>185void mgpuMemGetDevicePointer(T *hostPtr, T **devicePtr) {186 HIP_REPORT_IF_ERROR(hipSetDevice(0));187 HIP_REPORT_IF_ERROR(188 hipHostGetDevicePointer((void **)devicePtr, hostPtr, /*flags=*/0));189}190 191extern "C" StridedMemRefType<float, 1>192mgpuMemGetDeviceMemRef1dFloat(float *allocated, float *aligned, int64_t offset,193 int64_t size, int64_t stride) {194 float *devicePtr = nullptr;195 mgpuMemGetDevicePointer(aligned, &devicePtr);196 return {devicePtr, devicePtr, offset, {size}, {stride}};197}198 199extern "C" StridedMemRefType<int32_t, 1>200mgpuMemGetDeviceMemRef1dInt32(int32_t *allocated, int32_t *aligned,201 int64_t offset, int64_t size, int64_t stride) {202 int32_t *devicePtr = nullptr;203 mgpuMemGetDevicePointer(aligned, &devicePtr);204 return {devicePtr, devicePtr, offset, {size}, {stride}};205}206 207extern "C" void mgpuSetDefaultDevice(int32_t device) {208 defaultDevice = device;209 HIP_REPORT_IF_ERROR(hipSetDevice(device));210}211