[mlir][linalg][transform] Add TileOp to transform dialect

[springerm]

Summary

This commit adds a tiling op to the transform dialect as an external op.

[springerm]

Fix CMake build

[ftynse]

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h

1

Nit: wrong filename

13

Nit: wouldn't OpDefinition.h suffice in the header?

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

9

These aren't really linear algebra operations, but transforms thereof.

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

63

function_ref (std::function is only useful if the callee needs to outlive the caller)

66–67

Is it possible to reserve space for these?

70

I'd rather dyn_cast and report a user-visible error here.

81–83

Nit: drop trivial braces.

87–89

Nit: drop trivial braces.

117

Ultra-nit: we tend to use single-quotes rather than backticks in user-visible text.

120

Don't we have something like getSizesAttrName() to avoid magic strings?

123

Nit: there is no need to use formatv to insert one statically known string into another statically known string, just write the full string.

I suppose this could have been taken from some other place that used getSizesAttrName() instead of hardcoding the attribute name. Even in that case, something like parser.emitError(opLoc) << "expected " << getSizesAttrName() << " attribute"; is probably cheaper / easier for the compiler to optimize than formatv.

159

We also need Read/Write effect on PayloadIRResource, not attached to a particular value.

[springerm]

address comments

[springerm]

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h

13

Actually requires OpImplementation.h because OpAsmOpInterface is needed.

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

9

Ah, this was copy pasta...

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

120

We do, but getSizesAttrName is a member function and TileOp::parse() is static.

[ftynse]

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

120

Isn't there a static version that takes OperationName as argument? You can get OperationName from OperationState.

[springerm]

address comments

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

120

Nice I didn't know that

mlir/include/mlir/Dialect/Linalg/CMakeLists.txt

 add_subdirectory(IR)
+add_subdirectory(TransformOps)
 
 set(LLVM_TARGET_DEFINITIONS Passes.td)
 mlir_tablegen(Passes.h.inc -gen-pass-decls -name Linalg)
 mlir_tablegen(Passes.capi.h.inc -gen-pass-capi-header --prefix Linalg)
 mlir_tablegen(Passes.capi.cpp.inc -gen-pass-capi-impl --prefix Linalg)
 add_public_tablegen_target(MLIRLinalgPassIncGen)
 
 add_mlir_doc(Passes LinalgPasses ./ -gen-pass-doc)

mlir/include/mlir/Dialect/Linalg/TransformOps/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS LinalgTransformOps.td)
+mlir_tablegen(LinalgTransformOps.h.inc -gen-op-decls)
+mlir_tablegen(LinalgTransformOps.cpp.inc -gen-op-defs)
+add_public_tablegen_target(MLIRLinalgTransformOpsIncGen)

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h

+//===- Linalg.h - Linalg dialect --------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_LINALG_TRANSFORMOPS_LINALGTRANSFORMOPS_H
+#define MLIR_DIALECT_LINALG_TRANSFORMOPS_LINALGTRANSFORMOPS_H
+
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/IR/OpImplementation.h"
+
+//===----------------------------------------------------------------------===//
+// Linalg Transform Operations
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h.inc"
+
+namespace mlir {
+class DialectRegistry;
+
+namespace linalg {
+void registerTransformDialectExtension(DialectRegistry &registry);
+} // namespace linalg
+} // namespace mlir
+
+#endif // MLIR_DIALECT_LINALG_TRANSFORMOPS_LINALGTRANSFORMOPS_H

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

