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 2af04c34949..6e874ab5147 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -36,6 +36,8 @@
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Interfaces/DestinationStyleOpInterface.h"
 #include "mlir/Pass/PassManager.h"
+#include "mlir/Support/MathExtras.h"
+#include "mlir/include/mlir/Dialect/Utils/IndexingUtils.h"
 #include "tensorflow/compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h"
 #include "tensorflow/compiler/mlir/disc/tools/disc-transform/utils.h"
 
@@ -1204,11 +1206,316 @@ DiagnosedSilenceableFailure InlineReductionInitializerOp::apply(
   return DiagnosedSilenceableFailure(success());
 }
 
+//===---------------------------------------------------------------------===//
+// DecomposeVectorsOp
+//===---------------------------------------------------------------------===//
+
+namespace {
+
+Optional<SmallVector<int64_t>> computeShapeRatio(ArrayRef<int64_t> superShape,
+                                                 ArrayRef<int64_t> subShape) {
+  if (superShape.size() < subShape.size()) {
+    return Optional<SmallVector<int64_t>>();
+  }
+
+  // Starting from the end, compute the integer divisors.
+  std::vector<int64_t> result;
+  result.reserve(superShape.size());
+  for (auto [superSize, subSize] :
+       llvm::zip(llvm::reverse(superShape), llvm::reverse(subShape))) {
+    assert(superSize > 0 && "superSize must be > 0");
+    assert(subSize > 0 && "subSize must be > 0");
+
+    // If integral division does not occur, return and let the caller decide.
+    if (superSize % subSize != 0) return llvm::None;
+    result.push_back(superSize / subSize);
+  }
+
+  // At this point we computed the ratio (in reverse) for the common
+  // size. Fill with the remaining entries from the super-vector shape (still in
+  // reverse).
+  int commonSize = subShape.size();
+  std::copy(superShape.rbegin() + commonSize, superShape.rend(),
+            std::back_inserter(result));
+
+  assert(result.size() == superShape.size() &&
+         "super to sub shape ratio is not of the same size as the super rank");
+
+  // Reverse again to get it back in the proper order and return.
+  return SmallVector<int64_t>{result.rbegin(), result.rend()};
+}
+
+SmallVector<int64_t> computeStrides(ArrayRef<int64_t> shape,
+                                    ArrayRef<int64_t> sizes) {
+  int64_t rank = shape.size();
+  // Compute the count for each dimension.
+  SmallVector<int64_t> sliceDimCounts(rank);
+  for (int64_t r = 0; r < rank; ++r)
+    sliceDimCounts[r] = ceilDiv(shape[r], sizes[r]);
+  // Use that to compute the slice stride for each dimension.
+  SmallVector<int64_t> sliceStrides(rank);
+  sliceStrides[rank - 1] = 1;
+  for (int64_t r = rank - 2; r >= 0; --r)
+    sliceStrides[r] = sliceStrides[r + 1] * sliceDimCounts[r + 1];
+  return sliceStrides;
+}
+
+SmallVector<int64_t> computeElementOffsetsFromVectorSliceOffsets(
+    ArrayRef<int64_t> sizes, ArrayRef<int64_t> vectorOffsets) {
+  SmallVector<int64_t> result;
+  for (auto it : llvm::zip(vectorOffsets, sizes))
+    result.push_back(std::get<0>(it) * std::get<1>(it));
+  return result;
+}
+
+/// During unrolling from `originalShape` to `targetShape` return the offset for
+/// the slice `index`.
+static SmallVector<int64_t> getVectorOffset(ArrayRef<int64_t> originalShape,
+                                            ArrayRef<int64_t> targetShape,
+                                            int64_t index) {
+  SmallVector<int64_t> dstSliceStrides =
+      computeStrides(originalShape, targetShape);
+  SmallVector<int64_t> vectorOffsets = delinearize(dstSliceStrides, index);
+  SmallVector<int64_t> elementOffsets =
+      computeElementOffsetsFromVectorSliceOffsets(targetShape, vectorOffsets);
+  return elementOffsets;
+}
+
+struct DecomposeVectorInOutsOfForOp : public OpRewritePattern<scf::ForOp> {
+  explicit DecomposeVectorInOutsOfForOp(
+      MLIRContext* context, const vector::UnrollVectorOptions& options,
+      PatternBenefit benefit = 1)
+      : mlir::OpRewritePattern<scf::ForOp>(context, benefit),
+        options_(options) {}
+
+  struct VectorDecomposeInfo {
+    VectorType dstVecType;
+    int64_t startIdx;
+    int64_t numSubVectors;
+    SmallVector<SmallVector<int64_t>> strides;
+    SmallVector<SmallVector<int64_t>> offsets;
+  };
+
+  LogicalResult matchAndRewrite(scf::ForOp forOp,
+                                PatternRewriter& rewriter) const final {
+    Location loc = forOp.getLoc();
+    if (failed(options_.filterConstraint(forOp))) return failure();
+    auto maybeTargetShape = options_.nativeShape(forOp);
+    if (!maybeTargetShape) return failure();
+    auto& targetShape = *maybeTargetShape;
+    int numNewInitArgs = 0;
+    VectorType targetVectorType;
+    DenseMap<int, VectorDecomposeInfo> candidateValueMap;
+    for (const auto& en : llvm::enumerate(forOp.getInitArgs())) {
+      auto ty = en.value().getType().dyn_cast<VectorType>();
+      Optional<SmallVector<int64_t>> maybeShapeRatio;
+      if (!ty ||
+          !(maybeShapeRatio = computeShapeRatio(ty.getShape(), targetShape)) ||
+          llvm::all_of(*maybeShapeRatio, [](int64_t v) { return v == 1; })) {
+        ++numNewInitArgs;
+        continue;
+      }
+      // TODO(wyzero): support multiple iter args with different vector type.
+      if (targetVectorType && targetVectorType != ty) return failure();
+      targetVectorType = ty;
+      auto& item = candidateValueMap[en.index()];
+      item.dstVecType = ty;
+      item.numSubVectors = computeMaxLinearIndex(*maybeShapeRatio);
+      item.startIdx = numNewInitArgs;
+      SmallVector<int64_t> strides(targetShape.size(), 1);
+      for (int i = 0; i < item.numSubVectors; ++i) {
+        auto offsets = getVectorOffset(ty.getShape(), targetShape, i);
+        item.strides.push_back(strides);
+        item.offsets.push_back(offsets);
+        ++numNewInitArgs;
+      }
+    }
+    if (candidateValueMap.empty()) return failure();
+
+    SmallVector<Value> newIterArgs;
+    for (const auto& en : llvm::enumerate(forOp.getInitArgs())) {
+      auto it = candidateValueMap.find(en.index());
+      if (it == candidateValueMap.end()) {
+        newIterArgs.push_back(en.value());
+      } else {
+        auto& item = it->second;
+        for (int i = 0; i < item.numSubVectors; ++i) {
+          newIterArgs.push_back(rewriter.create<vector::ExtractStridedSliceOp>(
+              loc, en.value(), item.offsets[i], targetShape, item.strides[i]));
+        }
+      }
+    }
+
+    scf::ForOp newForOp = rewriter.create<scf::ForOp>(
+        forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(),
+        forOp.getStep(), newIterArgs);
+    newForOp->setAttrs(forOp->getAttrs());
+    Block& newBlock = newForOp.getRegion().front();
+
+    // 1, merge block args to restore the old vector type
+    rewriter.setInsertionPointToStart(&newBlock);
+    SmallVector<Value> newBlockArgs;
+    newBlockArgs.push_back(newForOp.getInductionVar());
+    // skip the first block arg: loop induction var.
+    size_t blockArgIdx = 1;
+    for (int i = 0; i < forOp->getNumResults(); ++i) {
+      auto it = candidateValueMap.find(i);
+      if (it == candidateValueMap.end()) {
+        newBlockArgs.push_back(newBlock.getArgument(blockArgIdx++));
+        continue;
+      }
+      auto& item = it->second;
+      Value mergedVec = rewriter.create<arith::ConstantOp>(
+          loc, item.dstVecType, rewriter.getZeroAttr(item.dstVecType));
+      for (int subIdx = 0; subIdx < item.numSubVectors; ++subIdx) {
+        mergedVec = rewriter.create<vector::InsertStridedSliceOp>(
+            loc, newBlock.getArgument(blockArgIdx++), mergedVec,
+            item.offsets[subIdx], item.strides[subIdx]);
+      }
+      newBlockArgs.push_back(mergedVec);
+    }
+
+    // 2, clone ops inside the region of old for op.
+    Block& oldBlock = forOp.getRegion().front();
+    rewriter.mergeBlocks(&oldBlock, &newBlock, newBlockArgs);
+    auto oldYieldOp = &newBlock.back();
+    rewriter.setInsertionPointAfter(oldYieldOp);
+
+    // 3, split the yield result of old for op
+    SmallVector<Value> newYieldValues;
+    for (const auto& en : llvm::enumerate(oldYieldOp->getOperands())) {
+      auto it = candidateValueMap.find(en.index());
+      if (it == candidateValueMap.end()) {
+        newYieldValues.push_back(en.value());
+        continue;
+      }
+      auto& item = it->second;
+      for (int subIdx = 0; subIdx < item.numSubVectors; ++subIdx) {
+        newYieldValues.push_back(rewriter.create<vector::ExtractStridedSliceOp>(
+            loc, en.value(), item.offsets[subIdx], targetShape,
+            item.strides[subIdx]));
+      }
+    }
+    rewriter.create<scf::YieldOp>(loc, newYieldValues);
+    rewriter.eraseOp(oldYieldOp);
+
+    // 4, merge return value of new for op.
+    rewriter.setInsertionPointAfter(newForOp);
+    size_t resultIdx = 0;
+    SmallVector<Value> newResults;
+    for (const auto& en : llvm::enumerate(forOp->getResults())) {
+      auto it = candidateValueMap.find(en.index());
+      if (it == candidateValueMap.end()) {
+        newResults.push_back(newForOp->getResult(resultIdx++));
+        continue;
+      }
+      auto& item = it->second;
+      Value mergedVec = rewriter.create<arith::ConstantOp>(
+          loc, item.dstVecType, rewriter.getZeroAttr(item.dstVecType));
+      for (int subIdx = 0; subIdx < item.numSubVectors; ++subIdx) {
+        mergedVec = rewriter.create<vector::InsertStridedSliceOp>(
+            loc, newForOp->getResult(resultIdx++), mergedVec,
+            item.offsets[subIdx], item.strides[subIdx]);
+      }
+      newResults.push_back(mergedVec);
+    }
+    rewriter.replaceOp(forOp, newResults);
+    return success();
+  }
+
+ private:
+  vector::UnrollVectorOptions options_;
+};
+
+}  // namespace
+
+void DecomposeVectorsOp::build(OpBuilder& builder, OperationState& result,
+                               Value target, int64_t vectorSize) {
+  MLIRContext* ctx = builder.getContext();
+  result.addOperands(target);
+  result.addAttribute(
+      DecomposeVectorsOp::getVectorSizeAttrName(result.name),
+      builder.getIntegerAttr(builder.getIntegerType(64), vectorSize));
+  result.addTypes({});
+}
+
+DiagnosedSilenceableFailure DecomposeVectorsOp::applyToOne(
+    Operation* target, SmallVectorImpl<Operation*>& results,
+    transform::TransformState& state) {
+  if (!target->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
+    return mlir::emitDefiniteFailure(target,
+                                     "applies only to isolated-from-above "
+                                     "targets because it needs to apply "
+                                     "patterns greedily");
+  }
+
+  MLIRContext* ctx = getContext();
+  RewritePatternSet patterns(ctx);
+  // decompose outerproduct to bcast + fma ops.
+  vector::populateVectorContractLoweringPatterns(patterns);
+
+  vector::UnrollVectorOptions options;
+  auto isTargetType = [&](Type ty) {
+    auto castedTy = ty.dyn_cast<VectorType>();
+    return castedTy && castedTy.getRank() > 0 &&
+           castedTy.getShape()[castedTy.getRank() - 1] %
+                   this->getVectorSize() ==
+               0;
+  };
+  auto getVectorTypeOfForOp = [&](Operation* op) -> VectorType {
+    if (!isa<scf::ForOp>(op)) return nullptr;
+    VectorType vectorTy;
+    for (auto ty : op->getResultTypes()) {
+      if (!isTargetType(ty)) continue;
+      if (vectorTy && vectorTy != ty.dyn_cast<VectorType>()) return nullptr;
+      vectorTy = ty.dyn_cast<VectorType>();
+    }
+    return vectorTy;
+  };
+  options.setFilterConstraint([&](Operation* op) {
+    if (getVectorTypeOfForOp(op)) return success();
+    if (isa<vector::TransferReadOp, vector::TransferWriteOp>(op))
+      return success();
+    if (op->getNumResults() != 1) return failure();
+    if (op->getDialect()->getTypeID() != TypeID::get<vector::VectorDialect>() &&
+        op->getDialect()->getTypeID() != TypeID::get<arith::ArithDialect>())
+      return failure();
+    return success(isTargetType(op->getResult(0).getType()));
+  });
+  vector::UnrollVectorOptions::NativeShapeFnType nativeShapeFn =
+      [&](Operation* op) -> Optional<SmallVector<int64_t, 4>> {
+    VectorType targetVectorTy;
+    if (auto vectorTy = getVectorTypeOfForOp(op)) {
+      targetVectorTy = vectorTy;
+    } else if (isa<vector::TransferWriteOp>(op)) {
+      targetVectorTy = op->getOperand(0).getType().cast<VectorType>();
+    } else {
+      targetVectorTy = op->getResult(0).getType().cast<VectorType>();
+    }
+    SmallVector<int64_t, 4> nativeShape(targetVectorTy.getRank(), 1);
+    nativeShape[targetVectorTy.getRank() - 1] = this->getVectorSize();
+    return nativeShape;
+  };
+  options.setNativeShapeFn(nativeShapeFn);
+  vector::populateVectorUnrollPatterns(patterns, options);
+  patterns.insert<DecomposeVectorInOutsOfForOp>(ctx, options);
+
+  // some clean up patterns.
+  vector::populateVectorToVectorCanonicalizationPatterns(patterns);
+
+  // Apply everything.
+  if (failed(applyPatternsAndFoldGreedily(target, std::move(patterns))))
+    return DiagnosedSilenceableFailure::definiteFailure();
+
+  results.assign({target});
+  return DiagnosedSilenceableFailure(success());
+}
+
 }  // namespace transform_dialect
 
 void registerTransformDialectCommonExtension(DialectRegistry& registry) {
   registry
-      .addExtensions< ::mlir::disc_ral::transform_dialect::CommonExtensions>();
+      .addExtensions<::mlir::disc_ral::transform_dialect::CommonExtensions>();
 }
 
 }  // namespace disc_ral
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 91709c9bb31..a75283fb153 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
@@ -313,4 +313,58 @@ def InlineReductionInitializerOp : Op<Transform_Dialect, "disc.inline_reduction_
   let cppNamespace = "::mlir::disc_ral::transform_dialect";
 }
 
