[MLIR] Adding gpu.host_register op and lower it to a runtime call.

[csigg]

Rebase.

[herhut]

Thanks!

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h

444

Nice!

mlir/include/mlir/Dialect/GPU/GPUOps.td

745

You can drop the Results

mlir/tools/mlir-rocm-runner/rocm-runtime-wrappers.cpp

79–87

nit: CUDA -> ROCm

[csigg]

Apply reviewer comments.

[csigg]

Thanks a lot for the review, Stephan.

[mehdi_amini]

(just two nits)

mlir/include/mlir/Dialect/GPU/GPUOps.td

749

Can you extend the description, it is quite terse for something that isn't that trivial I think.

In particular I think this should explain is this is just memory space mapping and what is the expected consistency model in terms of visibility / fences.

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

128

s/patter/pattern/

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h

public:
   Type getElementPtrType(MemRefType type) const;
 
   /// Determines sizes to be used in the memref descriptor.
   void getMemRefDescriptorSizes(Location loc, MemRefType memRefType,
                                 ArrayRef<Value> dynSizes,
                                 ConversionPatternRewriter &rewriter,
                                 SmallVectorImpl<Value> &sizes) const;
 
+  /// Computes the size of type in bytes.
+  Value getSizeInBytes(Location loc, Type type,
+                       ConversionPatternRewriter &rewriter) const;
+
   /// Computes total size in bytes of to store the given shape.
   Value getCumulativeSizeInBytes(Location loc, Type elementType,
                                  ArrayRef<Value> shape,
                                  ConversionPatternRewriter &rewriter) const;
 
 protected:
   /// Reference to the type converter, with potential extensions.
   LLVMTypeConverter &typeConverter;

mlir/include/mlir/Dialect/GPU/GPUOps.td

]> {
   let description = [{
     This op terminates the only block inside the only region of a `gpu.module`.
   }];
 
   let parser = [{ return success(); }];
   let printer = [{ p << getOperationName(); }];
 }
 
+def GPU_HostRegisterOp : GPU_Op<"host_register">,
+    Arguments<(ins AnyUnrankedMemRef:$value)> {
+  let summary = "Registers a memref for access from device.";
+  let description = [{
+    This op registers the host memory pointed to by a memref to be accessed from
+    a device.
+  }];
+
+  let assemblyFormat = "$value attr-dict `:` type($value)";
+  let verifier = [{ return success(); }];
+}
+
 #endif // GPU_OPS

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

FunctionCallBuilder launchKernelCallBuilder = {
           llvmPointerPointerType  /* void **extra */
       }};
   FunctionCallBuilder streamCreateCallBuilder = {
       "mgpuStreamCreate", llvmPointerType /* void *stream */, {}};
   FunctionCallBuilder streamSynchronizeCallBuilder = {
       "mgpuStreamSynchronize",
       llvmVoidType,
       {llvmPointerType /* void *stream */}};
+  FunctionCallBuilder hostRegisterCallBuilder = {
+      "mgpuMemHostRegisterMemRef",
+      llvmVoidType,
+      {llvmIntPtrType /* intptr_t rank */,
+       llvmPointerType /* void *memrefDesc */,
+       llvmIntPtrType /* intptr_t elementSizeBytes */}};
+};
+
+/// A rewrite patter to convert gpu.host_register operations into a GPU runtime
+/// call. Currently it supports CUDA and ROCm (HIP).
+class ConvertHostRegisterOpToGpuRuntimeCallPattern
+    : public ConvertOpToGpuRuntimeCallPattern<gpu::HostRegisterOp> {
+public:
+  ConvertHostRegisterOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
+      : ConvertOpToGpuRuntimeCallPattern<gpu::HostRegisterOp>(typeConverter) {}
+
+private:
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override;
 };
 
 /// A rewrite patter to convert gpu.launch_func operations into a sequence of
 /// GPU runtime calls. Currently it supports CUDA and ROCm (HIP).
 ///
 /// In essence, a gpu.launch_func operations gets compiled into the following
 /// sequence of runtime calls:
 ///
auto function = [&] {
     return OpBuilder(module.getBody()->getTerminator())
         .create<LLVM::LLVMFuncOp>(loc, functionName, functionType);
   }();
   return builder.create<LLVM::CallOp>(
       loc, const_cast<LLVM::LLVMType &>(functionType).getFunctionResultType(),
       builder.getSymbolRefAttr(function), arguments);
 }
 
