[mlir] use transform-interpreter in test passes (#70040)

Update most test passes to use the transform-interpreter pass instead of
the test-transform-dialect-interpreter-pass. The new "main" interpreter
pass has a named entry point instead of looking up the top-level op with
`PossibleTopLevelOpTrait`, which is arguably a more understandable
interface. The change is mechanical, rewriting an unnamed sequence into
a named one and wrapping the transform IR in to a module when necessary.

Add an option to the transform-interpreter pass to target a tagged
payload op instead of the root anchor op, which is also useful for repro
generation.

Only the test in the transform dialect proper and the examples have not
been updated yet. These will be updated separately after a more careful
consideration of testing coverage of the transform interpreter logic.

Author: ftynse

mlir/include/mlir/Dialect/Transform/Transforms/Passes.td

@@ -70,6 +70,14 @@ def InterpreterPass : Pass<"transform-interpreter"> {
   }];
   let dependentDialects = ["::mlir::transform::TransformDialect"];
   let options = [
+    Option<"debugPayloadRootTag", "debug-payload-root-tag", "std::string",
+           /*default=*/[{""}],
+           "Select the operation with 'transform.target_tag' attribute having "
+           "the given value as payload IR root. If empty select the pass "
+           "anchor operation as the payload IR root.">,
+    Option<"disableExpensiveChecks", "disable-expensive-checks", "bool",
+           "false",
+           "Disable expensive checks in the interpreter for a faster run.">,
     Option<"entryPoint", "entry-point", "std::string",
            /*default=*/[{"__transform_main"}],
            "Entry point of the pass pipeline.">,

mlir/include/mlir/Dialect/Transform/Transforms/TransformInterpreterUtils.h

@@ -85,21 +85,18 @@ LogicalResult mergeSymbolsInto(Operation *target,
                                OwningOpRef<Operation *> other);
 } // namespace detail
 
-/// Standalone util to apply the named sequence `entryPoint` to the payload.
-/// This is done in 3 steps:
-///   1. lookup the `entryPoint` symbol in `{payload, sharedTransformModule}` by
-///   calling detail::findTransformEntryPoint.
-///   2. if the entry point is found and not nested under
-///   `sharedTransformModule`, call `detail::defineDeclaredSymbols` to "link" in
-///   the `sharedTransformModule`. Note: this may modify the transform IR
-///   embedded with the payload IR.
-///   3. apply the transform IR to the payload IR, relaxing the requirement that
-///   the transform IR is a top-level transform op. We are applying a named
-///   sequence anyway.
-LogicalResult applyTransformNamedSequence(
-    Operation *payload, ModuleOp transformModule,
-    const TransformOptions &options,
-    StringRef entryPoint = TransformDialect::kTransformEntryPointSymbolName);
+/// Standalone util to apply the named sequence `transformRoot` to `payload` IR.
+/// This is done in 2 steps:
+///   1. If `transformModule` is provided and is not nested under
+///      `transformRoot`, it will be "linked into" the IR containing
+///      `transformRoot` to resolve undefined named sequences.
+///   2. The transforms specified in `transformRoot` are applied to `payload`,
+///      assuming the named sequence has a single argument handle that will be
+///      associated with `payload` on run.
+LogicalResult applyTransformNamedSequence(Operation *payload,
+                                          Operation *transformRoot,
+                                          ModuleOp transformModule,
+                                          const TransformOptions &options);
 
 } // namespace transform
 } // namespace mlir

mlir/lib/Dialect/Transform/Transforms/InterpreterPass.cpp

@@ -19,6 +19,39 @@ namespace transform {
 } // namespace transform
 } // namespace mlir
 
