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 b105d70bb50..7b6557340ea 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -96,8 +96,7 @@ LogicalResult comprehensiveBufferizeDeallocationFn(OpBuilder& builder,
 
 LogicalResult comprehensiveBufferizeCopyFn(OpBuilder& builder, Location loc,
                                            Value from, Value to) {
-  createLinalgCopyOp(builder, loc, from, to);
-  return success();
+  return success(createLinalgCopyOp(builder, loc, from, to) != nullptr);
 }
 
 OneShotBufferizationOptions getBufferizationOptions() {
@@ -557,6 +556,63 @@ struct FoldXferReadOfXferWriterPattern
   }
 };
 
+/// convert:
+///   %0 = vector.transfer_read %0, %cst : vector<8x12xf32>
+///   %1 = arith.select %pred, %0, %vector_cst : vector<8x12xf32>
+///   %2 = vector.transfer_write %1, %0[...]
+///   %3 = vector.transfer_read %0[...], %cst : vector<8x12xf32>
+///   use(%3)
+/// to:
+///   %0 = vector.transfer_read %0, %cst : vector<8x12xf32>
+///   %1 = arith.select %pred, %0, %vector_cst : vector<8x12xf32>
+///   use(%1)
+struct FoldXferReadOfXferWriterWithSelectPattern
+    : public OpRewritePattern<vector::TransferReadOp> {
+  using OpRewritePattern<vector::TransferReadOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::TransferReadOp op,
+                                PatternRewriter& rewriter) const override {
+    // Find xfer write op
+    auto writeOp = op.getSource().getDefiningOp<vector::TransferWriteOp>();
+    if (!writeOp || writeOp.getVector().getType() != op->getResult(0).getType())
+      return failure();
+
+    // Find select op
+    auto selectOp = writeOp.getVector().getDefiningOp<arith::SelectOp>();
+    if (!selectOp ||
+        selectOp->getResult(0).getType() != op->getResult(0).getType())
+      return failure();
+
+    // find const op and another xfer read op
+    auto prevReadOp =
+        selectOp->getOperand(1).getDefiningOp<vector::TransferReadOp>();
+    auto vectorConstOp =
+        selectOp->getOperand(2).getDefiningOp<arith::ConstantOp>();
+    if (!vectorConstOp || !prevReadOp) {
+      prevReadOp =
+          selectOp->getOperand(2).getDefiningOp<vector::TransferReadOp>();
+      vectorConstOp =
+          selectOp->getOperand(1).getDefiningOp<arith::ConstantOp>();
+    }
+    if (!vectorConstOp || !prevReadOp) return failure();
+    auto denseAttr = vectorConstOp.getValue().cast<DenseElementsAttr>();
+    if (!denseAttr.isSplat()) return failure();
+
+    // check the first read op and the second read have the same padding value
+    if (op.getPadding() != prevReadOp.getPadding()) return failure();
+
+    // check the padding value for xfer read op is a const
+    auto scalarConstOp = op.getPadding().getDefiningOp<arith::ConstantOp>();
+    if (!scalarConstOp) return failure();
+
+    if (*denseAttr.getValues<Attribute>().begin() != scalarConstOp.getValue())
+      return failure();
+
+    rewriter.replaceOp(op, selectOp->getResults());
+    return success();
+  }
+};
+
 static void addAllRegisteredCanonicalizationPatterns(
     RewritePatternSet& patterns) {
   MLIRContext* ctx = patterns.getContext();
@@ -574,6 +630,7 @@ static void addAllRegisteredCanonicalizationPatterns(
   patterns.insert<FoldMultiLevelPackOfConstantWrapperPattern>(ctx);
   patterns.insert<FoldTensorExtractOfConstantWrapperPattern>(ctx);
   patterns.insert<FoldXferReadOfXferWriterPattern>(ctx);
+  patterns.insert<FoldXferReadOfXferWriterWithSelectPattern>(ctx);
 }
 
 }  // namespace
@@ -2420,16 +2477,22 @@ DiagnosedSilenceableFailure ReductionInputFuseOp::apply(
         this->getOperation(),
         "expect the result of target is the iter arg of the loop.");
   Value iterArg = forOp.getRegionIterArg(argIt - iterOperands.begin());
