Files
clang-p2996/mlir/test/python/dialects/transform_interpreter.py
Oleksandr "Alex" Zinenko 91f1161133 [mlir] expose transform interpreter to Python (#82365)
Transform interpreter functionality can be used standalone without going
through the interpreter pass, make it available in Python.
2024-02-21 11:01:00 +01:00

57 lines
1.4 KiB
Python

# RUN: %PYTHON %s | FileCheck %s
from mlir import ir
from mlir.dialects.transform import interpreter as interp
def test_in_context(f):
with ir.Context(), ir.Location.unknown():
f()
return f
print_root_module = """
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%root: !transform.any_op) {
transform.print %root { name = \"from interpreter\" }: !transform.any_op
transform.yield
}
}"""
@test_in_context
def print_self():
m = ir.Module.parse(print_root_module.replace("from interpreter", "print_self"))
interp.apply_named_sequence(m, m.body.operations[0], m)
# CHECK-LABEL: print_self
# CHECK: transform.named_sequence @__transform_main
# CHECK: transform.print
# CHECK: transform.yield
@test_in_context
def print_other():
transform = ir.Module.parse(
print_root_module.replace("from interpreter", "print_other")
)
payload = ir.Module.parse("module attributes { this.is.payload } {}")
interp.apply_named_sequence(payload, transform.body.operations[0], transform)
# CHECK-LABEL: print_other
# CHECK-NOT: transform
# CHECK: this.is.payload
@test_in_context
def failed():
payload = ir.Module.parse("module attributes { this.is.payload } {}")
try:
interp.apply_named_sequence(payload, payload, payload)
except ValueError as e:
assert (
"must implement TransformOpInterface to be used as transform root" in str(e)
)