+/// Returns the payload operation to be used as payload root:
+///   - the operation nested under `passRoot` that has the given tag attribute,
+///     must be unique;
+///   - the `passRoot` itself if the tag is empty.
+static Operation *findPayloadRoot(Operation *passRoot, StringRef tag) {
+  // Fast return.
+  if (tag.empty())
+    return passRoot;
+
+  // Walk to do a lookup.
+  Operation *target = nullptr;
+  auto tagAttrName = StringAttr::get(
+      passRoot->getContext(), transform::TransformDialect::kTargetTagAttrName);
+  WalkResult walkResult = passRoot->walk([&](Operation *op) {
+    auto attr = op->getAttrOfType<StringAttr>(tagAttrName);
+    if (!attr || attr.getValue() != tag)
+      return WalkResult::advance();
+
+    if (!target) {
+      target = op;
+      return WalkResult::advance();
+    }
+
+    InFlightDiagnostic diag = op->emitError()
+                              << "repeated operation with the target tag '"
+                              << tag << "'";
+    diag.attachNote(target->getLoc()) << "previously seen operation";
+    return WalkResult::interrupt();
+  });
+
+  return walkResult.wasInterrupted() ? nullptr : target;
+}
+
 namespace {
 class InterpreterPass
     : public transform::impl::InterpreterPassBase<InterpreterPass> {
@@ -29,10 +62,22 @@ class InterpreterPass
     MLIRContext *context = &getContext();
     ModuleOp transformModule =
         transform::detail::getPreloadedTransformModule(context);
+    Operation *payloadRoot =
+        findPayloadRoot(getOperation(), debugPayloadRootTag);
+    Operation *transformEntryPoint = transform::detail::findTransformEntryPoint(
+        getOperation(), transformModule, entryPoint);
+    if (!transformEntryPoint) {
+      getOperation()->emitError()
+          << "could not find transform entry point: " << entryPoint
+          << " in either payload or transform module";
+      return signalPassFailure();
+    }
+
     if (failed(transform::applyTransformNamedSequence(
-            getOperation(), transformModule,
-            options.enableExpensiveChecks(true), entryPoint)))
+            payloadRoot, transformEntryPoint, transformModule,
+            options.enableExpensiveChecks(!disableExpensiveChecks)))) {
       return signalPassFailure();
+    }
   }
 
 private:

mlir/lib/Dialect/Transform/Transforms/TransformInterpreterUtils.cpp