+def DecomposeVectorsOp : Op<Transform_Dialect, "disc.decompose_vectors",
+    [FunctionalStyleTransformOpTrait,
+     MemoryEffectsOpInterface,
+     TransformEachOpTrait,
+     TransformOpInterface]> {
+  let description = [{
+    Decompose vector ops and related vector transform_read/write ops into fine-grained-size vector ops.
+
+    Without this decomposition, llvm backend will mis-allocate vector register for some vector ops
+    with specific shape (e.g. 8x12xf32 failed to map to hardware register). This is a workaround for this.
+
+    Examples:
+     convert
+      ```
+       %0 = vector.transfer_read %arg0[..., %c0] : vector<8xf32>
+       %1 = vector.transfer_read %arg1[..., %c0] : vector<8xf32>
+       %2 = vector.transfer_read %arg2[..., %c0] : vector<8xf32>
+       %3 = vector.fma %0, %1, %2 : vector<8xf32>
+       vector.transfer_write %3, %arg2[..., %c0] : vector<8xf32>
+      ```
+     to:
+      ```
+       %0_0 = vector.transfer_read %arg0[..., %c0] : vector<4xf32>
+       %0_1 = vector.transfer_read %arg0[..., %c4] : vector<4xf32>
+       %1_0 = vector.transfer_read %arg1[..., %c0] : vector<4xf32>
+       %1_1 = vector.transfer_read %arg1[..., %c4] : vector<4xf32>
+       %2_0 = vector.transfer_read %arg2[..., %c0] : vector<4xf32>
+       %2_1 = vector.transfer_read %arg2[..., %c4] : vector<4xf32>
+       %3_0 = vector.fma %0_0, %1_0, %2_0 : vector<4xf32>
+       %3_1 = vector.fma %0_1, %1_1, %2_1 : vector<4xf32>
+       vector.transfer_write %3_0, %arg2[..., %c0] : vector<4xf32>
+       vector.transfer_write %3_1, %arg2[..., %c4] : vector<4xf32>
+      ```
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                       I64Attr:$vector_size);
+  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":$vector_size)>
+  ];
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target,
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
 #endif // DISC_TRANSFORM_OPS_EXT
diff --git a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
index 633cd53ab3a..464983fb3c5 100644
--- a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
@@ -244,6 +244,13 @@ buildInlineReductionInitializerOp(OpBuilder& b, Location& loc, Value initOp,
       loc, pdlType, initOp, loopOp, readerOp);
 }
 