-  auto iterArgUsers = iterArg.getUsers();
-  auto userIt = llvm::find_if(iterArgUsers, [&](Operation* user) {
-    auto sliceOp = dyn_cast<tensor::ExtractSliceOp>(user);
-    return sliceOp && forOp->isProperAncestor(sliceOp);
-  });
-  if (userIt == iterArgUsers.end())
+  tensor::ExtractSliceOp sliceOp;
+  auto numCandidateSliceOp =
+      llvm::count_if(iterArg.getUsers(), [&](Operation* user) -> bool {
+        if (!forOp->isProperAncestor(user)) return false;
+        if (auto candidate = dyn_cast<tensor::ExtractSliceOp>(user)) {
+          return sliceOp = candidate;
+        }
+        return false;
+      });
+  if (numCandidateSliceOp != 1) {
     return mlir::emitDefiniteFailure(
-        this->getOperation(),
-        "could not find candidate extract_slice of target in the loop.");
-  auto sliceOp = cast<tensor::ExtractSliceOp>(*userIt);
+        this->getOperation(), (Twine("only support one candidate extract_slice "
+                                     "of target in the loop but got") +
+                               Twine(numCandidateSliceOp))
+                                  .str());
+  }
 
   OpBuilder b(sliceOp);
   forOp->setOperand(
@@ -2461,6 +2524,196 @@ DiagnosedSilenceableFailure ReductionInputFuseOp::apply(
   return DiagnosedSilenceableFailure(success());
 }
 
+//===---------------------------------------------------------------------===//
+// VectorizeConditionalGenericOp
+//===---------------------------------------------------------------------===//
+
+namespace {
+
+Operation* convertConditionalGenericOpToGenericOp(OpBuilder& b,
+                                                  linalg::LinalgOp linalgOp) {
+  SmallVector<Value> inputs = linalgOp.getDpsInputOperands();
+  inputs.erase(inputs.begin());
+  SmallVector<Value> outputs = linalgOp.getDpsInitOperands();
+  SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();
+  indexingMaps.erase(indexingMaps.begin());
+  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().drop_front(), operands);
+    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<linalg::YieldOp>(
+        linalgOp.getBlock()->getTerminator()->getLoc(), newResults);
+  };
+  auto genericOp = b.create<linalg::GenericOp>(linalgOp.getLoc(), resultTypes,
+                                               inputs, outputs, indexingMaps,
+                                               iterators, bodyBuilder, attrs);
+  return genericOp;
+}
+
+}  // namespace
+
+DiagnosedSilenceableFailure VectorizeConditionalGenericOp::applyToOne(
+    Operation* target, SmallVectorImpl<Operation*>& results,
+    transform::TransformState& state) {
+  auto genericOp = dyn_cast<disc_linalg_ext::ConditionalGenericOp>(target);
+  if (!genericOp) {
+    return mlir::emitDefiniteFailure(target,
+                                     "applies only to ConditionalGenericOp");
+  }
+
+  MLIRContext* ctx = target->getContext();
+  IRRewriter rewriter(ctx);
+  Location loc = genericOp->getLoc();
+  Value pred = genericOp->getOperand(0);
+  auto linalgOp = cast<linalg::LinalgOp>(target);
+
+  // 0, check and collect init tensor types
+  SmallVector<Type> resultTensorTypes;
+  for (auto init : linalgOp.getDpsInitOperands()) {
+    auto ty = init->get().getType().dyn_cast<RankedTensorType>();
+    if (!ty || !ty.hasStaticShape()) {
+      return mlir::emitDefiniteFailure(
+          this->getOperation(),
+          "expect the result of target op has ranked tensor type with static "
+          "shape\n");
+    }
+    resultTensorTypes.push_back(ty);
+  }
+
+  // 0, build init loop like following:
+  // ```
+  //   scf.if %pred {
+  //     disc_linalg_ext.conditional_generic ...
+  //   } else {
+  //     // simply forward init operands to outputs
+  //     linalg.generic ...
+  //   }
+  // ```
+  rewriter.setInsertionPoint(target);
+  auto ifOp = rewriter.create<scf::IfOp>(loc, resultTensorTypes, pred, true);
+
+  rewriter.setInsertionPointToStart(ifOp.thenBlock());
+  auto thenTarget = dyn_cast<linalg::LinalgOp>(
+      convertConditionalGenericOpToGenericOp(rewriter, linalgOp));
+  rewriter.create<scf::YieldOp>(loc, thenTarget->getResults());
+  rewriter.setInsertionPoint(thenTarget);
+  if (failed(linalg::vectorize(rewriter, thenTarget)))
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "failed to vectorize the then target\n");
+
+  rewriter.setInsertionPointToStart(ifOp.elseBlock());
+  auto elseTarget = dyn_cast<linalg::LinalgOp>(
+      convertConditionalGenericOpToGenericOp(rewriter, linalgOp));
+  elseTarget.getBlock()->getTerminator()->setOperands(
+      elseTarget.getRegionOutputArgs());
+  SmallVector<Operation*> toDeleteOps;
+  for (auto& op : llvm::reverse(elseTarget.getBlock()->without_terminator())) {
+    toDeleteOps.push_back(&op);
+  }
+  for (Operation* op : toDeleteOps) op->erase();
+  // simply forward init input tensors to corresponding outputs in the else
+  // branch.
+  rewriter.create<scf::YieldOp>(loc, elseTarget->getResults());
+  rewriter.setInsertionPoint(elseTarget);
+  if (failed(linalg::vectorize(rewriter, elseTarget)))
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "failed to vectorize the else target\n");
+
+  rewriter.setInsertionPointAfter(ifOp);
+
+  // 1, replace the return type of ifOp: from tensor type to vector type
+  SmallVector<Type> resultVectorTypes;
+  SmallVector<Value> resultVectorValues;
+  DenseMap<int, Operation*> resultIdx2writeOpMap;
+  DenseMap<Operation*, int> writeOp2resultIdxMap;
+  for (auto [idx, val] :
+       llvm::enumerate(ifOp.elseBlock()->getTerminator()->getOperands())) {
+    auto xferWriteOp = val.getDefiningOp<vector::TransferWriteOp>();
+    if (!xferWriteOp) {
+      return mlir::emitDefiniteFailure(
+          this->getOperation(),
+          "failed to find xfer write op after vectorization\n");
+    }
+    resultVectorValues.push_back(xferWriteOp.getVector());
+    resultVectorTypes.push_back(xferWriteOp.getVector().getType());
+    resultIdx2writeOpMap[idx] = xferWriteOp;
+    writeOp2resultIdxMap[xferWriteOp] = idx;
+  }
+
+  auto vectorIfOp =
+      rewriter.create<scf::IfOp>(loc, resultVectorTypes, pred, true);
+  DenseMap<Value, Value> initTensor2LoadedVectorMap;
+  Operation* moveAfterAnchor = vectorIfOp;
+  for (auto& op :
+       llvm::make_early_inc_range(ifOp.elseBlock()->without_terminator())) {
+    if (auto xferWriteOp = dyn_cast<vector::TransferWriteOp>(&op)) {
+      op.moveAfter(moveAfterAnchor);
+      op.replaceUsesOfWith(xferWriteOp.getVector(),
+                           vectorIfOp->getResult(writeOp2resultIdxMap[&op]));
+      moveAfterAnchor = &op;
+    } else {
+      auto initArgs = linalgOp.getDpsInitOperands();
+      auto xferReadOp = dyn_cast<vector::TransferReadOp>(&op);
+      bool readFromInitArg =
+          xferReadOp && llvm::find_if(initArgs, [&](OpOperand* operand) {
+                          return operand->get() == xferReadOp.getSource();
+                        }) != initArgs.end();
+      if (!xferReadOp || readFromInitArg) op.moveBefore(ifOp);
+      if (readFromInitArg) {
+        initTensor2LoadedVectorMap[xferReadOp.getSource()] =
+            xferReadOp->getResult(0);
+      }
+    }
+  }
+  rewriter.setInsertionPointToStart(vectorIfOp.elseBlock());
+  rewriter.create<scf::YieldOp>(loc, resultVectorValues);
+  rewriter.setInsertionPointToStart(vectorIfOp.thenBlock());
+  resultVectorValues.clear();
+  for (auto [idx, val] :
+       llvm::enumerate(ifOp.thenBlock()->getTerminator()->getOperands())) {
+    auto xferWriteOp = val.getDefiningOp<vector::TransferWriteOp>();
+    if (!xferWriteOp) {
+      return mlir::emitDefiniteFailure(
+          this->getOperation(),
+          "failed to find xfer write op after vectorization\n");
+    }
+    resultVectorValues.push_back(xferWriteOp.getVector());
+  }
+  auto yieldOp = rewriter.create<scf::YieldOp>(loc, resultVectorValues);
+  for (auto& op :
+       llvm::make_early_inc_range(ifOp.thenBlock()->without_terminator())) {
+    op.moveBefore(yieldOp);
+    auto xferReadOp = dyn_cast<vector::TransferReadOp>(&op);
+    if (!xferReadOp) continue;
+    auto it = initTensor2LoadedVectorMap.find(xferReadOp.getSource());
+    if (it != initTensor2LoadedVectorMap.end())
+      xferReadOp->getResult(0).replaceAllUsesWith(it->second);
+  }
+  for (auto& e : resultIdx2writeOpMap) {
+    target->getResult(e.first).replaceAllUsesWith(e.second->getResult(0));
+  }
+  ifOp->erase();
+
+  results.assign({vectorIfOp.getOperation()});
+  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 872ea1e6158..cec879893a2 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
@@ -729,4 +729,54 @@ def ReductionInputFuseOp : Op<Transform_Dialect, "disc.reduction_input_fuse",
   let cppNamespace = "::mlir::disc_ral::transform_dialect";
 }
 
