Add a new Op in the gpuruntime dialect for deallocating memory allocated by gpu.alloc (#187)

Author: nbpatel

dpcomp_gpu_runtime/lib/gpu_runtime.cpp

@@ -350,6 +350,8 @@ struct Stream {
     return {info, mem, event};
   }
 
+  void deallocBuffer(void *ptr) { zeMemFree(context.get(), ptr); }
+
   void suggestBlockSize(ze_kernel_handle_t kernel, const uint32_t *gridSize,
                         uint32_t *blockSize, size_t numDims) {
     assert(kernel);
@@ -483,6 +485,12 @@ dpcompGpuAlloc(void *stream, size_t size, size_t alignment, int shared,
   });
 }
 
+extern "C" DPCOMP_GPU_RUNTIME_EXPORT void dpcompGpuDeAlloc(void *stream,
+                                                           void *ptr) {
+  LOG_FUNC();
+  catchAll([&]() { static_cast<Stream *>(stream)->deallocBuffer(ptr); });
+}
+
 extern "C" DPCOMP_GPU_RUNTIME_EXPORT void
 dpcompGpuSuggestBlockSize(void *stream, void *kernel, const uint32_t *gridSize,
                           uint32_t *blockSize, size_t numDims) {

mlir/include/mlir-extensions/dialect/gpu_runtime/IR/GpuRuntimeOps.td

@@ -136,6 +136,12 @@ def GPUAllocOp
   }];
 }
 
+def GPUDeallocOp : GpuRuntime_Op<"gpu_dealloc"> {
+  let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
+                  Res<AnyMemRef, "", [MemAlloc]>:$memref,
+                  GpuRuntime_OpaqueType : $stream);
+}
+
 def GPUSuggestBlockSizeOp : GpuRuntime_Op<"suggest_block_size",
                                      [AttrSizedOperandSegments, NoSideEffect]> {
   let arguments = (ins Optional<GpuRuntime_OpaqueType>:$stream,

mlir/lib/Conversion/gpu_runtime_to_llvm.cpp

@@ -169,6 +169,14 @@ class ConvertOpToGpuRuntimeCallPattern
           llvmAllocResPtrType,    // result
       }};
 
+  FunctionCallBuilder deallocCallBuilder = {
+      "dpcompGpuDeAlloc",
+      llvmVoidType,
+      {
+          llvmPointerType, // stream
+          llvmPointerType, // memory pointer
+      }};
+
   FunctionCallBuilder suggestBlockSizeBuilder = {
       "dpcompGpuSuggestBlockSize",
       llvmVoidType,
@@ -638,6 +646,30 @@ class ConvertGpuAllocPattern
   }
 };
 
+class ConvertGpuDeAllocPattern
+    : public ConvertOpToGpuRuntimeCallPattern<gpu_runtime::GPUDeallocOp> {
+public:
+  ConvertGpuDeAllocPattern(mlir::LLVMTypeConverter &converter)
+      : ConvertOpToGpuRuntimeCallPattern<gpu_runtime::GPUDeallocOp>(converter) {
+  }
+
+private:
+  mlir::LogicalResult
+  matchAndRewrite(gpu_runtime::GPUDeallocOp op,
+                  gpu_runtime::GPUDeallocOp::Adaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    mlir::Value pointer =
+        mlir::MemRefDescriptor(adaptor.memref()).allocatedPtr(rewriter, loc);
+    auto casted =
+        rewriter.create<mlir::LLVM::BitcastOp>(loc, llvmPointerType, pointer);
+    mlir::Value params[] = {adaptor.stream(), casted};
+    auto res = deallocCallBuilder.create(loc, rewriter, params);
+    rewriter.replaceOp(op, res.getResults());
+    return mlir::success();
+  }
+};
+
 class ConvertGpuSuggestBlockSizePattern
     : public ConvertOpToGpuRuntimeCallPattern<
           gpu_runtime::GPUSuggestBlockSizeOp> {
@@ -761,6 +793,7 @@ struct GPUToLLVMPass
         gpu_runtime::DestroyGpuKernelOp,
         gpu_runtime::LaunchGpuKernelOp,
         gpu_runtime::GPUAllocOp,
+        gpu_runtime::GPUDeallocOp,
         gpu_runtime::GPUSuggestBlockSizeOp
         // clang-format on
         >();
@@ -780,6 +813,7 @@ struct GPUToLLVMPass
         ConvertGpuKernelDestroyPattern,
         ConvertGpuKernelLaunchPattern,
         ConvertGpuAllocPattern,
+        ConvertGpuDeAllocPattern,
         ConvertGpuSuggestBlockSizePattern,
         LowerUndef
         // clang-format on

mlir/lib/Conversion/gpu_to_gpu_runtime.cpp

@@ -350,8 +350,7 @@ struct InsertGPUAllocs
       if (access.hostRead && access.deviceWrite)
         builder.create<mlir::memref::CopyOp>(loc, allocResult, param);
 
-      // TODO: Add a memref dealloc or gpu dealloc
-      // builder.create<mlir::memref::DeallocOp>(loc, allocResult);
+      builder.create<mlir::gpu::DeallocOp>(loc, llvm::None, allocResult);
     }
   }
 };
@@ -1023,6 +1022,23 @@ struct ExpandAllocOp : public mlir::OpRewritePattern<mlir::gpu::AllocOp> {
   }
 };
 
+struct ExpandDeallocOp : public mlir::OpRewritePattern<mlir::gpu::DeallocOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mlir::gpu::DeallocOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    auto stream = getGpuStream(rewriter, op);
+    if (!stream)
+      return mlir::failure();
+
+    auto res = rewriter.replaceOpWithNewOp<gpu_runtime::GPUDeallocOp>(
+        op, op.asyncDependencies(), op.memref(), *stream);
+
+    return mlir::success();
+  }
+};
+
 struct ExpandSuggestBlockSizeOp
     : public mlir::OpRewritePattern<gpu_runtime::GPUSuggestBlockSizeOp> {
   using OpRewritePattern::OpRewritePattern;
@@ -1158,8 +1174,8 @@ struct GPUExPass
     auto *ctx = &getContext();
     mlir::RewritePatternSet patterns(ctx);
 
-    patterns.insert<ExpandLaunchOp, ExpandAllocOp, ExpandSuggestBlockSizeOp>(
-        ctx);
+    patterns.insert<ExpandLaunchOp, ExpandAllocOp, ExpandDeallocOp,
+                    ExpandSuggestBlockSizeOp>(ctx);
 
     (void)mlir::applyPatternsAndFoldGreedily(getOperation(),
                                              std::move(patterns));