+transform_dialect::DecomposeVectorsOp buildDecomposeVectors(
+    OpBuilder& b, Location& loc, Value target, int64_t vectorSize) {
+  auto pdlType = pdl::OperationType::get(b.getContext());
+  return b.create<transform_dialect::DecomposeVectorsOp>(loc, pdlType, target,
+                                                         vectorSize);
+}
+
 class ParsedFromFileScheduleFactory : public ScheduleFactoryWithNoGuard {
  public:
   explicit ParsedFromFileScheduleFactory(int64_t id, PatternKind kind,
@@ -498,8 +505,8 @@ LogicalResult Aarch64GEMMLargeKScheduleFactory::buildGuardCondition(
   }
   bool lhsTranspose = (lhsContractingDims[0] == lhsTy.getRank() - 2);
   Value dimK = b.create<memref::DimOp>(loc, lhs, lhsTranspose ? 0 : 1);
-  Value largeK = b.create<arith::ConstantIndexOp>(loc, 1024);
-  pred = b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sgt, dimK, largeK);
+  Value largeK = b.create<arith::ConstantIndexOp>(loc, 768);
+  pred = b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sge, dimK, largeK);
   return success();
 }
 
@@ -563,12 +570,14 @@ LogicalResult Aarch64GEMMLargeKScheduleFactory::assignSchedule(
   auto fuseIntoContainingOp =
       buildFuseIntoContainingOp(b, loc, fill, forEachThreadLoop);
 
-  // first/second level tile size for dimension m
-  int64_t M0 = 6;
+  // first level tile size for dimension m
+  int64_t M0 = 8;
   // first/second level tile size for dimension n
-  int64_t N0 = 208, N1 = 16;
-  // first level tile size for dimension k
+  int64_t N0 = 204, N1 = 12;
+  // first/second level tile size for dimension k
   int64_t K0 = 512, K1 = 1;
+  // TODO(wyzero): query cpuinfo.
+  int64_t hardwareVectorSizeInBytes = 4;
 
   auto tileOp0 = buildTileOp(b, loc, tiledMatmul, {0, N0, K0}, {0, 2, 1});
   auto tileOp1 =
@@ -696,6 +705,7 @@ LogicalResult Aarch64GEMMLargeKScheduleFactory::assignSchedule(
                     "linalg-copy", true, "eltwise", false);
   buildLowerVectors(b, loc, {5, 6, 7}, "outerproduct", "innerparallel",
                     "linalg-copy", true, "eltwise", false);
+  variant = buildDecomposeVectors(b, loc, variant, hardwareVectorSizeInBytes);
   b.create<transform::YieldOp>(loc);
   return success();
 }


