[MLIR][Linalg] Allow fusion of more than 2 linalg ops.

[pifon2a]

Summary

LinalgDependenceGraph was not updated after successful producer-consumer
fusion for linalg ops. In this patch it is fixed by reconstructing
LinalgDependenceGraph on every iteration. This is very ineffective and
should be improved by updating LDGraph only when it is necessary.

[pifon2a]

@nicolasvasilache This is not the best way of doing this update, I just wanted you to check if the idea is correct.

[merge_guards_bot]

Unit tests: pass. 62094 tests passed, 0 failed and 785 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[rriddle]

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

346

llvm::find(linalgOps, originalOp)?

[nicolasvasilache]

please make the commit message significantly more descriptive: what was the problem, how did you fix it etc

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

346

does @rriddle's solution work?

[nicolasvasilache]

Please split out the NFC changes in a separate diff.

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

341

Well that is a really dirty "fix" :)
It's true however that it is a problem with the current LinalgDependenceGraph impl that does not know how to update itself.
Still that problem already existed because the graph is recreated at each pattern invocation.

Let's let this fly for now but please adda big TODO that the graph needs to support mutation.

mlir/test/Dialect/Linalg/fusion.mlir

27

please avoid mixing in NFC formatting changes in the flow and send a separate CL for that.
I cannot tell what changed without spending too much time on it.

[pifon2a]

Addressed the comments.

[merge_guards_bot]

Unit tests: fail. 62343 tests passed, 1 failed and 839 were skipped.

failed: libc++.std/containers/sequences/array/array_creation/to_array.fail.cpp

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[nicolasvasilache]

Approved conditioned on beefing up the test and having no reordering surprises.

mlir/test/Dialect/Linalg/fusion.mlir

266

Sorry for being annoying but there is a risk of bad ordering here and I cannot tell whether this is correct.
Could we extend the test to capture the subviews into A, B, C, D

In the end we can only have:

linalg.matmul(%A, %B, %C) 
linalg.matmul(%C, %B, %D)
linalg.matmul(%5, %7, %8)

i.e. we need to make sure that there is no bypass.
It should be fine but I'd prefer to check this new behavior.

Thanks!

[nicolasvasilache]

Thanks for catching this!

[pifon2a]

Capture the indices.

[pifon2a]

@nicolasvasilache No problem, I captured the indices. But please, check if the fusion is correct.

[merge_guards_bot]

Unit tests: pass. 62414 tests passed, 0 failed and 839 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[nicolasvasilache]

The syntactic order of statements is good.
The fusion rules are too lenient and we end up recomputing things too many times.
This is a know current issue that I need to fix, filing a bug for it.

The TL;DR is do not count on fusion with reductions to compute the correct value atm.

[nicolasvasilache]

And thanks for improving the test!

[pifon2a]

https://github.com/llvm/llvm-project/commit/eda6b2e2b3f88e1e08328a229c2fc0f97b472ca4

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

