[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

345

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

345

does @rriddle's solution work?

[nicolasvasilache]

Please split out the NFC changes in a separate diff.

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

340

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

265

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);
   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_if(linalgOps.begin(), linalgOps.end(),
+                         [&](const Operation *op) { return op == 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 | FileCheck %s
+// RUN: mlir-opt %s -linalg-fusion -split-input-file | FileCheck %s
 
-#map0 = affine_map<(d0) -> (d0 + 2)>
-#map1 = affine_map<(d0) -> (d0 + 4)>
-#map2 = affine_map<(d0) -> (d0 + 3)>
-#map3 = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
-#map4 = affine_map<(d0) -> (d0)>
-#map5 = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
-#map6 = affine_map<(d0, d1) -> (d0, d1)>
-// CHECK-DAG: #[[strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s0 + d1 * s1)>
-
-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]> {
+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
   %c4 = constant 4 : index
   %c3 = constant 3 : index
   %c2 = constant 2 : index
   %0 = dim %A, 0 : memref<?x?xf32, offset: 0, strides: [?, 1]>
   %1 = dim %A, 1 : memref<?x?xf32, offset: 0, strides: [?, 1]>
   %2 = dim %B, 1 : memref<?x?xf32, offset: 0, strides: [?, 1]>
-  linalg.matmul(%A, %B, %C) : memref<?x?xf32, offset: 0, strides: [?, 1]>, memref<?x?xf32, offset: 0, strides: [?, 1]>, memref<?x?xf32, offset: 0, strides: [?, 1]>
+  linalg.matmul(%A, %B, %C) :
+    memref<?x?xf32, offset: 0, strides: [?, 1]>,
+    memref<?x?xf32, offset: 0, strides: [?, 1]>,
+    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: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %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: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
       }
     }
   }
   return %E : memref<?x?xf32, offset: 0, strides: [?, 1]>
 }
 // CHECK-LABEL: func @f1
 //       CHECK:   (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}})
 // No RAW dependences, the pass does not fuse RAR atm.
 //      CHECK: linalg.matmul
 //      CHECK: loop.for
 //      CHECK:   loop.for
 //      CHECK:     loop.for
 //      CHECK:       linalg.matmul
 
-func @f2(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>, %B: memref<?x?xf32, offset: 0, strides: [?, ?]>, %C: memref<?x?xf32, offset: 0, strides: [?, ?]>, %D: memref<?x?xf32, offset: 0, strides: [?, ?]>, %E: memref<?x?xf32, offset: 0, strides: [?, ?]>) -> memref<?x?xf32, offset: 0, strides: [?, ?]> {
+// -----
+
+// CHECK-DAG: #[[strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s0 + d1 * s1)>
+func @f2(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %B: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %C: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %D: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %E: memref<?x?xf32, offset: 0, strides: [?, ?]>
+        ) -> memref<?x?xf32, offset: 0, strides: [?, ?]> {
   %c1 = constant 1 : index
   %c0 = constant 0 : index
   %c4 = constant 4 : index
   %c3 = constant 3 : index
   %c2 = constant 2 : index
-  linalg.matmul(%A, %B, %C) : memref<?x?xf32, offset: 0, strides: [?, ?]>, memref<?x?xf32, offset: 0, strides: [?, ?]>, 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: [?, ?]>
   %0 = dim %C, 0 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   %1 = dim %C, 1 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   %2 = dim %D, 1 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   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 %C[%arg5, %arg7][%c2, %c4][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        %7 = std.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        %8 = std.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        linalg.matmul(%5, %7, %8) : memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %5 = std.subview %C[%arg5, %arg7][%c2, %c4][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %7 = std.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %8 = std.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        linalg.matmul(%5, %7, %8) :
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
       }
     }
   }
   return %E : memref<?x?xf32, offset: 0, strides: [?, ?]>
 }
 // CHECK-LABEL: func @f2
 //       CHECK:   (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}})
 //   CHECK-DAG:   %[[C_0:.*]] = dim %[[C]], 0 : memref<?x?xf32, #[[strided2D]]>
 //   CHECK-DAG:   %[[C_1:.*]] = dim %[[C]], 1 : memref<?x?xf32, #[[strided2D]]>
 //   CHECK-DAG:   %[[D_1:.*]] = dim %[[D]], 1 : memref<?x?xf32, #[[strided2D]]>
 //       CHECK:   loop.for %{{.*}} = %{{.*}} to %[[C_0]] step %{{.*}} {
 //       CHECK:     loop.for %{{.*}} = %{{.*}} to %[[D_1]] step %{{.*}} {
 //       CHECK:       loop.for %{{.*}} = %{{.*}} to %[[C_1]] step %{{.*}} {
 //       CHECK:         linalg.matmul
 //       CHECK:         linalg.matmul
 
