Repository: alibaba/BladeDISC (language: python)

Issue:
Introduce MLIR transform dialect to BladeDISC
We'll start to explore using MLIR [transform dialect ](https://mlir.llvm.org/docs/Dialects/Transform) to do codegen for (fused) compute-intensive pattern. The initial target is to support gemm codegen on ARM platform to address the dynamic shape problem of Arm Compute Library.

The initial plan is:

- [x]  Step 1,  enhance the fusion decision pass. We’ll add a new fusion kind `kTransform` for the transform-based fusion pattern.
- [x]  Step 2,  lower the lmhlo fusion op to linalg on tensor.
- [x]  Step 3,  transform the linalg computation to loops using transform dialect.
- [x]  Step 4,  refined the transformed loop to make it suitable for BladeDISC runtime.
- [x]  Step 5, add a new pass to the disc pass pipeline to drive the above process.
- [x]   Step 6, weight pre-packing support
   - [x]  add `disc_linalg.multi_level_pack` op, used for doing packing.
   - [x]  add  `transform.disc.cache_read` transform op, relying on `disc_linalg.multi_level_pack` op.
   - [x]  add folding support for `disc_linalg.multi_level_pack`.
   - [x]  lower `disc_linalg.multi_level_pack` to loop if it can not be folded.
   - [x]  fuse const weight op into the `kTransform` fusion pattern, lower it to linalg and then schedule it.
- [x]  Step 7,  assign a default schedule for each `kTransform` pattern.  
- [x]  Step 8,  schedule selection logic injection
- [x]  Step 9,  initial model level testing:  bert (albert).
- [x]  Step 10, support nt, tn, tt format GEMM.
- [ ]  Step 11, support batch matmul
- [x]  Step 12, support GEMM epilogue fusion.
- [ ]  Step 13, performance optimization


Discussion / hints:
e2e model test on: Bert Base (TF) and Albert (PyTorch),  on g6r, using single thread. Note that we only have one default schedule for all shape and the schedule is known to be less performant when n or k is large, thus the initial performance is supposed to be improved when we support schedule selection logic.

Bert Base (TF) 

input | TF 2.8(s) | DISC-ACL(s) | DISC-Transform(s) | speedup (DISC-transform / DISC-ACL)
-- | -- | -- | -- | --
(1, 128) | 0.742 | 0.638 | 0.656 | 97.3%
(2, 128) | 1.41 | 1.24 | 1.27 | 97.6%
(4, 128) | 2.85 | 2.36 | 2.55 | 92.5%
(8, 128) | 5.84 | 4.68 | 5.07 | 92.3%
(16, 128) | 11.9 | 9.55 | 10.2 | 93.6%
----


Albert (PyTorch)

input | TorchScript | OnnxRuntime | DISC-ACL | DISC-Transform
-- | -- | -- | -- | --
(2, 12) | 0.197 | 0.140 | 0.117 | 0.139



Fix patch:
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
index 408d609971e..4f55b0de279 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -2084,6 +2084,296 @@ DiagnosedSilenceableFailure DISCFuseIntoContainingOp::apply(
   return DiagnosedSilenceableFailure::success();
 }
 
