[mlir][memref] Add MemRef Lifetime Markers

[AlexandreSinger]

Summary

Add MemRef Lifetime Markers similar to those found within LLVM in order
to mark the lifetime of MemRefs within MLIR for use in analysis passes such
as buffer reause. This is further explained in an RFC found here:

https://discourse.llvm.org/t/rfc-lifetime-annotations-of-memory-within-mlir/72697

This patch is based on previous discussions on the LLVM lifetime markers with
the intent of making a more constrained version in MLIR using higher level
semantics. This old patch on the llvm repo does a great job explaining the ideal
assumptions the llvm lifetime markers should have:

https://reviews.llvm.org/D93376

[AlexandreSinger]

Updated the documentation of the proposed lifetime operations using points discussed in the RFC.

[mehdi_amini]

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

1194

Accessing a memref at this point should still be undefined behavior I believe, but you op description above isn't clear about this...

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let extraClassDeclaration = [{
      /// Returns the constant initial value if the memref.global is a constant,
      /// or null otherwise.
      ElementsAttr getConstantInitValue();
   }];
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// Lifetime[Start/End]Op
+//===----------------------------------------------------------------------===//
+
+class MemRef_LifetimeBaseOp<string mnemonic> : MemRef_Op<mnemonic> {
+  let arguments = (ins AnyMemRef:$memref);
+
+  let assemblyFormat = "$memref attr-dict `:` type($memref)";
+
+  let extraClassDeclaration = [{
+    MemRefType getMemRefType() {
+      return ::llvm::cast<MemRefType>(getMemref().getType());
+    }
+    int64_t getSize() {
+      MemRefType ty = getMemRefType();
+      if (ty.hasStaticShape())
+        return ty.getNumElements() * ty.getElementTypeBitWidth() / 8;
+      return -1;
+    }
+  }];
+}
+
+def MemRef_LifetimeStartOp : MemRef_LifetimeBaseOp<"lifetime_start"> {
+  let summary = "mark the start of a MemRef's lifetime";
+  let description = [{
+    The `lifetime_start` op marks the start of a MemRef's lifetime. Prior to
+    the execution of this op (if it exists), the MemRef is considered "dead"
+    and any memory access to it would yield undefined behavior.
+
+    Example:
+
+    ```mlir
+    %0 = memref.alloca() : memref<?xf32>
+    // No expected memory uses of %0
+    memref.lifetime_start %0 : memref<?xf32>
+    // Expected memory uses of %0
+    memref.lifetime_end %0 : memref<?xf32>
+    // No further expected memory uses of %0
+    ```
+
+    The lifetime semantics are not enforced by the op that allocated the MemRef,
+    and instead is used as an assumption for analysis passes used on the MemRef
+    This is useful for lifetime analysis (where the compiler can statically
+    analyze for memory out of bounds reads and writes). It is also useful for
+    transformation passes such as buffer reuse where allocations of MemRefs can
+    be statically analyzed for similar MemRefs with non-overlapping lifetimes to
+    be combined into a single allocation.
+
+    These operations are not expected to "execute" on a device and are used as
+    markers to define sections of execution where the writer of the program (or
+    the compiler) believes that the MemRef should be alive or dead.
+
+    Lifetime in the MemRef dialect uses the following assumptions, designed to
+    be a constrained subset of the assumptions used for
+    [LLVM Object Lifetime](https://llvm.org/docs/LangRef.html#object-lifetime):
+
+    1. Lifetime markers can only take MemRefs that have been allocated using
+       `alloc` or `alloca` ops. They cannot take subviewed MemRefs as operands.
+    2. A MemRef is considered alive at a point of execution if a
+       `lifetime_start` op is never called on this MemRef (within its use chain)
+       *OR* a `lifetime_start` op has executed prior to this point in execution
+       and a `lifetime_end` op has not.
+    3. Loading or storing to a dead MemRef is considered undefined behavior;
+       however, "Pure" ops that do not have memory effects can be called and
+       are expected to return valid results (such as `subview` ops).
+    4. A path of execution within the program cannot execute a `lifetime_start`
+       op on the same MemRef more than once.
+    5. It is possible for a path of execution to not execute a `lifetime_end`
+       op.
+  }];
+}
+
+def MemRef_LifetimeEndOp : MemRef_LifetimeBaseOp<"lifetime_end"> {
+  let summary = "mark the end of a MemRef's lifetime";
+  let description = [{
+    The `lifetime_end` op marks the end of a MemRef's lifetime. After this op is
+    executed, the MemRef is assumed to be "dead". Accessing the MemRef after
+    this point would yield undefined behavior.
+
+    See the  `lifetime_start` op for more information on lifetime within the
+    MemRef dialect.
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // LoadOp
 //===----------------------------------------------------------------------===//
 
 def LoadOp : MemRef_Op<"load",
      [TypesMatchWith<"result type matches element type of 'memref'",
                      "memref", "result",
                      "::llvm::cast<MemRefType>($_self).getElementType()">,
       MemRefsNormalizable,

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

 //===- MemRefToLLVM.cpp - MemRef to LLVM dialect conversion ---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
 
 #include "mlir/Analysis/DataLayoutAnalysis.h"
 #include "mlir/Conversion/ConvertToLLVM/ToLLVMInterface.h"
 #include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
+#include "mlir/Conversion/LLVMCommon/MemRefBuilder.h"
 #include "mlir/Conversion/LLVMCommon/Pattern.h"
 #include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 #include "mlir/Conversion/MemRefToLLVM/AllocLikeConversion.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/MathExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include <optional>
 
 namespace mlir {
 #define GEN_PASS_DEF_FINALIZEMEMREFTOLLVMCONVERSIONPASS
auto deadBeefPtr =
         rewriter.create<LLVM::IntToPtrOp>(loc, elementPtrType, deadBeefConst);
 
     // Both allocated and aligned pointers are same. We could potentially stash
     // a nullptr for the allocated pointer since we do not expect any dealloc.
     return std::make_tuple(deadBeefPtr, gep);
   }
 };
 
+// Lifetime marker operations are lowered to their associated intrinsic.
+template <typename MemRefOp, typename LLVMOp>
+struct LifetimeOpLowering : public ConvertOpToLLVMPattern<MemRefOp> {
+  using ConvertOpToLLVMPattern<MemRefOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(MemRefOp op, typename MemRefOp::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    MemRefDescriptor desc = MemRefDescriptor(adaptor.getMemref());
+    Value ptr = desc.alignedPtr(rewriter, op.getLoc());
+    IntegerAttr size = rewriter.getI64IntegerAttr(op.getSize());
+    rewriter.replaceOpWithNewOp<LLVMOp>(op, size, ptr);
+    return success();
+  }
+};
+
+using LifetimeStartOpLowering =
+    LifetimeOpLowering<memref::LifetimeStartOp, LLVM::LifetimeStartOp>;
+using LifetimeEndOpLowering =
+    LifetimeOpLowering<memref::LifetimeEndOp, LLVM::LifetimeEndOp>;
+
 // Load operation is lowered to obtaining a pointer to the indexed element
 // and loading it.
 struct LoadOpLowering : public LoadStoreOpLowering<memref::LoadOp> {
   using Base::Base;
 
   LogicalResult
   matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
patterns.add<
       AtomicRMWOpLowering,
       AssumeAlignmentOpLowering,
       ConvertExtractAlignedPointerAsIndex,
       DimOpLowering,
       ExtractStridedMetadataOpLowering,
       GenericAtomicRMWOpLowering,
       GlobalMemrefOpLowering,
       GetGlobalMemrefOpLowering,
+      LifetimeEndOpLowering,
+      LifetimeStartOpLowering,
       LoadOpLowering,
       MemRefCastOpLowering,
       MemRefCopyOpLowering,
       MemorySpaceCastOpLowering,
       MemRefReinterpretCastOpLowering,
       MemRefReshapeOpLowering,
       PrefetchOpLowering,
       RankOpLowering,

mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir

 // CHECK-LABEL: func @store_non_temporal(
 func.func @store_non_temporal(%input : memref<32xf32, affine_map<(d0) -> (d0)>>, %output : memref<32xf32, affine_map<(d0) -> (d0)>>) {
   %1 = arith.constant 7 : index
   %2 = memref.load %input[%1] {nontemporal = true} : memref<32xf32, affine_map<(d0) -> (d0)>>
   // CHECK: llvm.store %{{.*}}, %{{.*}}  {nontemporal} : f32, !llvm.ptr
   memref.store %2, %output[%1] {nontemporal = true} : memref<32xf32, affine_map<(d0) -> (d0)>>
   func.return
 }
+
+// -----
+
+// CHECK-LABEL: func @lifetime_start_end
+func.func @lifetime_start_end(%d : index) {
+  %0 = memref.alloca() : memref<17x42xf32>
+  // CHECK: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+  // CHECK: %[[BASE_PTR:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+  // CHECK: llvm.intr.lifetime.start 2856, %[[BASE_PTR]] : !llvm.ptr
+  memref.lifetime_start %0 : memref<17x42xf32>
+  // CHECK: %[[BASE_PTR_2:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+  // CHECK: llvm.intr.lifetime.end 2856, %[[BASE_PTR_2]] : !llvm.ptr
+  memref.lifetime_end %0 : memref<17x42xf32>
+  %1 = memref.alloca(%d) : memref<?xf32>
+  // CHECK: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+  // CHECK: %[[BASE_PTR_3:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+  // CHECK: llvm.intr.lifetime.start -1, %[[BASE_PTR_3]] : !llvm.ptr
+  memref.lifetime_start %1 : memref<?xf32>
+  // CHECK: %[[BASE_PTR_4:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+  // CHECK: llvm.intr.lifetime.end -1, %[[BASE_PTR_4]] : !llvm.ptr
+  memref.lifetime_end %1 : memref<?xf32>
+  func.return
+}

mlir/test/Dialect/MemRef/ops.mlir

func.func @memref_extract_aligned_pointer(%src : memref<?xf32>) -> index {
   return %0 : index
 }
 
 // CHECK-LABEL: func @memref_memory_space_cast
 func.func @memref_memory_space_cast(%src : memref<?xf32>) -> memref<?xf32, 1> {
   %dst = memref.memory_space_cast %src : memref<?xf32> to memref<?xf32, 1>
   return %dst : memref<?xf32, 1>
 }
+
+// CHECK-LABEL: func @memref_lifetime
+func.func @memref_lifetime(%d : index) {
+  %0 = memref.alloca() : memref<32xf32>
+  // CHECK: memref.lifetime_start %{{.*}} : memref<32xf32>
+  memref.lifetime_start %0 : memref<32xf32>
+  // CHECK-NEXT: memref.lifetime_end %{{.*}} : memref<32xf32>
+  memref.lifetime_end %0 : memref<32xf32>
+  return
+}