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 3547aecb7d8..f087cb1359c 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -1511,6 +1511,137 @@ DiagnosedSilenceableFailure DecomposeVectorsOp::applyToOne(
   return DiagnosedSilenceableFailure(success());
 }
 
+//===---------------------------------------------------------------------===//
+// LinalgFuseOperandOp
+//===---------------------------------------------------------------------===//
+
+void LinalgFuseOperandOp::build(OpBuilder& builder, OperationState& result,
+                                Value target, int64_t operandIdx) {
+  MLIRContext* ctx = builder.getContext();
+  result.addOperands(target);
+  result.addAttribute(
+      LinalgFuseOperandOp::getOperandIdxAttrName(result.name),
+      builder.getIntegerAttr(builder.getIntegerType(64), operandIdx));
+  result.addTypes({pdl::OperationType::get(ctx)});
+}
+
+DiagnosedSilenceableFailure LinalgFuseOperandOp::applyToOne(
+    Operation* target, SmallVectorImpl<Operation*>& results,
+    transform::TransformState& state) {
+  auto linalgOp = dyn_cast<linalg::GenericOp>(target);
+  if (!linalgOp) {
+    return mlir::emitDefiniteFailure(target,
+                                     "applies only to linalg::GenericOp");
+  }
+  if (getOperandIdx() >= linalgOp->getNumOperands()) {
+    return mlir::emitDefiniteFailure(target, "illegal operand_idx\n");
+  }
+  OpOperand& opOperand = linalgOp->getOpOperand(getOperandIdx());
+  if (!linalg::areElementwiseOpsFusable(&opOperand)) {
+    return mlir::emitDefiniteFailure(
+        target, "operand of linalg::GenericOp is not fusible");
+  }
+  SimplePatternRewriter rewriter(target->getContext());
+  rewriter.setInsertionPoint(target);
+  FailureOr<Operation*> fusedOp =
+      linalg::fuseElementwiseOps(rewriter, &opOperand);
+  if (!succeeded(fusedOp)) {
+    return mlir::emitDefiniteFailure(
+        target, "failed to fuse operand into linalg::GenericOp");
+  }
+  // copy custom attributes.
+  for (const auto& namedAttr : linalg::getPrunedAttributeList(linalgOp)) {
+    (*fusedOp)->setAttr(namedAttr.getName(), namedAttr.getValue());
+  }
+
+  auto replacements =
+      (*fusedOp)->getResults().take_back(linalgOp.getNumResults());
+  rewriter.replaceOp(linalgOp, replacements);
+
+  results.push_back(*fusedOp);
+  return DiagnosedSilenceableFailure(success());
+}
+
+//===---------------------------------------------------------------------===//
+// LinalgFuseProducersOp
+//===---------------------------------------------------------------------===//
+
+void LinalgFuseProducersOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance>& effects) {
+  transform::consumesHandle({this->getOperation()->getOperand(0)}, effects);
+  transform::onlyReadsHandle(this->getOperands().drop_front(), effects);
+  transform::producesHandle(this->getOperation()->getResults(), effects);
+  transform::modifiesPayload(effects);
+}
+
+DiagnosedSilenceableFailure LinalgFuseProducersOp::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 op but got ")
+           << targetOps.size();
+  Operation* targetOp = targetOps[0];
+  auto linalgOp = dyn_cast<linalg::GenericOp>(targetOp);
+  if (!linalgOp) {
+    return mlir::emitDefiniteFailure(targetOp,
+                                     "applies only to linalg::GenericOp");
+  }
+
+  // collect producers and assign an id for each producer.
+  // The id system will help to get a deterministic fused result.
+  DenseMap<Operation*, int> producerMap;
+  for (Value producerHandle : getProducers()) {
+    for (auto producer : state.getPayloadOps(producerHandle))
+      if (producerMap.find(producer) == producerMap.end()) {
+        producerMap[producer] = producerMap.size();
+      }
+  }
+
+  bool stop = false;
+  Operation* fusedOp = targetOp;
+  SimplePatternRewriter rewriter(targetOp->getContext());
+  rewriter.setInsertionPoint(targetOp);
+  do {
+    stop = true;
+    int currentMaxIdx = -1;
+    OpOperand* targetOpOperand = nullptr;
+    for (auto& opOperand : fusedOp->getOpOperands()) {
+      auto definingOp = opOperand.get().getDefiningOp();
+      auto it = producerMap.find(definingOp);
+      if (it == producerMap.end() ||
+          !linalg::areElementwiseOpsFusable(&opOperand))
+        continue;
+      if (it->second > currentMaxIdx) {
+        currentMaxIdx = it->second;
+        targetOpOperand = &opOperand;
+      }
+    }
+
+    if (!targetOpOperand) continue;
+    FailureOr<Operation*> newFusedOp =
+        linalg::fuseElementwiseOps(rewriter, targetOpOperand);
+    if (!succeeded(newFusedOp)) {
+      return mlir::emitDefiniteFailure(
+          targetOp, "failed to fuse producer into linalg::GenericOp");
+    }
+    for (const auto& namedAttr : linalg::getPrunedAttributeList(linalgOp)) {
+      (*newFusedOp)->setAttr(namedAttr.getName(), namedAttr.getValue());
+    }
+    fusedOp = *newFusedOp;
+    stop = false;
+  } while (!stop);
+
+  if (fusedOp != targetOp) {
+    auto replacements =
+        fusedOp->getResults().take_back(linalgOp.getNumResults());
+    rewriter.replaceOp(linalgOp, replacements);
+  }
+
+  results.set(getResult().cast<OpResult>(), {fusedOp});
+  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 a75283fb153..eef1a6f08a3 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
@@ -367,4 +367,118 @@ def DecomposeVectorsOp : Op<Transform_Dialect, "disc.decompose_vectors",
   }];
 }
 
