[mlir][linalg] Use memory effect to detecting allocation

[antiagainst]

Summary

This commit marks AllocLikeOp as MemAlloc in StandardOps.

Also in Linalg dependency analysis use memory effect to detect
allocation. This allows the dependency analysis to be more
general and recognize other allocation-like operations

[rriddle]

Thanks!

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

279

Can we add to AllocLikeOp instead? Ideally marking the returned value directly.

[mehdi_amini]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

279

Should this be rebased on top of https://reviews.llvm.org/D78619 ?

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

41–42

FYI, less operations:

Operation *defOp = v.getDefiningOp();
if (!defOp) return v;

44

The "MemAlloc" trait can be attached to individual results, I believe that the right way here is to look-up which result is v for the operation and query the interface for the MemAlloc trait for this particular result.

[antiagainst]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

279

Oh, thanks for the pointer! I'll wait that CL to land and then rebase. :)

[antiagainst]

Rebase and address comments

[mehdi_amini]

Can this be exercised by a test? (using a test.alloc op?)

[antiagainst]

Address comments

[antiagainst]

In D78705#2006112, @mehdi_amini wrote:

Can this be exercised by a test? (using a test.alloc op?)

Actually we have two allocation-like ops in std. I think leveraging them should suffice for a test over the memory effect here.

[antiagainst]

Rebase

[antiagainst]

@mehdi_amini: Any further comments?

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

 class FloatArithmeticOp<string mnemonic, list<OpTrait> traits = []> :
     ArithmeticOp<mnemonic, traits>,
     Arguments<(ins FloatLike:$lhs, FloatLike:$rhs)>;
 
 // Base class for memref allocating ops: alloca and alloc.
 //
 //   %0 = alloclike(%m)[%s] : memref<8x?xf32, (d0, d1)[s0] -> ((d0 + s0), d1)>
 //
-class AllocLikeOp<string mnemonic,
-                  list<OpVariableDecorator> resultDecorators = [],
-                  list<OpTrait> traits = []> :
-    Std_Op<mnemonic, traits> {
+class AllocLikeOp<string mnemonic, list<OpTrait> traits = []> :
+    Std_Op<mnemonic, !listconcat(traits, [MemoryEffects<[MemAlloc]>])> {
 
   let arguments = (ins Variadic<Index>:$value,
                    Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$alignment);
-  let results = (outs Arg<AnyMemRef, "", resultDecorators>);
+  let results = (outs Res<AnyMemRef, "", [MemAlloc]>);
 
   let builders = [OpBuilder<
     "Builder *builder, OperationState &result, MemRefType memrefType", [{
        result.types.push_back(memrefType);
      }]>,
     OpBuilder<
     "Builder *builder, OperationState &result, MemRefType memrefType, " #
     "ValueRange operands, IntegerAttr alignment = IntegerAttr()", [{
def AddIOp : IntArithmeticOp<"addi", [Commutative]> {
   }];
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // AllocOp
 //===----------------------------------------------------------------------===//
 
-def AllocOp : AllocLikeOp<"alloc", [MemAlloc], [MemoryEffects<[MemAlloc]>]> {
+def AllocOp : AllocLikeOp<"alloc"> {
   let summary = "memory allocation operation";
   let description = [{
     The `alloc` operation allocates a region of memory, as specified by its
     memref type.
 
     Example:
 
     ```mlir

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

Value Aliases::find(Value v) {
   if (it != aliases.end()) {
     assert(it->getSecond().getType().isa<MemRefType>() && "Memref expected");
     return it->getSecond();
   }
 
   while (true) {
     if (v.isa<BlockArgument>())
       return v;
+
     Operation *defOp = v.getDefiningOp();
-    if (auto alloc = dyn_cast_or_null<AllocOp>(defOp)) {
-      if (isStrided(alloc.getType()))
-        return alloc.getResult();
+    if (!defOp)
+      return v;
+
+    if (auto memEffect = dyn_cast<MemoryEffectOpInterface>(defOp)) {
+      // Collect all memory effects on `v`.
+      SmallVector<MemoryEffects::EffectInstance, 1> effects;
+      memEffect.getEffectsOnValue(v, effects);
+
+      // If we have the 'Allocate' memory effect on `v`, then `v` should be the
+      // original buffer.
+      if (llvm::any_of(
+              effects, [](const MemoryEffects::EffectInstance &instance) {
+                return isa<MemoryEffects::Allocate>(instance.getEffect());
+              }))
+        return v;
     }
-    if (auto viewLikeOp = dyn_cast_or_null<ViewLikeOpInterface>(defOp)) {
+
+    if (auto viewLikeOp = dyn_cast<ViewLikeOpInterface>(defOp)) {
       auto it =
           aliases.insert(std::make_pair(v, find(viewLikeOp.getViewSource())));
       return it.first->second;
     }
 
     llvm::errs() << "View alias analysis reduces to: " << v << "\n";
     llvm_unreachable("unsupported view alias case");
   }

mlir/test/Dialect/Linalg/fusion.mlir

func @fill_and_conv(%arg0: memref<1x4x5x1xf32>, %arg1: memref<2x3x1x1xf32>, %arg2: memref<1x4x5x1xf32>) {
   }
   return
 }
 // CHECK-LABEL: func @fill_and_conv
 // CHECK: loop.for
 // CHECK:   loop.for
 // CHECK:     linalg.fill
 // CHECK:     linalg.conv
+
+// -----
+
+// Test that different allocation-like ops are recognized and properly handled.
+func @accept_different_alloc_ops(%dim: index, %s0 : index, %s1: index) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c3 = constant 3 : index
+  %c4 = constant 4 : index
+
+  %A = alloca(%dim, %dim)[%s0, %s1] : memref<?x?xf32, offset: 0, strides: [?, ?]>
+  %B = alloca(%dim, %dim)[%s0, %s1] : memref<?x?xf32, offset: 0, strides: [?, ?]>
+  %C = alloc(%dim, %dim)[%s0, %s1]  : memref<?x?xf32, offset: 0, strides: [?, ?]>
+
+  linalg.matmul(%A, %B, %C) :
+    memref<?x?xf32, offset: 0, strides: [?, ?]>,
+    memref<?x?xf32, offset: 0, strides: [?, ?]>,
+    memref<?x?xf32, offset: 0, strides: [?, ?]>
+
+  loop.for %i = %c0 to %dim step %c2 {
+    loop.for %j = %c0 to %dim step %c3 {
+      loop.for %k = %c0 to %dim step %c4 {
+        %0 = std.subview %A[%i, %k][%c2, %c4][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %1 = std.subview %B[%k, %j][%c4, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %2 = std.subview %C[%i, %j][%c2, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        linalg.matmul(%0, %1, %2) :
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+      }
+    }
+  }
+  return
+}
+
+// CHECK-LABEL: func @accept_different_alloc_ops
+// CHECK-COUNT-3: loop.for
+// CHECK-COUNT-2:   linalg.matmul