Add gpu::HostUnregisterOp

[makslevental]

Summary

Without explicitly unregistering you will get

'cuMemHostRegister(ptr, sizeBytes, 0)' failed with 'CUDA_ERROR_HOST_MEMORY_ALREADY_REGISTERED'

in CUDA (for example) after repeated runs (e.g., during benchmarking the same kernel).

[makslevental]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

167 ↗

(On Diff #509825)

As far as I can tell, the calling convention requires that the first argument passed is the rank of the memref and thus the hostUnregisterCallBuilder must be exactly the same (even though rank isn't used in the actual runtime call).

[mehdi_amini]

Can you upload the patch with full context?

[mehdi_amini]

Otherwise it seems symmetrical to the host_register op, so that seems like an obvious good addition to me, but it'd be best if someone else focusing on GPUs could confirm!

[ftynse]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

472 ↗

(On Diff #509876)

Please expand auto unless there is a cast on the RHS or the type is difficult to spell (lambdas, iterators).

167 ↗

(On Diff #509825)

Unranked memref gets lowered to { i64, ptr } in LLVM, where the first element is the rank and the second is a pointer to the ranked descriptor. At function boundary, the struct is unpacked into two individual arguments, hence the interface you see.

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

     way to check at runtime whether this feature is supported.
 
     Writes from the host are guaranteed to be visible to device kernels that are
     launched afterwards. Writes from the device are guaranteed to be visible on
     the host after synchronizing with the device kernel completion.
   }];
 
   let assemblyFormat = "$value attr-dict `:` type($value)";
 }
 
+def GPU_HostUnregisterOp : GPU_Op<"host_unregister">,
+    Arguments<(ins AnyUnrankedMemRef:$value)> {
+  let summary = "Unregisters a memref for access from device.";
+  let description = [{
+      This op unmaps the provided host buffer from the device address space.
+
+      This operation may not be supported in every environment, there is not yet a
+          way to check at runtime whether this feature is supported.
+  }];
+
+  let assemblyFormat = "$value attr-dict `:` type($value)";
+}
+
 def GPU_WaitOp : GPU_Op<"wait", [GPU_AsyncOpInterface]> {
   let summary = "Wait for async gpu ops to complete.";
   let description = [{
     This op synchronizes the host or the device with a list of dependent ops.
 
     If the op contains the `async` keyword, it returns a new async token which
     is synchronized with the op arguments. This new token is merely a shortcut
     to the argument list, and one could replace the uses of the result with the
     arguments for the same effect. The async version of this op is primarily
     used to make each async token have a single use during lowering and

lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

   FunctionCallBuilder eventRecordCallBuilder = {
       "mgpuEventRecord",
       llvmVoidType,
       {llvmPointerType /* void *event */, llvmPointerType /* void *stream */}};
   FunctionCallBuilder hostRegisterCallBuilder = {
       "mgpuMemHostRegisterMemRef",
       llvmVoidType,
       {llvmIntPtrType /* intptr_t rank */,
        llvmPointerType /* void *memrefDesc */,
        llvmIntPtrType /* intptr_t elementSizeBytes */}};
+  FunctionCallBuilder hostUnregisterCallBuilder = {
+      "mgpuMemHostUnregisterMemRef",
+      llvmVoidType,
+      {llvmIntPtrType /* intptr_t rank */,
+       llvmPointerType /* void *memrefDesc */,
+       llvmIntPtrType /* intptr_t elementSizeBytes */}};
   FunctionCallBuilder allocCallBuilder = {
       "mgpuMemAlloc",
       llvmPointerType /* void * */,
       {llvmIntPtrType /* intptr_t sizeBytes */,
        llvmPointerType /* void *stream */}};
   FunctionCallBuilder deallocCallBuilder = {
       "mgpuMemFree",
       llvmVoidType,
       {llvmPointerType /* void *ptr */, llvmPointerType /* void *stream */}};
   FunctionCallBuilder memcpyCallBuilder = {

 public:
   ConvertHostRegisterOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
       : ConvertOpToGpuRuntimeCallPattern<gpu::HostRegisterOp>(typeConverter) {}
 
 private:
   LogicalResult
   matchAndRewrite(gpu::HostRegisterOp hostRegisterOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+class ConvertHostUnregisterOpToGpuRuntimeCallPattern
+    : public ConvertOpToGpuRuntimeCallPattern<gpu::HostUnregisterOp> {
+public:
+  ConvertHostUnregisterOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
+      : ConvertOpToGpuRuntimeCallPattern<gpu::HostUnregisterOp>(typeConverter) {}
+
+private:
+  LogicalResult
+  matchAndRewrite(gpu::HostUnregisterOp hostUnregisterOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 /// A rewrite pattern to convert gpu.alloc operations into a GPU runtime
 /// call. Currently it supports CUDA and ROCm (HIP).
 class ConvertAllocOpToGpuRuntimeCallPattern
     : public ConvertOpToGpuRuntimeCallPattern<gpu::AllocOp> {
 public:
   ConvertAllocOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
       : ConvertOpToGpuRuntimeCallPattern<gpu::AllocOp>(typeConverter) {}
 
 private:
   LogicalResult

 
   auto arguments = getTypeConverter()->promoteOperands(
       loc, op->getOperands(), adaptor.getOperands(), rewriter);
   arguments.push_back(elementSize);
   hostRegisterCallBuilder.create(loc, rewriter, arguments);
 
   rewriter.eraseOp(op);
   return success();
 }
 
+LogicalResult ConvertHostUnregisterOpToGpuRuntimeCallPattern::matchAndRewrite(
+    gpu::HostUnregisterOp hostUnregisterOp, OpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  auto *op = hostUnregisterOp.getOperation();
+  if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)))
+    return failure();
+
+  Location loc = op->getLoc();
+
+  auto memRefType = hostUnregisterOp.getValue().getType();
+  auto elementType = memRefType.cast<UnrankedMemRefType>().getElementType();
+  auto elementSize = getSizeInBytes(loc, elementType, rewriter);
+
+  auto arguments = getTypeConverter()->promoteOperands(
+      loc, op->getOperands(), adaptor.getOperands(), rewriter);
+  arguments.push_back(elementSize);
+  hostUnregisterCallBuilder.create(loc, rewriter, arguments);
+
+  rewriter.eraseOp(op);
+  return success();
+}
+
 LogicalResult ConvertAllocOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::AllocOp allocOp, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (adaptor.getHostShared())
     return rewriter.notifyMatchFailure(
         allocOp, "host_shared allocation is not supported");
 
   MemRefType memRefType = allocOp.getType();
 
   if (failed(areAllLLVMTypes(allocOp, adaptor.getOperands(), rewriter)) ||

                                                RewritePatternSet &patterns,
                                                StringRef gpuBinaryAnnotation,
                                                bool kernelBarePtrCallConv) {
   converter.addConversion([&converter](gpu::AsyncTokenType type) -> Type {
     return converter.getPointerType(
         IntegerType::get(&converter.getContext(), 8));
   });
   patterns.add<ConvertAllocOpToGpuRuntimeCallPattern,
                ConvertDeallocOpToGpuRuntimeCallPattern,
                ConvertHostRegisterOpToGpuRuntimeCallPattern,
+               ConvertHostUnregisterOpToGpuRuntimeCallPattern,
                ConvertMemcpyOpToGpuRuntimeCallPattern,
                ConvertMemsetOpToGpuRuntimeCallPattern,
                ConvertSetDefaultDeviceOpToGpuRuntimeCallPattern,
                ConvertWaitAsyncOpToGpuRuntimeCallPattern,
                ConvertWaitOpToGpuRuntimeCallPattern,
                ConvertAsyncYieldToGpuRuntimeCallPattern>(converter);
   patterns.add<ConvertLaunchFuncOpToGpuRuntimeCallPattern>(
       converter, gpuBinaryAnnotation, kernelBarePtrCallConv);
   patterns.add<EraseGpuModuleOpPattern>(&converter.getContext());
 }

lib/ExecutionEngine/CudaRuntimeWrappers.cpp

   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;
 }

lib/ExecutionEngine/RocmRuntimeWrappers.cpp

   // Only densely packed tensors are currently supported.
   std::rotate(denseStrides.begin(), denseStrides.begin() + 1,
               denseStrides.end());
   denseStrides.back() = 1;
   assert(strides == llvm::ArrayRef(denseStrides));
 
   auto ptr = descriptor->data + descriptor->offset * elementSizeBytes;
   mgpuMemHostRegister(ptr, sizeBytes);
 }
 
+// Allows to unregister byte array with the ROCM runtime. Helpful until we have
+// transfer functions implemented.
+extern "C" void mgpuMemHostUnregister(void *ptr) {
+  HIP_REPORT_IF_ERROR(hipHostUnregister(ptr));
+}
+
+// Allows to unregister a MemRef with the ROCm runtime. Helpful until we have
+// transfer functions implemented.
+extern "C" void
+mgpuMemHostUnregisterMemRef(int64_t rank, StridedMemRefType<char, 1> *descriptor,
+                          int64_t elementSizeBytes) {
+  auto ptr = descriptor->data + descriptor->offset * elementSizeBytes;
+  mgpuMemHostUnregister(ptr);
+}
+
 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));
 }
 
 extern "C" StridedMemRefType<float, 1>
 mgpuMemGetDeviceMemRef1dFloat(float *allocated, float *aligned, int64_t offset,
                               int64_t size, int64_t stride) {