+def LinalgFuseOperandOp : Op<Transform_Dialect, "disc.linalg.fuse_operand",
+    [FunctionalStyleTransformOpTrait,
+     MemoryEffectsOpInterface,
+     TransformEachOpTrait,
+     TransformOpInterface]> {
+  let description = [{
+    Given a linalg op, try to fuse its specified operand linalg op into it.
+
+    Returns the fused linalg.
+
+    Example use:
+     convert
+    ```
+      %3 = linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%0, %1 : tensor<?x768xf32>, tensor<?x768xf32>) outs(%2 : tensor<?x768xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "subtract"} {
+      ^bb0(%in: f32, %in_0: f32, %out: f32):
+        %6 = arith.subf %in, %in_0 : f32
+        linalg.yield %6 : f32
+      } -> tensor<?x768xf32>
+      %6 = linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%3, %4 : tensor<?x768xf32>, tensor<?x768xf32>) outs(%5 : tensor<?x768xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "multiply"} {
+      ^bb0(%in: f32, %in_0: f32, %out: f32):
+        %6 = arith.mulf %in, %in_0 : f32
+        linalg.yield %6 : f32
+      } -> tensor<?x768xf32>
+    ```
+     to
+    ```
+     %6 = linalg.generic {indexing_maps = [#map1, #map1, #map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%0, %1, %3 : tensor<?x768xf32>, tensor<?x768xf32>, tensor<?x768xf32>) outs(%5 : tensor<?x768xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "multiply"} {
+     ^bb0(%in: f32, %in_0: f32, %in_1: f32, %out: f32):
+       %6 = arith.subf %in, %in_0 : f32
+       %7 = arith.mulf %6, %in_1 : f32
+       linalg.yield %7 : f32
+     } -> tensor<?x768xf32>
+    ```
+
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                       I64Attr:$operand_idx);
+  let results = (outs PDL_Operation:$result);
+
+  let assemblyFormat = "$target attr-dict";
+  let cppNamespace = "::mlir::disc_ral::transform_dialect";
+
+  let builders = [
+    OpBuilder<(ins "Value":$target, "int64_t":$operand_idx)>
+  ];
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target,
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
+def LinalgFuseProducersOp : Op<Transform_Dialect, "disc.linalg.fuse_producers",
+    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let description = [{
+    Given a linalg op, try to recursively and greedily fuse its operands if they are
+    produced by the given op set.
+
+    Returns the fused linalg.
+
+    Example use. Suppose we have the following payload IR:
+    ```
+      %3 = linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%0, %1 : tensor<?x768xf32>, tensor<?x768xf32>) outs(%2 : tensor<?x768xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "subtract"} {
+      ^bb0(%in: f32, %in_0: f32, %out: f32):
+        %6 = arith.subf %in, %in_0 : f32
+        linalg.yield %6 : f32
+      } -> tensor<?x768xf32>
+      %6 = linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%3, %4 : tensor<?x768xf32>, tensor<?x768xf32>) outs(%5 : tensor<?x768xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "multiply"} {
+      ^bb0(%in: f32, %in_0: f32, %out: f32):
+        %6 = arith.mulf %in, %in_0 : f32
+        linalg.yield %6 : f32
+      } -> tensor<?x768xf32>
+      %8 = linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%6, %6 : tensor<?x768xf32>, tensor<?x768xf32>) outs(%7 : tensor<?x768xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "add"} {
+      ^bb0(%in: f32, %in_0: f32, %out: f32):
+        %6 = arith.addf %in, %in_0 : f32
+        linalg.yield %6 : f32
+      } -> tensor<?x768xf32>
+    ```
+     and following transform IR:
+    ```
+     %0 = transform.structured.match attributes {disc.transform.name = "subtract"} in %arg0
+     %1 = transform.structured.match attributes {disc.transform.name = "multiply"} in %arg0
+     %2 = transform.structured.match attributes {disc.transform.name = "add"} in %arg0
+     %3 = transform.disc.fuse_producers %0, %1 into %2
+    ```
+     finally we'll get the following transformed IR:
+    ```
+     %6 = linalg.generic {indexing_maps = [#map1, #map1, #map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%0, %1, %3 : tensor<?x768xf32>, tensor<?x768xf32>, tensor<?x768xf32>) outs(%5 : tensor<?x768xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "multiply"} {
+     ^bb0(%in: f32, %in_0: f32, %in_1: f32, %out: f32):
+       %6 = arith.subf %in, %in_0 : f32
+       %7 = arith.mulf %6, %in_1 : f32
+       %8 = arith.addf %7, %7 : f32
+       linalg.yield %8 : f32
+     } -> tensor<?x768xf32>
+    ```
+
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                       Variadic<PDL_Operation>:$producers);
+  let results = (outs PDL_Operation:$result);
+
+  let assemblyFormat = "$target attr-dict";
+  let assemblyFormat = [{
+    attr-dict
+    $producers `into` $target
+  }];
+  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/linalg-fuse-operand.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/linalg-fuse-operand.mlir
new file mode 100644
index 00000000000..844b698c638
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/linalg-fuse-operand.mlir
@@ -0,0 +1,48 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s
+
+#map = affine_map<(d0, d1) -> (d1)>
+#map1 = affine_map<(d0, d1) -> (d0, d1)>
+
+// CHECK-LABEL: @elemwise_fuse
+func.func @elemwise_fuse(%arg0: tensor<?x3072xf32>, %arg1: tensor<2xindex>, %arg2: tensor<3072x768xf32>, %arg3: tensor<768xf32>, %arg4: tensor<?x768xf32>, %arg5: tensor<?x768xf32>, %arg6: tensor<?x768xf32>, %arg7: tensor<?x768xf32>) -> tensor<?x768xf32> {
+  // CHECK: disc_linalg_ext.constant_wrapper
+  // CHECK-NEXT: disc_linalg_ext.constant_wrapper
+  // CHECK-NEXT: linalg.fill
+  // CHECK-NEXT: linalg.matmul
+  // CHECK-NEXT: %[[RES:.*]] = linalg.generic
+  // CHECK-NEXT: (%[[IN0:.*]]: f32, %[[IN1:.*]]: f32, %[[OUT:.*]]: f32)
+  // CHECK-NEXT: %[[T0:.*]] = arith.subf %[[IN0]], %[[IN1]] : f32
+  // CHECK-NEXT: %[[T1:.*]] = arith.addf %[[T0]], %[[T0]] : f32
+  // CHECK-NEXT: linalg.yield %[[T1]] : f32
+  // CHECK: return %[[RES]]
+  %0 = disc_linalg_ext.constant_wrapper {disc.transform.name = "constant"} dense<-8.000000e-01> : tensor<3072x768xf32>
+  %1 = disc_linalg_ext.constant_wrapper {disc.transform.name = "constant_1"} dense<-1.000000e-01> : tensor<768xf32>
+  %cst = arith.constant 0.000000e+00 : f32
+  %2 = linalg.fill {disc.transform.name = "dot_general"} ins(%cst : f32) outs(%arg4 : tensor<?x768xf32>) -> tensor<?x768xf32>
+  %3 = linalg.matmul {disc.transform.name = "dot_general"} ins(%arg0, %0 : tensor<?x3072xf32>, tensor<3072x768xf32>) outs(%2 : tensor<?x768xf32>) -> tensor<?x768xf32>
+  %4 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel"]} ins(%1 : tensor<768xf32>) outs(%arg5 : tensor<?x768xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "dynamic_broadcast_in_dim"} {
+  ^bb0(%in: f32, %out: f32):
+    linalg.yield %in : f32
+  } -> tensor<?x768xf32>
+  %5 = linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%3, %4 : tensor<?x768xf32>, tensor<?x768xf32>) outs(%arg6 : tensor<?x768xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "subtract"} {
+  ^bb0(%in: f32, %in_0: f32, %out: f32):
+    %6 = arith.subf %in, %in_0 : f32
+    linalg.yield %6 : f32
+  } -> tensor<?x768xf32>
+  %6 = linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%5, %5 : tensor<?x768xf32>, tensor<?x768xf32>) outs(%arg7 : tensor<?x768xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "add"} {
+  ^bb0(%in: f32, %in_0: f32, %out: f32):
+    %6 = arith.addf %in, %in_0 : f32
+    linalg.yield %6 : f32
+  } -> tensor<?x768xf32>
+  return %6 : tensor<?x768xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %0 = transform.structured.match attributes {disc.transform.name = "dynamic_broadcast_in_dim"} in %arg0
+  %1 = transform.structured.match attributes {disc.transform.name = "subtract"} in %arg0
+  %2 = transform.structured.match attributes {disc.transform.name = "add"} in %arg0
+  %3 = transform.disc.linalg.fuse_operand %2 {operand_idx = 0 : i64}
+  %4 = transform.structured.match attributes {disc.transform.name = "add"} in %arg0
+  %5 = transform.disc.linalg.fuse_operand %4 {operand_idx = 1 : i64}
+}
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/linalg-fuse-producers.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/linalg-fuse-producers.mlir
new file mode 100644
index 00000000000..c1ed134552d
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/linalg-fuse-producers.mlir
@@ -0,0 +1,46 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s
+
+#map = affine_map<(d0, d1) -> (d1)>
+#map1 = affine_map<(d0, d1) -> (d0, d1)>
+
+// CHECK-LABEL: @elemwise_fuse
+func.func @elemwise_fuse(%arg0: tensor<?x3072xf32>, %arg1: tensor<2xindex>, %arg2: tensor<3072x768xf32>, %arg3: tensor<768xf32>, %arg4: tensor<?x768xf32>, %arg5: tensor<?x768xf32>, %arg6: tensor<?x768xf32>, %arg7: tensor<?x768xf32>) -> tensor<?x768xf32> {
+  // CHECK: disc_linalg_ext.constant_wrapper
+  // CHECK-NEXT: disc_linalg_ext.constant_wrapper
+  // CHECK-NEXT: linalg.fill
+  // CHECK-NEXT: linalg.matmul
+  // CHECK-NEXT: %[[RES:.*]] = linalg.generic
+  // CHECK-NEXT: (%[[IN0:.*]]: f32, %[[IN1:.*]]: f32, %[[OUT:.*]]: f32)
+  // CHECK-NEXT: %[[T0:.*]] = arith.subf %[[IN0]], %[[IN1]] : f32
+  // CHECK-NEXT: %[[T1:.*]] = arith.addf %[[T0]], %[[T0]] : f32
+  // CHECK-NEXT: linalg.yield %[[T1]] : f32
+  // CHECK: return %[[RES]]
+  %0 = disc_linalg_ext.constant_wrapper {disc.transform.name = "constant"} dense<-8.000000e-01> : tensor<3072x768xf32>
+  %1 = disc_linalg_ext.constant_wrapper {disc.transform.name = "constant_1"} dense<-1.000000e-01> : tensor<768xf32>
+  %cst = arith.constant 0.000000e+00 : f32
+  %2 = linalg.fill {disc.transform.name = "dot_general"} ins(%cst : f32) outs(%arg4 : tensor<?x768xf32>) -> tensor<?x768xf32>
+  %3 = linalg.matmul {disc.transform.name = "dot_general"} ins(%arg0, %0 : tensor<?x3072xf32>, tensor<3072x768xf32>) outs(%2 : tensor<?x768xf32>) -> tensor<?x768xf32>
+  %4 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel"]} ins(%1 : tensor<768xf32>) outs(%arg5 : tensor<?x768xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "dynamic_broadcast_in_dim"} {
+  ^bb0(%in: f32, %out: f32):
+    linalg.yield %in : f32
+  } -> tensor<?x768xf32>
+  %5 = linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%3, %4 : tensor<?x768xf32>, tensor<?x768xf32>) outs(%arg6 : tensor<?x768xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "subtract"} {
+  ^bb0(%in: f32, %in_0: f32, %out: f32):
+    %6 = arith.subf %in, %in_0 : f32
+    linalg.yield %6 : f32
+  } -> tensor<?x768xf32>
+  %6 = linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%5, %5 : tensor<?x768xf32>, tensor<?x768xf32>) outs(%arg7 : tensor<?x768xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "add"} {
+  ^bb0(%in: f32, %in_0: f32, %out: f32):
+    %6 = arith.addf %in, %in_0 : f32
+    linalg.yield %6 : f32
+  } -> tensor<?x768xf32>
+  return %6 : tensor<?x768xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %0 = transform.structured.match attributes {disc.transform.name = "dynamic_broadcast_in_dim"} in %arg0
+  %1 = transform.structured.match attributes {disc.transform.name = "subtract"} in %arg0
+  %2 = transform.structured.match attributes {disc.transform.name = "add"} in %arg0
+  %3 = transform.disc.linalg.fuse_producers %0, %1 into %2
+}