[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

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

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

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/LLVMDialect.h

public:
   /// Constructs a GEPArg with a constant index.
   /*implicit*/ GEPArg(int32_t integer) : BaseT(integer) {}
 
   /// Constructs a GEPArg with a dynamic index.
   /*implicit*/ GEPArg(Value value) : BaseT(value) {}
 
   using BaseT::operator=;
 };
+
+//===----------------------------------------------------------------------===//
+// Operation Traits
+//===----------------------------------------------------------------------===//
+
+namespace OpTrait {
+namespace impl {
+LogicalResult verifyImmArg(Operation *op, unsigned index);
+} // namespace impl
+
+/// This trait verifies that the nth operand is an "immediate". That is, it is
+/// define by an `llvm.mlir.constant` op.
+template <unsigned N>
+class ImmArg {
+public:
+  template <typename ConcreteType>
+  class Impl : public mlir::OpTrait::TraitBase<ConcreteType, ImmArg<N>::Impl> {
+  public:
+    static LogicalResult verifyTrait(Operation *op) {
+      return impl::verifyImmArg(op, N);
+    }
+  };
+};
+} // namespace OpTrait
+
 } // namespace LLVM
 } // namespace mlir
 
 ///// Ops /////
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LLVMIR/LLVMOps.h.inc"
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LLVMIR/LLVMIntrinsicOps.h.inc"

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
+//===----------------------------------------------------------------------===//
+
+def LLVM_LifetimeStartOp
+  : LLVM_ZeroResultIntrOp<"lifetime.start", [1], [ImmArg<0>]> {
+  let arguments = (ins I64:$size, LLVM_AnyPointer:$ptr);
+}
+
+def LLVM_LifetimeEndOp
+  : LLVM_ZeroResultIntrOp<"lifetime.end", [1], [ImmArg<0>]> {
+  let arguments = (ins I64:$size, LLVM_AnyPointer:$ptr);
+}
+
 // 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/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

class LLVM_OpBase<Dialect dialect, string mnemonic, list<Trait> traits = []> :
   //   - $_hasResult - substituted with a check that a variadic-result op does
   //                   have a result (LLVM ops can have 0 or 1 result);
   //   - $_location - mlir::Location object of the instruction.
   // Additionally, `$$` can be used to produce the dollar character.
   string llvmBuilder = "";
 }
 
 //===----------------------------------------------------------------------===//
+// LLVM Operation Traits
+//===----------------------------------------------------------------------===//
+
+class ImmArg<int N> : ParamNativeOpTrait<"ImmArg", !cast<string>(N)> {
+  let cppNamespace = "::mlir::LLVM::OpTrait";
+}
+
+//===----------------------------------------------------------------------===//
 // Base classes for LLVM dialect operations.
 //===----------------------------------------------------------------------===//
 
 // Base class for LLVM operations. All operations get an "llvm." prefix in
 // their name automatically. LLVM operations have either zero or one result,
 // this class is specialized below for both cases and should not be used
 // directly.
 class LLVM_Op<string mnemonic, list<Trait> traits = []> :
class LLVM_IntrOpBase<Dialect dialect, string opName, string enumName,
                       list<int> overloadedResults, list<int> overloadedOperands,
                       list<Trait> traits, int numResults,
                       bit requiresAccessGroup = 0, bit requiresAliasScope = 0>
     : LLVM_OpBase<dialect, opName, traits>,
       Results<!if(!gt(numResults, 0), (outs LLVM_Type:$res), (outs))> {
   string resultPattern = !if(!gt(numResults, 1),
                              LLVM_IntrPatterns.structResult,
                              LLVM_IntrPatterns.result);
-  string llvmEnumName = enumName;                           
+  string llvmEnumName = enumName;
   let llvmBuilder = [{
     llvm::Module *module = builder.GetInsertBlock()->getModule();
     llvm::Function *fn = llvm::Intrinsic::getDeclaration(
         module,
         llvm::Intrinsic::}] # enumName # [{,
         { }] # !interleave(!listconcat(
             ListIntSubst<resultPattern, overloadedResults>.lst,
             ListIntSubst<LLVM_IntrPatterns.operand,

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

Value mlir::LLVM::createGlobalString(Location loc, OpBuilder &builder,
   // Get the pointer to the first character in the global string.
   Value globalPtr = builder.create<LLVM::AddressOfOp>(loc, global);
   return builder.create<LLVM::GEPOp>(
       loc, LLVM::LLVMPointerType::get(IntegerType::get(ctx, 8)), globalPtr,
       ArrayRef<GEPArg>{0, 0});
 }
 
 bool mlir::LLVM::satisfiesLLVMModule(Operation *op) {
-  return op->hasTrait<OpTrait::SymbolTable>() &&
-         op->hasTrait<OpTrait::IsIsolatedFromAbove>();
+  return op->hasTrait<mlir::OpTrait::SymbolTable>() &&
+         op->hasTrait<mlir::OpTrait::IsIsolatedFromAbove>();
 }
 
 void FMFAttr::print(AsmPrinter &printer) const {
   printer << "<";
   printer << stringifyFastmathFlags(this->getFlags());
   printer << ">";
 }
 
auto parseLoopOptions = [&]() -> ParseResult {
     return success();
   };
   if (parser.parseCommaSeparatedList(parseLoopOptions) || parser.parseGreater())
     return {};
 
   llvm::sort(options, llvm::less_first());
   return get(parser.getContext(), options);
 }
+
+//===----------------------------------------------------------------------===//
+// Operation Traits
+//===----------------------------------------------------------------------===//
+
+LogicalResult LLVM::OpTrait::impl::verifyImmArg(Operation *op, unsigned index) {
+  if (!op->getOperand(index).getDefiningOp<LLVM::ConstantOp>())
+    return op->emitOpError("expected operand #")
+           << index << " to be an immediate argument";
+  return success();
+}

mlir/test/Dialect/LLVMIR/invalid.mlir

 }
 
 // -----
 
 func.func @extract_scalable_from_fixed_length_vector(%arg0 : vector<16xf32>) {
   // expected-error@+1 {{op failed to verify that it is not extracting scalable from fixed-length vectors.}}
   %0 = llvm.intr.vector.extract %arg0[0] : vector<[8]xf32> from vector<16xf32>
 }
+
+// -----
+
+func.func @lifetime_imm_arg(%size: i64) {
+  %0 = llvm.alloca %size x i32 : (i64) -> !llvm.ptr<i32>
+  // expected-error @below {{expected operand #0 to be an immediate argument}}
+  "llvm.intr.lifetime.start"(%size, %0) : (i64, !llvm.ptr<i32>) -> ()
+  return
+}