[Codegen] Add PipelineAttrInterface and PassPipelineAttr (#23590)

This change adds an attribute interface for representing pass pipelines
and a single basic attribute that uses the string based pass interpreter
to populate a pipeline. The intent of this change is NOT to induce a
refactor of all the pass pipelines, instead it's primarily to make
testing structural pipeline changes with partially lowered inputs much
easier. Today if you want to work on a change that affects later stages
of a pass pipeline but will also require changes to earlier steps, it's
hard to stage those changes since there isn't a convenient way to jump
into the middle of a codegen pass pipeline (unlike the rest of the
compiler which offers distinct stages).

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: qedawkins

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/BUILD.bazel

@@ -115,6 +115,7 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:LinalgDialect",
         "@llvm-project//mlir:MemRefDialect",
         "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:SCFTransforms",
         "@llvm-project//mlir:Support",

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/CMakeLists.txt

@@ -66,6 +66,7 @@ iree_cc_library(
     MLIRLinalgDialect
     MLIRMemRefDialect
     MLIRParser
+    MLIRPass
     MLIRSCFDialect
     MLIRSCFTransforms
     MLIRSupport

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp

@@ -19,6 +19,50 @@
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/StorageUniquerSupport.h"
+#include "mlir/Pass/PassRegistry.h"
+
+// Custom parse/print directives for TranslationInfoAttr's pipeline field.
+// These must be defined before the generated .cpp.inc is included because
+// the ODS-generated parse/print methods call them.
+namespace mlir::iree_compiler::IREE::Codegen {
+
+/// Parses either a DispatchLoweringPassPipeline enum keyword (e.g.,
+/// `CPUDefault`) or a generic attribute implementing PipelineAttrInterface
+/// (e.g., `#iree_codegen.pass_pipeline<"canonicalize">`).
+static ParseResult parsePipelineAttr(AsmParser &parser, Attribute &result) {
+  StringRef keyword;
+  SMLoc loc = parser.getCurrentLocation();
+  if (succeeded(parser.parseOptionalKeyword(&keyword))) {
+    std::optional<DispatchLoweringPassPipeline> pipeline =
+        symbolizeDispatchLoweringPassPipeline(keyword);
+    if (!pipeline) {
+      parser.emitError(loc, "unknown pipeline keyword: ") << keyword;
+      return failure();
+    }
+    result =
+        DispatchLoweringPassPipelineAttr::get(parser.getContext(), *pipeline);
+    return success();
+  }
+  Attribute attr;
+  if (parser.parseAttribute(attr)) {
+    return failure();
+  }
+  result = attr;
+  return success();
+}
+
+/// Prints DispatchLoweringPassPipelineAttr as a bare keyword and other
+/// attributes (e.g., PipelineAttrInterface impls) via the generic printer.
+static void printPipelineAttr(AsmPrinter &printer, Attribute pipelineAttr) {
+  if (auto enumAttr =
+          dyn_cast<DispatchLoweringPassPipelineAttr>(pipelineAttr)) {
+    printer << stringifyEnum(enumAttr.getValue());
+    return;
+  }
+  printer.printAttribute(pipelineAttr);
+}
+
+} // namespace mlir::iree_compiler::IREE::Codegen
 
 #define GET_ATTRDEF_CLASSES
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp.inc"
@@ -64,6 +108,28 @@ ArrayAttr ExportConfigAttr::getWorkgroupSizeIndexArray() {
   return getIndexArrayAttr(getContext(), getWorkgroupSize());
 }
 
+//===----------------------------------------------------------------------===//
+// iree_codegen.pass_pipeline
+//===----------------------------------------------------------------------===//
+
+LogicalResult PassPipelineAttr::buildPipeline(OpPassManager &pm) const {
+  if (failed(parsePassPipeline(getPipeline(), pm))) {
+    return failure();
+  }
+  return success();
+}
+
+LogicalResult
+PassPipelineAttr::verify(function_ref<InFlightDiagnostic()> emitError,
+                         StringRef pipeline) {
+  OpPassManager pm("builtin.module");
+  if (failed(parsePassPipeline(pipeline, pm))) {
+    return emitError() << "invalid pass pipeline specification: '" << pipeline
+                       << "'";
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // iree_codegen.translation_info
 //===----------------------------------------------------------------------===//
@@ -72,7 +138,7 @@ TranslationInfoAttr TranslationInfoAttr::get(
     MLIRContext *context, DispatchLoweringPassPipeline passPipeline,
     SymbolRefAttr codegenSpec, ArrayRef<int64_t> workgroupSize,
     std::optional<int64_t> subgroupSize, DictionaryAttr configuration) {
-  auto pipelineAttr =
+  Attribute pipelineAttr =
       DispatchLoweringPassPipelineAttr::get(context, passPipeline);
   return get(context, pipelineAttr, codegenSpec, workgroupSize,
              subgroupSize.value_or(int64_t()), configuration);
@@ -82,36 +148,46 @@ TranslationInfoAttr TranslationInfoAttr::get(
     MLIRContext *context, DispatchLoweringPassPipeline passPipeline,
     ArrayRef<int64_t> workgroupSize, std::optional<int64_t> subgroupSize,
     DictionaryAttr configuration) {
-  auto pipelineAttr =
+  Attribute pipelineAttr =
       DispatchLoweringPassPipelineAttr::get(context, passPipeline);
   return get(context, pipelineAttr, /*codegenSpec=*/SymbolRefAttr(),
              workgroupSize, subgroupSize.value_or(int64_t()), configuration);
 }
 
 DispatchLoweringPassPipeline
 TranslationInfoAttr::getDispatchLoweringPassPipeline() {
-  return getPassPipeline().getValue();
+  if (auto enumAttr =
+          dyn_cast<DispatchLoweringPassPipelineAttr>(getPassPipeline())) {
+    return enumAttr.getValue();
+  }
+  return DispatchLoweringPassPipeline::None;
 }
 
 LogicalResult TranslationInfoAttr::verify(
-    function_ref<InFlightDiagnostic()> emitError,
-    IREE::Codegen::DispatchLoweringPassPipelineAttr passPipeline,
+    function_ref<InFlightDiagnostic()> emitError, Attribute passPipeline,
     SymbolRefAttr codegenSpec, ArrayRef<int64_t> workgroupSize,
     int64_t subgroupSize, DictionaryAttr configuration) {
   if (!passPipeline) {
     return emitError() << "missing pass pipeline specification";
   }
-  auto passPipelineValue = passPipeline.getValue();
-  if (passPipelineValue > IREE::Codegen::DispatchLoweringPassPipeline::None) {
-    return emitError() << "invalid pass pipeline value : "
-                       << stringifyEnum(passPipeline.getValue());
-  }
-  auto tdPassPipeline =
-      IREE::Codegen::DispatchLoweringPassPipeline::TransformDialectCodegen;
-  if (codegenSpec && passPipelineValue != tdPassPipeline) {
+  if (auto enumAttr =
+          dyn_cast<DispatchLoweringPassPipelineAttr>(passPipeline)) {
+    DispatchLoweringPassPipeline passPipelineValue = enumAttr.getValue();
+    if (passPipelineValue > IREE::Codegen::DispatchLoweringPassPipeline::None) {
+      return emitError() << "invalid pass pipeline value : "
+                         << stringifyEnum(passPipelineValue);
+    }
+    DispatchLoweringPassPipeline tdPassPipeline =
+        IREE::Codegen::DispatchLoweringPassPipeline::TransformDialectCodegen;
+    if (codegenSpec && passPipelineValue != tdPassPipeline) {
+      return emitError()
+             << "transform dialect codegen spec requires pass pipeline : "
+             << stringifyEnum(tdPassPipeline);
+    }
+  } else if (!isa<PipelineAttrInterface>(passPipeline)) {
     return emitError()
-           << "transform dialect codegen spec requires pass pipeline : "
-           << stringifyEnum(tdPassPipeline);
+           << "pass pipeline must be a DispatchLoweringPassPipelineAttr or "
+              "implement PipelineAttrInterface";
   }
   if (workgroupSize.size() > 3) {
     return emitError() << "workgroup size cannot have more than 3 entries";

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.td

@@ -250,6 +250,32 @@ def IREECodegen_SimpleTargetAttr :
 }
 
 
+//===---------------------------------------------------------------------===//
+// iree_codegen.pass_pipeline
+//===---------------------------------------------------------------------===//
+
+def IREECodegen_PassPipelineAttr :
+    AttrDef<IREECodegen_Dialect, "PassPipeline", [
+      DeclareAttrInterfaceMethods<IREECodegen_PipelineAttrInterface, [
+        "buildPipeline"
+      ]>
+    ]> {
+  let mnemonic = "pass_pipeline";
+  let summary = "An attribute carrying a textual pass pipeline string.";
+  let description = [{
+    Specifies a pass pipeline using MLIR's textual pass pipeline syntax.
+    The pipeline string is parsed and populated into an OpPassManager
+    when `buildPipeline` is called.
+  }];
+  let parameters = (ins
+    StringRefParameter<"The textual pass pipeline specification">:$pipeline
+  );
+  let assemblyFormat = [{
+    `<` $pipeline `>`
+  }];
+  let genVerifyDecl = 1;
+}
+
 //===---------------------------------------------------------------------===//
 // iree_codegen.translation_info
 //===---------------------------------------------------------------------===//
@@ -269,22 +295,15 @@ def IREECodegen_TranslationInfoAttr :
     dispatch region (like `linalg.matmul`/`linalg.*conv*`), this
     attribute gets propagated to the entry point function.
 
-    The fields are
-    - `passPipeline` : The pass pipeline to use.
-
-  }];
-
-  let assemblyFormat = [{
-    `<` `pipeline` `=` `` $passPipeline
-    (`codegen_spec` `=` $codegenSpec^)?
-    (`workgroup_size` `=` `[` $workgroupSize^ `]`)?
-    (`subgroup_size` `=` $subgroupSize^)?
-    (`,` $configuration^)? `>`
+    The `passPipeline` field can be either:
+    - A `DispatchLoweringPassPipelineAttr` (enum keyword like `CPUDefault`).
+    - Any attribute implementing `PipelineAttrInterface` (e.g.,
+      `#iree_codegen.pass_pipeline<"...">`).
   }];
 
   let parameters = (ins
-    AttrParameter<"IREE::Codegen::DispatchLoweringPassPipelineAttr",
-        "Name of the pipeline to be invoked on the translation unit.">:$passPipeline,
+    AttrParameter<"Attribute",
+        "Pass pipeline specification.">:$passPipeline,
     OptionalParameter<"SymbolRefAttr",
         "The symbol pointing to the transform dialect codegen spec to be used">:$codegenSpec,
     OptionalArrayRefParameter<"int64_t", "The workgroup size to use">:$workgroupSize,
@@ -304,9 +323,20 @@ def IREECodegen_TranslationInfoAttr :
         CArg<"DictionaryAttr", "{}">:$configuration)>
   ];
   let extraClassDeclaration = [{
-    // Returns the lowering pass pipeline set.
+    // Returns the lowering pass pipeline enum value. Returns None if the
+    // pipeline is not a DispatchLoweringPassPipelineAttr.
     DispatchLoweringPassPipeline getDispatchLoweringPassPipeline();
   }];
+
+  let assemblyFormat = [{
+    `<` `pipeline` `=` custom<PipelineAttr>($passPipeline)
+      (`codegen_spec` `=` $codegenSpec^)?
+      (`workgroup_size` `=` `[` $workgroupSize^ `]`)?
+      (`subgroup_size` `=` $subgroupSize^)?
+      (`,` $configuration^)?
+    `>`
+  }];
+
   let genVerifyDecl = 1;
 }
 

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h

@@ -14,6 +14,7 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Pass/PassManager.h"
 
 #include "llvm/ADT/STLExtras.h"
 

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.td

@@ -752,6 +752,26 @@ def IREECodegen_UKernelProviderInterface :
   ];
 }
 
+def IREECodegen_PipelineAttrInterface :
+    AttrInterface<"PipelineAttrInterface"> {
+  let cppNamespace = "::mlir::iree_compiler::IREE::Codegen";
+  let description = [{
+    Attribute interface for building a pass pipeline. Implementations populate
+    the provided OpPassManager with the desired pass pipeline.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{
+        Populates the given pass manager with a pass pipeline.
+      }],
+      /*retTy=*/"::mlir::LogicalResult",
+      /*methodName=*/"buildPipeline",
+      /*args=*/(ins "::mlir::OpPassManager &":$pm)
+    >
+  ];
+}
+
 def IREECodegen_TargetInfoAttrInterface :
     AttrInterface<"TargetInfoAttrInterface"> {
   let cppNamespace = "::mlir::iree_compiler::IREE::Codegen";

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/test/lowering_config_attr.mlir

@@ -125,3 +125,37 @@ module {
     }
   }
 }
