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

[Mogball]

Summary

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.

[Mogball]

rebase

[dcaballe]

Thanks, @Mogball! LGTM. Please wait for somebody else to chime in.

mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h

126 ↗

(On Diff #460118)

Nice! I wonder if we already have more LLVM ops with this requirement. Are we just not checking that invariant for them?

[rriddle]

mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h

105–126 ↗

(On Diff #460118)

This doesn't seem like a good constraint for MLIR. There are plenty of things that could easily break this (e.g. block merging). Why not just have the size be an Attribute?

[Mogball]

Switched to an attribute for imm arg instead of an op trait

PTAL :)

[Mogball]

mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h

126 ↗

(On Diff #460118)

It's possible. I've switch to an attribute like river suggested. PTAL @dcaballe

[Mogball]

remove extra newline

[rriddle]

Nice!

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

let arguments = (ins Arg<LLVM_Type,"",[MemWrite]>:$dst, Arg<LLVM_Type,"",[MemRead]>:$src, LLVM_Type:$len,
                    LLVM_Type:$isVolatile);
 }
 
 def LLVM_MemsetOp : LLVM_ZeroResultIntrOp<"memset", [0, 2]> {
   let arguments = (ins Arg<LLVM_Type,"",[MemWrite]>:$dst, LLVM_Type:$val, LLVM_Type:$len,
                    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
     : LLVM_IntrOp<"sadd.with.overflow", [0], [], [], 2> {
   let arguments = (ins LLVM_Type, LLVM_Type);
 }
 def LLVM_UAddWithOverflowOp
     : LLVM_IntrOp<"uadd.with.overflow", [0], [], [], 2> {

mlir/test/Dialect/LLVMIR/roundtrip.mlir

llvm.func @vararg_func(%arg0: i32, ...) {
   llvm.intr.vacopy %3 to %5
   // CHECK: llvm.intr.vaend %[[CAST1]]
   // CHECK: llvm.intr.vaend %[[CAST0]]
   llvm.intr.vaend %5
   llvm.intr.vaend %3
   // 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
+}

mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir

llvm.func @vector_predication_intrinsics(%A: vector<8xi32>, %B: vector<8xi32>,
   // CHECK: call <8 x i32> @llvm.vp.select.v8i32
   "llvm.intr.vp.select" (%mask, %A, %B, %evl) :
          (vector<8xi1>, vector<8xi32>, vector<8xi32>, i32) -> vector<8xi32>
   // CHECK: call <8 x i32> @llvm.vp.merge.v8i32
   "llvm.intr.vp.merge" (%mask, %A, %B, %evl) :
          (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) :
          (!llvm.ptr<i32>, i32, vector<8xi1>, i32) -> vector<8xi32>
 
   // CHECK: call <8 x i32> @llvm.vp.trunc.v8i32.v8i64
   "llvm.intr.vp.trunc" (%E, %mask, %evl) :
          (vector<8xi64>, vector<8xi1>, i32) -> vector<8xi32>
llvm.func @vector_insert_extract(%f256: vector<8xi32>, %f128: vector<4xi32>,
   %5 = llvm.intr.vector.extract %2[0] :
               vector<4xi32> from vector<[4]xi32>
   // CHECK: call <2 x i32> @llvm.vector.extract.v2i32.v8i32
   %6 = llvm.intr.vector.extract %f256[6] :
               vector<2xi32> from vector<8xi32>
   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
 // CHECK-DAG: declare float @llvm.fmuladd.f32(float, float, float)
 // CHECK-DAG: declare <8 x float> @llvm.fmuladd.v8f32(<8 x float>, <8 x float>, <8 x float>) #0
 // CHECK-DAG: declare void @llvm.prefetch.p0(ptr nocapture readonly, i32 immarg, i32 immarg, i32 immarg)
 // CHECK-DAG: declare float @llvm.exp.f32(float)
 // CHECK-DAG: declare <8 x float> @llvm.exp.v8f32(<8 x float>) #0

 // CHECK-DAG: declare <8 x i64> @llvm.vp.ptrtoint.v8i64.v8p0(<8 x ptr>, <8 x i1>, i32) #2
 // CHECK-DAG: declare <8 x ptr> @llvm.vp.inttoptr.v8p0.v8i64(<8 x i64>, <8 x i1>, i32) #2
 // CHECK-DAG: declare <vscale x 4 x i32> @llvm.vector.insert.nxv4i32.v8i32(<vscale x 4 x i32>, <8 x i32>, i64 immarg) #0
 // CHECK-DAG: declare <vscale x 4 x i32> @llvm.vector.insert.nxv4i32.v4i32(<vscale x 4 x i32>, <4 x i32>, i64 immarg) #0
 // CHECK-DAG: declare <8 x i32> @llvm.vector.insert.v8i32.v4i32(<8 x i32>, <4 x i32>, i64 immarg) #0
 // 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)