+// Returns whether value is of LLVM type.
+static bool isLLVMType(Value value) {
+  return value.getType().isa<LLVM::LLVMType>();
+}
+
+LogicalResult ConvertHostRegisterOpToGpuRuntimeCallPattern::matchAndRewrite(
+    Operation *op, ArrayRef<Value> operands,
+    ConversionPatternRewriter &rewriter) const {
+  if (!llvm::all_of(operands, isLLVMType))
+    return rewriter.notifyMatchFailure(
+        op, "Cannot convert if operands aren't of LLVM type.");
+
+  Location loc = op->getLoc();
+
+  auto memRefType = cast<gpu::HostRegisterOp>(op).value().getType();
+  auto elementType = memRefType.cast<UnrankedMemRefType>().getElementType();
+  auto elementSize = getSizeInBytes(loc, elementType, rewriter);
+
+  auto arguments =
+      typeConverter.promoteOperands(loc, op->getOperands(), operands, rewriter);
+  arguments.push_back(elementSize);
+  hostRegisterCallBuilder.create(loc, rewriter, arguments);
+
+  rewriter.eraseOp(op);
+  return success();
+}
+
 // Creates a struct containing all kernel parameters on the stack and returns
 // an array of type-erased pointers to the fields of the struct. The array can
 // then be passed to the CUDA / ROCm (HIP) kernel launch calls.
 // The generated code is essentially as follows:
 //
 // %struct = alloca(sizeof(struct { Parameters... }))
 // %array = alloca(NumParameters * sizeof(void *))
 // for (i : [0, NumParameters))
Value ConvertLaunchFuncOpToGpuRuntimeCallPattern::generateKernelNameConstant(
 
   std::string globalName =
       std::string(llvm::formatv("{0}_{1}_kernel_name", moduleName, name));
   return LLVM::createGlobalString(
       loc, builder, globalName, StringRef(kernelName.data(), kernelName.size()),
       LLVM::Linkage::Internal);
 }
 
-// Returns whether value is of LLVM type.
-static bool isLLVMType(Value value) {
-  return value.getType().isa<LLVM::LLVMType>();
-}
-
 // Emits LLVM IR to launch a kernel function. Expects the module that contains
 // the compiled kernel function as a cubin in the 'nvvm.cubin' attribute, or a
 // hsaco in the 'rocdl.hsaco' attribute of the kernel function in the IR.
 //
 // %0 = call %binarygetter
 // %1 = call %moduleLoad(%0)
 // %2 = <see generateKernelNameConstant>
 // %3 = call %moduleGetFunction(%1, %2)

 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
 mlir::createGpuToLLVMConversionPass(StringRef gpuBinaryAnnotation) {
   return std::make_unique<GpuToLLVMConversionPass>(gpuBinaryAnnotation);
 }
 
 void mlir::populateGpuToLLVMConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
     StringRef gpuBinaryAnnotation) {
+  patterns.insert<ConvertHostRegisterOpToGpuRuntimeCallPattern>(converter);
   patterns.insert<ConvertLaunchFuncOpToGpuRuntimeCallPattern>(
       converter, gpuBinaryAnnotation);
   patterns.insert<EraseGpuModuleOpPattern>(&converter.getContext());
 }

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

void ConvertToLLVMPattern::getMemRefDescriptorSizes(
   sizes.reserve(memRefType.getRank());
   unsigned i = 0;
   for (int64_t s : memRefType.getShape())
     sizes.push_back(s == ShapedType::kDynamicSize
                         ? dynSizes[i++]
                         : createIndexConstant(rewriter, loc, s));
 }
 