+//===- LinalgTransformOps.td - Linalg transform ops --------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This is the operation definition file for linear algebra operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LINALG_TRANSFORM_OPS
+#define LINALG_TRANSFORM_OPS
+
+include "mlir/Dialect/Transform/IR/TransformDialect.td"
+include "mlir/Dialect/Transform/IR/TransformEffects.td"
+include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
+include "mlir/Dialect/PDL/IR/PDLTypes.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/OpBase.td"
+
+def TileOp : Op<Transform_Dialect, "structured.tile",
+       [DeclareOpInterfaceMethods<TransformOpInterface>,
+        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let description = [{
+    Indicates that the given `target` op should be tiled with the options
+    provided as attributes. This transform generates a loop nest with a smaller
+    ("tiled") target operation in its body. Currently limited to LinalgOps.
+
+    `sizes` are the tile sizes. A tile size of `0` indicates that the
+    respective dimension should not be tiled. No loop will be generated for such
+    dimensions. If all tile sizes are `0`, this transform is effectively a
+    no-op.
+
+    This op returns handles to the tiled op (in the generated loop nest) and the
+    generated loops. The number of loops is the number of non-zero tile sizes.
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                   DefaultValuedAttr<I64ArrayAttr, "{}">:$sizes,
+                   DefaultValuedAttr<I64ArrayAttr, "{}">:$interchange);
+  let results = (outs PDL_Operation:$tiled_linalg_op,
+                      Variadic<PDL_Operation>:$loops);
+
+  let hasCustomAssemblyFormat = 1;
+}
+
+#endif // LINALG_TRANSFORM_OPS

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/EmitC/IR/EmitC.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h"
 #include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/NVGPU/NVGPUDialect.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
registry.insert<acc::OpenACCDialect,
                   ROCDL::ROCDLDialect,
                   shape::ShapeDialect,
                   sparse_tensor::SparseTensorDialect,
                   tensor::TensorDialect,
                   transform::TransformDialect,
                   tosa::TosaDialect,
                   x86vector::X86VectorDialect>();
   // clang-format on
+
+  // Register all dialect extensions.
+  linalg::registerTransformDialectExtension(registry);
+
+  // Register all external models.
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);
   shape::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerInferTypeOpInterfaceExternalModels(registry);

mlir/lib/Dialect/Linalg/CMakeLists.txt

 add_subdirectory(Analysis)
 add_subdirectory(IR)
+add_subdirectory(TransformOps)
 add_subdirectory(Transforms)
 add_subdirectory(Utils)

mlir/lib/Dialect/Linalg/TransformOps/CMakeLists.txt

+add_mlir_dialect_library(MLIRLinalgTransformOps
+  LinalgTransformOps.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg/TransformOps
+
+  DEPENDS
+  MLIRLinalgTransformOpsIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRLinalg
+  MLIRLinalgTransforms
+  MLIRParser
+  MLIRPDL
+  MLIRSideEffectInterfaces
+  MLIRTransformDialect
+  )

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