+//===---------------------------------------------------------------------===//
+// ReductionOutputFuseOp
+//===---------------------------------------------------------------------===//
+
+namespace {
+
+Operation* convertGenericOpToConditionalGenericOp(OpBuilder& b, Value pred,
+                                                  linalg::LinalgOp linalgOp) {
+  SmallVector<Value> inputs{pred};
+  inputs.append(linalgOp.getDpsInputOperands());
+  SmallVector<Value> outputs = linalgOp.getDpsInitOperands();
+  SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();
+  if (indexingMaps.empty()) return nullptr;
+  SmallVector<AffineMap> newIndexingMaps;
+  newIndexingMaps.push_back(AffineMap::get(indexingMaps.back().getNumDims(),
+                                           indexingMaps[0].getNumSymbols(), {},
+                                           b.getContext()));
+  newIndexingMaps.append(std::move(indexingMaps));
+  SmallVector<StringRef> iterators = linalgOp.getIteratorTypesArray();
+  SmallVector<Type> resultTypes = linalgOp.hasTensorSemantics()
+                                      ? TypeRange(ValueRange(outputs))
+                                      : TypeRange{};
+  SmallVector<NamedAttribute> attrs;
+  for (auto attr : linalgOp->getAttrs()) {
+    if (attr.getName().getValue().starts_with("disc"))
+      attrs.push_back(std::move(attr));
+  }
+  auto bodyBuilder = [&](OpBuilder& innerBuilder, Location loc,
+                         ValueRange operands) {
+    BlockAndValueMapping mapping;
+    mapping.map(linalgOp.getBlock()->getArguments(), operands.drop_front());
+    for (Operation& op : linalgOp.getBlock()->without_terminator()) {
+      innerBuilder.clone(op, mapping);
+    }
+    SmallVector<Value> newResults;
+    for (Value v : linalgOp.getBlock()->getTerminator()->getOperands()) {
+      newResults.push_back(mapping.lookupOrDefault(v));
+    }
+    innerBuilder.create<disc_linalg_ext::YieldOp>(
+        linalgOp.getBlock()->getTerminator()->getLoc(), newResults);
+  };
+  auto conditionalGenericOp = b.create<disc_linalg_ext::ConditionalGenericOp>(
+      linalgOp.getLoc(), resultTypes, inputs, outputs, newIndexingMaps,
+      iterators, bodyBuilder, attrs);
+  return conditionalGenericOp;
+}
+
+}  // namespace
+
+void ReductionOutputFuseOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance>& effects) {
+  transform::consumesHandle({this->getOperation()->getOperand(0)}, effects);
+  transform::consumesHandle({this->getOperation()->getOperand(1)}, effects);
+  transform::producesHandle(this->getOperation()->getResults(), effects);
+  transform::modifiesPayload(effects);
+}
+
+DiagnosedSilenceableFailure ReductionOutputFuseOp::apply(
+    transform::TransformResults& results, transform::TransformState& state) {
+  ArrayRef<Operation*> targetOps = state.getPayloadOps(getTarget());
+  if (targetOps.size() != 1)
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "expect only one target but got ")
+           << targetOps.size();
+  Operation* target = targetOps[0];
+  auto linalgOp = dyn_cast<linalg::LinalgOp>(target);
+  if (!linalgOp)
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "expect the target is a linalg op");
+
+  ArrayRef<Operation*> loopOps = state.getPayloadOps(getLoop());
+  if (loopOps.size() != 1)
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "expect only one loop but got ")
+           << loopOps.size();
+  auto forOp = dyn_cast<scf::ForOp>(loopOps[0]);
+  if (!forOp)
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "expect the loop to a for op");
+
+  // Check all the outputs of the linalg have the same affine map.
+  if (linalgOp->getNumResults() < 1)
+    return mlir::emitDefiniteFailure(
+        this->getOperation(), "expect the target has at least one result");
+  auto operands = linalgOp.getDpsInitOperands();
+  auto indexMap = linalgOp.getMatchingIndexingMap(operands.front());
+  AffineMap resultIndexMap;
+  for (auto [idx, operand] : llvm::enumerate(linalgOp.getDpsInitOperands())) {
+    auto indexMap = linalgOp.getMatchingIndexingMap(operand);
+    if (idx == 0) {
+      resultIndexMap = indexMap;
+    } else if (resultIndexMap != indexMap) {
+      return mlir::emitDefiniteFailure(
+          this->getOperation(),
+          "expect all the results of the target have the same shape");
+    }
+  }
+  // result of for op -> result idx.
+  DenseMap<Value, int> forOpResultMap;
+  for (auto [idx, result] : llvm::enumerate(forOp->getResults())) {
+    forOpResultMap[result] = idx;
+  }
+  // linalg input operand idx -> result idx of for op.
+  DenseMap<int, int> linalgInputOperandToForResultMap;
+  for (auto [idx, operand] : llvm::enumerate(linalgOp.getDpsInputOperands())) {
+    auto it = forOpResultMap.find(operand->get());
+    if (it == forOpResultMap.end()) continue;
+    auto indexMap = linalgOp.getMatchingIndexingMap(operand);
+    if (resultIndexMap != indexMap) {
+      return mlir::emitDefiniteFailure(
+          this->getOperation(),
+          "expect all the results of for loop which are consumed by the target "
+          "have the same shape as the results of linalg op");
+    }
+    linalgInputOperandToForResultMap[idx] = it->second;
+  }
+  // linalg init operand idx -> result idx of for op.
+  DenseMap<int, int> linalgInitOperandToForResultMap;
+  DenseMap<int, int> forResultToLinalgInitOperandMap;
+  for (auto [idx, operand] : llvm::enumerate(linalgOp.getDpsInitOperands())) {
+    auto it = forOpResultMap.find(operand->get());
+    if (it == forOpResultMap.end()) continue;
+    auto indexMap = linalgOp.getMatchingIndexingMap(operand);
+    if (resultIndexMap != indexMap) {
+      return mlir::emitDefiniteFailure(
+          this->getOperation(),
+          "expect all the results of for loop which are consumed by the target "
+          "have the same shape as the results of linalg op");
+    }
+    if (!llvm::hasSingleElement(operand->get().getUsers())) {
+      return mlir::emitDefiniteFailure(
+          this->getOperation(),
+          "expect the init operand of target which is produced by the loop op "
+          "only has one consumer");
+    }
+    linalgInitOperandToForResultMap[idx] = it->second;
+    forResultToLinalgInitOperandMap[it->second] = idx;
+  }
+
+  if (linalgInputOperandToForResultMap.empty() &&
+      linalgInitOperandToForResultMap.empty()) {
+    return mlir::emitDefiniteFailure(
+        this->getOperation(),
+        "none of the results of the reduction loop are consumed by target\n");
+  }
+
+  OpBuilder b(target);
+  auto newInitArgs = llvm::to_vector<>(forOp.getInitArgs());
+  for (auto [idx, operand] : llvm::enumerate(linalgOp.getDpsInitOperands())) {
+    if (linalgInitOperandToForResultMap.find(idx) !=
+        linalgInitOperandToForResultMap.end())
+      continue;
+    newInitArgs.push_back(operand->get());
+  }
+  auto newForOp =
+      b.create<scf::ForOp>(forOp.getLoc(), forOp.getLowerBound(),
+                           forOp.getUpperBound(), forOp.getStep(), newInitArgs);
+  BlockAndValueMapping mapping;
+  mapping.map(forOp.getInductionVar(), newForOp.getInductionVar());
+  for (auto [oldValue, newValue] :
+       llvm::zip(forOp.getRegionIterArgs(), newForOp.getRegionIterArgs())) {
+    mapping.map(oldValue, newValue);
+  }
+  int nonFusedInitOperandIdx = 0;
+  for (auto [idx, operand] : llvm::enumerate(linalgOp.getDpsInitOperands())) {
+    auto it = linalgInitOperandToForResultMap.find(idx);
+    if (linalgInitOperandToForResultMap.find(idx) !=
+        linalgInitOperandToForResultMap.end())
+      continue;
+    mapping.map(operand->get(),
+                newForOp.getRegionIterArgs()
+                    .drop_front(forOp.getRegionIterArgs().size() +
+                                nonFusedInitOperandIdx++)
+                    .front());
+  }
+
+  // move ops inside the old for op to the region of the new for op.
+  for (Value arg : forOp.getBody()->getArguments()) {
+    arg.replaceAllUsesWith(mapping.lookup(arg));
+  }
+  for (auto& op : llvm::make_early_inc_range(*forOp.getBody())) {
+    op.moveBefore(newForOp.getBody(), newForOp.getBody()->end());
+  }
+
+  int forResultIdx;
+  if (!linalgInputOperandToForResultMap.empty()) {
+    forResultIdx = linalgInputOperandToForResultMap.begin()->second;
+  } else if (!linalgInitOperandToForResultMap.empty()) {
+    forResultIdx = linalgInitOperandToForResultMap.begin()->second;
+  }
+
+  auto yieldOp = cast<scf::YieldOp>(newForOp.getBody()->getTerminator());
+  auto insertSliceOp =
+      yieldOp->getOperand(forResultIdx).getDefiningOp<tensor::InsertSliceOp>();
+  if (!insertSliceOp) {
+    return mlir::emitDefiniteFailure(
+        this->getOperation(),
+        "insert_slice op not found for target result of for op\n");
+  }
+
+  b.setInsertionPoint(yieldOp);
+  for (auto [idx, val] : llvm::enumerate(yieldOp->getOperands())) {
+    auto it = forResultToLinalgInitOperandMap.find(idx);
+    if (it == forResultToLinalgInitOperandMap.end()) {
+      mapping.map(forOp->getResult(idx), val);
+      continue;
+    }
+    auto insertOp = val.getDefiningOp<tensor::InsertSliceOp>();
+    if (!insertOp)
+      return mlir::emitDefiniteFailure(
+          this->getOperation(),
+          "insert_slice op not found for target result of for op\n");
+    mapping.map(forOp->getResult(idx), insertOp->getOperand(0));
+  }
+  auto clonedTarget = b.clone(*target, mapping);
+  auto tileableTarget = cast<TilingInterface>(clonedTarget);
+  FailureOr<Value> tiledTarget = tileableTarget.generateResultTileValue(
+      b, 0, insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes());
+  if (failed(tiledTarget)) {
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "failed to tile target op\n");
+  }
+  Operation* tiledOp = tiledTarget->getDefiningOp();
+  Value iv = newForOp.getInductionVar();
+  Value nextIv =
+      b.create<arith::AddIOp>(tiledOp->getLoc(), iv, newForOp.getStep());
+  Value pred =
+      b.create<arith::CmpIOp>(tiledOp->getLoc(), arith::CmpIPredicate::sge,
+                              nextIv, newForOp.getUpperBound());
+  auto conditionalGenericOp = convertGenericOpToConditionalGenericOp(
+      b, pred, cast<linalg::LinalgOp>(tiledOp));
+  if (!conditionalGenericOp) {
+    return mlir::emitDefiniteFailure(
+        this->getOperation(), "failed to build conditional_generic op\n");
+  }
+  SmallVector<Value> newResults;
+  for (auto [idx, val] : llvm::enumerate(yieldOp->getOperands())) {
+    auto it = forResultToLinalgInitOperandMap.find(idx);
+    if (it == forResultToLinalgInitOperandMap.end()) {
+      newResults.push_back(val);
+      continue;
+    }
+    Value initValue = newForOp.getRegionIterArg(idx);
+    Value updateValue = conditionalGenericOp->getResult(it->second);
+    newResults.push_back(b.create<tensor::InsertSliceOp>(
+                              tiledOp->getLoc(), updateValue, initValue,
+                              insertSliceOp.getMixedOffsets(),
+                              insertSliceOp.getMixedSizes(),
+                              insertSliceOp.getMixedStrides())
+                             ->getResult(0));
+  }
+  nonFusedInitOperandIdx = 0;
+  for (auto [idx, updateValue] :
+       llvm::enumerate(conditionalGenericOp->getResults())) {
+    auto it = linalgInitOperandToForResultMap.find(idx);
+    if (it != linalgInitOperandToForResultMap.end()) continue;
+    Value initValue = newForOp.getRegionIterArgs()
+                          .drop_front(forOp.getRegionIterArgs().size() +
+                                      nonFusedInitOperandIdx++)
+                          .front();
+    newResults.push_back(b.create<tensor::InsertSliceOp>(
+                              tiledOp->getLoc(), updateValue, initValue,
+                              insertSliceOp.getMixedOffsets(),
+                              insertSliceOp.getMixedSizes(),
+                              insertSliceOp.getMixedStrides())
+                             ->getResult(0));
+  }
+  b.create<scf::YieldOp>(yieldOp->getLoc(), newResults);
+  yieldOp->erase();
+
+  nonFusedInitOperandIdx = 0;
+  for (auto [idx, val] : llvm::enumerate(linalgOp->getResults())) {
+    Value newVal;
+    auto it = linalgInitOperandToForResultMap.find(idx);
+    if (it != linalgInitOperandToForResultMap.end()) {
+      newVal = newForOp->getResult(it->second);
+    } else {
+      newVal = newForOp->getResult(forOp->getNumResults() +
+                                   nonFusedInitOperandIdx++);
+    }
+    val.replaceAllUsesWith(newVal);
+  }
+  linalgOp->erase();
+  forOp->replaceAllUsesWith(
+      newForOp->getResults().take_front(forOp->getNumResults()));
+  results.set(getTiledTarget().cast<OpResult>(), {conditionalGenericOp});
+  results.set(getFusedLoop().cast<OpResult>(), {newForOp});
+  return DiagnosedSilenceableFailure(success());
+}
+
 }  // namespace transform_dialect
 
 void registerTransformDialectCommonExtension(DialectRegistry& registry) {
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
index 8ba3ef5259f..865848c17de 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
@@ -624,4 +624,63 @@ def DISCFuseIntoContainingOp :
   let cppNamespace = "::mlir::disc_ral::transform_dialect";
 }
 