static void fuseLinalgOpsGreedily(FuncOp f) {
 
   // Save original Linalg ops, we only want to make a pass over those.
   SmallVector<Operation *, 8> linalgOps;
   f.walk([&](LinalgOp op) {
     if (op.hasBufferSemantics())
       linalgOps.push_back(op);
   });
 
-  Aliases aliases;
-  LinalgDependenceGraph G(aliases, linalgOps);
+  // TODO(pifon, ntv): LinalgDependenceGraph should be able to update itself.
+  // The current naive and expensive reconstruction of the graph should be
+  // removed.
   for (auto *op : llvm::reverse(linalgOps)) {
-    for (unsigned consumerIdx = 0, e = LinalgOp(op).getNumInputs();
-         consumerIdx < e; ++consumerIdx) {
-      if (auto fusionInfo = fuseProducerOf(b, op, consumerIdx, G, &folder))
-        eraseSet.insert(fusionInfo->originalProducer.getOperation());
+    for (unsigned id = 0, e = LinalgOp(op).getNumInputs(); id < e; ++id) {
+      linalg::Aliases aliases;
+      linalg::LinalgDependenceGraph graph(aliases, linalgOps);
+      if (auto info = fuseProducerOf(b, op, id, graph, &folder)) {
+        auto *originalOp = info->originalProducer.getOperation();
+        eraseSet.insert(originalOp);
+        auto *originalOpInLinalgOpsVector =
+            std::find(linalgOps.begin(), linalgOps.end(), originalOp);
+        *originalOpInLinalgOpsVector = info->fusedProducer.getOperation();
+      }
     }
   }
-
   // The `fuseProducerOf` function performs structural checks and in particular
   // that no covering read or write exist between the consumer and the producer.
   // As a consequence, the only fusions that may occur preserve subsequent
   // dependences and are guaranteed by construction to produce the whole view.
   // We may thus erase the producer once it is fused.
   for (auto *e : eraseSet)
     e->erase();
   LLVM_DEBUG(f.print(dbgs() << "\nAfter linalg-fusion: \n"));

mlir/test/Dialect/Linalg/fusion.mlir

-// RUN: mlir-opt %s -linalg-fusion -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -linalg-fusion -split-input-file | FileCheck %s --dump-input-on-failure
 
 func @f1(%A: memref<?x?xf32, offset: 0, strides: [?, 1]>,
          %B: memref<?x?xf32, offset: 0, strides: [?, 1]>,
          %C: memref<?x?xf32, offset: 0, strides: [?, 1]>,
          %D: memref<?x?xf32, offset: 0, strides: [?, 1]>,
          %E: memref<?x?xf32, offset: 0, strides: [?, 1]>
         ) -> memref<?x?xf32, offset: 0, strides: [?, 1]> {
   %c0 = constant 0 : index
linalg.matmul(%A, %B, %C) :
     memref<?x?xf32, offset: 0, strides: [?, 1]>
   %c1 = constant 1 : index
   loop.for %arg5 = %c0 to %0 step %c2 {
     loop.for %arg6 = %c0 to %2 step %c3 {
       loop.for %arg7 = %c0 to %1 step %c4 {
         %5 = std.subview %A[%arg5, %arg7][%c2, %c4][%c1, %c1] :
           memref<?x?xf32, offset: 0, strides: [?, 1]> to
           memref<?x?xf32, offset: ?, strides: [?, ?]>
         %7 = std.subview %B[%arg7, %arg6][%c4, %c3][%c1, %c1] :
           memref<?x?xf32, offset: 0, strides: [?, 1]> to
           memref<?x?xf32, offset: ?, strides: [?, ?]>
         %8 = std.subview %C[%arg5, %arg6][%c2, %c3][%c1, %c1] :
           memref<?x?xf32, offset: 0, strides: [?, 1]> to
           memref<?x?xf32, offset: ?, strides: [?, ?]>
         linalg.matmul(%5, %7, %8) :
           memref<?x?xf32, offset: ?, strides: [?, ?]>,
           memref<?x?xf32, offset: ?, strides: [?, ?]>,
func @f5(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>,
   }
   return %E : memref<?x?xf32, offset: 0, strides: [?, ?]>
 }
 // CHECK-LABEL: func @f5
 // CHECK:  (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}})
 // CHECK-DAG:  %[[B_1:.*]] = dim %[[B]], 1 : memref<?x?xf32, #[[strided2D]]>
 // CHECK-DAG:  %[[D_0:.*]] = dim %[[D]], 0 : memref<?x?xf32, #[[strided2D]]>
 // CHECK-DAG:  %[[D_1:.*]] = dim %[[D]], 1 : memref<?x?xf32, #[[strided2D]]>
-// Don't fuse C due to false dependence, note that this is too conservative though.
-// CHECK:  linalg.matmul(%{{.*}}, %{{.*}}, %{{.*}})
-// CHECK:  loop.for %{{.*}} = %{{.*}} to %[[D_0]] step %{{.*}} {
-// CHECK:    loop.for %{{.*}} = %{{.*}} to %[[B_1]] step %{{.*}} {
-// CHECK:      loop.for %{{.*}} = %{{.*}} to %[[D_1]] step %{{.*}} {
-// CHECK:        linalg.matmul
-// CHECK:        linalg.matmul
+// CHECK:  loop.for %[[I:.*]] = %{{.*}} to %[[D_0]] step %{{.*}} {
+// CHECK:    loop.for %[[J:.*]] = %{{.*}} to %[[B_1]] step %{{.*}} {
+// CHECK:      loop.for %[[K:.*]] = %{{.*}} to %[[D_1]] step %{{.*}} {
+// CHECK-DAG:    %[[D_IK:.*]] = std.subview %[[D]][%[[I]], %[[K]]]
+// CHECK-DAG:    %[[B_KJ:.*]] = std.subview %[[B]][%[[K]], %[[J]]]
+// CHECK-DAG:    %[[E_IJ:.*]] = std.subview %[[E]][%[[I]], %[[J]]]
+// CHECK:        dim
+// CHECK-DAG:    %[[C_I0:.*]] = std.subview %[[C]][%[[I]], %{{.*}}]
+// CHECK-DAG:    %[[B_0K:.*]] = std.subview %[[B]][%{{.*}}, %[[K]]]
+// CHECK-DAG:    %[[D_IK_:.*]] = std.subview %[[D]][%[[I]], %[[K]]]
+// CHECK:        dim
+// CHECK-DAG:    %[[A_I0:.*]] = std.subview %[[A]][%[[I]], %{{.*}}]
+// CHECK-DAG:    %[[B_00:.*]] = std.subview %[[B]][%{{.*}}, %{{.*}}]
+// CHECK-DAG:    %[[C_I0_:.*]] = std.subview %[[C]][%[[I]], %{{.*}}]
+// CHECK:        linalg.matmul(%[[A_I0]], %[[B_00]], %[[C_I0_]])
+// CHECK:        linalg.matmul(%[[C_I0]], %[[B_0K]], %[[D_IK_]])
+// CHECK:        linalg.matmul(%[[D_IK]], %[[B_KJ]], %[[E_IJ]])
 
 // -----
 
 #map0 = affine_map<(d0) -> (d0 + 2)>
 #map1 = affine_map<(d0) -> (d0 + 4)>
 #map2 = affine_map<(d0) -> (d0 + 3)>
 
 func @f6(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>,

 // CHECK-LABEL: func @indexed_generic_test
 // CHECK:  loop.for
 // CHECK:    loop.for
 // CHECK-NOT:  loop.for
 // CHECK:      linalg.generic
 // CHECK:        addf
 // CHECK:      linalg.indexed_generic
 // CHECK:        index_cast
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d0)>
+#map1 = affine_map<(d0, d1) -> (d0, d1)>
+#map2 = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
+
+func @fusion_of_three(%arg0: memref<100x10xf32>,
+                      %arg1: memref<100xf32>,
+                      %arg2: memref<100x10xf32>) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %0 = alloc() {temp = true} : memref<100x10xf32>
+  linalg.generic {
+    args_in = 1 : i64,
+    args_out = 1 : i64,
+    indexing_maps = [#map0, #map1],
+    iterator_types = ["parallel", "parallel"]
+  } %arg1, %0 {
+      ^bb0(%arg3: f32, %arg4: f32): // no predecessors
+        linalg.yield %arg3 : f32
+      }: memref<100xf32>, memref<100x10xf32>
+  %1 = alloc() {temp = true} : memref<100x10xf32>
+  linalg.generic {
+    args_in = 2 : i64,
+    args_out = 1 : i64,
+    indexing_maps = [#map1, #map1, #map1],
+    iterator_types = ["parallel", "parallel"]
+  } %arg0, %0, %1 {
+      ^bb0(%arg3: f32, %arg4: f32, %arg5: f32): // no predecessors
+        %2 = subf %arg3, %arg4 : f32
+        linalg.yield %2 : f32
+      }: memref<100x10xf32>, memref<100x10xf32>, memref<100x10xf32>
+  dealloc %0 : memref<100x10xf32>
+  %2 = dim %1, 0 : memref<100x10xf32>
+  %3 = dim %1, 1 : memref<100x10xf32>
+  %4 = dim %arg2, 0 : memref<100x10xf32>
+  %5 = dim %arg2, 1 : memref<100x10xf32>
+  loop.for %i = %c0 to %2 step %c1 {
+    loop.for %j = %c0 to %3 step %c1 {
+      %6 = std.subview %1[%i, %j][%c1, %c1][%c1, %c1] :
+      memref<100x10xf32> to memref<?x?xf32, #map2>
+      %7 = std.subview %arg2[%i, %j][%c1, %c1][%c1, %c1] :
+      memref<100x10xf32> to memref<?x?xf32, #map2>
+      linalg.generic {
+        args_in = 1 : i64,
+        args_out = 1 : i64,
+        indexing_maps = [#map1, #map1],
+        iterator_types = ["parallel", "parallel"]
+      } %6, %7 {
+          ^bb0(%arg3: f32, %arg4: f32):     // no predecessors
+            %8 = exp %arg3 : f32
+            linalg.yield %8 : f32
+          }: memref<?x?xf32, #map2>,
+             memref<?x?xf32, #map2>
+    }
+  }
+ dealloc %1 : memref<100x10xf32>
+ return
+}
+// CHECK-LABEL: func @fusion
+// CHECK-NOT: linalg.generic
+// CHECK:     loop.for
+// CHECK:       loop.for
+// CHECK-NOT: loop.for
+// CHECK:       linalg.generic
+// CHECK:         linalg.yield
+// CHECK:       linalg.generic
+// CHECK:         subf
+// CHECK:         linalg.yield
+// CHECK:       linalg.generic
+// CHECK:         exp
+// CHECK:         linalg.yield