[Mlir-commits] [mlir] 87e07b4 - [mlir] Use memory effect to detecting allocation

[Lei Zhang]

Author: Lei Zhang
Date: 2020-04-30T09:20:53-04:00
New Revision: 87e07b4c643f5dccd8ce202be0e715ddd3a0998e

URL: https://github.com/llvm/llvm-project/commit/87e07b4c643f5dccd8ce202be0e715ddd3a0998e
DIFF: https://github.com/llvm/llvm-project/commit/87e07b4c643f5dccd8ce202be0e715ddd3a0998e.diff

LOG: [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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
    mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp
    mlir/test/Dialect/Linalg/fusion.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
index 87f8e629a5c6..1f1658f9695f 100644
--- a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
+++ b/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

diff  --git a/mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp b/mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp
index 342a909df88b..4c218503ba17 100644
--- a/mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp
+++ b/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;

diff  --git a/mlir/test/Dialect/Linalg/fusion.mlir b/mlir/test/Dialect/Linalg/fusion.mlir
index 14a12840d1d0..3f20fd842675 100644
--- a/mlir/test/Dialect/Linalg/fusion.mlir
+++ b/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 
diff erent allocation-like ops are recognized and properly handled.
+func @accept_
diff erent_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_
diff erent_alloc_ops
+// CHECK-COUNT-3: loop.for
+// CHECK-COUNT-2:   linalg.matmul