+def ReductionOutputFuseOp : Op<Transform_Dialect, "disc.reduction_output_fuse",
+    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let description = [{
+    Given a linalg op, output fuse it into the reduction loop.
+
+    Returns the fused linalg op and the new loop op.
+
+    Example:
+     convert from:
+     ```
+      %1 = scf.for %arg1 = %c0 to %c768 step %c512 iter_args(%arg0 = %0) -> (tensor<?x3072xf32>) {
+        ...
+        %t0 = tensor.extract_slice %arg0[0, 0] [...] [1, 1] : tensor<?x3072xf32> to tensor<?x3072xf32>
+        %t1 = linalg.matmul {disc.transform.name = "dot_general"} ins(...) outs(%t0 : tensor<?x3072xf32>) -> tensor<?x3072xf32>
+        %t2 = tensor.insert_slice %t1 into %arg0[0, 0] [...] [1, 1] : tensor<?x3072xf32> into tensor<?x3072xf32>
+        scf.yield %t2 : tensor<?x3072xf32>
+      }
+      %3 = tensor.extract_slice %2[...] [...] [1, 1] : tensor<?x3072xf32> to tensor<?x3072xf32>
+      %4 = linalg.generic {indexing_maps = [#map5, #map6, #map7, #map6], iterator_types = ["parallel", "parallel"]} ins(..., %1, ...) outs(%3 : tensor<?x3072xf32>) attrs =  {disc.device = "cpu", disc.trasform.name = "maximum"} {
+      ^bb0(%in: f32, %in_5: f32, %in_6: f32, %out: f32):
+        %12 = arith.addf %in_5, %in_6 : f32
+        %13 = arith.maxf %in, %12 : f32
+        linalg.yield %13 : f32
+      } -> tensor<?x3072xf32>
+      use(%4)
+     ```
+     to:
+     ```
+      %3 = tensor.extract_slice %2[...] [...] [1, 1] : tensor<?x3072xf32> to tensor<?x3072xf32>
+      %1, %4 = scf.for %arg1 = %c0 to %c768 step %c512 iter_args(%arg0 = %0, %arg1 = %3) -> (tensor<?x3072xf32>, tensor<?x3072xf32>) {
+        ...
+        %t0 = tensor.extract_slice %arg0[0, 0] [...] [1, 1] : tensor<?x3072xf32> to tensor<?x3072xf32>
+        %t1 = linalg.matmul {disc.transform.name = "dot_general"} ins(...) outs(%t0 : tensor<?x3072xf32>) -> tensor<?x3072xf32>
+        %t2 = tensor.insert_slice %t1 into %arg0[0, 0] [...] [1, 1] : tensor<?x3072xf32> into tensor<?x3072xf32>
+        %t3 = tensor.extract_slice %arg1[0, 0] [...] [1, 1] : tensor<?x3072xf32> to tensor<?x3072xf32>
+        %pred = ... // is last iteration?
+        %t4 = disc_linalg_ext.conditional_generic {indexing_maps = [#map5, #map6, #map7, #map6], iterator_types = ["parallel", "parallel"]} ins(%pred, ..., %t2, ...) outs(%t3 : tensor<?x3072xf32>) attrs {disc.device = "cpu", disc.transform.name = "maximum"} {
+        ^bb0(%in: f32, %in_5: f32, %in_6: f32, %out: f32):
+          %12 = arith.addf %in_5, %in_6 : f32
+          %13 = arith.maxf %in, %12 : f32
+          linalg.yield %13 : f32
+        } -> tensor<?x3072xf32>
+        scf.yield %t2, %t4 : tensor<?x3072xf32>, tensor<?x3072xf32>
+      }
+      use(%4)
+     ```
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                       PDL_Operation:$loop);
+  let results = (outs PDL_Operation:$tiled_target, PDL_Operation:$fused_loop);
+
+  let assemblyFormat = [{
+    attr-dict $target `into` $loop
+  }];
+  let cppNamespace = "::mlir::disc_ral::transform_dialect";
+}
+
 #endif // DISC_TRANSFORM_OPS_EXT