-func @f3(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>, %B: memref<?x?xf32, offset: 0, strides: [?, ?]>, %C: memref<?x?xf32, offset: 0, strides: [?, ?]>, %D: memref<?x?xf32, offset: 0, strides: [?, ?]>, %E: memref<?x?xf32, offset: 0, strides: [?, ?]>) -> memref<?x?xf32, offset: 0, strides: [?, ?]> {
+// -----
+
+func @f3(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %B: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %C: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %D: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %E: memref<?x?xf32, offset: 0, strides: [?, ?]>
+        ) -> memref<?x?xf32, offset: 0, strides: [?, ?]> {
   %c1 = constant 1 : index
   %c0 = constant 0 : index
   %c4 = constant 4 : index
   %c3 = constant 3 : index
   %c2 = constant 2 : index
-  linalg.matmul(%A, %B, %C) : memref<?x?xf32, offset: 0, strides: [?, ?]>, memref<?x?xf32, offset: 0, strides: [?, ?]>, 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: [?, ?]>
   %0 = dim %D, 0 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   %1 = dim %D, 1 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   %2 = dim %C, 1 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   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 %D[%arg5, %arg7][%c2, %c4][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        %7 = std.subview %C[%arg7, %arg6][%c4, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        %8 = std.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        linalg.matmul(%5, %7, %8) : memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %5 = std.subview %D[%arg5, %arg7][%c2, %c4][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %7 = std.subview %C[%arg7, %arg6][%c4, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %8 = std.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        linalg.matmul(%5, %7, %8) :
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
       }
     }
   }
   return %E : memref<?x?xf32, offset: 0, strides: [?, ?]>
 }
 // CHECK-LABEL: func @f3
 //       CHECK:   (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}})
 //          CHECK:   %[[D_0:.*]] = dim %[[D]], 0 : memref<?x?xf32, #[[strided2D]]>
 //          CHECK:   %[[D_1:.*]] = dim %[[D]], 1 : memref<?x?xf32, #[[strided2D]]>
 //          CHECK:   %[[C_1:.*]] = dim %[[C]], 1 : memref<?x?xf32, #[[strided2D]]>
 //          CHECK:   loop.for %{{.*}} = %{{.*}} to %[[D_0]] step %{{.*}} {
 //          CHECK:     loop.for %{{.*}} = %{{.*}} to %[[C_1]] step %{{.*}} {
 //          CHECK:       loop.for %{{.*}} = %{{.*}} to %[[D_1]] step %{{.*}} {
 //          CHECK:         linalg.matmul
 //          CHECK:         linalg.matmul
 
-func @f4(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>, %B: memref<?x?xf32, offset: 0, strides: [?, ?]>, %C: memref<?x?xf32, offset: 0, strides: [?, ?]>, %D: memref<?x?xf32, offset: 0, strides: [?, ?]>, %E: memref<?x?xf32, offset: 0, strides: [?, ?]>) -> memref<?x?xf32, offset: 0, strides: [?, ?]> {
+// -----
+
+func @f4(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %B: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %C: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %D: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %E: memref<?x?xf32, offset: 0, strides: [?, ?]>
+        ) -> memref<?x?xf32, offset: 0, strides: [?, ?]> {
   %c1 = constant 1 : index
   %c0 = constant 0 : index
   %c4 = constant 4 : index
   %c3 = constant 3 : index
   %c2 = constant 2 : index
-  linalg.matmul(%A, %B, %C) : memref<?x?xf32, offset: 0, strides: [?, ?]>, memref<?x?xf32, offset: 0, strides: [?, ?]>, memref<?x?xf32, offset: 0, strides: [?, ?]>
-  linalg.matmul(%A, %B, %D) : memref<?x?xf32, offset: 0, strides: [?, ?]>, memref<?x?xf32, offset: 0, strides: [?, ?]>, 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: [?, ?]>
+  linalg.matmul(%A, %B, %D) :
+    memref<?x?xf32, offset: 0, strides: [?, ?]>,
+    memref<?x?xf32, offset: 0, strides: [?, ?]>,
+    memref<?x?xf32, offset: 0, strides: [?, ?]>
   %0 = dim %C, 0 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   %1 = dim %C, 1 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   %2 = dim %D, 1 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   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 %C[%arg5, %arg7][%c2, %c4][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        %7 = std.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        %8 = std.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        linalg.matmul(%5, %7, %8) : memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %5 = std.subview %C[%arg5, %arg7][%c2, %c4][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %7 = std.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %8 = std.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        linalg.matmul(%5, %7, %8) :
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
       }
     }
   }
   return %E : memref<?x?xf32, offset: 0, strides: [?, ?]>
 }
 // CHECK-LABEL: func @f4
 //       CHECK:   (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}})
 //          CHECK:   %[[C_0:.*]] = dim %[[C]], 0 : memref<?x?xf32, #[[strided2D]]>
 //          CHECK:   %[[C_1:.*]] = dim %[[C]], 1 : memref<?x?xf32, #[[strided2D]]>
 //          CHECK:   %[[D_1:.*]] = dim %[[D]], 1 : memref<?x?xf32, #[[strided2D]]>
 //          CHECK:   loop.for %{{.*}} = %{{.*}} to %[[C_0]] step %{{.*}} {
 //          CHECK:     loop.for %{{.*}} = %{{.*}} to %[[D_1]] step %{{.*}} {
 //          CHECK:       loop.for %{{.*}} = %{{.*}} to %[[C_1]] step %{{.*}} {
 // Fuse D then fuse C, no false dependence prevent it.
 //          CHECK:         linalg.matmul
 //          CHECK:         linalg.matmul
 //          CHECK:         linalg.matmul
 
-func @f5(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>, %B: memref<?x?xf32, offset: 0, strides: [?, ?]>, %C: memref<?x?xf32, offset: 0, strides: [?, ?]>, %D: memref<?x?xf32, offset: 0, strides: [?, ?]>, %E: memref<?x?xf32, offset: 0, strides: [?, ?]>) -> memref<?x?xf32, offset: 0, strides: [?, ?]> {
+// -----
+
+// CHECK-DAG: #[[strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s0 + d1 * s1)>
+func @f5(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %B: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %C: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %D: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %E: memref<?x?xf32, offset: 0, strides: [?, ?]>
+        ) -> memref<?x?xf32, offset: 0, strides: [?, ?]> {
   %c1 = constant 1 : index
   %c0 = constant 0 : index
   %c4 = constant 4 : index
   %c3 = constant 3 : index
   %c2 = constant 2 : index
   %0 = dim %B, 1 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   %1 = dim %D, 0 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   %2 = dim %D, 1 : 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: [?, ?]>
-  linalg.matmul(%C, %B, %D) : memref<?x?xf32, offset: 0, strides: [?, ?]>, memref<?x?xf32, offset: 0, strides: [?, ?]>, 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: [?, ?]>
+  linalg.matmul(%C, %B, %D) :
+    memref<?x?xf32, offset: 0, strides: [?, ?]>,
+    memref<?x?xf32, offset: 0, strides: [?, ?]>,
+    memref<?x?xf32, offset: 0, strides: [?, ?]>
   loop.for %arg5 = %c0 to %1 step %c2 {
     loop.for %arg6 = %c0 to %0 step %c3 {
       loop.for %arg7 = %c0 to %2 step %c4 {
-        %5 = std.subview %D[%arg5, %arg7][%c2, %c4][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        %7 = std.subview %B[%arg7, %arg6][%c4, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        %8 = std.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        linalg.matmul(%5, %7, %8) : memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %5 = std.subview %D[%arg5, %arg7][%c2, %c4][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %7 = std.subview %B[%arg7, %arg6][%c4, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %8 = std.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        linalg.matmul(%5, %7, %8) :
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, 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:         linalg.matmul
 
-func @f6(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>, %B: memref<?x?xf32, offset: 0, strides: [?, ?]>, %C: memref<?x?xf32, offset: 0, strides: [?, ?]>, %D: memref<?x?xf32, offset: 0, strides: [?, ?]>, %E: memref<?x?xf32, offset: 0, strides: [?, ?]>) -> memref<?x?xf32, offset: 0, strides: [?, ?]> {
+// -----
+
+#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: [?, ?]>,
+         %B: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %C: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %D: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %E: memref<?x?xf32, offset: 0, strides: [?, ?]>
+        ) -> memref<?x?xf32, offset: 0, strides: [?, ?]> {
   %c1 = constant 1 : index
   %c0 = constant 0 : index
   %c4 = constant 4 : index
   %c3 = constant 3 : index
   %c2 = constant 2 : index
   %0 = dim %C, 1 : 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: [?, ?]>
-  linalg.matmul(%A, %C, %E) : memref<?x?xf32, offset: 0, strides: [?, ?]>, memref<?x?xf32, offset: 0, strides: [?, ?]>, 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: [?, ?]>
+  linalg.matmul(%A, %C, %E) :
+    memref<?x?xf32, offset: 0, strides: [?, ?]>,
+    memref<?x?xf32, offset: 0, strides: [?, ?]>,
+    memref<?x?xf32, offset: 0, strides: [?, ?]>
   %1 = dim %C, 0 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   %2 = dim %D, 1 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   loop.for %arg5 = %c0 to %1 step %c2 {
     loop.for %arg6 = %c0 to %2 step %c3 {
       loop.for %arg7 = %c0 to %0 step %c4 {
         %3 = affine.apply #map0(%arg5)
         %4 = affine.apply #map1(%arg7)
-        %5 = std.subview %C[%arg5, %arg7][%c2, %c4][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %5 = std.subview %C[%arg5, %arg7][%c2, %c4][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
         %6 = affine.apply #map2(%arg6)
-        %7 = std.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        %8 = std.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        linalg.matmul(%5, %7, %8) : memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %7 = std.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %8 = std.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        linalg.matmul(%5, %7, %8) :
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
       }
     }
   }
   return %E : memref<?x?xf32, offset: 0, strides: [?, ?]>
 }
 // CHECK-LABEL: func @f6
 //       CHECK:   (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}})
 // Cannot fuse C due to interleaved read of C that would be bypassed.
 // Cannot fuse E (WAW).
 //   CHECK:   linalg.matmul
 //   CHECK:   linalg.matmul
 //   CHECK:   loop.for
 //   CHECK:     loop.for
 //   CHECK:       loop.for
 //   CHECK:         linalg.matmul
 // CHECK-NOT:       linalg.matmul
 
-func @f7(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>, %B: memref<?x?xf32, offset: 0, strides: [?, ?]>, %C: memref<?x?xf32, offset: 0, strides: [?, ?]>, %D: memref<?x?xf32, offset: 0, strides: [?, ?]>, %E: memref<?x?xf32, offset: 0, strides: [?, ?]>) -> memref<?x?xf32, offset: 0, strides: [?, ?]> {
+// -----
+
+func @f7(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %B: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %C: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %D: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %E: memref<?x?xf32, offset: 0, strides: [?, ?]>
+        ) -> memref<?x?xf32, offset: 0, strides: [?, ?]> {
   %c1 = constant 1 : index
   %c0 = constant 0 : index
   %c4 = constant 4 : index
   %c3 = constant 3 : index
   %c2 = constant 2 : index
   %0 = dim %A, 0 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   %1 = dim %A, 1 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   %2 = dim %C, 1 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   %3 = dim %C, 0 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   %4 = dim %D, 1 : memref<?x?xf32, offset: 0, strides: [?, ?]>
-  linalg.matmul(%A, %C, %E) : memref<?x?xf32, offset: 0, strides: [?, ?]>, memref<?x?xf32, offset: 0, strides: [?, ?]>, 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: [?, ?]>
+  linalg.matmul(%A, %C, %E) :
+    memref<?x?xf32, offset: 0, strides: [?, ?]>,
+    memref<?x?xf32, offset: 0, strides: [?, ?]>,
+    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 %arg5 = %c0 to %0 step %c2 {
     loop.for %arg6 = %c0 to %2 step %c3 {
       loop.for %arg7 = %c0 to %1 step %c4 {
-        %7 = std.subview %A[%arg5, %arg7][%c2, %c4][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        %9 = std.subview %C[%arg7, %arg6][%c4, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        %10 = std.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        linalg.matmul(%7, %9, %10) : memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %7 = std.subview %A[%arg5, %arg7][%c2, %c4][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %9 = std.subview %C[%arg7, %arg6][%c4, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %10 = std.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        linalg.matmul(%7, %9, %10) :
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
       }
     }
   }
   loop.for %arg5 = %c0 to %3 step %c2 {
     loop.for %arg6 = %c0 to %4 step %c3 {
       loop.for %arg7 = %c0 to %2 step %c4 {
-        %7 = std.subview %C[%arg5, %arg7][%c2, %c4][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        %9 = std.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        %10 = std.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        linalg.matmul(%7, %9, %10) : memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %7 = std.subview %C[%arg5, %arg7][%c2, %c4][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %9 = std.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %10 = std.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        linalg.matmul(%7, %9, %10) :
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
       }
     }
   }
   return %E : memref<?x?xf32, offset: 0, strides: [?, ?]>
 }
 // CHECK-LABEL: func @f7
 //       CHECK:   (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}})
 //       CHECK:   %[[A_0:.*]] = dim %[[A]], 0 : memref<?x?xf32, #[[strided2D]]>
 //       CHECK:   %[[A_1:.*]] = dim %[[A]], 1 : memref<?x?xf32, #[[strided2D]]>
 //       CHECK:   %[[C_1:.*]] = dim %[[C]], 1 : memref<?x?xf32, #[[strided2D]]>
 //       CHECK:   %[[C_0:.*]] = dim %[[C]], 0 : memref<?x?xf32, #[[strided2D]]>
 //       CHECK:   %[[D_1:.*]] = dim %[[D]], 1 : memref<?x?xf32, #[[strided2D]]>
 //       CHECK:   linalg.matmul(%[[A]], %[[C]], %[[E]])
 //       CHECK:   loop.for %{{.*}} = %{{.*}} to %[[A_0]] step %{{.*}} {
 //       CHECK:     loop.for %{{.*}} = %{{.*}} to %[[C_1]] step %{{.*}} {
 //       CHECK:       loop.for %{{.*}} = %{{.*}} to %[[A_1]] step %{{.*}} {
 //       CHECK:         linalg.matmul
 //       CHECK:         linalg.matmul
 //       CHECK:   loop.for %{{.*}} = %{{.*}} to %[[C_0]] step %{{.*}} {
 //       CHECK:     loop.for %{{.*}} = %{{.*}} to %[[D_1]] step %{{.*}} {
 //       CHECK:       loop.for %{{.*}} = %{{.*}} to %[[C_1]] step %{{.*}} {
 //       CHECK:         linalg.matmul
 //   CHECK-NOT:         linalg.matmul
 
-func @f8(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>, %B: memref<?x?xf32, offset: 0, strides: [?, ?]>, %C: memref<?x?xf32, offset: 0, strides: [?, ?]>, %D: memref<?x?xf32, offset: 0, strides: [?, ?]>, %E: memref<?x?xf32, offset: 0, strides: [?, ?]>) -> memref<?x?xf32, offset: 0, strides: [?, ?]> {
+// -----
+
+#map0 = affine_map<(d0) -> (d0 + 2)>
+#map1 = affine_map<(d0) -> (d0 + 4)>
+#map2 = affine_map<(d0) -> (d0 + 3)>
+
+func @f8(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %B: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %C: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %D: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+         %E: memref<?x?xf32, offset: 0, strides: [?, ?]>
+        ) -> memref<?x?xf32, offset: 0, strides: [?, ?]> {
   %c1 = constant 1 : index
   %c0 = constant 0 : index
   %c4 = constant 4 : index
   %c3 = constant 3 : index
   %c2 = constant 2 : index
   %0 = dim %A, 0 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   %1 = dim %A, 1 : memref<?x?xf32, offset: 0, strides: [?, ?]>
-  linalg.matmul(%A, %C, %D) : memref<?x?xf32, offset: 0, strides: [?, ?]>, memref<?x?xf32, offset: 0, strides: [?, ?]>, 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: [?, ?]>
+  linalg.matmul(%A, %C, %D) :
+    memref<?x?xf32, offset: 0, strides: [?, ?]>,
+    memref<?x?xf32, offset: 0, strides: [?, ?]>,
+    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: [?, ?]>
   %2 = dim %D, 1 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   loop.for %arg5 = %c0 to %0 step %c2 {
     loop.for %arg6 = %c0 to %2 step %c3 {
       loop.for %arg7 = %c0 to %1 step %c4 {
         %3 = affine.apply #map0(%arg5)
         %4 = affine.apply #map1(%arg7)
-        %5 = std.subview %A[%arg5, %arg7][%c2, %c4][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %5 = std.subview %A[%arg5, %arg7][%c2, %c4][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
         %6 = affine.apply #map2(%arg6)
-        %7 = std.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        %8 = std.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-        linalg.matmul(%5, %7, %8) : memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %7 = std.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %8 = std.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] :
+          memref<?x?xf32, offset: 0, strides: [?, ?]> to
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
+        linalg.matmul(%5, %7, %8) :
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>,
+          memref<?x?xf32, offset: ?, strides: [?, ?]>
       }
     }
   }
   return %E : memref<?x?xf32, offset: 0, strides: [?, ?]>
 }
 // CHECK-LABEL: func @f8
 //       CHECK:   (%[[A:.*]]: memref{{.*}}, %[[B:.*]]: memref{{.*}}, %[[C:.*]]: memref{{.*}}, %[[D:.*]]: memref{{.*}}, %[[E:.*]]: memref{{.*}})
 //   CHECK:   linalg.matmul
 //   CHECK:   linalg.matmul
 //   CHECK:   loop.for
 //   CHECK:     loop.for
 //   CHECK:       loop.for
 //   CHECK:         linalg.matmul
 // CHECK-NOT:       linalg.matmul
 
+// -----
+
 #id_2d = affine_map<(i, j) -> (i, j)>
 #pointwise_2d_trait = {
   args_in = 2,
   args_out = 1,
   indexing_maps = [#id_2d, #id_2d, #id_2d],
   iterator_types = ["parallel", "parallel"]
 }
-func @pointwise(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>, %B: memref<?x?xf32, offset: 0, strides: [?, ?]>, %C: memref<?x?xf32, offset: 0, strides: [?, ?]>, %D: memref<?x?xf32, offset: 0, strides: [?, ?]>) {
+func @pointwise(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+                %B: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+                %C: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+                %D: memref<?x?xf32, offset: 0, strides: [?, ?]>) {
   %c1 = constant 1 : index
   %c0 = constant 0 : index
   %c3 = constant 3 : index
   %c2 = constant 2 : index
   linalg.generic #pointwise_2d_trait %A, %A, %B {
   ^bb0(%E: f32, %arg5: f32, %arg6: f32):   // no predecessors
     %2 = addf %E, %arg5 : f32
     linalg.yield %2 : f32
-  }: memref<?x?xf32, offset: 0, strides: [?, ?]>, memref<?x?xf32, offset: 0, strides: [?, ?]>, memref<?x?xf32, offset: 0, strides: [?, ?]>
+  }: memref<?x?xf32, offset: 0, strides: [?, ?]>,
+     memref<?x?xf32, offset: 0, strides: [?, ?]>,
+     memref<?x?xf32, offset: 0, strides: [?, ?]>
   %0 = dim %B, 0 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   %1 = dim %B, 1 : memref<?x?xf32, offset: 0, strides: [?, ?]>
   loop.for %arg4 = %c0 to %0 step %c2 {
     loop.for %arg5 = %c0 to %1 step %c3 {
-      %4 = std.subview %B[%arg4, %arg5][%c2, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-      %5 = std.subview %C[%arg4, %arg5][%c2, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-      %6 = std.subview %D[%arg4, %arg5][%c2, %c3][%c1, %c1] : memref<?x?xf32, offset: 0, strides: [?, ?]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+      %4 = std.subview %B[%arg4, %arg5][%c2, %c3][%c1, %c1] :
+        memref<?x?xf32, offset: 0, strides: [?, ?]> to
+        memref<?x?xf32, offset: ?, strides: [?, ?]>
+      %5 = std.subview %C[%arg4, %arg5][%c2, %c3][%c1, %c1] :
+        memref<?x?xf32, offset: 0, strides: [?, ?]> to
+        memref<?x?xf32, offset: ?, strides: [?, ?]>
+      %6 = std.subview %D[%arg4, %arg5][%c2, %c3][%c1, %c1] :
+        memref<?x?xf32, offset: 0, strides: [?, ?]> to
+        memref<?x?xf32, offset: ?, strides: [?, ?]>
       linalg.generic #pointwise_2d_trait %4, %5, %6 {
       ^bb0(%arg6: f32, %arg7: f32, %arg8: f32):       // no predecessors
         %7 = mulf %arg6, %arg7 : f32
         linalg.yield %7 : f32
-      }: memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>
+      }: memref<?x?xf32, offset: ?, strides: [?, ?]>,
+         memref<?x?xf32, offset: ?, strides: [?, ?]>,
+         memref<?x?xf32, offset: ?, strides: [?, ?]>
     }
   }
   return
 }
 // CHECK-LABEL: func @pointwise
 //       CHECK:   loop.for
 //       CHECK:     loop.for
 //   CHECK-NOT:   loop.for
 //       CHECK:       linalg.generic
 //       CHECK:         addf
 //       CHECK:       linalg.generic
 //       CHECK:         mulf
 
+// -----
+
+#id_2d = affine_map<(i, j) -> (i, j)>
+#pointwise_2d_trait = {
+  args_in = 2,
+  args_out = 1,
+  indexing_maps = [#id_2d, #id_2d, #id_2d],
+  iterator_types = ["parallel", "parallel"]
+}
 func @pointwise_no_view(%M: index, %N: index) {
   %c1 = constant 1 : index
   %c0 = constant 0 : index
   %c3 = constant 3 : index
   %c2 = constant 2 : index
   %A = alloc (%M, %N): memref<?x?xf32>
   %B = alloc (%M, %N): memref<?x?xf32>
   %C = alloc (%M, %N): memref<?x?xf32>
   %D = alloc (%M, %N): memref<?x?xf32>
   %E = alloc (%M, %N): memref<?x?xf32>
   linalg.generic #pointwise_2d_trait %A, %A, %B {
   ^bb0(%e: f32, %arg5: f32, %arg6: f32):   // no predecessors
     %2 = addf %e, %arg5 : f32
     linalg.yield %2 : f32
-  }: memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>
+  }: memref<?x?xf32>,
+     memref<?x?xf32>,
+     memref<?x?xf32>
   %0 = dim %B, 0 : memref<?x?xf32>
   %1 = dim %B, 1 : memref<?x?xf32>
   loop.for %arg4 = %c0 to %0 step %c2 {
     loop.for %arg5 = %c0 to %1 step %c3 {
-      %4 = std.subview %B[%arg4, %arg5][%c2, %c3][%c1, %c1] : memref<?x?xf32> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-      %5 = std.subview %C[%arg4, %arg5][%c2, %c3][%c1, %c1] : memref<?x?xf32> to memref<?x?xf32, offset: ?, strides: [?, ?]>
-      %6 = std.subview %D[%arg4, %arg5][%c2, %c3][%c1, %c1] : memref<?x?xf32> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+      %4 = std.subview %B[%arg4, %arg5][%c2, %c3][%c1, %c1] :
+        memref<?x?xf32> to
+        memref<?x?xf32, offset: ?, strides: [?, ?]>
+      %5 = std.subview %C[%arg4, %arg5][%c2, %c3][%c1, %c1] :
+        memref<?x?xf32> to
+        memref<?x?xf32, offset: ?, strides: [?, ?]>
+      %6 = std.subview %D[%arg4, %arg5][%c2, %c3][%c1, %c1] :
+        memref<?x?xf32> to
+        memref<?x?xf32, offset: ?, strides: [?, ?]>
       linalg.generic #pointwise_2d_trait %4, %5, %6 {
       ^bb0(%arg6: f32, %arg7: f32, %arg8: f32):       // no predecessors
         %7 = mulf %arg6, %arg7 : f32
         linalg.yield %7 : f32
-      }: memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>, memref<?x?xf32, offset: ?, strides: [?, ?]>
+      }: memref<?x?xf32, offset: ?, strides: [?, ?]>,
+         memref<?x?xf32, offset: ?, strides: [?, ?]>,
+         memref<?x?xf32, offset: ?, strides: [?, ?]>
     }
   }
   return
 }
 // CHECK-LABEL: func @pointwise_no_view
 //       CHECK:   loop.for
 //       CHECK:     loop.for
 //   CHECK-NOT:   loop.for
 //       CHECK:       linalg.generic
 //       CHECK:         addf
 //       CHECK:       linalg.generic
 //       CHECK:         mulf
 
+// -----
+
+#map5 = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
+#map6 = affine_map<(d0, d1) -> (d0, d1)>
+#id_2d = affine_map<(i, j) -> (i, j)>
+#pointwise_2d_trait = {
+  args_in = 2,
+  args_out = 1,
+  indexing_maps = [#id_2d, #id_2d, #id_2d],
+  iterator_types = ["parallel", "parallel"]
+}
 func @indexed_generic_test(%A: memref<?x?xf32>,
                            %B: memref<?x?xf32>,
                            %C: memref<?x?xf32>,
                            %D: memref<?x?xf32>) {
   linalg.generic #pointwise_2d_trait %A, %B, %C {
   ^bb0(%e: f32, %arg5: f32, %arg6: f32):   // no predecessors
     %2 = addf %e, %arg5 : f32
     linalg.yield %2 : f32

 // 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