[mlir] Use memory effect to detecting allocation

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.

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

Author: antiagainst

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

@@ -129,14 +129,12 @@ class FloatArithmeticOp<string mnemonic, list<OpTrait> traits = []> :
 //
 //   %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<
     "OpBuilder &builder, OperationState &result, MemRefType memrefType", [{
@@ -278,7 +276,7 @@ def AddIOp : IntArithmeticOp<"addi", [Commutative]> {
 // 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

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

@@ -37,12 +37,26 @@ Value Aliases::find(Value v) {
   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;

mlir/test/Dialect/Linalg/fusion.mlir

@@ -727,3 +727,48 @@ func @fill_and_conv(%arg0: memref<1x4x5x1xf32>, %arg1: memref<2x3x1x1xf32>, %arg
 // 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