-Value ConvertToLLVMPattern::getCumulativeSizeInBytes(
-    Location loc, Type elementType, ArrayRef<Value> sizes,
-    ConversionPatternRewriter &rewriter) const {
-  // Compute the total number of memref elements.
-  Value cumulativeSizeInBytes =
-      sizes.empty() ? createIndexConstant(rewriter, loc, 1) : sizes.front();
-  for (unsigned i = 1, e = sizes.size(); i < e; ++i)
-    cumulativeSizeInBytes = rewriter.create<LLVM::MulOp>(
-        loc, getIndexType(), ArrayRef<Value>{cumulativeSizeInBytes, sizes[i]});
-
+Value ConvertToLLVMPattern::getSizeInBytes(
+    Location loc, Type type, ConversionPatternRewriter &rewriter) const {
   // Compute the size of an individual element. This emits the MLIR equivalent
   // of the following sizeof(...) implementation in LLVM IR:
   //   %0 = getelementptr %elementType* null, %indexType 1
   //   %1 = ptrtoint %elementType* %0 to %indexType
   // which is a common pattern of getting the size of a type in bytes.
-  auto convertedPtrType = typeConverter.convertType(elementType)
-                              .cast<LLVM::LLVMType>()
-                              .getPointerTo();
+  auto convertedPtrType =
+      typeConverter.convertType(type).cast<LLVM::LLVMType>().getPointerTo();
   auto nullPtr = rewriter.create<LLVM::NullOp>(loc, convertedPtrType);
   auto gep = rewriter.create<LLVM::GEPOp>(
       loc, convertedPtrType,
       ArrayRef<Value>{nullPtr, createIndexConstant(rewriter, loc, 1)});
-  auto elementSize =
-      rewriter.create<LLVM::PtrToIntOp>(loc, getIndexType(), gep);
+  return rewriter.create<LLVM::PtrToIntOp>(loc, getIndexType(), gep);
+}
+
+Value ConvertToLLVMPattern::getCumulativeSizeInBytes(
+    Location loc, Type elementType, ArrayRef<Value> sizes,
+    ConversionPatternRewriter &rewriter) const {
+  // Compute the total number of memref elements.
+  Value cumulativeSizeInBytes =
+      sizes.empty() ? createIndexConstant(rewriter, loc, 1) : sizes.front();
+  for (unsigned i = 1, e = sizes.size(); i < e; ++i)
+    cumulativeSizeInBytes = rewriter.create<LLVM::MulOp>(
+        loc, getIndexType(), ArrayRef<Value>{cumulativeSizeInBytes, sizes[i]});
+  auto elementSize = this->getSizeInBytes(loc, elementType, rewriter);
   return rewriter.create<LLVM::MulOp>(
       loc, getIndexType(), ArrayRef<Value>{cumulativeSizeInBytes, elementSize});
 }
 
 /// Only retain those attributes that are not constructed by
 /// `LLVMFuncOp::build`. If `filterArgAttrs` is set, also filter out argument
 /// attributes.
 static void filterFuncAttributes(ArrayRef<NamedAttribute> attrs,

mlir/test/mlir-cuda-runner/all-reduce-and.mlir

func @main() {
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c4 = constant 4 : index
   %c5 = constant 5 : index
   %c6 = constant 6 : index
 
   %cast_data = memref_cast %data : memref<2x6xi32> to memref<*xi32>
-  call @mgpuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
+  gpu.host_register %cast_data : memref<*xi32>
   %cast_sum = memref_cast %sum : memref<2xi32> to memref<*xi32>
-  call @mgpuMemHostRegisterInt32(%cast_sum) : (memref<*xi32>) -> ()
+  gpu.host_register %cast_sum : memref<*xi32>
 
   store %cst0, %data[%c0, %c0] : memref<2x6xi32>
   store %cst1, %data[%c0, %c1] : memref<2x6xi32>
   store %cst2, %data[%c0, %c2] : memref<2x6xi32>
   store %cst4, %data[%c0, %c3] : memref<2x6xi32>
   store %cst8, %data[%c0, %c4] : memref<2x6xi32>
   store %cst16, %data[%c0, %c5] : memref<2x6xi32>
 
func @main() {
   }
 
   call @print_memref_i32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [0, 2]
 
   return
 }
 
-func @mgpuMemHostRegisterInt32(%ptr : memref<*xi32>)
 func @print_memref_i32(memref<*xi32>)
 

mlir/test/mlir-cuda-runner/all-reduce-max.mlir

func @main() {
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c4 = constant 4 : index
   %c5 = constant 5 : index
   %c6 = constant 6 : index
 
   %cast_data = memref_cast %data : memref<2x6xi32> to memref<*xi32>
-  call @mgpuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
+  gpu.host_register %cast_data : memref<*xi32>
   %cast_sum = memref_cast %sum : memref<2xi32> to memref<*xi32>
-  call @mgpuMemHostRegisterInt32(%cast_sum) : (memref<*xi32>) -> ()
+  gpu.host_register %cast_sum : memref<*xi32>
 
   store %cst0, %data[%c0, %c0] : memref<2x6xi32>
   store %cst1, %data[%c0, %c1] : memref<2x6xi32>
   store %cst2, %data[%c0, %c2] : memref<2x6xi32>
   store %cst4, %data[%c0, %c3] : memref<2x6xi32>
   store %cst8, %data[%c0, %c4] : memref<2x6xi32>
   store %cst16, %data[%c0, %c5] : memref<2x6xi32>
 
func @main() {
   }
 
   call @print_memref_i32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [16, 11]
 
   return
 }
 
-func @mgpuMemHostRegisterInt32(%ptr : memref<*xi32>)
 func @print_memref_i32(memref<*xi32>)
 

mlir/test/mlir-cuda-runner/all-reduce-min.mlir

func @main() {
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c4 = constant 4 : index
   %c5 = constant 5 : index
   %c6 = constant 6 : index
 
   %cast_data = memref_cast %data : memref<2x6xi32> to memref<*xi32>
-  call @mgpuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
+  gpu.host_register %cast_data : memref<*xi32>
   %cast_sum = memref_cast %sum : memref<2xi32> to memref<*xi32>
-  call @mgpuMemHostRegisterInt32(%cast_sum) : (memref<*xi32>) -> ()
+  gpu.host_register %cast_sum : memref<*xi32>
 
   store %cst0, %data[%c0, %c0] : memref<2x6xi32>
   store %cst1, %data[%c0, %c1] : memref<2x6xi32>
   store %cst2, %data[%c0, %c2] : memref<2x6xi32>
   store %cst4, %data[%c0, %c3] : memref<2x6xi32>
   store %cst8, %data[%c0, %c4] : memref<2x6xi32>
   store %cst16, %data[%c0, %c5] : memref<2x6xi32>
 
func @main() {
   }
 
   call @print_memref_i32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [0, 2]
 
   return
 }
 
-func @mgpuMemHostRegisterInt32(%ptr : memref<*xi32>)
 func @print_memref_i32(memref<*xi32>)
 

mlir/test/mlir-cuda-runner/all-reduce-op.mlir

 // RUN: mlir-cuda-runner %s --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
 
 // CHECK-COUNT-8: [{{(5356, ){12}5356}}]
 func @main() {
   %arg = alloc() : memref<2x4x13xf32>
   %dst = memref_cast %arg : memref<2x4x13xf32> to memref<?x?x?xf32>
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %sx = dim %dst, %c2 : memref<?x?x?xf32>
   %sy = dim %dst, %c1 : memref<?x?x?xf32>
   %sz = dim %dst, %c0 : memref<?x?x?xf32>
   %cast_dst = memref_cast %dst : memref<?x?x?xf32> to memref<*xf32>
-  call @mgpuMemHostRegisterFloat(%cast_dst) : (memref<*xf32>) -> ()
+  gpu.host_register %cast_dst : memref<*xf32>
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %sy, %block_z = %sz) {
     %t0 = muli %tz, %block_y : index
     %t1 = addi %ty, %t0 : index
     %t2 = muli %t1, %block_x : index
     %idx = addi %tx, %t2 : index
     %t3 = index_cast %idx : index to i32
     %val = sitofp %t3 : i32 to f32
     %sum = "gpu.all_reduce"(%val) ({}) { op = "add" } : (f32) -> (f32)
     store %sum, %dst[%tz, %ty, %tx] : memref<?x?x?xf32>
     gpu.terminator
   }
   call @print_memref_f32(%cast_dst) : (memref<*xf32>) -> ()
   return
 }
 
-func @mgpuMemHostRegisterFloat(%ptr : memref<*xf32>)
 func @print_memref_f32(%ptr : memref<*xf32>)

mlir/test/mlir-cuda-runner/all-reduce-or.mlir

func @main() {
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c4 = constant 4 : index
   %c5 = constant 5 : index
   %c6 = constant 6 : index
 
   %cast_data = memref_cast %data : memref<2x6xi32> to memref<*xi32>
-  call @mgpuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
+  gpu.host_register %cast_data : memref<*xi32>
   %cast_sum = memref_cast %sum : memref<2xi32> to memref<*xi32>
-  call @mgpuMemHostRegisterInt32(%cast_sum) : (memref<*xi32>) -> ()
+  gpu.host_register %cast_sum : memref<*xi32>
 
   store %cst0, %data[%c0, %c0] : memref<2x6xi32>
   store %cst1, %data[%c0, %c1] : memref<2x6xi32>
   store %cst2, %data[%c0, %c2] : memref<2x6xi32>
   store %cst4, %data[%c0, %c3] : memref<2x6xi32>
   store %cst8, %data[%c0, %c4] : memref<2x6xi32>
   store %cst16, %data[%c0, %c5] : memref<2x6xi32>
 
func @main() {
   }
 
   call @print_memref_i32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [31, 15]
 
   return
 }
 
-func @mgpuMemHostRegisterInt32(%ptr : memref<*xi32>)
 func @print_memref_i32(memref<*xi32>)
 

mlir/test/mlir-cuda-runner/all-reduce-region.mlir

 // RUN: mlir-cuda-runner %s --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
 
 // CHECK: [{{(35, ){34}35}}]
 func @main() {
   %arg = alloc() : memref<35xf32>
   %dst = memref_cast %arg : memref<35xf32> to memref<?xf32>
   %one = constant 1 : index
   %c0 = constant 0 : index
   %sx = dim %dst, %c0 : memref<?xf32>
   %cast_dst = memref_cast %dst : memref<?xf32> to memref<*xf32>
-  call @mgpuMemHostRegisterFloat(%cast_dst) : (memref<*xf32>) -> ()
+  gpu.host_register %cast_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) {
     %val = index_cast %tx : index to i32
     %xor = "gpu.all_reduce"(%val) ({
     ^bb(%lhs : i32, %rhs : i32):
       %xor = xor %lhs, %rhs : i32
       "gpu.yield"(%xor) : (i32) -> ()
     }) : (i32) -> (i32)
     %res = sitofp %xor : i32 to f32
     store %res, %dst[%tx] : memref<?xf32>
     gpu.terminator
   }
   call @print_memref_f32(%cast_dst) : (memref<*xf32>) -> ()
   return
 }
 