@@ -409,16 +409,8 @@ transform::detail::mergeSymbolsInto(Operation *target,
 }
 
 LogicalResult transform::applyTransformNamedSequence(
-    Operation *payload, ModuleOp transformModule,
-    const TransformOptions &options, StringRef entryPoint) {
-  Operation *transformRoot =
-      detail::findTransformEntryPoint(payload, transformModule, entryPoint);
-  if (!transformRoot) {
-    return payload->emitError()
-           << "could not find transform entry point: " << entryPoint
-           << " in either payload or transform module";
-  }
-
+    Operation *payload, Operation *transformRoot, ModuleOp transformModule,
+    const TransformOptions &options) {
   // `transformModule` may not be modified.
   if (transformModule && !transformModule->isAncestor(transformRoot)) {
     OwningOpRef<Operation *> clonedTransformModule(transformModule->clone());

mlir/test/Conversion/FuncToLLVM/func-memref-return.mlir

@@ -2,7 +2,7 @@
 
 // RUN: mlir-opt -convert-func-to-llvm='use-bare-ptr-memref-call-conv=1 use-opaque-pointers=1'  %s | FileCheck %s --check-prefix=BAREPTR
 
-// RUN: mlir-opt -test-transform-dialect-interpreter %s | FileCheck %s --check-prefix=BAREPTR
+// RUN: mlir-opt -transform-interpreter %s | FileCheck %s --check-prefix=BAREPTR
 
 // These tests were separated from func-memref.mlir because applying
 // -reconcile-unrealized-casts resulted in `llvm.extractvalue` ops getting
@@ -110,17 +110,20 @@ func.func @unranked_memref(%arg0:memref<*xi32>) {
 }
 func.func private @printMemrefI32(memref<*xi32>)
 
-transform.sequence failures(propagate) {
-^bb1(%toplevel_module: !transform.any_op):
-  %func = transform.structured.match ops{["func.func"]} in %toplevel_module
-    : (!transform.any_op) -> !transform.any_op
-  transform.apply_conversion_patterns to %func {
-    transform.apply_conversion_patterns.func.func_to_llvm
-  } with type_converter {
-    transform.apply_conversion_patterns.memref.memref_to_llvm_type_converter
-      {use_bare_ptr_call_conv = true, use_opaque_pointers = true}
-  } {
-    legal_dialects = ["llvm"],
-    partial_conversion
-  } : !transform.any_op
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(
+      %toplevel_module: !transform.any_op {transform.readonly}) {
+    %func = transform.structured.match ops{["func.func"]} in %toplevel_module
+      : (!transform.any_op) -> !transform.any_op
+    transform.apply_conversion_patterns to %func {
+      transform.apply_conversion_patterns.func.func_to_llvm
+    } with type_converter {
+      transform.apply_conversion_patterns.memref.memref_to_llvm_type_converter
+        {use_bare_ptr_call_conv = true, use_opaque_pointers = true}
+    } {
+      legal_dialects = ["llvm"],
+      partial_conversion
+    } : !transform.any_op
+    transform.yield
+  }
 }

mlir/test/Conversion/FuncToLLVM/func-to-llvm.mlir

@@ -2,7 +2,7 @@
 
 // RUN: mlir-opt -pass-pipeline="builtin.module(func.func(convert-math-to-llvm,convert-arith-to-llvm{index-bitwidth=32}),convert-func-to-llvm{index-bitwidth=32 use-opaque-pointers=1},reconcile-unrealized-casts)" %s | FileCheck --check-prefix=CHECK32 %s
 
-// RUN: mlir-opt -test-transform-dialect-interpreter %s | FileCheck --check-prefix=CHECK32 %s
+// RUN: mlir-opt -transform-interpreter %s | FileCheck --check-prefix=CHECK32 %s
 
 // Same below, but using the `ConvertToLLVMPatternInterface` entry point
 // and the generic `convert-to-llvm` pass.
@@ -537,20 +537,22 @@ func.func @switchi8(%arg0 : i8) -> i32 {
 // CHECK-NEXT:     llvm.return %[[E1]] : i32
 // CHECK-NEXT:   }
 
-transform.sequence failures(propagate) {
-^bb1(%toplevel_module: !transform.any_op):
-  %func = transform.structured.match ops{["func.func"]} in %toplevel_module
-    : (!transform.any_op) -> !transform.any_op
-  transform.apply_conversion_patterns to %func {
-    transform.apply_conversion_patterns.dialect_to_llvm "math"
-    transform.apply_conversion_patterns.dialect_to_llvm "arith"
-    transform.apply_conversion_patterns.dialect_to_llvm "cf"
-    transform.apply_conversion_patterns.func.func_to_llvm
-  } with type_converter {
-    transform.apply_conversion_patterns.memref.memref_to_llvm_type_converter
-      {index_bitwidth = 32, use_opaque_pointers = true}
-  } {
-    legal_dialects = ["llvm"], 
-    partial_conversion
-  } : !transform.any_op
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%toplevel_module: !transform.any_op {transform.readonly}) {
+    %func = transform.structured.match ops{["func.func"]} in %toplevel_module
+      : (!transform.any_op) -> !transform.any_op
+    transform.apply_conversion_patterns to %func {
+      transform.apply_conversion_patterns.dialect_to_llvm "math"
+      transform.apply_conversion_patterns.dialect_to_llvm "arith"
+      transform.apply_conversion_patterns.dialect_to_llvm "cf"
+      transform.apply_conversion_patterns.func.func_to_llvm
+    } with type_converter {
+      transform.apply_conversion_patterns.memref.memref_to_llvm_type_converter
+        {index_bitwidth = 32, use_opaque_pointers = true}
+    } {
+      legal_dialects = ["llvm"], 
+      partial_conversion
+    } : !transform.any_op
+    transform.yield
+  }
 }

mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm-32b.mlir

@@ -1,6 +1,6 @@
 // RUN: mlir-opt %s -convert-gpu-to-nvvm='index-bitwidth=32 use-opaque-pointers=1' -split-input-file | FileCheck %s
 
-// RUN: mlir-opt %s -test-transform-dialect-interpreter | FileCheck %s
+// RUN: mlir-opt %s -transform-interpreter | FileCheck %s
 
 gpu.module @test_module_0 {
   // CHECK-LABEL: func @gpu_index_ops()
@@ -48,30 +48,32 @@ gpu.module @test_module_1 {
   }
 }
 
-transform.sequence failures(propagate) {
-^bb1(%toplevel_module: !transform.any_op):
-  %gpu_module = transform.structured.match ops{["gpu.module"]} in %toplevel_module
-    : (!transform.any_op) -> !transform.any_op
-  transform.apply_conversion_patterns to %gpu_module {
-    transform.apply_conversion_patterns.dialect_to_llvm "arith"
-    transform.apply_conversion_patterns.dialect_to_llvm "cf"
-    transform.apply_conversion_patterns.vector.vector_to_llvm
-    transform.apply_conversion_patterns.func.func_to_llvm
-    transform.apply_conversion_patterns.dialect_to_llvm "memref"
-    transform.apply_conversion_patterns.gpu.gpu_to_nvvm
-    transform.apply_conversion_patterns.gpu.gpu_wmma_to_nvvm
-    transform.apply_conversion_patterns.gpu.gpu_subgroup_reduce_to_nvvm {has_redux = true}
-    transform.apply_conversion_patterns.nvgpu.nvgpu_to_nvvm
-  } with type_converter {
-    transform.apply_conversion_patterns.memref.memref_to_llvm_type_converter
-      {index_bitwidth = 32, use_opaque_pointers = true}
-  } {
-    legal_dialects = ["llvm", "memref", "nvvm"],
-    legal_ops = ["func.func", "gpu.module", "gpu.module_end", "gpu.yield"],
-    illegal_dialects = ["gpu"],
-    illegal_ops = ["llvm.cos", "llvm.exp", "llvm.exp2", "llvm.fabs", "llvm.fceil",
-                   "llvm.ffloor", "llvm.log", "llvm.log10", "llvm.log2", "llvm.pow",
-                   "llvm.sin", "llvm.sqrt"],
-    partial_conversion
-  } : !transform.any_op
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%toplevel_module: !transform.any_op {transform.readonly}) {
+    %gpu_module = transform.structured.match ops{["gpu.module"]} in %toplevel_module
+      : (!transform.any_op) -> !transform.any_op
+    transform.apply_conversion_patterns to %gpu_module {
+      transform.apply_conversion_patterns.dialect_to_llvm "arith"
+      transform.apply_conversion_patterns.dialect_to_llvm "cf"
+      transform.apply_conversion_patterns.vector.vector_to_llvm
+      transform.apply_conversion_patterns.func.func_to_llvm
+      transform.apply_conversion_patterns.dialect_to_llvm "memref"
+      transform.apply_conversion_patterns.gpu.gpu_to_nvvm
+      transform.apply_conversion_patterns.gpu.gpu_wmma_to_nvvm
+      transform.apply_conversion_patterns.gpu.gpu_subgroup_reduce_to_nvvm {has_redux = true}
+      transform.apply_conversion_patterns.nvgpu.nvgpu_to_nvvm
+    } with type_converter {
+      transform.apply_conversion_patterns.memref.memref_to_llvm_type_converter
+        {index_bitwidth = 32, use_opaque_pointers = true}
+    } {
+      legal_dialects = ["llvm", "memref", "nvvm"],
+      legal_ops = ["func.func", "gpu.module", "gpu.module_end", "gpu.yield"],
+      illegal_dialects = ["gpu"],
+      illegal_ops = ["llvm.cos", "llvm.exp", "llvm.exp2", "llvm.fabs", "llvm.fceil",
+                    "llvm.ffloor", "llvm.log", "llvm.log10", "llvm.log2", "llvm.pow",
+                    "llvm.sin", "llvm.sqrt"],
+      partial_conversion
+    } : !transform.any_op
+    transform.yield
+  }
 }

mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir

@@ -1,5 +1,5 @@
 // RUN: mlir-opt %s -convert-gpu-to-nvvm='has-redux=1 use-opaque-pointers=1' -split-input-file | FileCheck %s
-// RUN: mlir-opt %s -test-transform-dialect-interpreter | FileCheck %s
+// RUN: mlir-opt %s -transform-interpreter | FileCheck %s
 
 gpu.module @test_module_0 {
   // CHECK-LABEL: func @gpu_index_ops()
@@ -627,38 +627,40 @@ gpu.module @test_module_31 {
   }
 }
 
-transform.sequence failures(propagate) {
-^bb1(%toplevel_module: !transform.any_op):
-  %gpu_module = transform.structured.match ops{["gpu.module"]} in %toplevel_module
-    : (!transform.any_op) -> !transform.any_op
-
-  transform.apply_patterns to %gpu_module {
-    transform.apply_patterns.gpu.gpu_rewrite_patterns
-  } : !transform.any_op
-
-  transform.apply_conversion_patterns to %gpu_module {
-    transform.apply_conversion_patterns.dialect_to_llvm "arith"
-    transform.apply_conversion_patterns.dialect_to_llvm "cf"
-    transform.apply_conversion_patterns.vector.vector_to_llvm
-    transform.apply_conversion_patterns.func.func_to_llvm
-    transform.apply_conversion_patterns.dialect_to_llvm "memref"
-    transform.apply_conversion_patterns.gpu.gpu_to_nvvm
-    transform.apply_conversion_patterns.gpu.gpu_wmma_to_nvvm
-    transform.apply_conversion_patterns.gpu.gpu_subgroup_reduce_to_nvvm
-    transform.apply_conversion_patterns.nvgpu.nvgpu_to_nvvm
-  } with type_converter {
-    transform.apply_conversion_patterns.memref.memref_to_llvm_type_converter
-      {index_bitwidth = 64,
-       use_bare_ptr = true,
-       use_bare_ptr_memref_call_conv = true,
-       use_opaque_pointers = true}
-  } {
-    legal_dialects = ["llvm", "memref", "nvvm", "test"],
-    legal_ops = ["func.func", "gpu.module", "gpu.module_end", "gpu.yield"],
-    illegal_dialects = ["gpu"],
-    illegal_ops = ["llvm.cos", "llvm.exp", "llvm.exp2", "llvm.fabs", "llvm.fceil",
-                   "llvm.ffloor", "llvm.log", "llvm.log10", "llvm.log2","llvm.pow",
-                   "llvm.sin", "llvm.sqrt"],
-    partial_conversion
-  } : !transform.any_op
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%toplevel_module: !transform.any_op {transform.readonly}) {
+    %gpu_module = transform.structured.match ops{["gpu.module"]} in %toplevel_module
+      : (!transform.any_op) -> !transform.any_op
+
+    transform.apply_patterns to %gpu_module {
+      transform.apply_patterns.gpu.gpu_rewrite_patterns
+    } : !transform.any_op
+
+    transform.apply_conversion_patterns to %gpu_module {
+      transform.apply_conversion_patterns.dialect_to_llvm "arith"
+      transform.apply_conversion_patterns.dialect_to_llvm "cf"
+      transform.apply_conversion_patterns.vector.vector_to_llvm
+      transform.apply_conversion_patterns.func.func_to_llvm
+      transform.apply_conversion_patterns.dialect_to_llvm "memref"
+      transform.apply_conversion_patterns.gpu.gpu_to_nvvm
+      transform.apply_conversion_patterns.gpu.gpu_wmma_to_nvvm
+      transform.apply_conversion_patterns.gpu.gpu_subgroup_reduce_to_nvvm
+      transform.apply_conversion_patterns.nvgpu.nvgpu_to_nvvm
+    } with type_converter {
+      transform.apply_conversion_patterns.memref.memref_to_llvm_type_converter
+        {index_bitwidth = 64,
+        use_bare_ptr = true,
+        use_bare_ptr_memref_call_conv = true,
+        use_opaque_pointers = true}
+    } {
+      legal_dialects = ["llvm", "memref", "nvvm", "test"],
+      legal_ops = ["func.func", "gpu.module", "gpu.module_end", "gpu.yield"],
+      illegal_dialects = ["gpu"],
+      illegal_ops = ["llvm.cos", "llvm.exp", "llvm.exp2", "llvm.fabs", "llvm.fceil",
+                    "llvm.ffloor", "llvm.log", "llvm.log10", "llvm.log2","llvm.pow",
+                    "llvm.sin", "llvm.sqrt"],
+      partial_conversion
+    } : !transform.any_op
+    transform.yield
+  }
 }