+def VectorizeConditionalGenericOp : Op<Transform_Dialect, "disc.vectorize_conditional_generic",
+    [FunctionalStyleTransformOpTrait,
+     MemoryEffectsOpInterface,
+     TransformEachOpTrait,
+     TransformOpInterface]> {
+  let description = [{
+    Vectorize the given conditional generic op and return the generated if op.
+
+    Suppose the conditional generic op has static shape.
+
+    Example:
+     convert from:
+     ```
+       %pred = ... : i1
+       %in0 = ... : tensor<8x12xf32>
+       disc_linalg_ext.conditional_generic ins(%pred, %in0 : ...) outs(... : tensor<8x12xf32>) {
+       ^bb0(%arg0 : i1, %arg1 : f32, %arg2 : f32):
+         %t0 = arith.addf %arg1, %arg1 : f32
+         disc_linalg_ext.yield %t0 : f32
+       }
+     ```
+     to (replace v0 with v1 and remove v0):
+     ```
+      %pred = ... : i1
+      %in0 = ... : tensor<8x12xf32>
+      %v0 = vector.transfer_read %in0 ...
+      %out = scf.if %pred {
+        %v1 = arith.addf %v0, %v0 : vector<8x12xf32>
+        scf.yield %v1 : vector<8x12xf32>
+      } else {
+        scf.yield %v0 : vector<8x12xf32>
+      }
+     ```
+
+  }];
+
+  let arguments = (ins PDL_Operation:$target);
+  let results = (outs PDL_Operation:$result);
+
+  let assemblyFormat = "$target attr-dict";
+  let cppNamespace = "::mlir::disc_ral::transform_dialect";
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target,
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
 #endif // DISC_TRANSFORM_OPS_EXT


Test patch:
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/apply-patterns-canonicalization.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/apply-patterns-canonicalization.mlir
index ec162bf4079..8ced3f78e58 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/apply-patterns-canonicalization.mlir
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/apply-patterns-canonicalization.mlir
@@ -127,3 +127,25 @@ transform.structured.canonicalized_sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   transform.disc.apply_patterns %arg1 {canonicalization}
 }
