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:

Fix patch:
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc
index c19e4ccb3df..4be60fc98f5 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc
@@ -94,19 +94,6 @@ SmallVector<OpFoldResult> getDims(OpBuilder& builder, Location loc,
       [&](int64_t dim) { return getDim(builder, loc, shapedTypeValue, dim); }));
 }
 
-/// Returns a vector that interchanges `elements` starting at offset `offset`
-/// based on the indexes in `interchangeVector`.
-template <typename T>
-SmallVector<T> interchange(ArrayRef<T> elements,
-                           ArrayRef<int64_t> interchangeVector,
-                           int offset = 0) {
-  SmallVector<T> vec = llvm::to_vector(elements);
-  for (auto en : llvm::enumerate(interchangeVector)) {
-    vec[en.index() + offset] = elements[en.value() + offset];
-  }
-  return vec;
-}
-
 //===----------------------------------------------------------------------===//
 // PackOp
 //===----------------------------------------------------------------------===//
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h
index a93d6dfd22b..f2c41c4cc0b 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h
@@ -25,4 +25,29 @@
 #define GET_OP_CLASSES
 #include "tensorflow/compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h.inc"
 
+namespace mlir {
+namespace disc_ral {
+namespace disc_linalg_ext {
+
+// Returns dimension value vector for `shapedTypeValue`.
+SmallVector<OpFoldResult> getDims(OpBuilder& builder, Location loc,
+                                  Value shapedTypeValue);
+
+/// Returns a vector that interchanges `elements` starting at offset `offset`
+/// based on the indexes in `interchangeVector`.
+template <typename T>
+SmallVector<T> interchange(ArrayRef<T> elements,
+                           ArrayRef<int64_t> interchangeVector,
+                           int offset = 0) {
+  SmallVector<T> vec = llvm::to_vector(elements);
+  for (auto en : llvm::enumerate(interchangeVector)) {
+    vec[en.index() + offset] = elements[en.value() + offset];
+  }
+  return vec;
+}
+
+}  // namespace disc_linalg_ext
+}  // namespace disc_ral
+}  // namespace mlir
+
 #endif  // DISC_TOOLS_DISC_TRANSFORM_LINALGEXT_OPS_EXT_
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.td b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.td
index 58f9e54e70c..d0a74972dd9 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.td
@@ -82,6 +82,7 @@ def DISCLinalgExt_MultiLevelPackOp : DISCLinalgExt_Op<"multi_level_pack", [
     disc_linalg_ext.multi_level_pack %arg0 with padding_value(%pad : f32)
       tile_levels = [1, 1] tile_sizes = [8, 2] permutation = [0, 2, 1, 3]
       into %arg1 : (tensor<13x15xf32> tensor<2x8x8x2xf32>)
+    ```
 
     Example KC_to_KCc{1}k{1}k{2}c{2} with padding:
 
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 f9ddfeb2c36..09b0b4f3ddd 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,7 @@
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Interfaces/DestinationStyleOpInterface.h"
 #include "mlir/Pass/PassManager.h"
+#include "tensorflow/compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h"
 #include "tensorflow/compiler/mlir/disc/tools/disc-transform/utils.h"
 
 namespace mlir {
@@ -56,6 +57,7 @@ CommonExtensions::CommonExtensions() {
 using bufferization::BufferizationOptions;
 using bufferization::OneShotAnalysisState;
 using bufferization::OneShotBufferizationOptions;
+using disc_linalg_ext::MultiLevelPackOp;
 
 namespace {
 
@@ -575,6 +577,309 @@ DiagnosedSilenceableFailure FoldProducerExtractSliceOp::applyToOne(
   return DiagnosedSilenceableFailure(success());
 }
 
+//===---------------------------------------------------------------------===//
+// CacheReadOp
+//===---------------------------------------------------------------------===//
+
+namespace {
+
+// Suppose:
+//   - `di` is one dimension of the source tensor.
+//   - `dj` is one of the tiled dimension of `di`.
+// A dimension `dj` of the packed tensor is the inner most dim if
+//   - the tile level of `di` is 0, or
+//   - dj is corresponding to the last tile size of `di`.
+struct InnerMostDimsInfo {
+  // the dimension indices for the inner most dimensions.
+  SmallVector<int64_t> dims;
+  // the dimension sizes for the inner most dimensions.
+  SmallVector<int64_t> dimSizes;
+  // permutation for the inner most dimensions.
+  SmallVector<int64_t> permutation;
+  // reverse permutation for the inner most dimensions.
+  SmallVector<int64_t> reversePermutation;
+
+  // Returns true if the dimension `d` is the inner most.
+  bool isInnerMostDim(int d) { return llvm::find(dims, d) != dims.end(); }
+};
+
+void parseInnerMostDimsInfo(InnerMostDimsInfo& info, ShapedType sourceTy,
+                            ArrayRef<int64_t> tileLevelsVec,
+                            ArrayRef<int64_t> tileSizesVec,
+                            ArrayRef<int64_t> permutationVec) {
+  int64_t totalNumDims = -1;
+  for (const auto& en : llvm::enumerate(tileLevelsVec)) {
+    totalNumDims += en.value() + 1;
+    info.dims.push_back(totalNumDims);
+    if (en.value() == 0) {
+      info.dimSizes.push_back(sourceTy.getShape()[en.index()]);
+    } else {
+      info.dimSizes.push_back(tileSizesVec[totalNumDims - en.index() - 1]);
+    }
+  }
+
+  for (int64_t d : permutationVec) {
+    auto it = llvm::find(info.dims, d);
+    if (it == info.dims.end()) continue;
+    info.permutation.push_back(std::distance(info.dims.begin(), it));
+  }
+
+  info.reversePermutation = info.permutation;
+  for (int i = 0; i < tileLevelsVec.size(); ++i) {
+    info.reversePermutation[info.permutation[i]] = i;
+  }
+}
+
+linalg::GenericOp makeTransposeOp(OpBuilder& b, Location loc, Value inputTensor,
+                                  Value outputTensor,
+                                  ArrayRef<int64_t> transposeVector) {
+  auto resultTensorType = outputTensor.getType().cast<RankedTensorType>();
+  Type elementType = resultTensorType.getElementType();
+
+  // Compute the transpose and the indentity indexing maps.
+  SmallVector<AffineMap> indexingMaps = {
+      inversePermutation(AffineMap::getPermutationMap(
+          SmallVector<unsigned>(transposeVector.begin(), transposeVector.end()),
+          b.getContext())),
+      AffineMap::getMultiDimIdentityMap(transposeVector.size(),
+                                        b.getContext())};
+  SmallVector<llvm::StringRef> iteratorTypes(transposeVector.size(),
+                                             getParallelIteratorTypeName());
+
+  // Create a GenericOp to transpose `inputTensor` into `outputTensor`.
+  auto transposeOp = b.create<linalg::GenericOp>(
+      loc, resultTensorType, inputTensor, outputTensor,
+      b.getAffineMapArrayAttr(indexingMaps), b.getStrArrayAttr(iteratorTypes),
+      /*doc=*/nullptr,
+      /*library_call=*/nullptr);
+  Region& body = transposeOp.getRegion();
+  body.push_back(new Block());
+  body.front().addArguments({elementType, elementType}, {loc, loc});
+
+  // Create the body of the transpose operation.
+  OpBuilder::InsertionGuard g(b);
+  b.setInsertionPointToEnd(&body.front());
+  b.create<linalg::YieldOp>(loc,
+                            transposeOp.getRegion().front().getArgument(0));
+  return transposeOp;
+}
+
+LogicalResult readFromPackedValue(
+    OpBuilder& b, Location loc, Value packedValue, ShapedType targetTy,
+    ArrayRef<OpFoldResult> offsets, ArrayRef<int64_t> tileLevelsVec,
+    ArrayRef<int64_t> tileSizesVec, ArrayRef<int64_t> permutationVec,
+    InnerMostDimsInfo& info, Operation*& resultOp) {
+  auto packedTy = packedValue.getType().cast<ShapedType>();
+  int packedRank = packedTy.getRank();
+
+  auto const2IndexAttr = [&](int64_t val) {
+    return IntegerAttr::get(b.getIndexType(), val);
+  };
+
+  SmallVector<OpFoldResult> newOffsets(packedRank);
+  SmallVector<OpFoldResult> newSizes(packedRank);
+  SmallVector<OpFoldResult> newStrides(packedRank,
+                                       OpFoldResult{const2IndexAttr(1)});
+
+  AffineExpr s0, s1;
+  bindSymbols(b.getContext(), s0, s1);
+  AffineExpr floorDivExpr = s0.floorDiv(s1);
+  AffineExpr modExpr = s0 % s1;
+
+  int tileSizeIdx = 0;
+  int resultDimIdx = 0;
+  for (int dimIdx = 0; dimIdx < offsets.size(); ++dimIdx) {
+    int64_t level = tileLevelsVec[dimIdx];
+    OpFoldResult offset = offsets[dimIdx];
+    for (int localResultDimIdx = 0; localResultDimIdx <= level;
+         ++localResultDimIdx) {
+      int d = resultDimIdx + localResultDimIdx;
+      if (info.isInnerMostDim(d)) {
+        newOffsets[d] = const2IndexAttr(0);
+        newSizes[d] = const2IndexAttr(info.dimSizes[dimIdx]);
+      } else {
+        newSizes[d] = const2IndexAttr(1);
+        OpFoldResult tileSize =
+            const2IndexAttr(tileSizesVec[tileSizeIdx + localResultDimIdx]);
+        newOffsets[d] = makeComposedFoldedAffineApply(b, loc, floorDivExpr,
+                                                      {offset, tileSize});
+        offset =
+            makeComposedFoldedAffineApply(b, loc, modExpr, {offset, tileSize});
+      }
+    }
+    tileSizeIdx += level;
+    resultDimIdx += 1 + level;
+  }
+
+  if (!permutationVec.empty()) {
+    newOffsets =
+        disc_linalg_ext::interchange<OpFoldResult>(newOffsets, permutationVec);
+    newSizes =
+        disc_linalg_ext::interchange<OpFoldResult>(newSizes, permutationVec);
+  }
+
+  auto newShape = info.dimSizes;
+  newShape = disc_linalg_ext::interchange<int64_t>(newShape, info.permutation);
+  auto sliceTy = RankedTensorType::get(newShape, packedTy.getElementType());
+  auto sliceOp = b.create<tensor::ExtractSliceOp>(
+      loc, sliceTy, packedValue, newOffsets, newSizes, newStrides);
+  if (sliceTy != targetTy) {
+    Value emptyResult = b.create<tensor::EmptyOp>(loc, targetTy, ValueRange{});
+    resultOp =
+        makeTransposeOp(b, loc, sliceOp, emptyResult, info.reversePermutation);
+  } else {
+    resultOp = sliceOp.getOperation();
+  }
+  return success();
+}
+
+}  // namespace
+
+void CacheReadOp::build(OpBuilder& builder, OperationState& result,
+                        Value target, Value anchor,
+                        ArrayRef<int64_t> tileLevels,
+                        ArrayRef<int64_t> tileSizes, bool padded,
+                        ArrayRef<int64_t> permutation) {
+  SmallVector<int64_t> permutationVec;
+  int64_t expectedResultRank =
+      disc_linalg_ext::MultiLevelPackOp::getExpectedResultRank(tileLevels);
+  if (expectedResultRank > 0 && permutation.empty()) {
+    permutationVec = llvm::to_vector<>(
+        llvm::seq(static_cast<int64_t>(0), expectedResultRank));
+    permutation = permutationVec;
+  }
+  build(builder, result, pdl::OperationType::get(builder.getContext()), target,
+        anchor, builder.getI64ArrayAttr(tileLevels),
+        builder.getI64ArrayAttr(tileSizes),
+        builder.getI64ArrayAttr(permutation));
+  if (padded) {
+    result.addAttribute(CacheReadOp::getPaddedAttrName(result.name),
+                        builder.getUnitAttr());
+  }
+}
+
+void CacheReadOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance>& effects) {
+  transform::consumesHandle({this->getOperation()->getOperand(0)}, effects);
+  transform::onlyReadsHandle({this->getOperation()->getOperand(1)}, effects);
+  transform::producesHandle(this->getOperation()->getResults(), effects);
+  transform::modifiesPayload(effects);
+}
+
+DiagnosedSilenceableFailure CacheReadOp::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];
+  tensor::ExtractSliceOp sliceOp;
+  Optional<Value> paddingValue;
+  if (getPadded()) {
+    auto padOp = dyn_cast<tensor::PadOp>(targetOp);
+    if (!padOp)
+      return mlir::emitDefiniteFailure(targetOp,
+                                       "expect target to be pad op when "
+                                       "`padded` attr is set, but got ")
+             << *targetOp;
+    sliceOp = padOp->getOperand(0).getDefiningOp<tensor::ExtractSliceOp>();
+    if (!sliceOp)
+      return mlir::emitDefiniteFailure(
+          targetOp,
+          "expect the source of pad is a slice op when `padded` attr is set");
+    if (padOp.getBody()->getOperations().size() != 1)
+      return mlir::emitDefiniteFailure(
+          targetOp, "expect the padding region of pad op only has one op\n");
+    paddingValue =
+        cast<tensor::YieldOp>(&padOp.getBody()->getOperations().front())
+            ->getOperand(0);
+  } else {
+    sliceOp = dyn_cast<tensor::ExtractSliceOp>(targetOp);
+    if (!sliceOp)
+      return mlir::emitDefiniteFailure(
+                 targetOp,
+                 "expect target to be extract_slice op when `padded` attr is "
+                 "not set, but got ")
+             << *targetOp;
+  }
+
+  if (!strideAllOnes(sliceOp))
+    return mlir::emitDefiniteFailure(
+        sliceOp.getOperation(),
+        "expect the strides of target slice op are all ones\n");
+
+  // verify the target op have static shape
+  auto targetTy = targetOp->getResult(0).getType().cast<ShapedType>();
+  if (!targetTy.hasStaticShape() || targetTy.getRank() == 0)
+    return mlir::emitDefiniteFailure(
+        targetOp, "expect the targetOp has static shape with rank > 0\n");
+
+  ArrayRef<Operation*> anchorOps = state.getPayloadOps(getAnchor());
+  if (anchorOps.size() != 1)
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "expect only one anchor op but got ")
+           << anchorOps.size();
+
+  Operation* anchorOp = anchorOps[0];
+  OpBuilder b(anchorOp);
+  Location loc = anchorOp->getLoc();
+  Value source = sliceOp.getSource();
+  auto sourceTy = source.getType().cast<ShapedType>();
+  auto tileLevelsVec =
+      MultiLevelPackOp::convertI64ArrayAttrToVec(getTileLevels());
+  auto tileSizesVec =
+      MultiLevelPackOp::convertI64ArrayAttrToVec(getTileSizes());
+  auto permutationVec =
+      MultiLevelPackOp::convertI64ArrayAttrToVec(getPermutation());
+  auto packedType = MultiLevelPackOp::getPackedType(
+      sourceTy, tileLevelsVec, tileSizesVec, permutationVec);
+  if (!packedType)
+    return mlir::emitDefiniteFailure(targetOp,
+                                     "failed to infer the packed type\n");
+
+  // Verify that: the inner most tile size (or the dim size if the dimension is
+  // not tiled at all) for each dimension equals to the dimension size of the
+  // (padded) slice correspondingly.
+  InnerMostDimsInfo innerMostDimsInfo;
+  parseInnerMostDimsInfo(innerMostDimsInfo, sourceTy, tileLevelsVec,
+                         tileSizesVec, permutationVec);
+  for (const auto& z :
+       llvm::zip(targetTy.getShape(), innerMostDimsInfo.dimSizes)) {
+    if (std::get<0>(z) != std::get<1>(z))
+      return mlir::emitDefiniteFailure(this->getOperation(),
+                                       "expect the inner most tile size match "
+                                       "the shape of the result of slice: ")
+             << std::get<0>(z) << " vs " << std::get<1>(z) << "\n";
+  }
+
+  SmallVector<OpFoldResult> sourceDims =
+      disc_linalg_ext::getDims(b, loc, source);
+  auto resultDims = MultiLevelPackOp::getResultShape(
+      b, loc, sourceDims, tileLevelsVec, tileSizesVec, permutationVec);
+  SmallVector<Value> resultDynDims;
+  for (auto r : resultDims)
+    if (auto v = r.dyn_cast<Value>()) resultDynDims.push_back(v);
+  auto empty = b.create<tensor::EmptyOp>(loc, packedType, resultDynDims);
+  auto packedOp =
+      b.create<MultiLevelPackOp>(loc, source, empty, tileLevelsVec,
+                                 tileSizesVec, permutationVec, paddingValue);
+
+  b.setInsertionPoint(targetOp);
+  Operation* resultOp;
+  if (failed(readFromPackedValue(b, loc, packedOp->getResult(0), targetTy,
+                                 sliceOp.getMixedOffsets(), tileLevelsVec,
+                                 tileSizesVec, permutationVec,
+                                 innerMostDimsInfo, resultOp)))
+    return mlir::emitDefiniteFailure(
+        this->getOperation(),
+        "failed to create new extract_slice op for the packed value\n");
+  targetOp->getResult(0).replaceAllUsesWith(resultOp->getResult(0));
+  results.set(getResult().cast<OpResult>(), {resultOp});
+  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 2af4f1acc9f..a0584345112 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
@@ -139,4 +139,99 @@ def FoldProducerExtractSliceOp : Op<Transform_Dialect, "disc.fold_producer_extra
   }];
 }
 
+def CacheReadOp : Op<Transform_Dialect, "disc.cache_read",
+    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let description = [{
+    Given a target `tensor.ExtractSliceOp` named 's0', first tile and pack source of the `s0`,
+    and then replace `s0` with a new `tensor.ExtractSliceOp` named `s1`. `s1` will read from
+    the packed and tiled source to increase the cache hit ratio. the transformed source tensor
+    will be placed right before the `anthor` op.
+
+    Returns the new tensor.ExtractSliceOp.
+
+    The following attributes need to be set:
+      - tileLevels: please see the document of `disc_ral::disc_linalg_ext::MultiLevelPackOp`
+      - tileSizes: please see the document of `disc_ral::disc_linalg_ext::MultiLevelPackOp`
+      - permutation: please see the document of `disc_ral::disc_linalg_ext::MultiLevelPackOp`
+      - padded (optional): indicates that the slice op is padded and target tile is the padded
+        version.
+
+    Example #0:
+     convert from:
+     ```
+      for (i, j) {
+        %0 = tensor.extract_slice %arg0[i, j][32, 32][1, 1] : tensor<256x256xf32> to tensor<32x32xf32>
+        use(%0)
+      }
+     ```
+     to:
+     ```
+      %0 = tensor.empty() : tensor<8x8x32x32xf32>
+      %packed = disc_linalg_ext.multi_level_pack %arg0 with
+          tile_levels = [1, 1] tile_sizes = [32, 32] permutation = [0, 3, 1, 3] into %0
+          (tensor<256x256xf32> tensor<8x8x32x32xf32>) -> tensor<8x8x32x32xf32>
+      for (i, j) {
+        i', j' = f(i, j) // index mapping
+        %0 = tensor.extract_slice %packed[i', j', 0, 0][1, 1, 32, 32][1, 1, 1, 1] : tensor<8x8x32x32xf32> to tensor<32x32xf32>
+        use(%0)
+      }
+     ```
+
+    Example #1:
+     convert from:
+     ```
+      #map = affine_map<(d0)[s0] -> (-d0 + s0, 32)>
+      %cst0 = arith.constant 0.000000e+00 : f32
+      for (i, j) {
+        %s0 = affine.min #map(%i)[%d0]
+        %s1 = affine.min #map(%j)[%d1]
+        %0 = tensor.extract_slice %arg0[i, j][%s0, %s1][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+        %1 = tensor.pad %0 low[0, 0] high[32, 32] {
+        ^bb0(%arg12: index, %arg13: index):
+          tensor.yield %cst0 : f32
+        } : tensor<?x?xf32> to tensor<32x32xf32>
+        use(%1)
+      }
+     ```
+     to:
+     ```
+      %0 = tensor.empty() : tensor<?x?x32x32xf32>
+      %packed = disc_linalg_ext.multi_level_pack %arg0 with
+          %cst0 tile_levels = [1, 1] tile_sizes = [32, 32] permutation = [0, 3, 1, 3] into %0
+          (tensor<?x?xf32> tensor<?x?x32x32xf32>) -> tensor<?x?x32x32xf32>
+      for (i, j) {
+        i', j' = f(i, j) // index mapping
+        %0 = tensor.extract_slice %packed[i', j', 0, 0][1, 1, 32, 32][1, 1, 1, 1] : tensor<?x?x32x32xf32> to tensor<32x32xf32>
+        use(%0)
+      }
+     ```
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                       PDL_Operation:$anchor,
+                       I64ArrayAttr:$tile_levels,
+                       I64ArrayAttr:$tile_sizes,
+                       I64ArrayAttr:$permutation,
+                       UnitAttr:$padded);
+  let results = (outs PDL_Operation:$result);
+
+  let assemblyFormat = [{
+    attr-dict
+    $target `at` $anchor `with`
+    `tile_levels` `=` $tile_levels
+    `tile_sizes` `=` $tile_sizes
+    `permutation` `=` $permutation
+  }];
+  let cppNamespace = "::mlir::disc_ral::transform_dialect";
+
+  let builders = [
+    OpBuilder<(ins "Value":$target, "Value":$anchor,
+                   "ArrayRef<int64_t>":$tileLevels,
+                   "ArrayRef<int64_t>":$tileSizes,
+                   CArg<"bool", "false">:$padded,
+                   CArg<"ArrayRef<int64_t>", "{}">:$permutation)>
+  ];
+}
+
 #endif // DISC_TRANSFORM_OPS_EXT


Test patch:
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/cache-read.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/cache-read.mlir
new file mode 100644
index 00000000000..59131464e04
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/cache-read.mlir
@@ -0,0 +1,95 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s
+
+#map = affine_map<()[s0] -> (0, s0)>
+#map1 = affine_map<(d0)[s0] -> (d0 * s0)>
+#map2 = affine_map<(d0)[s0] -> (-d0 + s0, 6)>
+#map3 = affine_map<(d0) -> (-d0 + 6)>
+#map4 = affine_map<(d0)[s0] -> (-d0 + s0, 16)>
+#map5 = affine_map<(d0) -> (-d0 + 16)>
+#map6 = affine_map<(d0)[s0] -> (-d0 + s0, 2)>
+#map7 = affine_map<(d0) -> (-d0 + 2)>
+
+module {
+  // CHECK-LABEL: @matmul_nn
+  // CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: tensor<?x?xf32>, %[[ARG2:.*]]: tensor<?x?xf32>)
+  func.func @matmul_nn(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
+    %c6 = arith.constant 6 : index
+    %c16 = arith.constant 16 : index
+    %cst = arith.constant 0.000000e+00 : f32
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %c2 = arith.constant 2 : index
+    %dim = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+    %dim_0 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
+    %dim_1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
+    // CHECK: %[[PACKED:.*]] = disc_linalg_ext.multi_level_pack %[[ARG1]]
+    // CHECK-SAME: tile_levels = [1, 1] tile_sizes = [2, 16] permutation = [0, 2, 3, 1]
+    // CHECK-SAME: (tensor<?x?xf32> tensor<?x?x16x2xf32>) -> tensor<?x?x16x2xf32>
+    // CHECK: %[[RES:.*]] = scf.foreach_thread
+    // CHECK: scf.for
+    // CHECK: scf.for
+    // CHECK: scf.for
+    // CHECK: %[[SLICE_FROM_PACKED:.*]] = tensor.extract_slice %[[PACKED]]
+    // CHECK: %[[TRANSPOSE:.*]] = linalg.generic
+    // CHECK-SAME: %[[SLICE_FROM_PACKED]]
+    // CHECK: linalg.matmul
+    // CHECK-SAME: %[[TRANSPOSE]]
+    %0 = scf.foreach_thread (%arg3, %arg4) in (%c1, %c1) shared_outs(%arg5 = %arg2) -> (tensor<?x?xf32>) {
+      %1 = affine.max #map()[%dim]
+      %2 = affine.max #map()[%dim_1]
+      %3 = affine.apply #map1(%arg3)[%dim]
+      %4 = affine.apply #map1(%arg4)[%dim_1]
+      %extracted_slice = tensor.extract_slice %arg0[%3, 0] [%1, %dim_0] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+      %extracted_slice_2 = tensor.extract_slice %arg5[%3, %4] [%1, %2] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+      %5 = scf.for %arg6 = %c0 to %1 step %c6 iter_args(%arg7 = %extracted_slice_2) -> (tensor<?x?xf32>) {
+        %6 = affine.min #map2(%arg6)[%1]
+        %7 = affine.apply #map3(%6)
+        %8 = scf.for %arg8 = %c0 to %2 step %c16 iter_args(%arg9 = %arg7) -> (tensor<?x?xf32>) {
+          %9 = affine.min #map4(%arg8)[%2]
+          %extracted_slice_3 = tensor.extract_slice %arg9[%arg6, %arg8] [%6, %9] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+          %10 = linalg.fill ins(%cst : f32) outs(%extracted_slice_3 : tensor<?x?xf32>) -> tensor<?x?xf32>
+          %11 = arith.addi %arg8, %4 : index
+          %12 = affine.apply #map5(%9)
+          %13 = scf.for %arg10 = %c0 to %dim_0 step %c2 iter_args(%arg11 = %10) -> (tensor<?x?xf32>) {
+            %14 = affine.min #map6(%arg10)[%dim_0]
+            %extracted_slice_4 = tensor.extract_slice %extracted_slice[%arg6, %arg10] [%6, %14] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+            %extracted_slice_5 = tensor.extract_slice %arg1[%arg10, %11] [%14, %9] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+            %extracted_slice_6 = tensor.extract_slice %arg11[0, 0] [%6, %9] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+            %15 = affine.apply #map7(%14)
+            %padded = tensor.pad %extracted_slice_4 nofold low[%c0, %c0] high[%7, %15] {
+            ^bb0(%arg12: index, %arg13: index):
+              tensor.yield %cst : f32
+            } : tensor<?x?xf32> to tensor<6x2xf32>
+            %padded_7 = tensor.pad %extracted_slice_5 low[%c0, %c0] high[%15, %12] {
+            ^bb0(%arg12: index, %arg13: index):
+              tensor.yield %cst : f32
+            } : tensor<?x?xf32> to tensor<2x16xf32>
+            %padded_8 = tensor.pad %extracted_slice_6 low[%c0, %c0] high[%7, %12] {
+            ^bb0(%arg12: index, %arg13: index):
+              tensor.yield %cst : f32
+            } : tensor<?x?xf32> to tensor<6x16xf32>
+            %16 = linalg.matmul ins(%padded, %padded_7 : tensor<6x2xf32>, tensor<2x16xf32>) outs(%padded_8 : tensor<6x16xf32>) -> tensor<6x16xf32>
+            %extracted_slice_9 = tensor.extract_slice %16[0, 0] [%6, %9] [1, 1] : tensor<6x16xf32> to tensor<?x?xf32>
+            %inserted_slice_10 = tensor.insert_slice %extracted_slice_9 into %arg11[0, 0] [%6, %9] [1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
+            scf.yield %inserted_slice_10 : tensor<?x?xf32>
+          }
+          %inserted_slice = tensor.insert_slice %13 into %arg9[%arg6, %arg8] [%6, %9] [1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
+          scf.yield %inserted_slice : tensor<?x?xf32>
+        }
+        scf.yield %8 : tensor<?x?xf32>
+      }
+      scf.foreach_thread.perform_concurrently {
+        tensor.parallel_insert_slice %5 into %arg5[%3, %4] [%1, %2] [1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
+      }
+    } {thread_dim_mapping = []}
+    return %0 : tensor<?x?xf32>
+  }
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
+  %foreach_op = transform.structured.match ops{["scf.foreach_thread"]} in %arg1
+  %pad_for_weight = get_producer_of_operand %matmul[1] : (!pdl.operation) -> !pdl.operation
+  transform.disc.cache_read {padded} %pad_for_weight at %foreach_op with tile_levels = [1, 1] tile_sizes = [2, 16] permutation = [0, 2, 3, 1]
+}
\ No newline at end of file