[Preprocessing] Apply a pass-pipeline based on attributes on func-like ops. (iree-org#20287)

This PR adds a new mechanism for running some preprocessing passes on
func-like ops. This is done by use of a new attribute
`preprocessing-passpipeline` which has the textual representation of the
pass-pipeline to run. It is expected that the textual representation is
for a pass-manager already nested on a func-like op.

---------

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

Author: MaheshRavishankar

compiler/src/iree/compiler/Dialect/Util/IR/UtilAttrs.td

@@ -207,6 +207,33 @@ def Util_NullAttr : AttrDef<Util_Dialect, "Null", [
 // #util.uninitialized
 //===----------------------------------------------------------------------===//
 
+def Util_PreprocessingPipelineAttr : AttrDef<Util_Dialect,
+    "PreprocessingPassPipeline"> {
+  let mnemonic = "preprocessing_pipeline";
+  let summary = "[{preprocessing pass-pipeline to run on the function}]";
+  let description = [{
+    Textual description of the preprocessing pass-pipeline to run on the
+    function-like object that the attribute is specified on. Note that the
+    string pipeline is implicitly nested to `builtin.module(func-lik(..))`
+    operation. For example to run the transpose convolution pipeline, the
+    attribute is specified as
+
+    ```preprocess_pipeline = #util.preprocess_pipeline<
+        "iree-preprocessing-transpose-convolution-pipeline">```
+  }];
+  let parameters = (ins
+    AttrParameter<"std::string", "">:$pipelineString
+  );
+
+  let assemblyFormat = [{
+    `<` $pipelineString `>`
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// #util.uninitialized
+//===----------------------------------------------------------------------===//
+
 def Util_UninitializedAttr : AttrDef<Util_Dialect, "Uninitialized", [
   TypedAttrInterface,
   DeclareAttrInterfaceMethods<Util_SizedStorageAttr, [

compiler/src/iree/compiler/Dialect/Util/IR/UtilTypes.h

@@ -12,6 +12,7 @@
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/MathExtras.h"
 #include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Location.h"

compiler/src/iree/compiler/Preprocessing/Common/AttrBasedPipelinePass.cpp

@@ -0,0 +1,90 @@
+// 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/Preprocessing/Common/Passes.h"
+
+#include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
+#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
+#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
+#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
+#include "llvm/Support/Debug.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
+
+#define DEBUG_TYPE "iree-attr-based-pipeline"
+
+namespace mlir::iree_compiler::Preprocessing {
+
+#define GEN_PASS_DEF_ATTRBASEDPIPELINEPASS
+#include "iree/compiler/Preprocessing/Common/Passes.h.inc" // IWYU pragma: export
+
+namespace {
+
+struct AttrBasedPipelinePass
+    : public iree_compiler::Preprocessing::impl::AttrBasedPipelinePassBase<
+          AttrBasedPipelinePass> {
+  void runOnOperation() override;
+};
+} // namespace
+
+static const char kPreprocessingPipelineAttrName[] = "preprocessing_pipeline";
+
+// Method to get the pass manager nested on a particular operation. There does
+// not seem to be a way to do this without specializing on the op itself.
+// When possible to do so, this method could be deleted.
+static std::optional<OpPassManager>
+getFunctionOpInterfacePassManager(FunctionOpInterface interfaceOp) {
+  return TypeSwitch<Operation *, std::optional<OpPassManager>>(
+             interfaceOp.getOperation())
+      .Case<func::FuncOp, IREE::Util::FuncOp>(
+          [&](auto funcOp) { return OpPassManager(funcOp.getOperationName()); })
+      .Default([&](Operation *op) { return std::nullopt; });
+}
+
+void AttrBasedPipelinePass::runOnOperation() {
+  auto op = getOperation();
+  SmallVector<FunctionOpInterface> funcLikeOps;
+  op->walk([&](FunctionOpInterface funcLikeOp) {
+    funcLikeOps.push_back(funcLikeOp);
+  });
+
+  for (auto funcLikeOp : funcLikeOps) {
+    auto attr = funcLikeOp->getAttr(kPreprocessingPipelineAttrName);
+    if (!attr) {
+      continue;
+    }
+    auto passPipelineAttr =
+        dyn_cast<IREE::Util::PreprocessingPassPipelineAttr>(attr);
+    if (!passPipelineAttr) {
+      funcLikeOp.emitRemark(
+          "expected preprocessing_pipeline attribute to be a `StringAttr` that "
+          "specifies the pass pipeline to apply");
+      continue;
+    }
+    LLVM_DEBUG({ llvm::dbgs() << "Parsed Attribute : " << passPipelineAttr; });
+
+    std::optional<OpPassManager> passManager =
+        getFunctionOpInterfacePassManager(funcLikeOp);
+    if (!passManager) {
+      continue;
+    }
+
+    std::string pipelineStr = passPipelineAttr.getPipelineString();
+    if (failed(parsePassPipeline(pipelineStr, *passManager))) {
+      funcLikeOp->emitOpError("failed to populate pass manager specified by : ")
+          << pipelineStr;
+      return signalPassFailure();
+    }
+
+    if (failed(runPipeline(*passManager, funcLikeOp))) {
+      funcLikeOp->emitOpError("failed to run pass specified as attribute : ")
+          << pipelineStr;
+      return signalPassFailure();
+    }
+  }
+}
+
+} // namespace mlir::iree_compiler::Preprocessing

compiler/src/iree/compiler/Preprocessing/Common/BUILD.bazel

@@ -31,6 +31,7 @@ iree_compiler_cc_library(
     name = "Transforms",
     srcs = [
         "ApplyPDLPatterns.cpp",
+        "AttrBasedPipelinePass.cpp",
         "ConvertConv2DToImg2Col.cpp",
         "ConvertConvFilterToChannelsLast.cpp",
         "ConvertConvToChannelsLast.cpp",

compiler/src/iree/compiler/Preprocessing/Common/CMakeLists.txt

@@ -27,6 +27,7 @@ iree_cc_library(
     "Passes.h.inc"
   SRCS
     "ApplyPDLPatterns.cpp"
+    "AttrBasedPipelinePass.cpp"
     "ConvertConv2DToImg2Col.cpp"
     "ConvertConvFilterToChannelsLast.cpp"
     "ConvertConvToChannelsLast.cpp"

compiler/src/iree/compiler/Preprocessing/Common/Passes.td

@@ -27,6 +27,19 @@ def ApplyPDLPatternsPass : Pass<"iree-preprocessing-apply-pdl-patterns", "Module
   ];
 }
 
+def AttrBasedPipelinePass : Pass<"iree-preprocessing-attr-based-pipeline", ""> {
+  let summary = "Run a pass pipeline specified as an attribute on any function-like operation";
+  let description = [{
+    The textual representation of the pass pipeline specified in the attribute
+    `preprocessing_pipeline=#util.preprocessing_pipeline<...>` is to be run on the
+    function-like operation.
+  }];
+  let dependentDialects = [
+    "iree_compiler::IREE::Flow::FlowDialect",
+    "iree_compiler::IREE::Util::UtilDialect",
+  ];
+}
+
 def ConvertConv2DToImg2ColPass :
     Pass<"iree-preprocessing-convert-conv2d-to-img2col", ""> {
   let summary = "Convert linalg convolution ops to matmul img2col based implementation";

compiler/src/iree/compiler/Preprocessing/Common/test/BUILD.bazel

@@ -16,6 +16,7 @@ iree_lit_test_suite(
     name = "lit",
     srcs = enforce_glob(
         [
+            "attr_based_pipeline.mlir",
             "conv2d_to_img2col.mlir",
             "conv_filter_to_channels_last.mlir",
             "conv_to_channels_last.mlir",

compiler/src/iree/compiler/Preprocessing/Common/test/CMakeLists.txt

@@ -14,6 +14,7 @@ iree_lit_test_suite(
   NAME
     lit
   SRCS
+    "attr_based_pipeline.mlir"
     "conv2d_to_img2col.mlir"
     "conv_filter_to_channels_last.mlir"
     "conv_to_channels_last.mlir"

compiler/src/iree/compiler/Preprocessing/Common/test/attr_based_pipeline.mlir

@@ -0,0 +1,70 @@
+// RUN: iree-opt --iree-preprocessing-attr-based-pipeline --mlir-print-local-scope --split-input-file --verify-diagnostics %s | FileCheck %s
+
+func.func @test(%lhs : tensor<10x20xf16>, %rhs : tensor<20x40xf16>,
+    %outs : tensor<10x40xf16>) -> tensor<10x40xf16> attributes {
+    preprocessing_pipeline = #util.preprocessing_pipeline<"iree-preprocessing-make-single-dispatch">} {
+  %matmul = linalg.matmul ins(%lhs, %rhs : tensor<10x20xf16>, tensor<20x40xf16>)
+      outs(%outs : tensor<10x40xf16>) -> tensor<10x40xf16>
+  return %matmul : tensor<10x40xf16>
+}
+// CHECK-LABEL: test
+//       CHECK:   %[[DISPATCH:.+]] = flow.dispatch.region
+//       CHECK:     %[[MATMUL:.+]] = linalg.matmul
+//       CHECK:     flow.return %[[MATMUL]]
+//       CHECK:   return %[[DISPATCH]]
+
+// -----
+
+module {
+func.func @function1(%lhs : tensor<10x20xf16>, %rhs : tensor<20x40xf16>,
+    %outs : tensor<10x40xf16>) -> tensor<10x40xf16> attributes {
+    preprocessing_pipeline = #util.preprocessing_pipeline<"iree-preprocessing-generalize-linalg-matmul-experimental">} {
+  %matmul = linalg.matmul ins(%lhs, %rhs : tensor<10x20xf16>, tensor<20x40xf16>)
+      outs(%outs : tensor<10x40xf16>) -> tensor<10x40xf16>
+  return %matmul : tensor<10x40xf16>
+}
+func.func @function2(%lhs : tensor<10x20xf16>, %rhs : tensor<20x40xf16>,
+    %outs : tensor<10x40xf16>) -> tensor<10x40xf16> attributes {
+    preprocessing_pipeline = #util.preprocessing_pipeline<"iree-preprocessing-pad-linalg-ops">} {
+  %matmul = linalg.matmul ins(%lhs, %rhs : tensor<10x20xf16>, tensor<20x40xf16>)
+      outs(%outs : tensor<10x40xf16>) -> tensor<10x40xf16>
+  return %matmul : tensor<10x40xf16>
+}
+}
+// CHECK-LABEL: func @function1
+//       CHECK:   %[[GENERIC:.+]] = linalg.generic
+//       CHECK:   return %[[GENERIC]]
+// CHECK-LABEL: func @function2
+//   CHECK-DAG:   %[[PAD1:.+]] = tensor.pad
+//   CHECK-DAG:   %[[PAD2:.+]] = tensor.pad
+//       CHECK:   %[[MATMUL:.+]] = linalg.matmul
+//  CHECK-SAME:       ins(%[[PAD1]],
+//  CHECK-SAME:       outs(%[[PAD2]]
+//       CHECK:   %[[SLICE:.+]] = tensor.extract_slice %[[MATMUL]]
+//       CHECK:   return %[[SLICE]]
+
+// -----
+
+func.func @function(%lhs : tensor<10x20xf16>, %rhs : tensor<20x40xf16>,
+    %outs : tensor<10x40xf16>) -> tensor<10x40xf16> attributes {
+    preprocessing_pipeline = #util.preprocessing_pipeline<"iree-preprocessing-pad-linalg-ops,iree-preprocessing-generalize-linalg-matmul-experimental">} {
+  %matmul = linalg.matmul ins(%lhs, %rhs : tensor<10x20xf16>, tensor<20x40xf16>)
+      outs(%outs : tensor<10x40xf16>) -> tensor<10x40xf16>
+  return %matmul : tensor<10x40xf16>
+}
+// CHECK-LABEL: func @function
+//   CHECK-DAG:   %[[PAD1:.+]] = tensor.pad
+//   CHECK-DAG:   %[[PAD2:.+]] = tensor.pad
+//       CHECK:   %[[MATMUL:.+]] = linalg.generic
+//  CHECK-SAME:       ins(%[[PAD1]],
+//  CHECK-SAME:       outs(%[[PAD2]]
+//       CHECK:   %[[SLICE:.+]] = tensor.extract_slice %[[MATMUL]]
+//       CHECK:   return %[[SLICE]]
+
+// -----
+
+// expected-remark@+1 {{expected preprocessing_pipeline attribute to be a `StringAttr` that specifies the pass pipeline to apply}}
+func.func @function() attributes {
+    preprocessing_pipeline = "iree-preprocessing-pad-linalg-ops"} {
+  return
+}

compiler/src/iree/compiler/Preprocessing/Passes.cpp

@@ -108,6 +108,10 @@ void buildPreprocessingPassPipeline(
   if (pipelineExtensions) {
     pipelineExtensions->extendPreprocessingPassPipeline(passManager);
   }
+
+  // 3. Run any pass pipelines specified through the use of
+  //    `preprocessing_pipeline` attribute.
+  FunctionLikeNest(passManager).addPass(createAttrBasedPipelinePass);
 }
 
 static void
@@ -124,11 +128,20 @@ buildTransposeConvolutionPassPipeline(OpPassManager &passManager,
   passManager.addPass(createCSEPass());
 }
 
+/// Pass pipeline to make the computation within a function a single dispatch.
+/// Note that this expects a `OpPassManager` nested on `FunctionOpInterface`
+/// ops.
+static void
+buildMakeSingleDispatchPassPipeline(OpPassManager &passManager,
+                                    const TransformOptions &options) {
+  passManager.addPass(createMakeSingleDispatchForFunctionPass());
+}
+
 void registerPreprocessingPasses() {
   registerCommonPreprocessingPasses();
 
   PassPipelineRegistration<TransformOptions>
-      globalOptimizationTransformPassPipeline(
+      preprocessingTransposeConvolutionPassPipeline(
           "iree-preprocessing-transpose-convolution-pipeline",
           "Runs a pass pipeline for transposing and canonicalizing "
           "convolutions",
@@ -137,6 +150,15 @@ void registerPreprocessingPasses() {
             buildTransposeConvolutionPassPipeline(passManager,
                                                   transformOptions);
           });
+
+  PassPipelineRegistration<TransformOptions>
+      preprocessingMakeSingleDispatchPassPipeline(
+          "iree-preprocessing-make-single-dispatch",
+          "Runs passes to get a single dispatch for a function",
+          [](OpPassManager &passManager,
+             const TransformOptions &transformOptions) {
+            buildMakeSingleDispatchPassPipeline(passManager, transformOptions);
+          });
 }
 
 } // namespace mlir::iree_compiler::Preprocessing