+//===- LinalgTransformOps.cpp - Implementation of Linalg transform ops ----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h"
+
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/PDL/IR/PDL.h"
+#include "mlir/Dialect/PDL/IR/PDLTypes.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Parser/Parser.h"
+#include "llvm/Support/FormatVariadic.h"
+
+using namespace mlir;
+using namespace mlir::linalg;
+using namespace mlir::transform;
+
+/// Extracts a vector of int64_t from an array attribute. Asserts if the
+/// attribute contains values other than integers.
+static SmallVector<int64_t> extractI64Array(ArrayAttr attr) {
+  SmallVector<int64_t> result;
+  result.reserve(attr.size());
+  for (APInt value : attr.getAsValueRange<IntegerAttr>())
+    result.push_back(value.getSExtValue());
+  return result;
+}
+
+/// Extracts a vector of unsigned from an array attribute. Asserts if the
+/// attribute contains values other than intergers. May truncate.
+static SmallVector<unsigned> extractUIntArray(ArrayAttr attr) {
+  SmallVector<unsigned> result;
+  result.reserve(attr.size());
+  for (APInt value : attr.getAsValueRange<IntegerAttr>())
+    result.push_back(value.getZExtValue());
+  return result;
+}
+
+namespace {
+/// A simple pattern rewriter that implements no special logic.
+class SimpleRewriter : public PatternRewriter {
+public:
+  SimpleRewriter(MLIRContext *context) : PatternRewriter(context) {}
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// TileOp
+//===----------------------------------------------------------------------===//
+
+/// Apply a tiling transformation to all payload ops and store both the
+/// tiled operation as well as the created tile loops.
+static LogicalResult
+applyTilingToAll(Operation *transformOp, Value target,
+                 ArrayRef<int64_t> tileSizes,
+                 transform::TransformResults &transformResults,
+                 transform::TransformState &state,
+                 std::function<FailureOr<TiledLinalgOp>(LinalgOp)> applyFn) {
+  size_t numLoops = tileSizes.size() - llvm::count(tileSizes, 0);
+
+  SmallVector<Operation *> tiledLinalgOps;
+  SmallVector<SmallVector<Operation *>> loopOps(numLoops);
+
+  for (Operation *target : state.getPayloadOps(target)) {
+    auto linalgOp = cast<linalg::LinalgOp>(target);
+    FailureOr<TiledLinalgOp> tiled = applyFn(linalgOp);
+    if (failed(tiled))
+      return failure();
+
+    tiledLinalgOps.push_back(tiled->op);
+    if (tiled->loops.size() != numLoops)
+      // Not enough loops were generated. This usually means that the input size
+      // was smaller than the tiling size.
+      // TODO: LinalgTilingPattern should return failure().
+      return failure();
+    for (unsigned int i = 0; i < numLoops; ++i) {
+      loopOps[i].push_back(tiled->loops[i]);
+    }
+  }
+
+  transformResults.set(transformOp->getOpResult(0), tiledLinalgOps);
+  for (unsigned int i = 0; i < numLoops; ++i) {
+    transformResults.set(transformOp->getOpResult(i + 1), loopOps[i]);
+  }
+  return success();
+}
+
+LogicalResult transform::TileOp::apply(TransformResults &transformResults,
+                                       TransformState &state) {
+  LinalgTilingOptions tilingOptions;
+  SmallVector<int64_t> tileSizes = extractI64Array(getSizes());
+
+  if (!tileSizes.empty())
+    tilingOptions.setTileSizes(tileSizes);
+  tilingOptions.setInterchange(extractUIntArray(getInterchange()));
+  LinalgTilingPattern pattern(getContext(), tilingOptions);
+
+  return applyTilingToAll(getOperation(), getTarget(), tileSizes,
+                          transformResults, state, [&](LinalgOp linalgOp) {
+                            SimpleRewriter rewriter(linalgOp.getContext());
+                            return pattern.returningMatchAndRewrite(linalgOp,
+                                                                    rewriter);
+                          });
+}
+
+ParseResult transform::TileOp::parse(OpAsmParser &parser,
+                                     OperationState &result) {
+  OpAsmParser::UnresolvedOperand targetOperand;
+  SMLoc opLoc;
+  parser.getCurrentLocation(&opLoc);
+  if (parser.parseOperand(targetOperand))
+    return parser.emitError(opLoc, "expected `target` operand");
+  if (parser.parseOptionalAttrDict(result.attributes))
+    return failure();
+  Attribute sizesAttr = result.attributes.get("sizes");
+  if (!sizesAttr) {
+    return parser.emitError(opLoc,
+                            llvm::formatv("expected `{0}` attribute", "sizes"));
+  }
+  auto sizesArrayAttr = sizesAttr.dyn_cast<ArrayAttr>();
+  if (!sizesArrayAttr) {
+    return parser.emitError(
+        opLoc, llvm::formatv("`{0}` attribute must be an array", "sizes"));
+  }
+  Type pdlOpType = parser.getBuilder().getType<pdl::OperationType>();
+  size_t numExpectedLoops =
+      sizesArrayAttr.size() - llvm::count(extractI64Array(sizesArrayAttr), 0);
+  result.addTypes(SmallVector<Type>(numExpectedLoops + 1, pdlOpType));
+  if (parser.resolveOperand(targetOperand, pdlOpType, result.operands))
+    return failure();
+  return success();
+}
+
+void TileOp::print(OpAsmPrinter &p) {
+  p << ' ';
+  p << getTarget();
+  p.printOptionalAttrDict((*this)->getAttrs());
+}
+
+void TileOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+        &effects) {
+  // `target` arg is consumed and can no longer be used.
+  effects.emplace_back(MemoryEffects::Read::get(), getTarget(),
+                       TransformMappingResource::get());
+  effects.emplace_back(MemoryEffects::Free::get(), getTarget(),
+                       TransformMappingResource::get());
+
+  for (Value r : getResults()) {
+    effects.emplace_back(MemoryEffects::Write::get(), r,
+                         TransformMappingResource::get());
+    effects.emplace_back(MemoryEffects::Allocate::get(), r,
+                         TransformMappingResource::get());
+  }
+}
+
+//===----------------------------------------------------------------------===//
+// Transform op registration
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// Registers new ops and declares PDL as dependent dialect since the additional
+/// ops are using PDL types for operands and results.
+class LinalgTransformDialectExtension
+    : public transform::TransformDialectExtension<
+          LinalgTransformDialectExtension> {
+public:
+  LinalgTransformDialectExtension() {
+    declareDependentDialect<pdl::PDLDialect>();
+    declareDependentDialect<scf::SCFDialect>();
+    registerTransformOps<
+#define GET_OP_LIST
+#include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp.inc"
+        >();
+  }
+};
+} // namespace
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp.inc"
+
+void mlir::linalg::registerTransformDialectExtension(
+    DialectRegistry &registry) {
+  registry.addExtensions<LinalgTransformDialectExtension>();
+}

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

