shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

Make type and rank explicit in mcuMemHostRegister function.

Browse Source 
Fix registered size of indirect MemRefType kernel arguments.

PiperOrigin-RevId: 281362940
This commit is contained in:
Christian Sigg 2019-11-19 13:12:19 -08:00 committed by A. Unique TensorFlower
parent ee95f6f259
commit f868adafee
5 changed files with 40 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
mlir/lib/Conversion/GPUToCUDA/ConvertLaunchFuncToCudaCalls.cpp
View File

@ -49,7 +49,7 @@ static constexpr const char *cuModuleGetFunctionName = "mcuModuleGetFunction";
static constexpr const char *cuLaunchKernelName = "mcuLaunchKernel";
static constexpr const char *cuGetStreamHelperName = "mcuGetStreamHelper";
static constexpr const char *cuStreamSynchronizeName = "mcuStreamSynchronize";
static constexpr const char *kMcuMemHostRegisterPtr = "mcuMemHostRegisterPtr";
static constexpr const char *kMcuMemHostRegister = "mcuMemHostRegister";
static constexpr const char *kCubinAnnotation = "nvvm.cubin";
static constexpr const char *kCubinStorageSuffix = "_cubin_cst";
@ -228,13 +228,13 @@ void GpuLaunchFuncToCudaCallsPass::declareCudaFunctions(Location loc) {
getPointerType() /* CUstream stream */,
/*isVarArg=*/false));
}
if (!module.lookupSymbol(kMcuMemHostRegisterPtr)) {
if (!module.lookupSymbol(kMcuMemHostRegister)) {
builder.create<LLVM::LLVMFuncOp>(
loc, kMcuMemHostRegisterPtr,
loc, kMcuMemHostRegister,
LLVM::LLVMType::getFunctionTy(getVoidType(),
{
getPointerType(), /* void *ptr */
getInt32Type() /* int32 flags*/
getInt64Type() /* int64 sizeBytes*/
},
/*isVarArg=*/false));
}
@ -277,12 +277,14 @@ GpuLaunchFuncToCudaCallsPass::setupParamsArray(gpu::LaunchFuncOp launchOp,
// the descriptor pointer is registered via @mcuMemHostRegisterPtr
if (llvmType.isStructTy()) {
auto registerFunc =
getModule().lookupSymbol<LLVM::LLVMFuncOp>(kMcuMemHostRegisterPtr);
auto zero = builder.create<LLVM::ConstantOp>(
loc, getInt32Type(), builder.getI32IntegerAttr(0));
getModule().lookupSymbol<LLVM::LLVMFuncOp>(kMcuMemHostRegister);
auto nullPtr = builder.create<LLVM::NullOp>(loc, llvmType.getPointerTo());
auto gep = builder.create<LLVM::GEPOp>(loc, llvmType.getPointerTo(),
ArrayRef<Value *>{nullPtr, one});
auto size = builder.create<LLVM::PtrToIntOp>(loc, getInt64Type(), gep);
builder.create<LLVM::CallOp>(loc, ArrayRef<Type>{},
builder.getSymbolRefAttr(registerFunc),
ArrayRef<Value *>{casted, zero});
ArrayRef<Value *>{casted, size});
Value *memLocation = builder.create<LLVM::AllocaOp>(
loc, getPointerPointerType(), one, /*alignment=*/1);
builder.create<LLVM::StoreOp>(loc, casted, memLocation);

5
mlir/test/mlir-cuda-runner/all-reduce-op.mlir
View File

@ -4,12 +4,11 @@
func @main() {
%arg = alloc() : memref<13x4x2xf32>
%dst = memref_cast %arg : memref<13x4x2xf32> to memref<?x?x?xf32>
%zero = constant 0 : i32
%one = constant 1 : index
%sx = dim %dst, 0 : memref<?x?x?xf32>
%sy = dim %dst, 1 : memref<?x?x?xf32>
%sz = dim %dst, 2 : memref<?x?x?xf32>
call @mcuMemHostRegister(%dst, %zero) : (memref<?x?x?xf32>, i32) -> ()
call @mcuMemHostRegisterMemRef3dFloat(%dst) : (memref<?x?x?xf32>) -> ()
gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %one, %grid_y = %one, %grid_z = %one)
threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %sy, %block_z = %sz)
args(%kernel_dst = %dst) : memref<?x?x?xf32> {
@ -27,5 +26,5 @@ func @main() {
return
}
func @mcuMemHostRegister(%ptr : memref<?x?x?xf32>, %flags : i32)
func @mcuMemHostRegisterMemRef3dFloat(%ptr : memref<?x?x?xf32>)
func @mcuPrintFloat(%ptr : memref<?x?x?xf32>)

5
mlir/test/mlir-cuda-runner/all-reduce-region.mlir
View File

@ -4,10 +4,9 @@
func @main() {
%arg = alloc() : memref<35xf32>
%dst = memref_cast %arg : memref<35xf32> to memref<?xf32>
%zero = constant 0 : i32
%one = constant 1 : index
%sx = dim %dst, 0 : memref<?xf32>
call @mcuMemHostRegister(%dst, %zero) : (memref<?xf32>, i32) -> ()
call @mcuMemHostRegisterMemRef1dFloat(%dst) : (memref<?xf32>) -> ()
gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %one, %grid_y = %one, %grid_z = %one)
threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %one, %block_z = %one)
args(%kernel_dst = %dst) : memref<?xf32> {
@ -25,5 +24,5 @@ func @main() {
return
}
func @mcuMemHostRegister(%ptr : memref<?xf32>, %flags : i32)
func @mcuMemHostRegisterMemRef1dFloat(%ptr : memref<?xf32>)
func @mcuPrintFloat(%ptr : memref<?xf32>)

5
mlir/test/mlir-cuda-runner/gpu-to-cubin.mlir
View File

@ -15,10 +15,9 @@ func @other_func(%arg0 : f32, %arg1 : memref<?xf32>) {
// CHECK: [1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00]
func @main() {
%arg0 = alloc() : memref<5xf32>
%20 = constant 0 : i32
%21 = constant 5 : i32
%22 = memref_cast %arg0 : memref<5xf32> to memref<?xf32>
call @mcuMemHostRegister(%22, %20) : (memref<?xf32>, i32) -> ()
call @mcuMemHostRegisterMemRef1dFloat(%22) : (memref<?xf32>) -> ()
call @mcuPrintFloat(%22) : (memref<?xf32>) -> ()
%24 = constant 1.0 : f32
call @other_func(%24, %22) : (f32, memref<?xf32>) -> ()
@ -26,5 +25,5 @@ func @main() {
return
}
func @mcuMemHostRegister(%ptr : memref<?xf32>, %flags : i32)
func @mcuMemHostRegisterMemRef1dFloat(%ptr : memref<?xf32>)
func @mcuPrintFloat(%ptr : memref<?xf32>)

42
mlir/tools/mlir-cuda-runner/cuda-runtime-wrappers.cpp
View File

@ -21,8 +21,8 @@
//
//===----------------------------------------------------------------------===//
#include <assert.h>
#include <memory.h>
#include <cassert>
#include <numeric>
#include "llvm/Support/raw_ostream.h"
@ -80,6 +80,13 @@ extern "C" int32_t mcuStreamSynchronize(void *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" void mcuMemHostRegister(void *ptr, uint64_t sizeBytes) {
reportErrorIfAny(cuMemHostRegister(ptr, sizeBytes, /*flags=*/0),
"MemHostRegister");
}
// A struct that corresponds to how MLIR represents memrefs.
template <typename T, int N> struct MemRefType {
T *basePtr;
@ -89,23 +96,22 @@ template <typename T, int N> struct MemRefType {
int64_t strides[N];
};
// Allows to register a pointer with the CUDA runtime. Helpful until
// we have transfer functions implemented.
extern "C" void mcuMemHostRegister(const MemRefType<float, 1> *arg,
int32_t flags) {
reportErrorIfAny(
cuMemHostRegister(arg->data, arg->sizes[0] * sizeof(float), flags),
"MemHostRegister");
for (int pos = 0; pos < arg->sizes[0]; pos++) {
arg->data[pos] = 1.23f;
}
// Allows to register a MemRef with the CUDA runtime. Initializes array with
// value. Helpful until we have transfer functions implemented.
template <typename T, int N>
void mcuMemHostRegisterMemRef(const MemRefType<T, N> *arg, T value) {
auto count = std::accumulate(arg->sizes, arg->sizes + N, 1,
std::multiplies<int64_t>());
std::fill_n(arg->data, count, value);
mcuMemHostRegister(arg->data, count * sizeof(T));
}
// Allows to register a pointer with the CUDA runtime. Helpful until
// we have transfer functions implemented.
extern "C" void mcuMemHostRegisterPtr(void *ptr, int32_t flags) {
reportErrorIfAny(cuMemHostRegister(ptr, sizeof(void *), flags),
"MemHostRegister");
extern "C" void
mcuMemHostRegisterMemRef1dFloat(const MemRefType<float, 1> *arg) {
mcuMemHostRegisterMemRef(arg, 1.23f);
}
extern "C" void
mcuMemHostRegisterMemRef3dFloat(const MemRefType<float, 3> *arg) {
mcuMemHostRegisterMemRef(arg, 1.23f);
}
/// Prints the given float array to stderr.