-func @mgpuMemHostRegisterFloat(%ptr : memref<*xf32>)
 func @print_memref_f32(memref<*xf32>)

mlir/test/mlir-cuda-runner/all-reduce-xor.mlir

func @main() {
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c4 = constant 4 : index
   %c5 = constant 5 : index
   %c6 = constant 6 : index
 
   %cast_data = memref_cast %data : memref<2x6xi32> to memref<*xi32>
-  call @mgpuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
+  gpu.host_register %cast_data : memref<*xi32>
   %cast_sum = memref_cast %sum : memref<2xi32> to memref<*xi32>
-  call @mgpuMemHostRegisterInt32(%cast_sum) : (memref<*xi32>) -> ()
+  gpu.host_register %cast_sum : memref<*xi32>
 
   store %cst0, %data[%c0, %c0] : memref<2x6xi32>
   store %cst1, %data[%c0, %c1] : memref<2x6xi32>
   store %cst2, %data[%c0, %c2] : memref<2x6xi32>
   store %cst4, %data[%c0, %c3] : memref<2x6xi32>
   store %cst8, %data[%c0, %c4] : memref<2x6xi32>
   store %cst16, %data[%c0, %c5] : memref<2x6xi32>
 
func @main() {
   }
 
   call @print_memref_i32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [31, 1]
 
   return
 }
 