Test patch:
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/reduction-output-fuse.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/reduction-output-fuse.mlir
new file mode 100644
index 00000000000..a8573122722
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/reduction-output-fuse.mlir
@@ -0,0 +1,107 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (512, d0 - d1)>
+
+// CHECK-LABEL: @test_reduction_output_fuse
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: tensor<?x?xf32>, %[[ARG2:.*]]: tensor<?x?xf32>, %[[ARG3:.*]]: tensor<?x?xf32>)
+func.func @test_reduction_output_fuse(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>, %arg3 : tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c512 = arith.constant 512 : index
+  // CHECK: %[[M:.*]] = tensor.dim %[[ARG0]], %c0 : tensor<?x?xf32>
+  // CHECK: %[[K:.*]] = tensor.dim %[[ARG0]], %c1 : tensor<?x?xf32>
+  // CHECK: %[[N:.*]] = tensor.dim %[[ARG1]], %c1 : tensor<?x?xf32>
+  %m = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+  %k = tensor.dim %arg0, %c1 : tensor<?x?xf32>
+  %n = tensor.dim %arg1, %c1 : tensor<?x?xf32>
+  // CHECK: scf.for %[[IV:.*]] = %c0 to %[[K]] step %c512 iter_args(%[[V0:.*]] = %[[ARG0]], %[[V1:.*]] = %[[ARG3]])
+  %t0 = scf.for %iv = %c0 to %k step %c512 iter_args(%v0 = %arg0) -> (tensor<?x?xf32>) {
+    %r0 = affine.min #map1(%k, %iv)
+    // CHECK: %[[R1:.*]] = tensor.extract_slice %[[ARG0]]
+    // CHECK: %[[R2:.*]] = tensor.extract_slice %[[ARG1]]
+    // CHECK: %[[R3:.*]] = tensor.extract_slice %[[V0]]
+    // CHECK: %[[R4:.*]] = linalg.matmul
+    // CHECK-SAME: ins(%[[R1]], %[[R2]] : tensor<?x?xf32>, tensor<?x?xf32>)
+    // CHECK-SAME: outs(%[[R3]] : tensor<?x?xf32>)
+    // CHECK: %[[R5:.*]] = tensor.insert_slice %[[R4]] into %[[V0]]
+    // CHECK: tensor.extract_slice %[[V1]]
+    // CHECK: %[[R6:.*]] = disc_linalg_ext.conditional_generic
+    // CHECK: %[[R7:.*]] = tensor.insert_slice %[[R6]] into %[[V1]]
+    // CHECK: scf.yield %[[R5]], %[[R7]]
+    %r1 = tensor.extract_slice %arg0[0, %iv] [%m, %r0] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+    %r2 = tensor.extract_slice %arg1[%iv, 0] [%r0, %n] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+    %r3 = tensor.extract_slice %v0[0, 0] [%m, %n] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+    %r4 = linalg.matmul {disc.transform.name = "dot_general"} ins(%r1, %r2 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%r3 : tensor<?x?xf32>) -> tensor<?x?xf32>
+    %r5 = tensor.insert_slice %r4 into %v0[0, 0] [%m, %n] [1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
+    scf.yield %r5 : tensor<?x?xf32>
+  }
+  %t1 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%t0 : tensor<?x?xf32>) outs(%arg3 : tensor<?x?xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "maximum"} {
+  ^bb0(%in: f32, %out: f32):
+    %t1 = arith.addf %in, %in : f32
+    linalg.yield %t1 : f32
+  } -> tensor<?x?xf32>
+  return %t1 : tensor<?x?xf32>
+}
+
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %0 = transform.structured.match attributes {disc.transform.name = "maximum"} in %arg0
+  %1 = transform.structured.match ops{["scf.for"]} in %arg0
+  %2, %3 = transform.disc.reduction_output_fuse %0 into %1
+}
+
+// -----
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (512, d0 - d1)>
+
+// CHECK-LABEL: @test_reduction_output_fuse2
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: tensor<?x?xf32>, %[[ARG2:.*]]: tensor<?x?xf32>, %[[ARG3:.*]]: tensor<?x?xf32>)
+func.func @test_reduction_output_fuse2(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>, %arg3 : tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c512 = arith.constant 512 : index
+  // CHECK: %[[M:.*]] = tensor.dim %[[ARG0]], %c0 : tensor<?x?xf32>
+  // CHECK: %[[K:.*]] = tensor.dim %[[ARG0]], %c1 : tensor<?x?xf32>
+  // CHECK: %[[N:.*]] = tensor.dim %[[ARG1]], %c1 : tensor<?x?xf32>
+  %m = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+  %k = tensor.dim %arg0, %c1 : tensor<?x?xf32>
+  %n = tensor.dim %arg1, %c1 : tensor<?x?xf32>
+  // CHECK: scf.for %[[IV:.*]] = %c0 to %[[K]] step %c512 iter_args(%[[V0:.*]] = %[[ARG0]])
+  %t0 = scf.for %iv = %c0 to %k step %c512 iter_args(%v0 = %arg0) -> (tensor<?x?xf32>) {
+    %r0 = affine.min #map1(%k, %iv)
+    // CHECK: %[[R1:.*]] = tensor.extract_slice %[[ARG0]]
+    // CHECK: %[[R2:.*]] = tensor.extract_slice %[[ARG1]]
+    // CHECK: %[[R3:.*]] = tensor.extract_slice %[[V0]]
+    // CHECK: %[[R4:.*]] = linalg.matmul
+    // CHECK-SAME: ins(%[[R1]], %[[R2]] : tensor<?x?xf32>, tensor<?x?xf32>)
+    // CHECK-SAME: outs(%[[R3]] : tensor<?x?xf32>)
+    // CHECK: %[[R5:.*]] = disc_linalg_ext.conditional_generic
+    // CHECK-SAME: outs(%[[R4]] : tensor<?x?xf32>)
+    // CHECK: %[[R7:.*]] = tensor.insert_slice %[[R5]] into %[[V0]]
+    // CHECK: scf.yield %[[R7]]
+    %r1 = tensor.extract_slice %arg0[0, %iv] [%m, %r0] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+    %r2 = tensor.extract_slice %arg1[%iv, 0] [%r0, %n] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+    %r3 = tensor.extract_slice %v0[0, 0] [%m, %n] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+    %r4 = linalg.matmul {disc.transform.name = "dot_general"} ins(%r1, %r2 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%r3 : tensor<?x?xf32>) -> tensor<?x?xf32>
+    %r5 = tensor.insert_slice %r4 into %v0[0, 0] [%m, %n] [1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
+    scf.yield %r5 : tensor<?x?xf32>
+  }
+  %t1 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel"]} outs(%t0 : tensor<?x?xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "maximum"} {
+  ^bb0(%out: f32):
+    %t1 = arith.addf %out, %out : f32
+    linalg.yield %t1 : f32
+  } -> tensor<?x?xf32>
+  return %t1 : tensor<?x?xf32>
+}
+
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %0 = transform.structured.match attributes {disc.transform.name = "maximum"} in %arg0
+  %1 = transform.structured.match ops{["scf.for"]} in %arg0
+  %2, %3 = transform.disc.reduction_output_fuse %0 into %1
+    %arg1 = transform.disc.apply_patterns %arg0 {canonicalization}
+}