[mlir] use transform-interpreter in test passes

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
+  }
 }