[Mlir-commits] [mlir] e63b574 - [mlir][LLVMIR] Add lifetime start and end marker instrinsics

[Jeff Niu]

Author: Jeff Niu
Date: 2022-09-14T13:15:14-07:00
New Revision: e63b574ff6a827ac64128490f117c94a54859992

URL: https://github.com/llvm/llvm-project/commit/e63b574ff6a827ac64128490f117c94a54859992
DIFF: https://github.com/llvm/llvm-project/commit/e63b574ff6a827ac64128490f117c94a54859992.diff

LOG: [mlir][LLVMIR] Add lifetime start and end marker instrinsics

This patch adds the `llvm.intr.lifetime.start` and `llvm.intr.lifetime.end`
intrinsics which are used to indicate to LLVM the lifetimes of allocated
memory.

These ops have the requirement that the first argument (the size) be an
"immediate argument". I added an OpTrait to check this, but it is
possible that an approach like GEPArg would work too.

Reviewed By: rriddle, dcaballe

Differential Revision: https://reviews.llvm.org/D133867

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
    mlir/test/Dialect/LLVMIR/roundtrip.mlir
    mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
index 8324408f0ef73..3577447f0a040 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
@@ -98,6 +98,32 @@ def LLVM_MemsetOp : LLVM_ZeroResultIntrOp<"memset", [0, 2]> {
                    LLVM_Type:$isVolatile);
 }
 
+//===----------------------------------------------------------------------===//
+// Lifetime Markers
+//===----------------------------------------------------------------------===//
+
+/// Base operation for lifetime markers. The LLVM intrinsics require the size
+/// operand to be an immediate. In MLIR it is encoded as an attribute.
+class LLVM_LifetimeBaseOp<string opName> : LLVM_ZeroResultIntrOp<opName> {
+  let arguments = (ins I64Attr:$size, LLVM_AnyPointer:$ptr);
+
+  // Custom builder to convert the size attribute to an integer.
+  let llvmBuilder = [{
+    llvm::Module *module = builder.GetInsertBlock()->getModule();
+    llvm::Function *fn = llvm::Intrinsic::getDeclaration(
+        module, llvm::Intrinsic::}] # llvmEnumName # [{, {}] #
+        !interleave(ListIntSubst<LLVM_IntrPatterns.operand, [0]>.lst, ", ")
+        # [{});
+    builder.CreateCall(fn, {builder.getInt64(op.getSizeAttr().getInt()),
+                            moduleTranslation.lookupValue(op.getPtr())});
+  }];
+
+  let assemblyFormat = "$size `,` $ptr attr-dict `:` type($ptr)";
+}
+
+def LLVM_LifetimeStartOp : LLVM_LifetimeBaseOp<"lifetime.start">;
+def LLVM_LifetimeEndOp : LLVM_LifetimeBaseOp<"lifetime.end">;
+
 // Intrinsics with multiple returns.
 
 def LLVM_SAddWithOverflowOp

diff  --git a/mlir/test/Dialect/LLVMIR/roundtrip.mlir b/mlir/test/Dialect/LLVMIR/roundtrip.mlir
index ad0050e251397..37a7905f9c584 100644
--- a/mlir/test/Dialect/LLVMIR/roundtrip.mlir
+++ b/mlir/test/Dialect/LLVMIR/roundtrip.mlir
@@ -522,3 +522,13 @@ llvm.func @vararg_func(%arg0: i32, ...) {
   // CHECK: llvm.return
   llvm.return
 }
+
+// CHECK-LABEL: @lifetime
+// CHECK-SAME: %[[P:.*]]: !llvm.ptr
+llvm.func @lifetime(%p: !llvm.ptr) {
+  // CHECK: llvm.intr.lifetime.start 16, %[[P]]
+  llvm.intr.lifetime.start 16, %p : !llvm.ptr
+  // CHECK: llvm.intr.lifetime.end 16, %[[P]]
+  llvm.intr.lifetime.end 16, %p : !llvm.ptr
+  llvm.return
+}

diff  --git a/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir b/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
index c87ff8551002b..b8bf5abf0e568 100644
--- a/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
@@ -635,13 +635,13 @@ llvm.func @vector_predication_intrinsics(%A: vector<8xi32>, %B: vector<8xi32>,
          (vector<8xi1>, vector<8xi32>, vector<8xi32>, i32) -> vector<8xi32>
 
   // CHECK: call void @llvm.vp.store.v8i32.p0
-  "llvm.intr.vp.store" (%A, %iptr, %mask, %evl) : 
+  "llvm.intr.vp.store" (%A, %iptr, %mask, %evl) :
          (vector<8xi32>, !llvm.ptr<i32>, vector<8xi1>, i32) -> ()
   // CHECK: call <8 x i32> @llvm.vp.load.v8i32.p0
   "llvm.intr.vp.load" (%iptr, %mask, %evl) :
          (!llvm.ptr<i32>, vector<8xi1>, i32) -> vector<8xi32>
   // CHECK: call void @llvm.experimental.vp.strided.store.v8i32.p0.i32
-  "llvm.intr.experimental.vp.strided.store" (%A, %iptr, %i, %mask, %evl) : 
+  "llvm.intr.experimental.vp.strided.store" (%A, %iptr, %i, %mask, %evl) :
          (vector<8xi32>, !llvm.ptr<i32>, i32, vector<8xi1>, i32) -> ()
   // CHECK: call <8 x i32> @llvm.experimental.vp.strided.load.v8i32.p0.i32
   "llvm.intr.experimental.vp.strided.load" (%iptr, %i, %mask, %evl) :
@@ -707,6 +707,15 @@ llvm.func @vector_insert_extract(%f256: vector<8xi32>, %f128: vector<4xi32>,
   llvm.return
 }
 
+// CHECK-LABEL: @lifetime
+llvm.func @lifetime(%p: !llvm.ptr) {
+  // CHECK: call void @llvm.lifetime.start
+  llvm.intr.lifetime.start 16, %p : !llvm.ptr
+  // CHECK: call void @llvm.lifetime.end
+  llvm.intr.lifetime.end 16, %p : !llvm.ptr
+  llvm.return
+}
+
 // Check that intrinsics are declared with appropriate types.
 // CHECK-DAG: declare float @llvm.fma.f32(float, float, float)
 // CHECK-DAG: declare <8 x float> @llvm.fma.v8f32(<8 x float>, <8 x float>, <8 x float>) #0
@@ -814,3 +823,5 @@ llvm.func @vector_insert_extract(%f256: vector<8xi32>, %f128: vector<4xi32>,
 // CHECK-DAG: declare <8 x i32> @llvm.vector.extract.v8i32.nxv4i32(<vscale x 4 x i32>, i64 immarg) #0
 // CHECK-DAG: declare <4 x i32> @llvm.vector.extract.v4i32.nxv4i32(<vscale x 4 x i32>, i64 immarg) #0
 // CHECK-DAG: declare <2 x i32> @llvm.vector.extract.v2i32.v8i32(<8 x i32>, i64 immarg) #0
+// CHECK-DAG: declare void @llvm.lifetime.start.p0(i64 immarg, ptr nocapture)
+// CHECK-DAG: declare void @llvm.lifetime.end.p0(i64 immarg, ptr nocapture)