-func @mgpuMemHostRegisterInt32(%ptr : memref<*xi32>)
 func @print_memref_i32(memref<*xi32>)
 

mlir/test/mlir-cuda-runner/gpu-to-cubin.mlir

 }
 
 // CHECK: [1, 1, 1, 1, 1]
 func @main() {
   %arg0 = alloc() : memref<5xf32>
   %21 = constant 5 : i32
   %22 = memref_cast %arg0 : memref<5xf32> to memref<?xf32>
   %23 = memref_cast %22 : memref<?xf32> to memref<*xf32>
-  call @mgpuMemHostRegisterFloat(%23) : (memref<*xf32>) -> ()
+  gpu.host_register %23 : memref<*xf32>
   call @print_memref_f32(%23) : (memref<*xf32>) -> ()
   %24 = constant 1.0 : f32
   call @other_func(%24, %22) : (f32, memref<?xf32>) -> ()
   call @print_memref_f32(%23) : (memref<*xf32>) -> ()
   return
 }
 
-func @mgpuMemHostRegisterFloat(%ptr : memref<*xf32>)
 func @print_memref_f32(%ptr : memref<*xf32>)

mlir/test/mlir-cuda-runner/multiple-all-reduce.mlir

func @main() {
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c4 = constant 4 : index
   %c5 = constant 5 : index
   %c6 = constant 6 : index
 
   %cast_data = memref_cast %data : memref<2x6xf32> to memref<*xf32>
-  call @mgpuMemHostRegisterFloat(%cast_data) : (memref<*xf32>) -> ()
+  gpu.host_register %cast_data : memref<*xf32>
   %cast_sum = memref_cast %sum : memref<2xf32> to memref<*xf32>
-  call @mgpuMemHostRegisterFloat(%cast_sum) : (memref<*xf32>) -> ()
+  gpu.host_register %cast_sum : memref<*xf32>
   %cast_mul = memref_cast %mul : memref<2xf32> to memref<*xf32>
-  call @mgpuMemHostRegisterFloat(%cast_mul) : (memref<*xf32>) -> ()
+  gpu.host_register %cast_mul : memref<*xf32>
 
   store %cst0, %data[%c0, %c0] : memref<2x6xf32>
   store %cst1, %data[%c0, %c1] : memref<2x6xf32>
   store %cst2, %data[%c0, %c2] : memref<2x6xf32>
   store %cst4, %data[%c0, %c3] : memref<2x6xf32>
   store %cst8, %data[%c0, %c4] : memref<2x6xf32>
   store %cst16, %data[%c0, %c5] : memref<2x6xf32>
 
func @main() {
   // CHECK: [31, 39]
 
   call @print_memref_f32(%cast_mul) : (memref<*xf32>) -> ()
   // CHECK: [0, 27720]
 
   return
 }
 
-func @mgpuMemHostRegisterFloat(%ptr : memref<*xf32>)
 func @print_memref_f32(memref<*xf32>)

mlir/test/mlir-cuda-runner/shuffle.mlir

 // RUN: mlir-cuda-runner %s --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
 
 // CHECK: [4, 5, 6, 7, 0, 1, 2, 3, 12, -1, -1, -1, 8]
 func @main() {
   %arg = alloc() : memref<13xf32>
   %dst = memref_cast %arg : memref<13xf32> to memref<?xf32>
   %one = constant 1 : index
   %c0 = constant 0 : index
   %sx = dim %dst, %c0 : memref<?xf32>
-  %cast_dest = memref_cast %dst : memref<?xf32> to memref<*xf32>
-  call @mgpuMemHostRegisterFloat(%cast_dest) : (memref<*xf32>) -> ()
+  %cast_dst = memref_cast %dst : memref<?xf32> to memref<*xf32>
+  gpu.host_register %cast_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) {
     %t0 = index_cast %tx : index to i32
     %val = sitofp %t0 : i32 to f32
     %width = index_cast %block_x : index to i32
     %offset = constant 4 : i32
     %shfl, %valid = gpu.shuffle %val, %offset, %width xor : f32
     cond_br %valid, ^bb1(%shfl : f32), ^bb0
   ^bb0:
     %m1 = constant -1.0 : f32
     br ^bb1(%m1 : f32)
   ^bb1(%value : f32):
     store %value, %dst[%tx] : memref<?xf32>
     gpu.terminator
   }
-  call @print_memref_f32(%cast_dest) : (memref<*xf32>) -> ()
+  call @print_memref_f32(%cast_dst) : (memref<*xf32>) -> ()
   return
 }
 