+
+// -----
+
+module {
+  /// Pass pipeline attribute round-trips correctly.
+  func.func @test_pass_pipeline() attributes {
+      translation_info = #iree_codegen.translation_info<pipeline = #iree_codegen.pass_pipeline<"canonicalize">>} {
+    return
+  }
+}
+// CHECK: #translation = #iree_codegen.translation_info<pipeline = #iree_codegen.pass_pipeline<"canonicalize">>
+
+// -----
+
+module {
+  /// Pass pipeline attribute with workgroup size and subgroup size round-trips.
+  func.func @test_pass_pipeline_with_config() attributes {
+      translation_info = #iree_codegen.translation_info<pipeline = #iree_codegen.pass_pipeline<"canonicalize"> workgroup_size = [64, 1, 1] subgroup_size = 32>} {
+    return
+  }
+}
+// CHECK: #translation = #iree_codegen.translation_info<pipeline = #iree_codegen.pass_pipeline<"canonicalize"> workgroup_size = [64, 1, 1] subgroup_size = 32>
+
+// -----
+
+module {
+  /// Invalid pass pipeline string should be caught at verify time.
+  func.func @invalid_pass_pipeline() attributes {
+    // expected-error @+1 {{invalid pass pipeline specification: 'not_a_real_pass'}}
+    translation_info = #iree_codegen.translation_info<pipeline = #iree_codegen.pass_pipeline<"not_a_real_pass">>
+  } {
+    return
+  }
+}

compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp

@@ -3976,9 +3976,10 @@ adjustTileSizesForRootUnPackOp(mlir::FunctionOpInterface entryPointFn,
     }
   }
 
-  auto tInfo = getTranslationInfo(entryPointFn);
-  auto pipeline = tInfo.getPassPipeline().getValue();
-  auto pipelineConfig = tInfo.getConfiguration();
+  IREE::Codegen::TranslationInfoAttr tInfo = getTranslationInfo(entryPointFn);
+  DispatchLoweringPassPipeline pipeline =
+      tInfo.getDispatchLoweringPassPipeline();
+  DictionaryAttr pipelineConfig = tInfo.getConfiguration();
   if (isOptEnabled(entryPointFn, getEnableLoopPeelingStr())) {
     // See #16406
     LDBG() << "unpack fusion does not work with peeling, falling back to "
@@ -4167,7 +4168,8 @@ setTranslationInfoAndRootConfig(mlir::FunctionOpInterface entryPointFn,
 
   // The transform dialect codegen has different logics and codegen flow.
   // Ignore the tile sizes adjustment.
-  auto pipeline = getTranslationInfo(entryPointFn).getPassPipeline().getValue();
+  DispatchLoweringPassPipeline pipeline =
+      getTranslationInfo(entryPointFn).getDispatchLoweringPassPipeline();
   if (pipeline != DispatchLoweringPassPipeline::TransformDialectCodegen) {
     if (failed(adjustTileSizesForRootUnPackOp(entryPointFn, rootOperation))) {
       return failure();