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-opt/disc-opt.cc b/tao_compiler/mlir/disc/tools/disc-opt/disc-opt.cc
index 9faa39155b9..878927a201b 100644
--- a/tao_compiler/mlir/disc/tools/disc-opt/disc-opt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-opt/disc-opt.cc
@@ -33,6 +33,7 @@ limitations under the License.
 #include "mlir/disc/IR/hlo_disc_ops.h"
 #include "mlir/disc/IR/lhlo_disc_ops.h"
 #include "mlir/disc/tools/disc-transform/LinalgExt/LinalgExtDialect.h"
+#include "mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h"
 #include "mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.h"
 #include "mlir/disc/tools/disc-transform/transforms/register_passes.h"
 #include "mlir/disc/transforms/register_passes.h"
@@ -67,6 +68,8 @@ int main(int argc, char** argv) {
   registry.addExtensions<
       mlir::iree_compiler::IREE::LinalgExt::LinalgExtTransformOpsExtension,
       transform_ext::StructuredTransformOpsExtension>();
+  mlir::disc_ral::disc_linalg_ext::registerTilingInterfaceExternalModels(
+      registry);
 
   return failed(mlir::MlirOptMain(argc, argv, "MLIR HLO pass driver\n",
                                   registry,
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/BUILD b/tao_compiler/mlir/disc/tools/disc-transform/BUILD
index edd6fa6f124..8874885de38 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/BUILD
+++ b/tao_compiler/mlir/disc/tools/disc-transform/BUILD
@@ -314,6 +314,7 @@ cc_library(
     name = "transform_dialect_interpreter",
     srcs = ["transforms/transform_dialect_interpreter.cc"],
     deps = [
+        ":DISCLinalgExtDialect",
         ":disc_transform_ops",
         ":pass_details",
         "@iree-dialects//:IREELinalgExtDialect",
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtDialect.h b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtDialect.h
index 03456f9cb00..ac85276629f 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtDialect.h
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtDialect.h
@@ -12,6 +12,7 @@
 #ifndef DISC_TOOLS_DISC_TRANSFORM_LINALGEXT_DIALECT_EXT_
 #define DISC_TOOLS_DISC_TRANSFORM_LINALGEXT_DIALECT_EXT_
 
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/disc/tools/disc-transform/LinalgExt/LinalgExtDialect.h.inc"
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 e85c187f314..b08b6a354d9 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc
@@ -14,6 +14,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/SMLoc.h"
@@ -23,6 +24,7 @@
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
@@ -332,6 +334,573 @@ LogicalResult MultiLevelPackOp::fold(ArrayRef<Attribute> operands,
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// ConditionalGenericOp
+//===----------------------------------------------------------------------===//
+
+/// Common parsing used for both named structured ops created by ods-gen and by
+/// manually defined C++ ops. Does not handle regions.
+static ParseResult parseCommonStructuredOpParts(
+    OpAsmParser& parser, OperationState& result,
+    SmallVectorImpl<Type>& inputTypes, SmallVectorImpl<Type>& outputTypes,
+    bool addOperandSegmentSizes = true) {
+  SMLoc inputsOperandsLoc, outputsOperandsLoc;
+  SmallVector<OpAsmParser::UnresolvedOperand, 4> inputsOperands,
+      outputsOperands;
+
+  if (parser.parseOptionalAttrDict(result.attributes)) return failure();
+
+  if (succeeded(parser.parseOptionalKeyword("ins"))) {
+    if (parser.parseLParen()) return failure();
+
+    inputsOperandsLoc = parser.getCurrentLocation();
+    if (parser.parseOperandList(inputsOperands) ||
+        parser.parseColonTypeList(inputTypes) || parser.parseRParen())
+      return failure();
+  }
+
+  if (succeeded(parser.parseOptionalKeyword("outs"))) {
+    outputsOperandsLoc = parser.getCurrentLocation();
+    if (parser.parseLParen() || parser.parseOperandList(outputsOperands) ||
+        parser.parseColonTypeList(outputTypes) || parser.parseRParen())
+      return failure();
+  }
+
+  if (parser.resolveOperands(inputsOperands, inputTypes, inputsOperandsLoc,
+                             result.operands) ||
+      parser.resolveOperands(outputsOperands, outputTypes, outputsOperandsLoc,
+                             result.operands))
+    return failure();
+
+  if (addOperandSegmentSizes) {
+    result.addAttribute("operand_segment_sizes",
+                        parser.getBuilder().getDenseI32ArrayAttr(
+                            {static_cast<int32_t>(inputsOperands.size()),
+                             static_cast<int32_t>(outputsOperands.size())}));
+  }
+  return success();
+}
+
+static void printCommonStructuredOpParts(OpAsmPrinter& p, ValueRange inputs,
+                                         ValueRange outputs) {
+  if (!inputs.empty())
+    p << " ins(" << inputs << " : " << inputs.getTypes() << ")";
+  if (!outputs.empty())
+    p << " outs(" << outputs << " : " << outputs.getTypes() << ")";
+}
+
+static void printNamedStructuredOpResults(OpAsmPrinter& p,
+                                          TypeRange resultTypes) {
+  if (resultTypes.empty()) return;
+  p.printOptionalArrowTypeList(resultTypes);
+}
+
+static ParseResult parseNamedStructuredOpResults(
+    OpAsmParser& parser, SmallVectorImpl<Type>& resultTypes) {
+  if (parser.parseOptionalArrowTypeList(resultTypes)) return failure();
+  return success();
+}
+
+void ConditionalGenericOp::getAsmBlockArgumentNames(
+    Region& region, OpAsmSetValueNameFn setNameFn) {
+  for (Value v : getRegionInputArgs()) setNameFn(v, "in");
+  for (Value v : getRegionOutputArgs()) setNameFn(v, "out");
+}
+
+void ConditionalGenericOp::build(
+    OpBuilder& builder, OperationState& result, TypeRange resultTensorTypes,
+    ValueRange inputs, ValueRange outputs, ArrayAttr indexingMaps,
+    ArrayAttr iteratorTypes, StringAttr doc, StringAttr libraryCall,
+    function_ref<void(OpBuilder&, Location, ValueRange)> bodyBuild,
+    ArrayRef<NamedAttribute> attributes) {
+  build(builder, result, resultTensorTypes, inputs, outputs, indexingMaps,
+        iteratorTypes, doc, libraryCall);
+  result.addAttributes(attributes);
+  if (!bodyBuild) return;
+
+  SmallVector<Type, 4> blockArgTypes;
+  SmallVector<Location, 4> blockArgLocs;
+  for (ValueRange container : {inputs, outputs}) {
+    for (Value v : container) {
+      blockArgTypes.push_back(getElementTypeOrSelf(v));
+      blockArgLocs.push_back(v.getLoc());
+    }
+  }
+
+  OpBuilder::InsertionGuard guard(builder);
+  auto& region = *result.regions.front();
+  Block* bodyBlock =
+      builder.createBlock(&region, region.end(), blockArgTypes, blockArgLocs);
+  bodyBuild(builder, result.location, bodyBlock->getArguments());
+}
+
+void ConditionalGenericOp::build(
+    OpBuilder& builder, OperationState& result, TypeRange resultTensorTypes,
+    ValueRange inputs, ValueRange outputs, ArrayRef<AffineMap> indexingMaps,
+    ArrayRef<StringRef> iteratorTypes, StringRef doc, StringRef libraryCall,
+    function_ref<void(OpBuilder&, Location, ValueRange)> bodyBuild,
+    ArrayRef<NamedAttribute> attributes) {
+  build(builder, result, resultTensorTypes, inputs, outputs,
+        builder.getAffineMapArrayAttr(indexingMaps),
+        builder.getStrArrayAttr(iteratorTypes),
+        doc.empty() ? StringAttr() : builder.getStringAttr(doc),
+        libraryCall.empty() ? StringAttr() : builder.getStringAttr(libraryCall),
+        bodyBuild, attributes);
+}
+
+void ConditionalGenericOp::build(
+    OpBuilder& builder, OperationState& result, ValueRange inputs,
+    ValueRange outputs, ArrayRef<AffineMap> indexingMaps,
+    ArrayRef<StringRef> iteratorTypes, StringRef doc, StringRef libraryCall,
+    function_ref<void(OpBuilder&, Location, ValueRange)> bodyBuild,
+    ArrayRef<NamedAttribute> attributes) {
+  build(builder, result, TypeRange{}, inputs, outputs, indexingMaps,
+        iteratorTypes, doc, libraryCall, bodyBuild, attributes);
+}
+
+void ConditionalGenericOp::build(
+    OpBuilder& builder, OperationState& result, ValueRange inputs,
+    ValueRange outputs, ArrayRef<AffineMap> indexingMaps,
+    ArrayRef<StringRef> iteratorTypes,
+    function_ref<void(OpBuilder&, Location, ValueRange)> bodyBuild,
+    ArrayRef<NamedAttribute> attributes) {
+  build(builder, result, inputs, outputs, indexingMaps, iteratorTypes,
+        /*doc=*/"",
+        /*libraryCall=*/"", bodyBuild, attributes);
+}
+
+void ConditionalGenericOp::build(
+    OpBuilder& builder, OperationState& result, TypeRange resultTensorTypes,
+    ValueRange inputs, ValueRange outputs, ArrayRef<AffineMap> indexingMaps,
+    ArrayRef<StringRef> iteratorTypes,
+    function_ref<void(OpBuilder&, Location, ValueRange)> bodyBuild,
+    ArrayRef<NamedAttribute> attributes) {
+  build(builder, result, resultTensorTypes, inputs, outputs, indexingMaps,
+        iteratorTypes,
+        /*doc=*/"",
+        /*libraryCall=*/"", bodyBuild, attributes);
+}
+
+void ConditionalGenericOp::print(OpAsmPrinter& p) {
+  p << " ";
+
+  // Print extra attributes.
+  auto genericAttrNames = linalgTraitAttrNames();
+
+  llvm::StringSet<> genericAttrNamesSet;
+  genericAttrNamesSet.insert(genericAttrNames.begin(), genericAttrNames.end());
+  SmallVector<NamedAttribute, 8> genericAttrs;
+  for (auto attr : (*this)->getAttrs())
+    if (genericAttrNamesSet.count(attr.getName().strref()) > 0)
+      genericAttrs.push_back(attr);
+  if (!genericAttrs.empty()) {
+    auto genericDictAttr = DictionaryAttr::get(getContext(), genericAttrs);
+    p << genericDictAttr;
+  }
+
+  // Printing is shared with named ops, except for the region and attributes
+  printCommonStructuredOpParts(p, SmallVector<Value>(getDpsInputOperands()),
+                               SmallVector<Value>(getDpsInitOperands()));
+
+  genericAttrNames.push_back("operand_segment_sizes");
+  genericAttrNamesSet.insert(genericAttrNames.back());
+
+  bool hasExtraAttrs = false;
+  for (NamedAttribute n : (*this)->getAttrs()) {
+    if ((hasExtraAttrs = !genericAttrNamesSet.contains(n.getName().strref())))
+      break;
+  }
+  if (hasExtraAttrs) {
+    p << " attrs = ";
+    p.printOptionalAttrDict((*this)->getAttrs(),
+                            /*elidedAttrs=*/genericAttrNames);
+  }
+
+  // Print region.
+  if (!getRegion().empty()) {
+    p << ' ';
+    p.printRegion(getRegion());
+  }
+
+  // Print results.
+  printNamedStructuredOpResults(p, getResultTensors().getTypes());
+}
+
+ParseResult ConditionalGenericOp::parse(OpAsmParser& parser,
+                                        OperationState& result) {
+  DictionaryAttr dictAttr;
+  // Parse the core linalg traits that must check into a dictAttr.
+  // The name is unimportant as we will overwrite result.attributes.
+  // The core linalg traits must contain the information necessary to pass the
+  // verifier.
+  if (parser.parseAttribute(dictAttr, "_", result.attributes)) return failure();
+  result.attributes.assign(dictAttr.getValue().begin(),
+                           dictAttr.getValue().end());
+
+  // Parsing is shared with named ops, except for the region.
+  SmallVector<Type, 1> inputTypes, outputTypes;
+  if (parseCommonStructuredOpParts(parser, result, inputTypes, outputTypes))
+    return failure();
+
+  // Optional attributes may be added.
+  if (succeeded(parser.parseOptionalKeyword("attrs")))
+    if (failed(parser.parseEqual()) ||
+        failed(parser.parseOptionalAttrDict(result.attributes)))
+      return failure();
+
+  std::unique_ptr<Region> region = std::make_unique<Region>();
+  if (parser.parseRegion(*region, {})) return failure();
+  result.addRegion(std::move(region));
+
+  // Generic ops may specify that a subset of its outputs are tensors. Such
+  // outputs are specified in the result type.
+  // TODO: may need to move output parsing before region parsing.
+  // Need to wait for declarative assembly resolution to decide.
+  SmallVector<Type, 1> outputTensorsTypes;
+  if (parseNamedStructuredOpResults(parser, outputTensorsTypes))
+    return failure();
+  result.addTypes(outputTensorsTypes);
+
+  return success();
+}
+
+static void getGenericEffectsImpl(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>&
+        effects,
+    ValueRange results, OpOperandVector inputOperands,
+    OpOperandVector outputOperands) {
+  for (auto* operand : inputOperands) {
+    if (!operand->get().getType().isa<MemRefType>()) continue;
+    effects.emplace_back(MemoryEffects::Read::get(), operand->get(),
+                         SideEffects::DefaultResource::get());
+  }
+  for (auto* operand : outputOperands) {
+    if (!operand->get().getType().isa<MemRefType>()) continue;
+    effects.emplace_back(MemoryEffects::Read::get(), operand->get(),
+                         SideEffects::DefaultResource::get());
+    effects.emplace_back(MemoryEffects::Write::get(), operand->get(),
+                         SideEffects::DefaultResource::get());
+  }
+}
+
+void ConditionalGenericOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>&
+        effects) {
+  getGenericEffectsImpl(effects, getOperation()->getResults(),
+                        getDpsInputOperands(), getDpsInitOperands());
+}
+
+LogicalResult ConditionalGenericOp::verify() {
+  if (getInputs().size() < 1)
+    return emitOpError("expected the first input is the pred");
+  if (!getInputs().front().getType().isInteger(1))
+    return emitOpError("expected the pred having i1 type");
+  return success();
+}
+
+void ConditionalGenericOp::getCanonicalizationPatterns(
+    RewritePatternSet& results, MLIRContext* context) {}
+
+//===----------------------------------------------------------------------===//
+// Utility methods for implementation of Tiling Interface for Conditional
+// Generic Op
+//===----------------------------------------------------------------------===//
+
+/// Return the SSA values that represent the data point accessed using a given
+/// `indexingMap` for a given point in the iteration space represented by `ivs`.
+static SmallVector<Value> getIndicesForAccess(OpBuilder& b, Location loc,
+                                              AffineMap indexingMap,
+                                              ValueRange ivs) {
+  SmallVector<Value> indices;
+  indices.reserve(indexingMap.getNumResults());
+  for (auto result : indexingMap.getResults()) {
+    AffineMap m = AffineMap::get(indexingMap.getNumDims(),
+                                 indexingMap.getNumSymbols(), result);
+    Value v = b.create<AffineApplyOp>(loc, m, ivs);
+    indices.push_back(v);
+  }
+  return indices;
+}
+
+/// Method to inline the payload of a `linalgOp` given the iteration space
+/// point and values for the arguments of the payload.
+static LogicalResult inlinePayload(OpBuilder& b, LinalgOp linalgOp,
+                                   ValueRange ivs, ValueRange argValues) {
+  Block* body = linalgOp.getBlock();
+  BlockAndValueMapping map;
+  map.map(body->getArguments(), argValues);
+  for (auto& op : body->without_terminator()) {
+    if (auto indexOp = dyn_cast<IndexOp>(&op)) {
+      map.map(indexOp.getResult(), ivs[indexOp.getDim()]);
+      continue;
+    }
+    b.clone(op, map);
+  }
+
+  Operation* terminator = body->getTerminator();
+  Location loc = terminator->getLoc();
+  for (const auto& operand : llvm::enumerate(terminator->getOperands())) {
+    Value toStore = map.lookupOrDefault(operand.value());
+    OpOperand* storeInto = linalgOp.getDpsInitOperand(operand.index());
+    auto indices = getIndicesForAccess(
+        b, loc, linalgOp.getMatchingIndexingMap(storeInto), ivs);
+    b.create<memref::StoreOp>(
+        loc, toStore, linalgOp.getDpsInitOperand(operand.index())->get(),
+        indices);
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// External Model for implementing `TilingInterface` for `LinalgOp`s.
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// External model implementation of TilingInterface for LinalgOps. An external
+/// model implementation is used for now till the use of `TilingInterface` is
+/// on-par with the current Linalg tiling + fusion patterns. Once it is
+/// maybe possible to move this into the op-definition (though there are
+/// advantages to leaving it as an external model)
+template <typename LinalgOpTy>
+struct LinalgOpTilingInterface
+    : public TilingInterface::ExternalModel<LinalgOpTilingInterface<LinalgOpTy>,
+                                            LinalgOpTy> {
+  /// Return the loop iterator type.
+  SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation* op) const {
+    LinalgOpTy concreteOp = cast<LinalgOpTy>(op);
+    return llvm::to_vector(llvm::map_range(
+        concreteOp.getIteratorTypesArray(), [](StringRef iteratorType) {
+          return utils::symbolizeIteratorType(iteratorType).value();
+        }));
+  }
+
+  /// Return the iteration domain range.
+  SmallVector<Range> getIterationDomain(Operation* op, OpBuilder& b) const {
+    OpBuilder::InsertionGuard g(b);
+    b.setInsertionPoint(op);
+    Location loc = op->getLoc();
+    LinalgOp linalgOp = cast<LinalgOp>(op);
+    SmallVector<OpFoldResult> allShapesSizes =
+        linalgOp.createFlatListOfOperandDims(b, loc);
+    AffineMap map = linalgOp.getShapesToLoopsMap();
+
+    return llvm::to_vector(
+        llvm::map_range(map.getResults(), [&](AffineExpr loopExpr) {
+          OpFoldResult ofr =
+              makeComposedFoldedAffineApply(b, loc, loopExpr, allShapesSizes);
+          return Range{b.getIndexAttr(0), ofr, b.getIndexAttr(1)};
+        }));
+  }
+
+  // Instantiate the tiled implementation of the operation.
+  SmallVector<Operation*> getTiledImplementation(
+      Operation* op, OpBuilder& b, ArrayRef<OpFoldResult> offsets,
+      ArrayRef<OpFoldResult> sizes) const {
+    // Leave the `sizeBounds` value empty. That is only needed when the `sizes`
+    // specified could lead to out of bounds accesses.
+    Location loc = op->getLoc();
+    LinalgOp linalgOp = cast<LinalgOp>(op);
+    SmallVector<Value> valuesToTile = linalgOp->getOperands();
+    SmallVector<Value, 4> tiledOperands = linalg::makeTiledShapes(
+        b, loc, linalgOp, valuesToTile, offsets, sizes, {}, true);
+
+    SmallVector<Type> resultTensorTypes =
+        getTensorOutputTypes(linalgOp, tiledOperands);
+
+    Operation* tiledOp =
+        linalgOp.clone(b, loc, resultTensorTypes, tiledOperands);
+    offsetIndices(b, cast<LinalgOp>(tiledOp), offsets);
+
+    return {tiledOp};
+  }
+
+  // Return the details of the output tile generated by the tiled
+  // implementation.
+  LogicalResult getResultTilePosition(
+      Operation* op, OpBuilder& b, unsigned resultNumber,
+      ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
+      SmallVector<OpFoldResult>& resultOffsets,
+      SmallVector<OpFoldResult>& resultSizes) const {
+    Location loc = op->getLoc();
+    LinalgOp linalgOp = cast<LinalgOp>(op);
+
+    AffineExpr d0;
+    bindDims(b.getContext(), d0);
+    SmallVector<OpFoldResult> subShapeSizes =
+        llvm::to_vector(llvm::map_range(sizes, [&](OpFoldResult ofr) {
+          return makeComposedFoldedAffineApply(b, loc, d0 - 1, ofr);
+        }));
+
+    OpOperand* outOperand = linalgOp.getDpsInitOperand(resultNumber);
+    linalg::SliceParameters sliceParams = linalg::computeSliceParameters(
+        b, loc, outOperand->get(), sizes,
+        linalgOp.getMatchingIndexingMap(outOperand), offsets,
+        /*ubs*/ {}, subShapeSizes, true);
+    resultOffsets = sliceParams.offsets;
+    resultSizes = sliceParams.sizes;
+    return success();
+  }
+
+  FailureOr<Value> generateResultTileValue(Operation* op, OpBuilder& b,
+                                           unsigned resultNumber,
+                                           ArrayRef<OpFoldResult> offsets,
+                                           ArrayRef<OpFoldResult> sizes) const {
+    auto linalgOp = cast<LinalgOp>(op);
+
+    // Check that the indexing map used for the output is a projected
+    // permutation. This could be relaxed with a more general approach that can
+    // map the offsets and sizes from the result to iteration space tiles
+    // (filling in full extent for dimensions not used to access the result).
+    AffineMap indexingMap =
+        linalgOp.getIndexingMapMatchingResult(op->getResult(resultNumber));
+    if (!indexingMap.isProjectedPermutation()) {
+      return op->emitOpError(
+          "unhandled tiled implementation generation when result is not "
+          "accessed using a permuted projection");
+    }
+
+    auto numLoops = linalgOp.getNumLoops();
+    auto tilingInterfaceOp = cast<TilingInterface>(op);
+    SmallVector<OpFoldResult> iterationTileOffsets(numLoops),
+        iterationTileSizes(numLoops);
+    if (!indexingMap.isPermutation()) {
+      SmallVector<Range> iterationDomain =
+          tilingInterfaceOp.getIterationDomain(b);
+      for (const auto& range : llvm::enumerate(iterationDomain)) {
+        iterationTileOffsets[range.index()] = range.value().offset;
+        iterationTileSizes[range.index()] = range.value().size;
+      }
+    }
+    for (const auto& resultExpr : llvm::enumerate(indexingMap.getResults())) {
+      unsigned dimPosition =
+          resultExpr.value().cast<AffineDimExpr>().getPosition();
+      iterationTileOffsets[dimPosition] = offsets[resultExpr.index()];
+      iterationTileSizes[dimPosition] = sizes[resultExpr.index()];
+    }
+
+    SmallVector<Operation*> tiledOp = tilingInterfaceOp.getTiledImplementation(
+        b, iterationTileOffsets, iterationTileSizes);
+    if (tiledOp.size() != 1)
+      return op->emitOpError("failed to generate tiled implementation");
+
+    return tiledOp[0]->getResult(resultNumber);
+  }
+
+  LogicalResult generateScalarImplementation(Operation* op, OpBuilder& builder,
+                                             Location loc,
+                                             ValueRange ivs) const {
+    auto linalgOp = cast<LinalgOp>(op);
+    if (!linalgOp.hasBufferSemantics())
+      return op->emitOpError("expected operation to have buffer semantics");
+
+    SmallVector<Value> indexedValues;
+    indexedValues.reserve(linalgOp->getNumOperands());
+    Location linalgOpLoc = op->getLoc();
+    /// Load the data corresponding to the block arguments that
+    /// represent input operands.
+    for (OpOperand& operand : linalgOp->getOpOperands()) {
+      if (!linalgOp.payloadUsesValueFromOperand(&operand)) {
+        indexedValues.push_back(nullptr);
+        continue;
+      }
+      if (linalgOp.isScalar(&operand)) {
+        indexedValues.push_back(operand.get());
+        continue;
+      }
+      SmallVector<Value> indices = getIndicesForAccess(
+          builder, linalgOpLoc, linalgOp.getMatchingIndexingMap(&operand), ivs);
+      Value load =
+          builder.create<memref::LoadOp>(linalgOpLoc, operand.get(), indices);
+      indexedValues.push_back(load);
+    }
+
+    /// Inline the op payload and store the result.
+    return inlinePayload(builder, linalgOp, ivs, indexedValues);
+  }
+};
+
+template <typename OpType>
+static void registerOne(MLIRContext* ctx) {
+  OpType::template attachInterface<LinalgOpTilingInterface<OpType>>(*ctx);
+}
+
+}  // namespace
+
+void registerTilingInterfaceExternalModels(DialectRegistry& registry) {
+  registry.addExtension(+[](MLIRContext* ctx, DISCLinalgExtDialect* dialect) {
+    registerOne<ConditionalGenericOp>(ctx);
+  });
+}
+
+//===----------------------------------------------------------------------===//
+// YieldOp
+//===----------------------------------------------------------------------===//
+
+void YieldOp::print(OpAsmPrinter& p) {
+  if (getNumOperands() > 0) p << ' ' << getOperands();
+  p.printOptionalAttrDict((*this)->getAttrs());
+  if (getNumOperands() > 0) p << " : " << getOperandTypes();
+}
+
+ParseResult YieldOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::UnresolvedOperand, 2> opInfo;
+  SmallVector<Type, 2> types;
+  SMLoc loc = parser.getCurrentLocation();
+  return failure(parser.parseOperandList(opInfo) ||
+                 parser.parseOptionalAttrDict(result.attributes) ||
+                 (!opInfo.empty() && parser.parseColonTypeList(types)) ||
+                 parser.resolveOperands(opInfo, types, loc, result.operands));
+}
+
+// Check the operand number and types must match the element types of the
+// LinalgOp interface's shaped operands.
+static LogicalResult verifyYield(YieldOp op, LinalgOp linalgOp) {
+  if (op.getNumOperands() != linalgOp.getNumDpsInits())
+    return op.emitOpError("expected number of yield values (")
+           << linalgOp.getNumDpsInits()
+           << ") to match the number of operands of the enclosing "
+           << "LinalgOp (" << op.getNumOperands() << ")";
+
+  for (OpOperand& opOperand : op->getOpOperands()) {
+    OpOperand* outputOperand =
+        linalgOp.getDpsInitOperand(opOperand.getOperandNumber());
+    Type elementType = getElementTypeOrSelf(outputOperand->get().getType());
+    if (opOperand.get().getType() != elementType)
+      return op.emitOpError("type of yield operand ")
+             << (opOperand.getOperandNumber() + 1) << " ("
+             << opOperand.get().getType() << ") doesn't match "
+             << "the element type of the enclosing linalg.generic op ("
+             << elementType << ")";
+  }
+  return success();
+}
+
+LogicalResult YieldOp::verify() {
+  auto* parentOp = (*this)->getParentOp();
+  if (parentOp->getNumRegions() != 1 || parentOp->getRegion(0).empty())
+    return emitOpError("expected single non-empty parent region");
+
+  if (auto linalgOp = dyn_cast<LinalgOp>(parentOp))
+    return verifyYield(*this, linalgOp);
+
+  return emitOpError("expected parent op with LinalgOp interface");
+}
+
+//===----------------------------------------------------------------------===//
+// IndexOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult IndexOp::verify() {
+  auto linalgOp = dyn_cast<LinalgOp>((*this)->getParentOp());
+  if (!linalgOp)
+    return emitOpError("expected parent op with LinalgOp interface");
+  if (linalgOp.getNumLoops() <= getDim())
+    return emitOpError("expected dim (")
+           << getDim() << ") to be lower than the number of loops ("
+           << linalgOp.getNumLoops() << ") of the enclosing LinalgOp";
+  return success();
+}
+
 namespace {
 /// This is derived from mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp without any
 /// changes.
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 01b4f10ffda..d94a116cbc4 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h
@@ -12,6 +12,7 @@
 #ifndef DISC_TOOLS_DISC_TRANSFORM_LINALGEXT_OPS_EXT_
 #define DISC_TOOLS_DISC_TRANSFORM_LINALGEXT_OPS_EXT_
 
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
@@ -23,6 +24,16 @@
 #include "mlir/disc/tools/disc-transform/LinalgExt/LinalgExtEnums.h.inc"
 #include "mlir/disc/tools/disc-transform/LinalgExt/LinalgExtInterfaces.h"
 
+namespace mlir {
+namespace disc_ral {
+namespace disc_linalg_ext {
+
+using linalg::LinalgOp;
+
+}  // namespace disc_linalg_ext
+}  // namespace disc_ral
+}  // namespace mlir
+
 #define GET_OP_CLASSES
 #include "mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h.inc"
 
@@ -47,6 +58,8 @@ SmallVector<T> interchange(ArrayRef<T> elements,
   return vec;
 }
 
+void registerTilingInterfaceExternalModels(DialectRegistry& registry);
+
 }  // namespace disc_linalg_ext
 }  // namespace disc_ral
 }  // namespace mlir
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 a6d6b1316b6..3879235ad96 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.td
@@ -15,9 +15,11 @@
 include "external/org_tensorflow/tensorflow/compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtBase.td"
 include "external/org_tensorflow/tensorflow/compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtInterfaces.td"
 
+include "mlir/Dialect/Linalg/IR/LinalgInterfaces.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/BuiltinAttributeInterfaces.td"
 include "mlir/IR/EnumAttr.td"
+include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/OpBase.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/DestinationStyleOpInterface.td"
@@ -30,6 +32,9 @@ include "mlir/Interfaces/ViewLikeInterface.td"
 // Base class.
 //===----------------------------------------------------------------------===//
 
+class DISCLinalgExt_BaseOp<string mnemonic, list<Trait> traits = []> :
+    Op<DISCLinalgExt_Dialect, mnemonic, traits>;
+
 class DISCLinalgExt_Op<string mnemonic, list<Trait> traits = []> :
     Op<DISCLinalgExt_Dialect, mnemonic, !listconcat(traits,
         [AttrSizedOperandSegments,
@@ -358,4 +363,230 @@ def DISCLinalgExt_MultiLevelPackOp : DISCLinalgExt_Op<"multi_level_pack", [
   }];
 }
 
+// Copied from linalg dialect
+def DISCLinalgExt_YieldOp : DISCLinalgExt_BaseOp<"yield", [Pure, ReturnLike, Terminator]>,
+    Arguments<(ins Variadic<AnyType>:$values)> {
+  let summary = "disc_linalg_ext yield operation";
+  let description = [{
+    `disc_linalg_ext.yield` is a special terminator operation for blocks inside regions
+    in `disc_linalg_ext` condition_generic ops. It returns values to the immediately enclosing
+    `disc_linalg_ext` generic op.
+
+    Example:
+
+    ```mlir
+    disc_linalg_ext.yield %f0, %f1 : f32, f32
+    ```
+  }];
+  let builders = [OpBuilder<(ins), [{ /* nothing to do */ }]>];
+  let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
+}
+
+// Copied from linalg dialect
+def DISCLinalgExt_IndexOp : DISCLinalgExt_BaseOp<"index", [Pure]>,
+    Arguments<(ins ConfinedAttr<I64Attr, [IntMinValue<0>]>:$dim)>,
+    Results<(outs Index:$result)> {
+  let summary = "disc_linalg_ext index operation";
+  let description = [{
+    The `disc_linalg_ext.index` operation returns the iteration index of the immediately
+    enclosing disc_linalg_ext structured operation for the iteration dimension `dim`. The
+    `dim` attribute specifies the position of the accessed dimension in the
+    indexing map domain.
+
+    Example:
+
+    ```mlir
+    #map = affine_map<(i, j) -> (i, j)>
+    disc_linalg_ext.condition_generic {indexing_maps = [#map, #map],
+                    iterator_types = ["parallel", "parallel"]}
+      outs(%I, %J : memref<?x?xindex>, memref<?x?xindex>) {
+      ^bb0(%arg0 : index, %arg1 : index):
+      // Access the outer iteration dimension i
+      %i = disc_linalg_ext.index 0 : index
+      // Access the inner iteration dimension j
+      %j = disc_linalg_ext.index 1 : index
+      disc_linalg_ext.yield %i, %j : index, index
+    }
+    ```
+
+    This may lower to IR resembling:
+
+    ```mlir
+    %0 = dim %I, %c0 : memref<?x?xindex>
+    %1 = dim %I, %c1 : memref<?x?xindex>
+    scf.for %i = %c0 to %0 step %c1 {
+      scf.for %j = %c0 to %1 step %c1 {
+        store %i, %I[%i, %j] : memref<?x?xindex>
+        store %j, %J[%i, %j] : memref<?x?xindex>
+      }
+    }
+    ```
+  }];
+
+  let assemblyFormat = [{ $dim attr-dict `:` type($result) }];
+  let hasVerifier = 1;
+}
+
+// Copied from linalg dialect
+// Base Tablegen class for Linalg ops.
+// Linalg ops that correspond to library calls operate on ShapedType as their
+// first operands. These may be optionally followed by non-view operands
+// depending on the specific Linalg op.
+class DISCLinalgExt_StructuredOp<string mnemonic, list<Trait> props>
+  : DISCLinalgExt_BaseOp<mnemonic, !listconcat([
+       SingleBlockImplicitTerminator<"YieldOp">,
+       DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+       DestinationStyleOpInterface,
+       LinalgStructuredInterface,
+       RegionBranchOpInterface,
+       ReifyRankedShapedTypeOpInterface], props)> {
+  code structuredOpsBaseDecls = [{
+    // Return whether the op accesses the iteration indices.
+    bool hasIndexSemantics() {
+      return !this->getBody()->getOps<IndexOp>().empty();
+    }
+
+    LogicalResult reifyResultShapes(OpBuilder &b,
+        ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
+      return llvm::cast<LinalgOp>(getOperation()).reifyResultShapes(b,
+          reifiedReturnShapes);
+    }
+
+    void getSuccessorRegions(
+        Optional<unsigned> index, ArrayRef<Attribute> operands,
+        SmallVectorImpl<RegionSuccessor> &regions) {
+      // Op has a region, but conceptually the control flow does not enter the
+      // region.
+    }
+  }];
+}
+
+def DISCLinalgExt_ConditionalGenericOp : DISCLinalgExt_StructuredOp<"conditional_generic", [
+    DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>,
+    AttrSizedOperandSegments]> {
+  let description = [{
+    Conditional Generic Linalg op form where the key properties of the computation are
+    specified as attributes. In pretty form, a `disc_linalg_ext.conditional_generic` op
+    is written as:
+
+      ```mlir
+      disc_linalg_ext.conditional_generic #trait_attribute
+          ins(%pred, %A, %B : i1, memref<?x?xf32>, memref<?x?xf32>)
+          outs(%C : memref<?x?xf32>)
+          attrs = {other-optional-attributes}
+          {region}
+      ```
+
+    Where #trait_attributes is an alias of a dictionary attribute containing:
+      - doc [optional]: a documentation string
+      - indexing_maps: a list of AffineMapAttr, one AffineMapAttr per each input
+        and output view. Such AffineMapAttr specifies the mapping between the
+        loops and the indexing within each view.
+      - library_call [optional]: a StringAttr containing the name of an
+        external library function that the linalg.generic operation maps to.
+        The external library is assumed to be dynamically linked and no strong
+        compile-time guarantees are provided. In the absence of such a library
+        call, linalg.generic will always lower to loops.
+      - iterator_types: an ArrayAttr specifying the type of the enclosing loops.
+        Each element of the list represents and iterator of one of the following
+        types:
+          parallel, reduction, window
+
+    Example:
+    Defining a elemment-wise epilogue ops
+      ```mlir
+      disc_linalg_ext.conditional_generic #trait_attribute
+        ins(%pred, %A, %B : i1, memref<?x?xf32>, memref<?x?xf32>)
+        outs(%C : memref<?x?xf32>)
+        {other-optional-attributes} {
+        ^bb0(%a: f32, %b: f32, %c: f32) :
+          %d = arith.mulf %a, %b: f32
+          linalg.yield %d : f32
+      }
+      ```
+
+    This may lower to:
+    ```mlir
+    scf.if (%pred) {
+      scf.for %m = %c0 to %M step %c1 {
+        scf.for %n = %c0 to %N step %c1 {
+            %a = load %A[%m, %n] : memref<?x?xf32>
+            %b = load %B[%m, %n] : memref<?x?xf32>
+            %d = arith.mulf %a, %b: f32
+            store %d, %C[%m, %n] : memref<?x?x?xf32>
+          }
+        }
+      }
+    }
+    ```
+  }];
+
+  let arguments = (ins Variadic<AnyType>:$inputs,
+                       Variadic<AnyShaped>:$outputs,
+                       AffineMapArrayAttr:$indexing_maps,
+                       ArrayAttr:$iterator_types,
+                       OptionalAttr<StrAttr>:$doc,
+                       OptionalAttr<StrAttr>:$library_call);
+  let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
+  let regions = (region AnyRegion:$region);
+
+  let builders = [
+    OpBuilder<(ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
+      "ValueRange":$outputs, "ArrayAttr":$indexingMaps,
+      "ArrayAttr":$iteratorTypes, "StringAttr":$doc,
+      "StringAttr":$libraryCall,
+      "function_ref<void(OpBuilder &, Location, ValueRange)>",
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
+    OpBuilder<(ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
+      "ValueRange":$outputs, "ArrayRef<AffineMap>":$indexingMaps,
+      "ArrayRef<StringRef>":$iteratorTypes, "StringRef":$doc,
+      "StringRef":$libraryCall,
+      CArg<"function_ref<void(OpBuilder &, Location, ValueRange)>", "nullptr">,
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
+    OpBuilder<(ins "ValueRange":$inputs, "ValueRange":$outputBuffers,
+      "ArrayRef<AffineMap>":$indexingMaps, "ArrayRef<StringRef>":$iteratorTypes,
+      "StringRef":$doc, "StringRef":$libraryCall,
+      CArg<"function_ref<void(OpBuilder &, Location, ValueRange)>", "nullptr">,
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
+    OpBuilder<(ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
+      "ValueRange":$outputs, "ArrayRef<AffineMap>":$indexingMaps,
+      "ArrayRef<StringRef>":$iteratorTypes,
+      CArg<"function_ref<void(OpBuilder &, Location, ValueRange)>", "nullptr">,
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
+    OpBuilder<(ins "ValueRange":$inputs, "ValueRange":$outputBuffers,
+      "ArrayRef<AffineMap>":$indexingMaps, "ArrayRef<StringRef>":$iteratorTypes,
+      CArg<"function_ref<void(OpBuilder &, Location, ValueRange)>", "nullptr">,
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>
+  ];
+
+  let extraClassDeclaration = structuredOpsBaseDecls # [{
+    SmallVector<StringRef, 8> linalgTraitAttrNames() {
+      return SmallVector<StringRef, 8>{
+        getDocAttrName(),
+        getIndexingMapsAttrName(), getLibraryCallAttrName(),
+        getIteratorTypesAttrName(),
+      };
+    }
+    std::string getLibraryCallName() {
+      return getLibraryCall() ?
+        getLibraryCall()->str() : "op_has_no_registered_library_name";
+    }
+
+    static std::function<void(ImplicitLocOpBuilder &,
+                              Block &, ArrayRef<NamedAttribute>)>
+    getRegionBuilder() {
+      return nullptr;
+    }
+    std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
+      int64_t getNumOperands = this->getNumOperands();
+      return {getNumOperands - getOutputs().size(), getNumOperands};
+    }
+  }];
+
+  let hasCanonicalizer = 1;
+  let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
+}
+
 #endif  // DISC_LINALGEXT_OPS
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/transform_dialect_interpreter.cc b/tao_compiler/mlir/disc/tools/disc-transform/transforms/transform_dialect_interpreter.cc
index 3ba1ac2acc2..d16d00ef0d4 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/transforms/transform_dialect_interpreter.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/transform_dialect_interpreter.cc
@@ -39,6 +39,7 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/Passes.h"
 #include "mlir/disc/tools/disc-transform/LinalgExt/LinalgExtDialect.h"
+#include "mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h"
 #include "mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.h"
 #include "mlir/disc/tools/disc-transform/transforms/PassDetail.h"
 #include "mlir/disc/tools/disc-transform/utils.h"
@@ -90,6 +91,7 @@ void addTransformDialectDependentDialects(DialectRegistry& registry) {
   linalg::registerTransformDialectExtension(registry);
   scf::registerTransformDialectExtension(registry);
   registerTransformDialectCommonExtension(registry);
+  disc_linalg_ext::registerTilingInterfaceExternalModels(registry);
 }
 
 namespace {


Test patch:
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/tests/ops.mlir b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/tests/ops.mlir
index cf342ebf7a0..89510e260e1 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/tests/ops.mlir
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/tests/ops.mlir
@@ -24,3 +24,20 @@ func.func @padding_value_placeholder() -> f32 {
   %2 = arith.addf %0, %1 : f32
   return %2 : f32
 }
+
+// -----
+
+// CHECK-LABEL: @conditional_generic
+#map0 = affine_map<(d0, d1, d2) -> ()>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+func.func @conditional_generic(%pred : i1, %arg0 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
+  %out = disc_linalg_ext.conditional_generic  {indexing_maps = [#map0, #map1],
+                   iterator_types = ["parallel", "parallel", "parallel"]}
+                  ins(%pred : i1)
+                  outs(%arg0 : tensor<?x?x?xf32>) {
+  ^bb0(%arg4: i1, %arg3: f32):
+    %cst = arith.constant 0.000000e+00 : f32
+    disc_linalg_ext.yield %cst : f32
+  } -> tensor<?x?x?xf32>
+  return %out : tensor<?x?x?xf32>
+}
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/conditional-generic.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/conditional-generic.mlir
new file mode 100644
index 00000000000..19e2bdf4330
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/conditional-generic.mlir
@@ -0,0 +1,95 @@
+// 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) -> (d0, d1)>
+
+// CHECK-LABEL: @tile_conditional_generic
+func.func @tile_conditional_generic(%pred : i1, %arg0: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  // CHECK: scf.for
+  // CHECK: scf.for
+  // CHECK: tensor.extract_slice
+  // CHECK-NEXT: disc_linalg_ext.conditional_generic
+  %out = disc_linalg_ext.conditional_generic {indexing_maps = [#map0, #map1],
+      iterator_types = ["parallel", "parallel"]}
+      ins(%pred : i1)
+      outs(%arg0 : tensor<?x?xf32>) {
+  ^bb0(%arg4: i1, %arg3: f32):
+    disc_linalg_ext.yield %cst : f32
+  } -> tensor<?x?xf32>
+  return %out : tensor<?x?xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %0 = transform.structured.match ops{["disc_linalg_ext.conditional_generic"]} in %arg0
+  %1, %loops:2 = transform.structured.tile %0 [288, 512]
+}
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> ()>
+#map1 = affine_map<(d0, d1) -> (d0, d1)>
+
+// CHECK-LABEL: @fuse_into_containing_op_conditional_generic
+func.func @fuse_into_containing_op_conditional_generic(
+    %pred : i1, %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %cst = arith.constant 1.000000e+00 : f32
+  %0 = disc_linalg_ext.conditional_generic {indexing_maps = [#map0, #map1],
+      iterator_types = ["parallel", "parallel"]}
+      ins(%pred : i1)
+      outs(%arg0 : tensor<?x?xf32>) {
+  ^bb0(%arg4: i1, %arg3: f32):
+    disc_linalg_ext.yield %cst : f32
+  } -> tensor<?x?xf32>
+  // CHECK: scf.for
+  // CHECK:   %[[T0:.*]] = disc_linalg_ext.conditional_generic
+  // CHECK:   scf.for
+  // CHECK:     linalg.matmul
+  // CHECK-SAME:  %[[T0]]
+  %1 = linalg.matmul ins(%0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %1 : tensor<?x?xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+  ^bb0(%arg0: !pdl.operation):
+    %0 = transform.structured.match ops{["disc_linalg_ext.conditional_generic"]} in %arg0
+    %1 = transform.structured.match ops{["linalg.matmul"]} in %arg0
+    %2, %loops:2 = transform.structured.tile %1 [2, 3]
+    transform.structured.fuse_into_containing_op %0 into %loops#0
+}
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> ()>
+#map1 = affine_map<(d0, d1) -> (d0, d1)>
+#map2 = affine_map<(d0) -> (6, d0)>
+#map3 = affine_map<(d0) -> (16, d0)>
+
+// CHECK-LABEL: @pad_conditional_generic
+func.func @pad_conditional_generic(%pred : i1, %arg1 : index, %arg2 : index, %arg_out : tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %d0 = affine.min #map2(%arg1)
+  %d1 = affine.min #map2(%arg2)
+  %1 = tensor.extract_slice %arg_out[0, 0] [%d0, %d1] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+  // CHECK: %[[T0:.*]] = tensor.pad
+  // CHECK: %[[T1:.*]] = disc_linalg_ext.conditional_generic
+  // CHECK-SAME: outs(%[[T0]] : tensor<6x6xf32>)
+  // CHECK: tensor.extract_slice
+  // CHECK-SAME: %[[T1]]
+  %out = disc_linalg_ext.conditional_generic {indexing_maps = [#map0, #map1],
+      iterator_types = ["parallel", "parallel"]}
+      ins(%pred : i1)
+      outs(%1 : tensor<?x?xf32>) {
+  ^bb0(%arg4: i1, %arg3: f32):
+    disc_linalg_ext.yield %cst : f32
+  } -> tensor<?x?xf32>
+  return %out : tensor<?x?xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %0 = transform.structured.match ops{["disc_linalg_ext.conditional_generic"]} in %arg0
+  %1 = transform.structured.pad %0 {padding_dimensions = [0, 1], padding_values = [0 : i1, 0.000000e+00 : f32]}
+}
+