[DT] Set encodings if iree.opt.data_tiling unit attribute is attached. (#21676)

The revision moves the filter to AnnotateDataTilingHints pass that adds
the unit attribute to target gemms.

If any operation has the data-tiling hint, it does nothing. It allows
users to perform data-tiling selectively in preprocessing. E.g., user
can run a transform dialect script that matches linalg ops and attaches
the unit attribute.

E.g., the below transform script only attaches `iree.opt.data_tiling` to
the first matmul.

```mlir
func.func @matmuls(%lhs: tensor<?x?xi8>, %rhs: tensor<?x?xi8>, %acc: tensor<?x?xi32>) -> (tensor<?x?xi32>, tensor<?x?xi32>) {
  %res0 = linalg.matmul_transpose_b
        ins(%lhs, %lhs : tensor<?x?xi8>, tensor<?x?xi8>)
        outs(%acc : tensor<?x?xi32>) -> tensor<?x?xi32>
  %res1 = linalg.matmul_transpose_b
        ins(%lhs, %rhs : tensor<?x?xi8>, tensor<?x?xi8>)
        outs(%acc : tensor<?x?xi32>) -> tensor<?x?xi32>
  return %res0, %res1 : tensor<?x?xi32>, tensor<?x?xi32>
}

module attributes {transform.with_named_sequence} {
  transform.named_sequence @match_matmul_repeated_operand(%arg0: !transform.any_op {transform.readonly}) -> !transform.any_op {
    %inputs, %outputs = transform.iree.match.cast_compatible_dag_from_root %arg0 {
    ^bb0(%arg1: tensor<?x?xi8>, %arg2: tensor<?x?xi32>):
      %1 = linalg.matmul_transpose_b
          ins(%arg1, %arg1 : tensor<?x?xi8>, tensor<?x?xi8>)
          outs(%arg2 : tensor<?x?xi32>) -> tensor<?x?xi32>
    } : (!transform.any_op) -> (!transform.any_value, !transform.any_value)
    transform.yield %arg0 : !transform.any_op
  }

  transform.named_sequence @Annotate(%generic: !transform.any_op {transform.readonly}) {
    transform.annotate %generic "iree.opt.data_tiling" : !transform.any_op
    transform.yield
  }

  transform.named_sequence @__transform_main(%module: !transform.any_op) {
    %func = transform.structured.match ops{["func.func"]} in %module : (!transform.any_op) -> !transform.any_op
    transform.foreach_match in %module
        @match_matmul_repeated_operand -> @Annotate
      : (!transform.any_op) -> (!transform.any_op)
    transform.yield
  }
}
```

Fixes #21246

---------

Signed-off-by: hanhanW <[email protected]>

Author: hanhanW

compiler/src/iree/compiler/Dialect/Encoding/Utils/Utils.h

@@ -9,10 +9,23 @@
 
 #include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
 #include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 
 namespace mlir::iree_compiler::IREE::Encoding {
 
+constexpr char kDataTilingHint[] = "iree.opt.data_tiling";
+
+/// Returns true if the operation has data-tiling hint attribute.
+inline bool hasDataTilingHint(Operation *op) {
+  return op->getAttr(kDataTilingHint) ? true : false;
+}
+
+/// Adds an unit attribute with `kDataTilingHint` key to the operation.
+inline void setDataTilingHint(Operation *op) {
+  op->setAttr(kDataTilingHint, UnitAttr::get(op->getContext()));
+}
+
 /// Returns the encoding attribute from the type if there is an encoding that
 /// implements SerializableAttr. Otherwise, returns null.
 SerializableAttr getSerializableAttr(RankedTensorType type);

compiler/src/iree/compiler/DispatchCreation/AnnotateDataTilingHints.cpp

@@ -0,0 +1,151 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed 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 "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
+#include "iree/compiler/Dialect/Encoding/Utils/Utils.h"
+#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
+#include "iree/compiler/DispatchCreation/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Support/WalkResult.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+#define DEBUG_TYPE "iree-dispatch-creation-annotate-data-tiling-hints"
+
+namespace mlir::iree_compiler::DispatchCreation {
+#define GEN_PASS_DEF_ANNOTATEDATATILINGHINTSPASS
+#include "iree/compiler/DispatchCreation/Passes.h.inc"
+
+namespace {
+struct AnnotateDataTilingHintsPass final
+    : impl::AnnotateDataTilingHintsPassBase<AnnotateDataTilingHintsPass> {
+  using Base::Base;
+  void runOnOperation() override;
+};
+} // namespace
+
+/// Returns true iff the linalgOp has a body like a regular matmul, i.e.
+/// yield(add(out, mul(cast(in0), cast(in1))))
+static bool hasMatmulLikeBody(linalg::LinalgOp linalgOp) {
+  auto outBlockArg =
+      linalgOp.getMatchingBlockArgument(linalgOp.getDpsInitOperand(0));
+  auto yieldOp =
+      dyn_cast<linalg::YieldOp>(outBlockArg.getParentBlock()->getTerminator());
+  if (!yieldOp) {
+    return false;
+  }
+  Operation *addOp = yieldOp->getOperand(0).getDefiningOp();
+  if (!addOp || !isa<arith::AddIOp, arith::AddFOp>(addOp)) {
+    return false;
+  }
+  Value addLhs = addOp->getOperand(0);
+  Value addRhs = addOp->getOperand(1);
+  Operation *addLhsOp = addLhs.getDefiningOp();
+  Operation *addRhsOp = addRhs.getDefiningOp();
+  if (!(addLhsOp && addRhs == outBlockArg) &&
+      !(addRhsOp && addLhs == outBlockArg)) {
+    return false;
+  }
+  Operation *mulOp = addLhsOp ? addLhsOp : addRhsOp;
+  if (!isa<arith::MulFOp, arith::MulIOp>(mulOp)) {
+    return false;
+  }
+  Value mulLhs = mulOp->getOperand(0);
+  Value mulRhs = mulOp->getOperand(1);
+  auto mulLhsOp = mulLhs.getDefiningOp<CastOpInterface>();
+  auto mulRhsOp = mulRhs.getDefiningOp<CastOpInterface>();
+  if (!isa<BlockArgument>(mulLhs) && !mulLhsOp && !isa<BlockArgument>(mulRhs) &&
+      !mulRhsOp) {
+    return false;
+  }
+  if ((mulLhsOp && !isa<BlockArgument>(mulLhsOp->getOperand(0))) ||
+      (mulRhsOp && !isa<BlockArgument>(mulRhsOp->getOperand(0)))) {
+    return false;
+  }
+  return true;
+}
+
+/// Not all contractions are supported by data tiling, so return true if:
+///   1) linalgOp has pure tensor semantics.
+///   2) linalgOp does not have a preset compilation info.
+///   3) The workgroup count is not present if linalgOp is wrapped within
+///      Flow::DispatchRegionOp.
+///   4) All the operands do not have encodings.
+///   5) linalgOp has contraction indexingMaps.
+///   6) There are not more than one of each contraction dimension.
+///   7) There is an M or N dimension, and there is a K dimension.
+///   8) linalgOp has the same body as an ordinary int or float matmul.
+///
+/// These restrictions are required because data tiling currently creates
+/// an Mmt4DOp or BatchMmt4DOp on the packed inputs.
+///
+/// TODO(#16176): Loosen restrictions on contraction ops once data tiling
+/// can support more cases.
+static bool isSupportedContractionOp(linalg::LinalgOp linalgOp) {
+  if (!linalgOp.hasPureTensorSemantics()) {
+    return false;
+  }
+  if (getCompilationInfo(linalgOp)) {
+    return false;
+  }
+  auto hasWorkgroupCounts = [](Operation *op) -> bool {
+    auto parentDispatchOp = op->getParentOfType<IREE::Flow::DispatchRegionOp>();
+    return parentDispatchOp && !parentDispatchOp.getWorkgroupCount().empty();
+  };
+  if (hasWorkgroupCounts(linalgOp)) {
+    return false;
+  }
+  auto hasEncoding = [](Value operand) -> bool {
+    auto type = llvm::dyn_cast<RankedTensorType>(operand.getType());
+    return type && type.getEncoding();
+  };
+  if (llvm::any_of(linalgOp.getDpsInputs(), hasEncoding) ||
+      llvm::any_of(linalgOp.getDpsInits(), hasEncoding)) {
+    return false;
+  }
+
+  if (!linalg::isaContractionOpInterface(linalgOp)) {
+    return false;
+  }
+  auto cDims = linalg::inferContractionDims(linalgOp);
+  if (failed(cDims) || cDims->batch.size() > 1 || cDims->m.size() > 1 ||
+      cDims->n.size() > 1 || cDims->k.size() > 1) {
+    return false;
+  }
+  if ((cDims->n.empty() && cDims->m.empty()) || cDims->k.empty()) {
+    return false;
+  }
+  if (!hasMatmulLikeBody(linalgOp)) {
+    return false;
+  }
+  return true;
+}
+
+void AnnotateDataTilingHintsPass::runOnOperation() {
+  FunctionOpInterface funcOp = getOperation();
+  SmallVector<Operation *> candidates;
+  WalkResult result = funcOp.walk([&](Operation *op) -> WalkResult {
+    if (IREE::Encoding::hasDataTilingHint(op)) {
+      return WalkResult::interrupt();
+    }
+    auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
+    if (linalgOp && isSupportedContractionOp(linalgOp)) {
+      candidates.push_back(op);
+      return WalkResult::advance();
+    }
+    return WalkResult::advance();
+  });
+  if (result.wasInterrupted()) {
+    return;
+  }
+  for (Operation *op : candidates) {
+    IREE::Encoding::setDataTilingHint(op);
+  }
+}
+
+} // namespace mlir::iree_compiler::DispatchCreation

compiler/src/iree/compiler/DispatchCreation/BUILD.bazel

@@ -20,6 +20,7 @@ package(
 iree_compiler_cc_library(
     name = "DispatchCreation",
     srcs = [
+        "AnnotateDataTilingHints.cpp",
         "BitcastUnsupportedElementTypes.cpp",
         "BubbleUpExpandShapes.cpp",
         "CloneProducersIntoDispatchRegions.cpp",
@@ -57,6 +58,7 @@ iree_compiler_cc_library(
         ":PassesIncGen",
         "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
         "//compiler/src/iree/compiler/Dialect/Encoding/IR",
+        "//compiler/src/iree/compiler/Dialect/Encoding/Utils",
         "//compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow",
         "//compiler/src/iree/compiler/Dialect/Flow/IR",
         "//compiler/src/iree/compiler/Dialect/Flow/Transforms",

compiler/src/iree/compiler/DispatchCreation/CMakeLists.txt

@@ -17,6 +17,7 @@ iree_cc_library(
     "FusionUtils.h"
     "Passes.h"
   SRCS
+    "AnnotateDataTilingHints.cpp"
     "BitcastUnsupportedElementTypes.cpp"
     "BubbleUpExpandShapes.cpp"
     "CloneProducersIntoDispatchRegions.cpp"
@@ -83,6 +84,7 @@ iree_cc_library(
     MLIRTransforms
     iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
     iree::compiler::Dialect::Encoding::IR
+    iree::compiler::Dialect::Encoding::Utils
     iree::compiler::Dialect::Flow::Conversion::TensorToFlow
     iree::compiler::Dialect::Flow::IR
     iree::compiler::Dialect::Flow::Transforms

compiler/src/iree/compiler/DispatchCreation/Passes.cpp

@@ -271,6 +271,7 @@ addDispatchRegionCreationPasses(OpPassManager &passManager,
           options.cseConstants = false;
           return IREE::Flow::createCanonicalizePass(options);
         })
+        .addPass(createAnnotateDataTilingHintsPass)
         // Set encodings on all eligible ops. All ops should be in compiler
         // formed dispatch regions, so encodings will be placed inside of the
         // dispatch regions with the data-tiled op.

compiler/src/iree/compiler/DispatchCreation/Passes.td

@@ -326,6 +326,16 @@ def PropagateEncodingsPass :
   ];
 }
 
+def AnnotateDataTilingHintsPass :
+    InterfacePass<"iree-dispatch-creation-annotate-data-tiling-hints", "mlir::FunctionOpInterface"> {
+  let summary = "Adds data-tiling hint attribute to linalg operations.";
+  let description = [{
+    The pass does nothing, if any operation already has the data-tiling hint
+    attribute. Otherwise, it aggressively filters linalg ops and adds the
+    data-tiling hint attribute to the operations.
+  }];
+}
+
 def SetEncodingPass : InterfacePass<"iree-dispatch-creation-set-encoding",
                                     "mlir::FunctionOpInterface"> {
   let summary = "Introduces tensor encoding for flow dispatch regions.";

compiler/src/iree/compiler/DispatchCreation/SetEncoding.cpp

@@ -7,6 +7,7 @@
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
 #include "iree/compiler/Dialect/Encoding/IR/EncodingDialect.h"
 #include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
+#include "iree/compiler/Dialect/Encoding/Utils/Utils.h"
 #include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
 #include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
@@ -73,108 +74,11 @@ static Type getContractionInputTypeWithSignedness(OpBuilder &builder,
   return elemType;
 }
 
-/// Returns true iff the linalgOp has a body like a regular matmul, i.e.
-/// yield(add(out, mul(cast(in0), cast(in1))))
-static bool hasMatmulLikeBody(linalg::LinalgOp linalgOp) {
-  auto outBlockArg =
-      linalgOp.getMatchingBlockArgument(linalgOp.getDpsInitOperand(0));
-  auto yieldOp =
-      dyn_cast<linalg::YieldOp>(outBlockArg.getParentBlock()->getTerminator());
-  if (!yieldOp) {
-    return false;
-  }
-  Operation *addOp = yieldOp->getOperand(0).getDefiningOp();
-  if (!addOp || !isa<arith::AddIOp, arith::AddFOp>(addOp)) {
-    return false;
-  }
-  Value addLhs = addOp->getOperand(0);
-  Value addRhs = addOp->getOperand(1);
-  Operation *addLhsOp = addLhs.getDefiningOp();
-  Operation *addRhsOp = addRhs.getDefiningOp();
-  if (!(addLhsOp && addRhs == outBlockArg) &&
-      !(addRhsOp && addLhs == outBlockArg)) {
-    return false;
-  }
-  Operation *mulOp = addLhsOp ? addLhsOp : addRhsOp;
-  if (!isa<arith::MulFOp, arith::MulIOp>(mulOp)) {
-    return false;
-  }
-  Value mulLhs = mulOp->getOperand(0);
-  Value mulRhs = mulOp->getOperand(1);
-  auto mulLhsOp = mulLhs.getDefiningOp<CastOpInterface>();
-  auto mulRhsOp = mulRhs.getDefiningOp<CastOpInterface>();
-  if (!isa<BlockArgument>(mulLhs) && !mulLhsOp && !isa<BlockArgument>(mulRhs) &&
-      !mulRhsOp) {
-    return false;
-  }
-  if ((mulLhsOp && !isa<BlockArgument>(mulLhsOp->getOperand(0))) ||
-      (mulRhsOp && !isa<BlockArgument>(mulRhsOp->getOperand(0)))) {
-    return false;
-  }
-  return true;
-}
-
-/// Not all contractions are supported by data tiling, so return true if:
-///   1) linalgOp has pure tensor semantics.
-///   2) linalgOp does not have a preset compilation info.
-///   3) The workgroup count is not present if linalgOp is wrapped within
-///      Flow::DispatchRegionOp.
-///   4) All the operands do not have encodings.
-///   5) linalgOp has contraction indexingMaps.
-///   6) There are not more than one of each contraction dimension.
-///   7) There is an M or N dimension, and there is a K dimension.
-///   8) linalgOp has the same body as an ordinary int or float matmul.
-///
-/// These restrictions are required because data tiling currently creates
-/// an Mmt4DOp or BatchMmt4DOp on the packed inputs.
-///
-/// TODO(#16176): Loosen restrictions on contraction ops once data tiling
-/// can support more cases.
-static bool isSupportedContractionOp(linalg::LinalgOp linalgOp) {
-  if (!linalgOp.hasPureTensorSemantics()) {
-    return false;
-  }
-  if (getCompilationInfo(linalgOp)) {
-    return false;
-  }
-  auto hasWorkgroupCounts = [](Operation *op) -> bool {
-    auto parentDispatchOp = op->getParentOfType<IREE::Flow::DispatchRegionOp>();
-    return parentDispatchOp && !parentDispatchOp.getWorkgroupCount().empty();
-  };
-  if (hasWorkgroupCounts(linalgOp)) {
-    return false;
-  }
-  auto hasEncoding = [](Value operand) -> bool {
-    auto type = llvm::dyn_cast<RankedTensorType>(operand.getType());
-    return type && type.getEncoding();
-  };
-  if (llvm::any_of(linalgOp.getDpsInputs(), hasEncoding) ||
-      llvm::any_of(linalgOp.getDpsInits(), hasEncoding)) {
-    return false;
-  }
-
-  if (!linalg::isaContractionOpInterface(linalgOp)) {
-    return false;
-  }
-  auto cDims = linalg::inferContractionDims(linalgOp);
-  if (failed(cDims) || cDims->batch.size() > 1 || cDims->m.size() > 1 ||
-      cDims->n.size() > 1 || cDims->k.size() > 1) {
-    return false;
-  }
-  if ((cDims->n.empty() && cDims->m.empty()) || cDims->k.empty()) {
-    return false;
-  }
-  if (!hasMatmulLikeBody(linalgOp)) {
-    return false;
-  }
-  return true;
-}
-
 static SmallVector<linalg::LinalgOp>
 getDataTilingCandidates(FunctionOpInterface funcOp) {
   SmallVector<linalg::LinalgOp> result;
   funcOp.walk([&](linalg::LinalgOp op) {
-    if (!isSupportedContractionOp(op)) {
+    if (!IREE::Encoding::hasDataTilingHint(op)) {
       return;
     }
     result.push_back(op);

compiler/src/iree/compiler/DispatchCreation/test/BUILD.bazel

@@ -16,6 +16,7 @@ iree_lit_test_suite(
     name = "lit",
     srcs = enforce_glob(
         [
+            "annotate_data_tiling_hints.mlir",
             "bitcast_unsupported_element_types.mlir",
             "clone_producers_into_dispatch_regions.mlir",
             "collapse_dimensions.mlir",

compiler/src/iree/compiler/DispatchCreation/test/CMakeLists.txt

@@ -14,6 +14,7 @@ iree_lit_test_suite(
   NAME
     lit
   SRCS
+    "annotate_data_tiling_hints.mlir"
     "bitcast_unsupported_element_types.mlir"
     "bubble_up_expand_shapes.mlir"
     "bubble_up_extract_slice.mlir"

compiler/src/iree/compiler/DispatchCreation/test/annotate_data_tiling_hints.mlir

@@ -0,0 +1,27 @@
+// RUN: iree-opt --pass-pipeline="builtin.module(util.func(iree-dispatch-creation-annotate-data-tiling-hints))" --split-input-file %s | FileCheck %s
+
+util.func public @matmul(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %0 = linalg.matmul
+         ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+         outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  util.return %0 : tensor<?x?xf32>
+}
+// CHECK-LABEL: @matmul(
+// CHECK:         linalg.matmul
+// CHECK-SAME:      iree.opt.data_tiling
+
+// -----
+
+util.func public @matmul_with_preset_hints(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
+  %0 = linalg.matmul {"iree.opt.data_tiling"}
+         ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+         outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  %1 = linalg.matmul
+         ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+         outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  util.return %0, %1 : tensor<?x?xf32>, tensor<?x?xf32>
+}
+// CHECK-LABEL: @matmul_with_preset_hints(
+// CHECK:         linalg.matmul
+// CHECK-SAME:      iree.opt.data_tiling
+// CHECK-NOT:       iree.opt.data_tiling