-func @mgpuMemHostRegisterFloat(%ptr : memref<*xf32>)
 func @print_memref_f32(%ptr : memref<*xf32>)

mlir/test/mlir-cuda-runner/two-modules.mlir

 // RUN: mlir-cuda-runner %s --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
 
 // CHECK: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
 func @main() {
   %arg = alloc() : memref<13xi32>
   %dst = memref_cast %arg : memref<13xi32> to memref<?xi32>
   %one = constant 1 : index
   %c0 = constant 0 : index
   %sx = dim %dst, %c0 : memref<?xi32>
   %cast_dst = memref_cast %dst : memref<?xi32> to memref<*xi32>
-  call @mgpuMemHostRegisterInt32(%cast_dst) : (memref<*xi32>) -> ()
+  gpu.host_register %cast_dst : memref<*xi32>
   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) {
     %t0 = index_cast %tx : index to i32
     store %t0, %dst[%tx] : memref<?xi32>
     gpu.terminator
   }
   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) {
     %t0 = index_cast %tx : index to i32
     store %t0, %dst[%tx] : memref<?xi32>
     gpu.terminator
   }
   call @print_memref_i32(%cast_dst) : (memref<*xi32>) -> ()
   return
 }
 
-func @mgpuMemHostRegisterInt32(%memref : memref<*xi32>)
 func @print_memref_i32(%memref : memref<*xi32>)

mlir/test/mlir-rocm-runner/gpu-to-hsaco.mlir

 }
 
 // CHECK: [1, 1, 1, 1, 1]
 func @main() {
   %arg0 = alloc() : memref<5xf32>
   %21 = constant 5 : i32
   %22 = memref_cast %arg0 : memref<5xf32> to memref<?xf32>
   %cast = memref_cast %22 : memref<?xf32> to memref<*xf32>
-  call @mgpuMemHostRegisterFloat(%cast) : (memref<*xf32>) -> ()
+  gpu.host_register %cast : memref<*xf32>
   %23 = memref_cast %22 : memref<?xf32> to memref<*xf32>
   call @print_memref_f32(%23) : (memref<*xf32>) -> ()
   %24 = constant 1.0 : f32
   %25 = call @mgpuMemGetDeviceMemRef1dFloat(%22) : (memref<?xf32>) -> (memref<?xf32>)
   call @other_func(%24, %25) : (f32, memref<?xf32>) -> ()
   call @print_memref_f32(%23) : (memref<*xf32>) -> ()
   return
 }
 
-func @mgpuMemHostRegisterFloat(%ptr : memref<*xf32>)
 func @mgpuMemGetDeviceMemRef1dFloat(%ptr : memref<?xf32>) -> (memref<?xf32>)
 func @print_memref_f32(%ptr : memref<*xf32>)

mlir/test/mlir-rocm-runner/two-modules.mlir

 // RUN: mlir-rocm-runner %s --shared-libs=%rocm_wrapper_library_dir/librocm-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
 
 // CHECK: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
 func @main() {
   %arg = alloc() : memref<13xi32>
   %dst = memref_cast %arg : memref<13xi32> to memref<?xi32>
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %sx = dim %dst, %c0 : memref<?xi32>
   %cast_dst = memref_cast %dst : memref<?xi32> to memref<*xi32>
-  call @mgpuMemHostRegisterInt32(%cast_dst) : (memref<*xi32>) -> ()
+  gpu.host_register %cast_dst : memref<*xi32>
   %dst_device = call @mgpuMemGetDeviceMemRef1dInt32(%dst) : (memref<?xi32>) -> (memref<?xi32>)
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %c1, %block_z = %c1) {
     %t0 = index_cast %tx : index to i32
     store %t0, %dst_device[%tx] : memref<?xi32>
     gpu.terminator
   }
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %c1, %block_z = %c1) {
     %t0 = index_cast %tx : index to i32
     store %t0, %dst_device[%tx] : memref<?xi32>
     gpu.terminator
   }
   call @print_memref_i32(%cast_dst) : (memref<*xi32>) -> ()
   return
 }
 
-func @mgpuMemHostRegisterInt32(%ptr : memref<*xi32>)
 func @mgpuMemGetDeviceMemRef1dInt32(%ptr : memref<?xi32>) -> (memref<?xi32>)
 func @print_memref_i32(%ptr : memref<*xi32>)

mlir/test/mlir-rocm-runner/vecadd.mlir

