[MLIR] Fix fusion of linalg.indexed_generic producer into tiled (Indexed)GenericOp.

[nicolasvasilache]

Looks great, thanks Alex!

mlir/test/Dialect/Linalg/fusion_indexed_generic.mlir

59

Can we add a

// CHECK-NOT:  addi

above this line?

[bondhugula]

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

107

You don't need region().getBlocks().front(), region().front() will work.

[bondhugula]

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

114

Please add a comment here on the replacement.

115–119

Could you just use replaceAllUsesExcept(oldIndex, newIndex, {newIndexOp} for this entire six lines? It's just needs to be exposed.

[ftynse]

A slight concern on this being ultimately used within the pattern rewriting infrastructure that does interact well with brutal editing of uses. There are other aspects in this code that make me think it shouldn't be called from that infrastructure anyway, so not blocking.

[pifon2a]

Addressed the comments and exposed replaceAllUsesExcept in LoopUtils.h.

[pifon2a]

Thanks Nicolas, Alex, Uday!

[bondhugula]

mlir/include/mlir/Transforms/LoopUtils.h

290 ↗

(On Diff #257994)

Why is this added to LoopUtils.h? It has nothing to do with loops in particular. This could have been in Value.h.

[ftynse]

mlir/include/mlir/Transforms/LoopUtils.h

290 ↗

(On Diff #257994)

It sounds like moving that across can be a separate NFC change

[pifon2a]

mlir/include/mlir/Transforms/LoopUtils.h

290 ↗

(On Diff #257994)

Yes, here it is https://reviews.llvm.org/D78426

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

for (auto r : viewRanges) {
       sizes.push_back(r.size);
       strides.push_back(r.stride);
     }
     clonedViews.push_back(
         b.create<SubViewOp>(loc, view, offsets, sizes, strides));
   }
   auto operands = getAssumedNonViewOperands(op);
   clonedViews.append(operands.begin(), operands.end());
-  return op.clone(b, loc, clonedViews);
+
+  Operation *clonedOp = op.clone(b, loc, clonedViews);
+  // When the producer is an IndexedGenercOp, we have to transform its block
+  // IV arguments according to the tiling of the consumer, i.e. offset them by
+  // the values computed in `loopRanges`.
+  if (auto indexedGenericOp = dyn_cast<IndexedGenericOp>(clonedOp)) {
+    auto &block = indexedGenericOp.region().getBlocks().front();
+
+    OpBuilder::InsertionGuard g(b);
+    b.setInsertionPointToStart(&block);
+    for (unsigned i = 0, e = indexedGenericOp.getNumLoops(); i < e; ++i) {
+      Value oldIndex = block.getArgument(i);
+      Value newIndex = b.create<AddIOp>(indexedGenericOp.getLoc(), oldIndex,
+                                        loopRanges[i].offset);
+      for (auto &use : oldIndex.getUses()) {
+        if (use.getOwner() == newIndex.getDefiningOp())
+          continue;
+        use.set(newIndex);
+      }
+    }
+  }
+  return clonedOp;
 }
 
 struct ViewDimension {
   Value view;
   unsigned dimension;
 };
 
 // Given an `op`, returns the first (`view`, `dimension`) pair that identifies
fuseProducerOfDep(OpBuilder &b, LinalgOp consumer, unsigned consumerIdx,
   assert(consumer.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
   LLVM_DEBUG(dbgs() << "\nStart examining consumer: "
                     << *consumer.getOperation());
   for (auto dependence : graph.getDependencesInto(consumer, depType)) {
     LLVM_DEBUG(dbgs() << "\n***Consider producer:\t"
                       << *dependence.dependentOpView.op << "\n");
     auto producer = cast<LinalgOp>(dependence.dependentOpView.op);
-    if (isa<linalg::IndexedGenericOp>(dependence.dependentOpView.op)) {
-      LLVM_DEBUG(dbgs() << "Not fusing indexed_generic producer");
-      continue;
-    }
 
     // Check that the dependence is indeed on the input `consumerIdx` view.
     auto consumedView = dependence.indexingView;
     if (consumer.getBuffer(consumerIdx) != consumedView)
       continue;
 
     // Consumer consumes this view, `isStructurallyFusableProducer` also checks
     // whether it is a strict subview of the producer view.

mlir/test/Dialect/Linalg/fusion.mlir

 // 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
-  }: memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>
-  %c1 = constant 1 : index
-  %c0 = constant 0 : index
-  %c25 = constant 25 : index
-  %c10 = constant 10 : index
-  %0 = dim %C, 0 : memref<?x?xf32>
-  %1 = dim %C, 1 : memref<?x?xf32>
-  %2 = dim %D, 0 : memref<?x?xf32>
-  %3 = dim %D, 1 : memref<?x?xf32>
-  loop.for %arg2 = %c0 to %0 step %c10 {
-    loop.for %arg3 = %c0 to %1 step %c25 {
-      %4 = std.subview %C[%arg2, %arg3][%c10, %c25][%c1, %c1] :
-          memref<?x?xf32> to memref<?x?xf32, #map5>
-      %5 = std.subview %D[%arg2, %arg3][%c10, %c25][%c1, %c1] :
-          memref<?x?xf32> to memref<?x?xf32, #map5>
-      linalg.indexed_generic {
-        indexing_maps = [#map6, #map6],
-        iterator_types = ["parallel", "parallel"],
-        args_in = 1,
-        args_out = 1
-      } %4, %5 {
-      ^bb0(%arg4: index, %arg5: index, %arg6: f32, %arg7: f32):
-        %6 = addi %arg4, %arg2 : index
-        %7 = addi %arg5, %arg3 : index
-        %8 = index_cast %6 : index to i32
-        %9 = sitofp %8 : i32 to f32
-        %10 = index_cast %7 : index to i32
-        %11 = sitofp %10 : i32 to f32
-        %12 = addf %9, %11 : f32
-        linalg.yield %12 : f32
-      }: memref<?x?xf32, #map5>, memref<?x?xf32, #map5>
-    }
-  }
-  return
-}
-// 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
-
-// -----
-
-//
-// We should not be fusing indexed_generic into a generic yet.
-// https://bugs.llvm.org/show_bug.cgi?id=44875
-//
-
-#map0 = affine_map<(d0)[s0,s1] -> (d0 * s1 + s0)>
-#pointwise_map = affine_map<(d0) -> (d0)>
-#pointwise_1d_trait = {
-  args_in = 1,
-  args_out = 1,
-  indexing_maps = [#pointwise_map, #pointwise_map],
-  iterator_types = ["parallel"]
-}
-
-func @nofuse_indexed_generic(%A: memref<?xf32>, %B: memref<?xf32>, %C: memref<?xf32>) {
-  linalg.indexed_generic #pointwise_1d_trait %A, %B {
-  ^bb0(%i: index, %a: f32, %b: f32):
-    linalg.yield %a : f32
-  }: memref<?xf32>, memref<?xf32>
-
-  %c0 = constant 0 : index
-  %c1 = constant 1 : index
-  %c10 = constant 10 : index
-  %dB = dim %B, 0 : memref<?xf32>
-  loop.for %i = %c0 to %dB step %c10 {
-    %subB = subview %B[%i][%c10][%c1] : memref<?xf32> to memref<?xf32, #map0>
-    %subC = subview %C[%i][%c10][%c1] : memref<?xf32> to memref<?xf32, #map0>
-    linalg.generic #pointwise_1d_trait %subB, %subC {
-    ^bb0(%b: f32, %c: f32):
-      linalg.yield %b : f32
-    }: memref<?xf32, #map0>, memref<?xf32, #map0>
-  }
-  return
-}
-// CHECK-LABEL: func @nofuse_indexed_generic
-// CHECK-NOT: loop.for
-// CHECK:     linalg.indexed_generic
-// CHECK:     loop.for
-// CHECK-NOT:   linalg.indexed_generic
-// CHECK:       linalg.generic
 
 // -----
 
 #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>,

mlir/test/Dialect/Linalg/fusion_indexed_generic.mlir

+// RUN: mlir-opt %s -linalg-fusion -split-input-file | FileCheck %s --dump-input-on-failure
+
+#map = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
+#id_2d = affine_map<(d0, d1) -> (d0, d1)>
+#pointwise_2d_trait = {
+  args_in = 2,
+  args_out = 1,
+  indexing_maps = [#id_2d, #id_2d, #id_2d],
+  iterator_types = ["parallel", "parallel"]
+}
+func @fuse_indexed_generic_consumer(%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
+  }: memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %c25 = constant 25 : index
+  %c10 = constant 10 : index
+  %0 = dim %C, 0 : memref<?x?xf32>
+  %1 = dim %C, 1 : memref<?x?xf32>
+  %2 = dim %D, 0 : memref<?x?xf32>
+  %3 = dim %D, 1 : memref<?x?xf32>
+  loop.for %arg2 = %c0 to %0 step %c10 {
+    loop.for %arg3 = %c0 to %1 step %c25 {
+      %4 = std.subview %C[%arg2, %arg3][%c10, %c25][%c1, %c1] :
+          memref<?x?xf32> to memref<?x?xf32, #map>
+      %5 = std.subview %D[%arg2, %arg3][%c10, %c25][%c1, %c1] :
+          memref<?x?xf32> to memref<?x?xf32, #map>
+      linalg.indexed_generic {
+        indexing_maps = [#id_2d, #id_2d],
+        iterator_types = ["parallel", "parallel"],
+        args_in = 1,
+        args_out = 1
+      } %4, %5 {
+      ^bb0(%arg4: index, %arg5: index, %arg6: f32, %arg7: f32):
+        %6 = addi %arg4, %arg2 : index
+        %7 = addi %arg5, %arg3 : index
+        %8 = index_cast %6 : index to i32
+        %9 = sitofp %8 : i32 to f32
+        %10 = index_cast %7 : index to i32
+        %11 = sitofp %10 : i32 to f32
+        %12 = addf %9, %11 : f32
+        linalg.yield %12 : f32
+      }: memref<?x?xf32, #map>, memref<?x?xf32, #map>
+    }
+  }
+  return
+}
+// CHECK-LABEL: func @fuse_indexed_generic_consumer
+// CHECK:  loop.for
+// CHECK:    loop.for
+// CHECK-NOT:  loop.for
+// CHECK:      linalg.generic
+// CHECK:        addf
+// CHECK:      linalg.indexed_generic
+// CHECK:        index_cast
+
+// -----
+
+#map = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
+#id_2d = affine_map<(d0, d1) -> (d0, d1)>
+#pointwise_2d_trait = {
+  args_in = 2,
+  args_out = 1,
+  indexing_maps = [#id_2d, #id_2d, #id_2d],
+  iterator_types = ["parallel", "parallel"]
+}
+func @fuse_indexed_generic_producer(%A: memref<?x?xf32>,
+                                    %B: memref<?x?xf32>,
+                                    %C: memref<?x?xf32>,
+                                    %D: memref<?x?xf32>) {
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %c25 = constant 25 : index
+  %c10 = constant 10 : index
+  linalg.indexed_generic #pointwise_2d_trait %A, %B, %C {
+    ^bb0(%i: index, %j: index, %a: f32, %b: f32, %c: f32): // no predecessors
+      %i_int = index_cast %i: index to i32
+      %i_float = sitofp %i_int : i32 to f32
+      %ab = addf %a, %b : f32
+      %out = addf %ab, %i_float : f32
+      linalg.yield %out : f32
+  }: memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>
+  %C_X = dim %C, 0 : memref<?x?xf32>
+  %C_Y = dim %C, 1 : memref<?x?xf32>
+  %D_X = dim %D, 0 : memref<?x?xf32>
+  %D_Y = dim %D, 1 : memref<?x?xf32>
+  loop.parallel (%arg2, %arg3) = (%c0, %c0) to (%C_X, %C_Y) step (%c10, %c25) {
+    %C_view = std.subview %C[%arg2, %arg3][%c10, %c25][%c1, %c1] :
+        memref<?x?xf32> to memref<?x?xf32, #map>
+    %D_view = std.subview %D[%arg2, %arg3][%c10, %c25][%c1, %c1] :
+        memref<?x?xf32> to memref<?x?xf32, #map>
+    linalg.generic {
+      indexing_maps = [#id_2d, #id_2d],
+      iterator_types = ["parallel", "parallel"],
+      args_in = 1,
+      args_out = 1
+    } %C_view, %D_view {
+    ^bb0( %a: f32, %b: f32):
+      %ab = addf %a, %b : f32
+      linalg.yield %ab : f32
+    }: memref<?x?xf32, #map>, memref<?x?xf32, #map>
+  }
+  return
+}
+// CHECK-LABEL: func @fuse_indexed_generic_producer
+// CHECK:  loop.parallel ([[I:%.*]], [[J:%.*]]) =
+// CHECK-NOT:  loop.parallel
+// CHECK:      linalg.indexed_generic
+// CHECK:        ^bb0([[i:%.*]]: index, [[j:%.*]]: index
+// CHECK:          [[i_new:%.*]] = addi [[i]], [[I]] : index
+// CHECK:          [[j_new:%.*]] = addi [[j]], [[J]] : index
+// CHECK:          {{.*}} = index_cast [[i_new]] : index to i32
+// CHECK:      linalg.generic
+// CHECK:          addf
+
+// -----
+
+#map = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
+#id_2d = affine_map<(d0, d1) -> (d0, d1)>
+#pointwise_2d_trait = {
+  args_in = 2,
+  args_out = 1,
+  indexing_maps = [#id_2d, #id_2d, #id_2d],
+  iterator_types = ["parallel", "parallel"]
+}
+func @fuse_indexed_generic_producer_tile_second_dim_only(%A: memref<?x?xf32>,
+                                                         %B: memref<?x?xf32>,
+                                                         %C: memref<?x?xf32>,
+                                                         %D: memref<?x?xf32>) {
+  %c1 = constant 1 : index
+  %c3 = constant 3 : index
+  %c0 = constant 0 : index
+  linalg.indexed_generic #pointwise_2d_trait %A, %B, %C {
+    ^bb0(%i: index, %j: index, %a: f32, %b: f32, %c: f32): // no predecessors
+      %j_int = index_cast %j: index to i32
+      %j_float = sitofp %j_int : i32 to f32
+      %ab = addf %a, %b : f32
+      %out = addf %ab, %j_float : f32
+      linalg.yield %out : f32
+  }: memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>
+  %C_X = dim %C, 0 : memref<?x?xf32>
+  %C_Y = dim %C, 1 : memref<?x?xf32>
+  %D_X = dim %D, 0 : memref<?x?xf32>
+  %D_Y = dim %D, 1 : memref<?x?xf32>
+  %3 = linalg.range %c0 : %C_Y : %c3 : !linalg.range
+  loop.parallel (%j) = (%c0) to (%C_Y) step (%c3) {
+    %0 = affine.min affine_map<(d0, d1, d2) -> (d0, d1 - d2)>(%c3, %C_Y, %j)
+    %C_view = subview %C[%c0, %j] [%C_X, %0] [%c1, %c1] :
+      memref<?x?xf32> to memref<?x?xf32, #map>
+
+    %1 = affine.min affine_map<(d0, d1, d2) -> (d0, d1 - d2)>(%c3, %D_Y, %j)
+    %D_view = subview %D[%c0, %j] [%D_X, %1] [%c1, %c1] :
+      memref<?x?xf32> to memref<?x?xf32, #map>
+
+    linalg.generic {
+      indexing_maps = [#id_2d, #id_2d],
+      iterator_types = ["parallel", "parallel"],
+      args_in = 1,
+      args_out = 1
+    } %C_view, %D_view {
+    ^bb0( %a: f32, %b: f32):
+      %ab = addf %a, %b : f32
+      linalg.yield %ab : f32
+    }: memref<?x?xf32, #map>, memref<?x?xf32, #map>
+    loop.yield
+  }
+  return
+}
+// CHECK-LABEL: func @fuse_indexed_generic_producer_tile_second_dim_only
+// CHECK:  [[C0:%.*]] = constant 0 : index
+// CHECK:  loop.parallel ([[J:%.*]]) =
+// CHECK-NOT:  loop.parallel
+// CHECK:      linalg.indexed_generic
+// CHECK:        ^bb0([[i:%.*]]: index, [[j:%.*]]: index
+// CHECK:          [[i_new:%.*]] = addi [[i]], [[C0]] : index
+// CHECK:          [[j_new:%.*]] = addi [[j]], [[J]] : index
+// CHECK:          {{.*}} = index_cast [[j_new]] : index to i32
+// CHECK:      linalg.generic
+// CHECK:          addf