Test patch:
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/decompose-vectors.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/decompose-vectors.mlir
new file mode 100644
index 00000000000..a89712d61d4
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/decompose-vectors.mlir
@@ -0,0 +1,130 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: @test_fma_xfer_ops
+// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: memref<?xf32>, %[[ARG2:.*]]: memref<?xf32>)
+func.func @test_fma_xfer_ops(%arg0: memref<?xf32>, %arg1: memref<?xf32>, %arg2: memref<?xf32>) -> memref<?xf32> {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 0.000000e+00 : f32
+  // CHECK: vector.transfer_read
+  // CHECK-SAME: %[[ARG0]][%c0]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: vector.transfer_read
+  // CHECK-SAME: %[[ARG0]][%c4]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: vector.transfer_read
+  // CHECK-SAME: %[[ARG0]][%c8]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: vector.transfer_read
+  // CHECK-SAME: %[[ARG0]][%c12]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: vector.transfer_read
+  // CHECK-SAME: %[[ARG1]][%c0]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: vector.transfer_read
+  // CHECK-SAME: %[[ARG1]][%c4]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: vector.transfer_read
+  // CHECK-SAME: %[[ARG1]][%c8]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: vector.transfer_read
+  // CHECK-SAME: %[[ARG1]][%c12]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: vector.transfer_read
+  // CHECK-SAME: %[[ARG2]][%c0]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: vector.transfer_read
+  // CHECK-SAME: %[[ARG2]][%c4]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: vector.transfer_read
+  // CHECK-SAME: %[[ARG2]][%c8]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: vector.transfer_read
+  // CHECK-SAME: %[[ARG2]][%c12]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: vector.fma
+  // CHECK: vector.fma
+  // CHECK: vector.fma
+  // CHECK: vector.fma
+
+  // CHECK: vector.transfer_write
+  // CHECK-SAME: %[[ARG2]][%c0]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: vector.transfer_write
+  // CHECK-SAME: %[[ARG2]][%c4]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: vector.transfer_write
+  // CHECK-SAME: %[[ARG2]][%c8]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: vector.transfer_write
+  // CHECK-SAME: %[[ARG2]][%c12]
+  // CHECK-SAME: vector<4xf32>
+
+  %0 = vector.transfer_read %arg0[%c0], %cst : memref<?xf32>, vector<16xf32>
+  %1 = vector.transfer_read %arg1[%c0], %cst : memref<?xf32>, vector<16xf32>
+  %2 = vector.transfer_read %arg2[%c0], %cst : memref<?xf32>, vector<16xf32>
+  %3 = vector.fma %0, %1, %2 : vector<16xf32>
+  vector.transfer_write %3, %arg2[%c0]: vector<16xf32>, memref<?xf32>
+  return %arg2 : memref<?xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  transform.disc.decompose_vectors %arg0 {vector_size = 4}
+}
+
+// -----
+
+// CHECK-LABEL: @test_for
+// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: memref<?xf32>)
+func.func @test_for(%arg0: memref<?xf32>, %arg1: memref<?xf32>) -> memref<?xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %c0 = arith.constant 0 : index
+  %c8 = arith.constant 8 : index
+  %d0 = memref.dim %arg0, %c0 : memref<?xf32>
+  // CHECK: %[[T0:.*]] = vector.transfer_read
+  // CHECK-SAME: %[[ARG1]][%c0]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: %[[T1:.*]] = vector.transfer_read
+  // CHECK-SAME: %[[ARG1]][%c4]
+  // CHECK-SAME: vector<4xf32>
+
+  // CHECK: %[[T2:.*]]:2 = scf.for %[[IV:.*]] = %c0 to %[[DIM0:.*]] step %c8
+  // CHECK-SAME: iter_args(%[[ITER_ARG0:.*]] = %[[T0]], %[[ITER_ARG1:.*]] = %[[T1]]) -> (vector<4xf32>, vector<4xf32>)
+  %0 = vector.transfer_read %arg1[%c0], %cst : memref<?xf32>, vector<8xf32>
+  %3= scf.for %iv0 = %c0 to %d0 step %c8 iter_args(%arg2 = %0) -> (vector<8xf32>) {
+      %1 = vector.transfer_read %arg0[%iv0], %cst : memref<?xf32>, vector<8xf32>
+      // CHECK: arith.addf
+      // CHECK-SAME: vector<4xf32>
+      // CHECK: arith.addf
+      // CHECK-SAME: vector<4xf32>
+      %2 = arith.addf %1, %arg2 : vector<8xf32>
+      scf.yield %2 : vector<8xf32>
+  }
+  // CHECK: vector.transfer_write
+  // CHECK-SAME: vector<4xf32>
+  // CHECK: vector.transfer_write
+  // CHECK-SAME: vector<4xf32>
+  vector.transfer_write %3, %arg1[%c0] : vector<8xf32>, memref<?xf32>
+  return %arg1: memref<?xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  transform.disc.decompose_vectors %arg0 {vector_size = 4}
+}