func @main() {
   %1 = alloc() : memref<5xf32>
   %2 = alloc() : memref<5xf32>
   %3 = memref_cast %0 : memref<5xf32> to memref<?xf32>
   %4 = memref_cast %1 : memref<5xf32> to memref<?xf32>
   %5 = memref_cast %2 : memref<5xf32> to memref<?xf32>
   %6 = memref_cast %3 : memref<?xf32> to memref<*xf32>
   %7 = memref_cast %4 : memref<?xf32> to memref<*xf32>
   %8 = memref_cast %5 : memref<?xf32> to memref<*xf32>
-  call @mgpuMemHostRegisterFloat(%6) : (memref<*xf32>) -> ()
-  call @mgpuMemHostRegisterFloat(%7) : (memref<*xf32>) -> ()
-  call @mgpuMemHostRegisterFloat(%8) : (memref<*xf32>) -> ()
+  gpu.host_register %6 : memref<*xf32>
+  gpu.host_register %7 : memref<*xf32>
+  gpu.host_register %8 : memref<*xf32>
   %9 = call @mgpuMemGetDeviceMemRef1dFloat(%3) : (memref<?xf32>) -> (memref<?xf32>)
   %10 = call @mgpuMemGetDeviceMemRef1dFloat(%4) : (memref<?xf32>) -> (memref<?xf32>)
   %11 = call @mgpuMemGetDeviceMemRef1dFloat(%5) : (memref<?xf32>) -> (memref<?xf32>)
 
   call @vecadd(%9, %10, %11) : (memref<?xf32>, memref<?xf32>, memref<?xf32>) -> ()
   call @print_memref_f32(%8) : (memref<*xf32>) -> ()
   return
 }
 
-func @mgpuMemHostRegisterFloat(%ptr : memref<*xf32>)
 func @mgpuMemGetDeviceMemRef1dFloat(%ptr : memref<?xf32>) -> (memref<?xf32>)
 func @print_memref_f32(%ptr : memref<*xf32>)

mlir/test/mlir-rocm-runner/vector-transferops.mlir

func @main() {
   %arg1 = alloc() : memref<4xf32>
 
   %22 = memref_cast %arg0 : memref<4xf32> to memref<?xf32>
   %23 = memref_cast %arg1 : memref<4xf32> to memref<?xf32>
 
   %cast0 = memref_cast %22 : memref<?xf32> to memref<*xf32>
   %cast1 = memref_cast %23 : memref<?xf32> to memref<*xf32>
 
-  call @mgpuMemHostRegisterFloat(%cast0) : (memref<*xf32>) -> ()
-  call @mgpuMemHostRegisterFloat(%cast1) : (memref<*xf32>) -> ()
+  gpu.host_register %cast0 : memref<*xf32>
+  gpu.host_register %cast1 : memref<*xf32>
 
   %24 = call @mgpuMemGetDeviceMemRef1dFloat(%22) : (memref<?xf32>) -> (memref<?xf32>)
   %26 = call @mgpuMemGetDeviceMemRef1dFloat(%23) : (memref<?xf32>) -> (memref<?xf32>)
 
   // CHECK: [1.23, 2.46, 2.46, 1.23]
   call @vectransferx2(%24, %26) : (memref<?xf32>,  memref<?xf32>) -> ()
   call @print_memref_f32(%cast1) : (memref<*xf32>) -> ()
 
   // CHECK: [2.46, 2.46, 2.46, 2.46]
   call @vectransferx4(%24, %26) : (memref<?xf32>,  memref<?xf32>) -> ()
   call @print_memref_f32(%cast1) : (memref<*xf32>) -> ()
   return
 }
 
-func @mgpuMemHostRegisterFloat(%ptr : memref<*xf32>)
 func @mgpuMemGetDeviceMemRef1dFloat(%ptr : memref<?xf32>) -> (memref<?xf32>)
 func @print_memref_f32(%ptr : memref<*xf32>)

mlir/tools/mlir-cuda-runner/cuda-runtime-wrappers.cpp

 /// 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 mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) {
   CUDA_REPORT_IF_ERROR(cuMemHostRegister(ptr, sizeBytes, /*flags=*/0));
 }
 