+
+// -----
+
+// CHECK-LABEL: @FoldXferReadOfXferWriterWithSelectPattern
+// CHECK-SAME: (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: tensor<?x?xf32>)
+func.func @FoldXferReadOfXferWriterWithSelectPattern(%pred: i1, %arg0: tensor<?x?xf32>) -> vector<8x12xf32> {
+  // CHECK: %[[T0:.*]] = arith.select %[[ARG0]]
+  // CHECK: return %[[T0]]
+  %cst = arith.constant 0.000000e+00 : f32
+  %c0 = arith.constant 0 : index
+  %cst_0 = arith.constant dense<0.000000e+00> : vector<8x12xf32>
+  %0 = vector.transfer_read %arg0[%c0, %c0], %cst : tensor<?x?xf32>, vector<8x12xf32>
+  %1 = arith.select %pred, %cst_0, %0 : vector<8x12xf32>
+  %2 = vector.transfer_write %1, %arg0[%c0, %c0] : vector<8x12xf32>, tensor<?x?xf32>
+  %3 = vector.transfer_read %2[%c0, %c0], %cst : tensor<?x?xf32>, vector<8x12xf32>
+  return %3 : vector<8x12xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  transform.disc.apply_patterns %arg1 {canonicalization}
+}
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/vectorize-conditional-generic.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/vectorize-conditional-generic.mlir
new file mode 100644
index 00000000000..82ef600f9ce
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/vectorize-conditional-generic.mlir
@@ -0,0 +1,37 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s --dump-input=always
+
+#map0 = affine_map<(d0, d1) -> ()>
+#map1 = affine_map<(d0, d1) -> (d1)>
+#map2 = affine_map<(d0, d1) -> (d0, d1)>
+
+// CHECK-LABEL: @vectorize_conditional_generic
+// CHECK-SAME: (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: tensor<f32>, %[[ARG2:.*]]: tensor<12xf32>, %[[ARG3:.*]]: tensor<8x12xf32>)
+func.func @vectorize_conditional_generic(
+    %pred : i1, %arg0: tensor<f32>, %arg1 : tensor<12xf32>, %arg2 : tensor<8x12xf32>) -> tensor<8x12xf32> {
+  // CHECK: %[[T0:.*]] = vector.transfer_read %[[ARG3]]
+  // CHECK: %[[OUT:.*]] = scf.if %[[ARG0:.*]] -> (vector<8x12xf32>)
+  // CHECK-NEXT:   %[[T1:.*]] = vector.transfer_read %[[ARG1]]
+  // CHECK-NEXT:   %[[T2:.*]] = vector.transfer_read %[[ARG2]]
+  // CHECK-NEXT:   %[[T3:.*]] = arith.addf %[[T0]], %[[T2]] : vector<8x12xf32>
+  // CHECK-NEXT:   %[[T4:.*]] = arith.maxf %[[T1]], %[[T3]] : vector<8x12xf32>
+  // CHECK-NEXT:   scf.yield %[[T4]] : vector<8x12xf32>
+  // CHECK-NEXT: } else {
+  // CHECK-NEXT:   scf.yield %0 : vector<8x12xf32>
+  // CHECK-NEXT: }
+  // CHECK-NEXT: %[[UPDATE_OUT:.*]] = vector.transfer_write %[[OUT]], %[[ARG3]]
+  // CHECK-NEXT: return %[[UPDATE_OUT]]
+  %out = disc_linalg_ext.conditional_generic {indexing_maps = [#map0, #map0, #map1, #map2], iterator_types = ["parallel", "parallel"]}
+      ins(%pred, %arg0, %arg1 : i1, tensor<f32>, tensor<12xf32>) outs(%arg2 : tensor<8x12xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "maximum"} {
+  ^bb0(%in: i1, %in_1: f32, %in_2: f32, %out: f32):
+    %t0 = arith.addf %out, %in_2 : f32
+    %t1 = arith.maxf %in_1, %t0 : f32
+    disc_linalg_ext.yield %t1 : f32
+  } -> tensor<8x12xf32>
+  return %out : tensor<8x12xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %0 = transform.structured.match ops{["disc_linalg_ext.conditional_generic"]} in %arg0
+  transform.disc.vectorize_conditional_generic %0
+}
\ No newline at end of file