[mlir][linalg] Fix tile and fuse for outermost reduction.

[gysit]

Summary

Tile and fuse failed if the outermost tile loop is a reduction dimension. Add the necessary check to handle outermost reductions and introduce a test case to verify the change.

[gysit]

Rebase.

[gysit]

Rebase.

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 class TileLoopNest {
 public:
   TileLoopNest(LinalgOp rootOp) : rootOp(rootOp) {}
 
   /// Tile the root operation using the given `tileSizes` and `tileInterchange`.
   LogicalResult tileRootOp(OpBuilder &b, ArrayRef<int64_t> tileSizes,
                            ArrayRef<int64_t> tileInterchange);
 
-  /// Fuse the producer of `rootOpOperand` into the tile loop nest. Returns the
-  /// fused producer of fails if fusion is not possible.
-  FailureOr<LinalgOp> fuseProducer(OpBuilder &b, OpOperand *rootOpOperand);
+  /// Fuse the producer of `consumerOpOperand` into the tile loop nest. Returns
+  /// the fused producer or fails if fusion is not possible.
+  FailureOr<LinalgOp> fuseProducer(OpBuilder &b, OpOperand *consumerOpOperand);
 
   /// Returns the replacement results for the original untiled root operation.
   ValueRange getRootOpReplacementResults();
 
   /// Returns the tiled root operation.
   LinalgOp getRootOp() { return rootOp; }
 
   /// Returns the loop ops generated from tiling.

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

for (auto en : enumerate(tileSizes)) {
     tiledRootAndFusedOpsLoops[rootOp].push_back(tileInterchange[en.index()]);
   }
   assert(isValid() && "expect tile loop nest to be valid after tiling");
   return success();
 }
 
 FailureOr<LinalgOp> TileLoopNest::fuseProducer(OpBuilder &b,
                                                OpOperand *consumerOpOperand) {
-  assert(tiledRootAndFusedOpsLoops.count(consumerOpOperand->getOwner()) != 0 &&
-         "expect the operand owner is the root operation or a fused producer");
+  // Check if the consumer has been tiled before. For example, it may not have
+  // been tiled if the outermost tile loop is a reduction loop.
+  if (tiledRootAndFusedOpsLoops.count(consumerOpOperand->getOwner()) == 0)
+    return failure();
+
   assert(this->isValid() &&
          "expect the tile loop nest to satisfy all invariants");
 
   // Check the tile loop nest is non-empty.
   if (isEmpty())
     return failure();
 
   // Check `consumerOpOperand` is defined by an ExtractSliceOp.

mlir/test/Dialect/Linalg/tile-and-fuse-on-tensors.mlir

builtin.func @fuse_indexed(%arg0: tensor<24x12xi32>,
   //      CHECK:  %[[IDX1_SHIFTED:.*]] = affine.apply #[[MAP0]](%[[IDX1]], %[[IV2]])
   //      CHECK:  %{{.*}} = arith.addi %[[IDX0_SHIFTED]], %[[IDX1_SHIFTED]]
   %1 = linalg.matmul ins(%arg0, %0 : tensor<24x12xi32>, tensor<12x25xi32>) outs(%arg2 : tensor<24x25xi32>) -> tensor<24x25xi32>
   return %1 : tensor<24x25xi32>
 }
 
 // -----
 
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d0)>
+
+//      CHECK:  fuse_outermost_reduction
+// CHECK-SAME:    %[[ARG0:[0-9a-zA-Z]*]]: tensor<10x17xf32>
+// CHECK-SAME:    %[[ARG1:[0-9a-zA-Z]*]]: tensor<10xf32>
+func @fuse_outermost_reduction(%arg0: tensor<10x17xf32>,
+                               %arg1: tensor<10xf32>) -> tensor<10xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = linalg.fill(%cst, %arg0) : f32, tensor<10x17xf32> -> tensor<10x17xf32>
+
+  // Cannot fuse the output fill since the reduction loop is the outermost loop.
+  //      CHECK:      %[[T0:.*]] = linalg.fill(%{{.*}}, %[[ARG1]])
+  %1 = linalg.fill(%cst, %arg1) : f32, tensor<10xf32> -> tensor<10xf32>
+
+  //      CHECK:  scf.for %[[IV0:[0-9a-zA-Z]*]] = {{.*}} iter_args(%[[ARG2:.*]] = %[[T0]]
+  //      CHECK:    scf.for %[[IV1:[0-9a-zA-Z]*]] = {{.*}} iter_args(%[[ARG3:.*]] = %[[ARG2]]
+
+  // Check the input fill has been fused.
+  //      CHECK:      %[[T1:.*]] = tensor.extract_slice %[[ARG0]]
+  // CHECK-SAME:                                        %[[IV1]], %[[IV0]]
+  //      CHECK:      %[[T2:.*]] = linalg.fill(%{{.*}}, %[[T1]])
+  //      CHECK:      %[[T3:.*]] = tensor.extract_slice %[[ARG3]]
+  // CHECK-SAME:                                        %[[IV1]]
+  //      CHECK:  linalg.generic {{.*}} ins(%[[T2]] {{.*}} outs(%[[T3]]
+  %2 = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "reduction"]} ins(%0 : tensor<10x17xf32>) outs(%1 : tensor<10xf32>) {
+  ^bb0(%arg2: f32, %arg3: f32):  // no predecessors
+    %3 = arith.addf %arg2, %arg3 : f32
+    linalg.yield %3 : f32
+  } -> tensor<10xf32>
+  return %2 : tensor<10xf32>
+}
+
+// -----
+
 //  CHECK-DAG:  #[[MAP0:.*]] = affine_map<(d0, d1) -> (d0 + d1)>
 //  CHECK-DAG:  #[[MAP1:.*]] = affine_map<(d0, d1) -> (8, -d0 - d1 + 17)>
 //  CHECK-DAG:  #[[MAP2:.*]] = affine_map<(d0, d1, d2) -> (d0, -d1 - d2 + 17)>
 #map0 = affine_map<(d0, d1) -> (d0, d0 + d1)>
 #map1 = affine_map<(d0, d1) -> (d0, d1)>
 
 //      CHECK:  fuse_non_rectangular
 // CHECK-SAME:    %[[ARG0:[0-9a-zA-Z]*]]: tensor<10x17xf32>