-// Allows to register a MemRef with the CUDA runtime. Initializes array with
-// value. Helpful until we have transfer functions implemented.
-template <typename T>
-void mgpuMemHostRegisterMemRef(const DynamicMemRefType<T> &memRef, T value) {
-  llvm::SmallVector<int64_t, 4> denseStrides(memRef.rank);
-  llvm::ArrayRef<int64_t> sizes(memRef.sizes, memRef.rank);
-  llvm::ArrayRef<int64_t> strides(memRef.strides, memRef.rank);
+// Allows to register a MemRef with the CUDA runtime. Helpful until we have
+// transfer functions implemented.
+extern "C" void
+mgpuMemHostRegisterMemRef(int64_t rank, StridedMemRefType<char, 1> *descriptor,
+                          int64_t elementSizeBytes) {
+
+  llvm::SmallVector<int64_t, 4> denseStrides(rank);
+  llvm::ArrayRef<int64_t> sizes(descriptor->sizes, rank);
+  llvm::ArrayRef<int64_t> strides(sizes.end(), rank);
 
   std::partial_sum(sizes.rbegin(), sizes.rend(), denseStrides.rbegin(),
                    std::multiplies<int64_t>());
-  auto count = denseStrides.front();
+  auto sizeBytes = denseStrides.front() * elementSizeBytes;
 
   // Only densely packed tensors are currently supported.
   std::rotate(denseStrides.begin(), denseStrides.begin() + 1,
               denseStrides.end());
   denseStrides.back() = 1;
   assert(strides == llvm::makeArrayRef(denseStrides));
 
-  auto *pointer = memRef.data + memRef.offset;
-  std::fill_n(pointer, count, value);
-  mgpuMemHostRegister(pointer, count * sizeof(T));
-}
-
-extern "C" void mgpuMemHostRegisterFloat(int64_t rank, void *ptr) {
-  UnrankedMemRefType<float> memRef = {rank, ptr};
-  mgpuMemHostRegisterMemRef(DynamicMemRefType<float>(memRef), 1.23f);
-}
-
-extern "C" void mgpuMemHostRegisterInt32(int64_t rank, void *ptr) {
-  UnrankedMemRefType<int32_t> memRef = {rank, ptr};
-  mgpuMemHostRegisterMemRef(DynamicMemRefType<int32_t>(memRef), 123);
+  auto ptr = descriptor->data + descriptor->offset * elementSizeBytes;
+  mgpuMemHostRegister(ptr, sizeBytes);
 }

mlir/tools/mlir-rocm-runner/rocm-runtime-wrappers.cpp

 /// Helper functions for writing mlir example code
 
 // Allows to register byte array with the ROCM runtime. Helpful until we have
 // transfer functions implemented.
 extern "C" void mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) {
   HIP_REPORT_IF_ERROR(hipHostRegister(ptr, sizeBytes, /*flags=*/0));
 }
 
-// Allows to register a MemRef with the ROCM runtime. Initializes array with
-// value. Helpful until we have transfer functions implemented.
-template <typename T>
-void mgpuMemHostRegisterMemRef(T *pointer, llvm::ArrayRef<int64_t> sizes,
-                               llvm::ArrayRef<int64_t> strides, T value) {
-  assert(sizes.size() == strides.size());
-  llvm::SmallVector<int64_t, 4> denseStrides(strides.size());
+// Allows to register a MemRef with the ROCm runtime. Helpful until we have
+// transfer functions implemented.
+extern "C" void
+mgpuMemHostRegisterMemRef(int64_t rank, StridedMemRefType<char, 1> *descriptor,
+                          int64_t elementSizeBytes) {
+
+  llvm::SmallVector<int64_t, 4> denseStrides(rank);
+  llvm::ArrayRef<int64_t> sizes(descriptor->sizes, rank);
+  llvm::ArrayRef<int64_t> strides(sizes.end(), rank);
 
   std::partial_sum(sizes.rbegin(), sizes.rend(), denseStrides.rbegin(),
                    std::multiplies<int64_t>());
-  auto count = denseStrides.front();
+  auto sizeBytes = denseStrides.front() * elementSizeBytes;
 
   // Only densely packed tensors are currently supported.
   std::rotate(denseStrides.begin(), denseStrides.begin() + 1,
               denseStrides.end());
   denseStrides.back() = 1;
   assert(strides == llvm::makeArrayRef(denseStrides));
 
-  std::fill_n(pointer, count, value);
-  mgpuMemHostRegister(pointer, count * sizeof(T));
-}
-
-extern "C" void mgpuMemHostRegisterFloat(int64_t rank, void *ptr) {
-  auto *desc = static_cast<StridedMemRefType<float, 1> *>(ptr);
-  auto sizes = llvm::ArrayRef<int64_t>(desc->sizes, rank);
-  auto strides = llvm::ArrayRef<int64_t>(desc->sizes + rank, rank);
-  mgpuMemHostRegisterMemRef(desc->data + desc->offset, sizes, strides, 1.23f);
-}
-
-extern "C" void mgpuMemHostRegisterInt32(int64_t rank, void *ptr) {
-  auto *desc = static_cast<StridedMemRefType<int32_t, 1> *>(ptr);
-  auto sizes = llvm::ArrayRef<int64_t>(desc->sizes, rank);
-  auto strides = llvm::ArrayRef<int64_t>(desc->sizes + rank, rank);
-  mgpuMemHostRegisterMemRef(desc->data + desc->offset, sizes, strides, 123);
+  auto ptr = descriptor->data + descriptor->offset * elementSizeBytes;
+  mgpuMemHostRegister(ptr, sizeBytes);
 }
 
 template <typename T>
 void mgpuMemGetDevicePointer(T *hostPtr, T **devicePtr) {
   HIP_REPORT_IF_ERROR(hipSetDevice(0));
   HIP_REPORT_IF_ERROR(
       hipHostGetDevicePointer((void **)devicePtr, hostPtr, /*flags=*/0));
 }