resultTypes.push_back(tensor::ExtractSliceOp::inferResultType(
         t.cast<RankedTensorType>(), staticOffsetsVector, staticSizesVector,
         staticStridesVector));
   }
 
   Operation *clonedOp = producer.clone(b, loc, resultTypes, clonedShapes);
 
   // Shift all IndexOp results by the tile offset.
   SmallVector<Value> allIvs;
-  transform(loopRanges, std::back_inserter(allIvs),
+  llvm::transform(loopRanges, std::back_inserter(allIvs),
                   [](Range range) { return range.offset; });
   addTileLoopIvsToIndexOpResults(b, clonedOp, allIvs);
 
   return clonedOp;
 }
 
 /// Get the loop range for a dimension `dim` based on the `shapedOperand`. It is
 /// expected to be defined by a subview op or an extract_slice op.
 static Range getRangeFromOperandShape(OpBuilder &b, Location loc,

mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp

getTiledProducerLoops(OpResult producerResult,
   // projected permutations, `tiledProducerIndexingSubMap` has to be a
   // projected permutation as well. We can thus obtain the producer loop indices
   // by getting the positions of the result dimensions.
   // Example:
   // (d0, d1, d2) -> (d0, d2) has the result positions [0, 2].
   assert(tiledProducerIndexingSubMap.isProjectedPermutation() &&
          "expect slice and producer loop dimensions map one-to-one");
   SmallVector<int64_t> tiledProducerLoopIndices;
-  transform(llvm::seq<unsigned>(0, tiledProducerIndexingSubMap.getNumResults()),
-            std::back_inserter(tiledProducerLoopIndices), [&](unsigned idx) {
-              return tiledProducerIndexingSubMap.getDimPosition(idx);
-            });
+  llvm::transform(
+      llvm::seq<unsigned>(0, tiledProducerIndexingSubMap.getNumResults()),
 
   return tiledProducerLoopIndices;
 }
 
 /// Returns the producer fused in place of `sliceOp`. Tile the producer operands
 /// along the `tiledSliceDimIndices` and clone the producer. Consider the case
 /// of fusion of an output tensor:
 /// ```
static LinalgOp getTiledProducer(OpBuilder &b, OpResult producerResult,
   b.setInsertionPointAfter(sliceOp);
 
   // Get the producer.
   LinalgOp producerOp = producerResult.getOwner();
   Location loc = producerOp.getLoc();
 
   // Obtain the `producerOp` loop bounds and the `sliceOp` ranges.
   SmallVector<Value> producerLoopBounds;
-  transform(producerOp.createLoopRanges(b, loc),
+  llvm::transform(producerOp.createLoopRanges(b, loc),
                   std::back_inserter(producerLoopBounds),
                   [](Range range) { return range.size; });
   SmallVector<Range> sliceOpRanges = sliceOp.getOrCreateRanges(b, loc);
 
   // Tile the producer operands given the `sliceOp` ranges. Iterate the
   // `tiledSliceDimIndices` and store the tile offset and size for the tiled
   // slice dimension.
   auto zero = b.create<arith::ConstantIndexOp>(loc, 0);
   SmallVector<Value> tileIvs(producerOp.getNumLoops(), nullptr);
   SmallVector<Value> tileSizes(producerOp.getNumLoops(), zero);

mlir/test/Dialect/Linalg/transform-ops.mlir

+// RUN: mlir-opt --test-transform-dialect-interpreter %s | FileCheck %s
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  sequence %arg0 {
+    ^bb0(%arg1: !pdl.operation):
+      %0 = pdl_match @pdl_target in %arg1
+      %1, %loops:3 = transform.structured.tile %0 {sizes = [4, 4, 4]}
+  }
+
+  pdl.pattern @pdl_target : benefit(1) {
+    %args = operands
+    %results = types
+    %0 = operation "linalg.matmul"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
+    rewrite %0 with "transform.dialect"
+  }
+}
+
+// CHECK-LABEL: func @tile_linalg_matmul(
+// CHECK-SAME:    %[[TA:[0-9a-z]+]]: tensor<128x128xf32>
+// CHECK-SAME:    %[[TB:[0-9a-z]+]]: tensor<128x128xf32>
+// CHECK-SAME:    %[[TC:[0-9a-z]+]]: tensor<128x128xf32>
+// CHECK-SAME:  -> tensor<128x128xf32> {
+func @tile_linalg_matmul(
+  %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32> {
+//      CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<128x128xf32>) {
+//      CHECK:   %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<128x128xf32>) {
+//      CHECK:     %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<128x128xf32>) {
+//      CHECK:       %[[sTA:.*]] = tensor.extract_slice %[[TA]][{{.*}}] : tensor<128x128xf32> to tensor<4x4xf32>
+//      CHECK:       %[[sTB:.*]] = tensor.extract_slice %[[TB]][{{.*}}] : tensor<128x128xf32> to tensor<4x4xf32>
+//      CHECK:       %[[sTC:.*]] = tensor.extract_slice %[[TC2]][{{.*}}] : tensor<128x128xf32> to tensor<4x4xf32>
+//      CHECK:       %[[sTD:.*]] = linalg.matmul ins(%[[sTA]], %[[sTB]] : tensor<4x4xf32>, tensor<4x4xf32>)
+// CHECK-SAME:                                   outs(%[[sTC]] : tensor<4x4xf32>)  -> tensor<4x4xf32>
+//      CHECK:       %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}]  : tensor<4x4xf32> into tensor<128x128xf32>
+//      CHECK:       scf.yield %[[TD]] : tensor<128x128xf32>
+//      CHECK:     scf.yield %[[TD2]] : tensor<128x128xf32>
+//      CHECK:   scf.yield %[[TD1]] : tensor<128x128xf32>
+  %0 = linalg.matmul  ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32>
+
+//      CHECK: return %[[TD0]] : tensor<128x128xf32>
+  return %0 : tensor<128x128xf32>
+}
+

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":LLVMDialect",
         ":LLVMIRTransforms",
         ":LLVMPassIncGen",
         ":LinalgOps",
         ":LinalgPassIncGen",
         ":LinalgToLLVM",
         ":LinalgToSPIRV",
         ":LinalgToStandard",
+        ":LinalgTransformOps",
         ":LinalgTransforms",
         ":MLProgramDialect",
         ":MathDialect",
         ":MathToLLVM",
         ":MathToLibm",
         ":MathToSPIRV",
         ":MathTransforms",
         ":MemRefDialect",
deps = [
         ":LoopLikeInterfaceTdFiles",
         ":OpBaseTdFiles",
         ":SideEffectInterfacesTdFiles",
         ":TilingInterfaceTdFiles",
         ":ViewLikeInterfaceTdFiles",
     ],
 )
 
+td_library(
+    name = "LinalgTransformOpsTdFiles",
+    srcs = [
+        "include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td",
+    ],
+    includes = ["include"],
+    deps = [
+        ":PDLDialectTdFiles",
+        ":TransformDialectTdFiles",
+    ],
+)
+
 gentbl_cc_library(
     name = "LinalgOpsIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
         (
             ["-gen-op-decls"],
             "include/mlir/Dialect/Linalg/IR/LinalgOps.h.inc",
         ),
tbl_outs = [
             "include/mlir/Dialect/Linalg/IR/LinalgOpsAttrDefs.cpp.inc",
         ),
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/Linalg/IR/LinalgOps.td",
     deps = [":LinalgOpsTdFiles"],
 )
 
+gentbl_cc_library(
+    name = "LinalgTransformOpsIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            ["-gen-op-decls"],
+            "include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h.inc",
+        ),
+        (
+            ["-gen-op-defs"],
+            "include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td",
+    deps = [
+        ":LinalgTransformOpsTdFiles",
+    ],
+)
+
 genlinalg(
     name = "LinalgNamedStructuredOpsYamlIncGen",
     src = "include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml",
     linalg_outs = [
         (
             "-o-impl=$@",
             "include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yamlgen.cpp.inc",
         ),
deps = [
         ":Support",
         ":TensorDialect",
         ":TilingInterface",
         ":ViewLikeInterface",
         "//llvm:Support",
     ],
 )
 
+cc_library(
+    name = "LinalgTransformOps",
+    srcs = [
+        "lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp",
+    ],
+    hdrs = [
+        "include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h",
+    ],
+    includes = ["include"],
+    deps = [
+        ":IR",
+        ":LinalgOps",
+        ":LinalgTransformOpsIncGen",
+        ":LinalgTransforms",
+        ":PDLDialect",
+        ":Parser",
+        ":SideEffectInterfaces",
+        ":TransformDialect",
+        "//llvm:Support",
+    ],
+)
+
 gentbl_cc_library(
     name = "LinalgPassIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
         (
             [
                 "-gen-pass-decls",
                 "-name=Linalg",