c1a2360918
Addresses #53169. Mirrors the Python bindings pattern used for Attribute subclasses so that Location discrimination uses `isinstance`, and fills two small gaps at the same time. ### Approach Previously `Location` was a single nanobind class with `is_a_file`, `is_a_name`, etc. predicates, plus field accessors for every kind defined on the base class. This PR introduces a `PyConcreteLocation<T>` CRTP template (parallel to `PyConcreteAttribute<T>`) and registers one subclass per `LocationAttr` kind: `UnknownLoc`, `FileLineColLoc`, `NameLoc`, `CallSiteLoc`, `FusedLoc`. TypeID-based downcasting is implemented in `PyLocation::maybeDownCast` (using `mlirAttributeGetTypeID(mlirLocationGetAttribute(...))`) and called at the boundaries that return Location objects: `op.location`, `value.location`, `Location.from_attr`, and the subclass getters themselves. ### Example ```python from mlir.ir import Context, Module, FileLineColLoc, NameLoc, FusedLoc with Context(): module = Module.parse("...") for op in module.body.operations: loc = op.location if isinstance(loc, FileLineColLoc): report(f"{loc.filename}:{loc.start_line}:{loc.start_col}") elif isinstance(loc, NameLoc): report(loc.name_str) elif isinstance(loc, FusedLoc) and loc.metadata is not None: report_fused(loc.locations, loc.metadata) ``` `FusedLoc.metadata` is also newly exposed (the underlying C API `mlirLocationFusedGetMetadata` already existed but was not bound). ### Custom/downstream Location classes `PyConcreteLocation<DerivedTy>` is marked `MLIR_PYTHON_API_EXPORTED` and is used identically to `PyConcreteAttribute<DerivedTy>`. Downstream projects that define their own `LocationAttr` in C++ expose a matching Python class by declaring a subclass with the standard statics (`isaFunction`, `pyClassName`, optional `getTypeIdFunction`) and calling its `bind(m)` at module init; `isinstance` and automatic downcasting then work with no further changes in the base bindings. ### Backward compatibility Existing factories — `Location.unknown()`, `Location.file(...)`, `Location.name(...)`, `Location.callsite(...)`, `Location.fused(...)` — remain as aliases and now return the concrete subclass instance (so `isinstance(Location.file(...), FileLineColLoc)` is true). No user migration is required for construction. The old `is_a_*` predicate methods are removed; downstream consumers should switch to `isinstance(...)`. No in-tree users of the removed predicates besides the existing Location tests, which have been rewritten accordingly. ### Scope note — OpaqueLoc OpaqueLoc is intentionally not included. It is keyed on a C++ `TypeID::get<T>()` tag that has no natural Python counterpart, and I could not find a downstream consumer asking for Python exposure. An earlier revision of this PR exposed it with an `id(obj)` + `ctypes.cast` example, but that pattern is lifetime-unsafe (MLIR does not INCREF the underlying pointer) and the C++ `getUnderlyingLocation<T>()` semantics don't translate cleanly. Happy to add it in a follow-up once a real use case surfaces. ### Changes C API (`mlir/include/mlir-c/IR.h`, `mlir/lib/CAPI/IR/IR.cpp`): - `mlirLocationIsAUnknown`, `mlirLocationUnknownGetTypeID` — other Location kinds already had their `IsA*`/`*GetTypeID` pair. Python (`mlir/include/mlir/Bindings/Python/IRCore.h`, `mlir/lib/Bindings/Python/IRCore.cpp`): - `PyConcreteLocation<T>` template and five concrete classes - `PyLocation::maybeDownCast` - Subclass `.get()` constructors, `.metadata` on `FusedLoc` Tests: - `mlir/test/python/ir/location.py` rewritten around `isinstance` and subclass construction; adds explicit-cast-failure coverage. - `mlir/test/CAPI/ir.c` adds `testLocation` covering the new `mlirLocationIsAUnknown` predicate. ### Tool-use disclosure This change was drafted with Claude (Anthropic) assistance per LLVM's [AI Tool Use Policy](https://llvm.org/docs/AIToolPolicy.html). All code was read, reviewed, and tested locally by me; the commit carries `Assisted-by:` and `Co-Authored-By:` trailers accordingly. --------- Co-authored-by: Claude (Anthropic) <noreply@anthropic.com>
295 lines
9.5 KiB
Python
295 lines
9.5 KiB
Python
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
import gc
|
|
from mlir.ir import *
|
|
|
|
|
|
def run(f):
|
|
print("\nTEST:", f.__name__)
|
|
f()
|
|
gc.collect()
|
|
assert Context._get_live_count() == 0
|
|
|
|
|
|
# CHECK-LABEL: TEST: testUnknown
|
|
def testUnknown():
|
|
with Context() as ctx:
|
|
loc = UnknownLoc.get()
|
|
assert loc.context is ctx
|
|
ctx = None
|
|
gc.collect()
|
|
# CHECK: unknown str: loc(unknown)
|
|
print("unknown str:", str(loc))
|
|
# CHECK: unknown repr: UnknownLoc(loc(unknown))
|
|
print("unknown repr:", repr(loc))
|
|
|
|
assert isinstance(loc, UnknownLoc)
|
|
assert isinstance(loc, Location)
|
|
assert not isinstance(loc, FileLineColLoc)
|
|
assert not isinstance(loc, NameLoc)
|
|
assert not isinstance(loc, CallSiteLoc)
|
|
assert not isinstance(loc, FusedLoc)
|
|
|
|
|
|
run(testUnknown)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testLocationAttr
|
|
def testLocationAttr():
|
|
with Context() as ctxt:
|
|
loc = UnknownLoc.get()
|
|
attr = loc.attr
|
|
clone = Location.from_attr(attr)
|
|
gc.collect()
|
|
# CHECK: loc: loc(unknown)
|
|
print("loc:", str(loc))
|
|
# CHECK: clone: loc(unknown)
|
|
print("clone:", str(clone))
|
|
assert loc == clone
|
|
assert isinstance(clone, UnknownLoc)
|
|
|
|
|
|
run(testLocationAttr)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testFileLineCol
|
|
def testFileLineCol():
|
|
with Context() as ctx:
|
|
loc = FileLineColLoc.get("foo1.txt", 123, 56)
|
|
range = FileLineColLoc.get("foo2.txt", 123, 56, 124, 100)
|
|
|
|
ctx = None
|
|
gc.collect()
|
|
|
|
# CHECK: file str: loc("foo1.txt":123:56)
|
|
print("file str:", str(loc))
|
|
# CHECK: file repr: FileLineColLoc(loc("foo1.txt":123:56))
|
|
print("file repr:", repr(loc))
|
|
# CHECK: file range str: loc("foo2.txt":123:56 to 124:100)
|
|
print("file range str:", str(range))
|
|
# CHECK: file range repr: FileLineColLoc(loc("foo2.txt":123:56 to 124:100))
|
|
print("file range repr:", repr(range))
|
|
|
|
assert isinstance(loc, FileLineColLoc)
|
|
assert not isinstance(loc, NameLoc)
|
|
assert not isinstance(loc, CallSiteLoc)
|
|
assert not isinstance(loc, FusedLoc)
|
|
|
|
# CHECK: file filename: foo1.txt
|
|
print("file filename:", loc.filename)
|
|
# CHECK: file start_line: 123
|
|
print("file start_line:", loc.start_line)
|
|
# CHECK: file start_col: 56
|
|
print("file start_col:", loc.start_col)
|
|
# CHECK: file end_line: 123
|
|
print("file end_line:", loc.end_line)
|
|
# CHECK: file end_col: 56
|
|
print("file end_col:", loc.end_col)
|
|
|
|
assert isinstance(range, FileLineColLoc)
|
|
# CHECK: file filename: foo2.txt
|
|
print("file filename:", range.filename)
|
|
# CHECK: file start_line: 123
|
|
print("file start_line:", range.start_line)
|
|
# CHECK: file start_col: 56
|
|
print("file start_col:", range.start_col)
|
|
# CHECK: file end_line: 124
|
|
print("file end_line:", range.end_line)
|
|
# CHECK: file end_col: 100
|
|
print("file end_col:", range.end_col)
|
|
|
|
with Context() as ctx:
|
|
ctx.allow_unregistered_dialects = True
|
|
loc = FileLineColLoc.get("foo3.txt", 127, 61)
|
|
with loc:
|
|
i32 = IntegerType.get_signless(32)
|
|
module = Module.create()
|
|
with InsertionPoint(module.body):
|
|
op = Operation.create("custom.op1", results=[i32])
|
|
new_value = op.result
|
|
# CHECK: new_value location: loc("foo3.txt":127:61)
|
|
print("new_value location: ", new_value.location)
|
|
# `op.location` and `value.location` both downcast to the
|
|
# concrete subclass.
|
|
assert isinstance(op.location, FileLineColLoc)
|
|
assert isinstance(new_value.location, FileLineColLoc)
|
|
assert op.location.typeid == FileLineColLoc.static_typeid
|
|
|
|
|
|
run(testFileLineCol)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testName
|
|
def testName():
|
|
with Context() as ctx:
|
|
loc = NameLoc.get("nombre")
|
|
loc_with_child_loc = NameLoc.get("naam", loc)
|
|
|
|
ctx = None
|
|
gc.collect()
|
|
|
|
# CHECK: name str: loc("nombre")
|
|
print("name str:", str(loc))
|
|
# CHECK: name repr: NameLoc(loc("nombre"))
|
|
print("name repr:", repr(loc))
|
|
# CHECK: name str: loc("naam"("nombre"))
|
|
print("name str:", str(loc_with_child_loc))
|
|
# CHECK: name repr: NameLoc(loc("naam"("nombre")))
|
|
print("name repr:", repr(loc_with_child_loc))
|
|
|
|
assert isinstance(loc, NameLoc)
|
|
# CHECK: name name_str: nombre
|
|
print("name name_str:", loc.name_str)
|
|
# CHECK: name child_loc: loc(unknown)
|
|
print("name child_loc:", loc.child_loc)
|
|
assert isinstance(loc.child_loc, UnknownLoc)
|
|
|
|
assert isinstance(loc_with_child_loc, NameLoc)
|
|
# CHECK: name name_str: naam
|
|
print("name name_str:", loc_with_child_loc.name_str)
|
|
# CHECK: name child_loc_with_child_loc: loc("nombre")
|
|
print("name child_loc_with_child_loc:", loc_with_child_loc.child_loc)
|
|
assert isinstance(loc_with_child_loc.child_loc, NameLoc)
|
|
|
|
|
|
run(testName)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testCallSite
|
|
def testCallSite():
|
|
with Context() as ctx:
|
|
loc = CallSiteLoc.get(
|
|
FileLineColLoc.get("foo.text", 123, 45),
|
|
[
|
|
FileLineColLoc.get("util.foo", 379, 21),
|
|
FileLineColLoc.get("main.foo", 100, 63),
|
|
],
|
|
)
|
|
ctx = None
|
|
# CHECK: callsite str: loc(callsite("foo.text":123:45 at callsite("util.foo":379:21 at "main.foo":100:63))
|
|
print("callsite str:", str(loc))
|
|
# CHECK: callsite repr: CallSiteLoc(loc(callsite("foo.text":123:45 at callsite("util.foo":379:21 at "main.foo":100:63)))
|
|
print("callsite repr:", repr(loc))
|
|
|
|
assert isinstance(loc, CallSiteLoc)
|
|
# CHECK: callsite callee: loc("foo.text":123:45)
|
|
print("callsite callee:", loc.callee)
|
|
assert isinstance(loc.callee, FileLineColLoc)
|
|
# CHECK: callsite caller: loc(callsite("util.foo":379:21 at "main.foo":100:63))
|
|
print("callsite caller:", loc.caller)
|
|
assert isinstance(loc.caller, CallSiteLoc)
|
|
|
|
|
|
run(testCallSite)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testFused
|
|
def testFused():
|
|
with Context() as ctx:
|
|
loc_single = Location.fused([NameLoc.get("apple")])
|
|
loc_empty = Location.fused([])
|
|
loc = FusedLoc.get([NameLoc.get("apple"), NameLoc.get("banana")])
|
|
attr = Attribute.parse('"sauteed"')
|
|
loc_attr = FusedLoc.get([NameLoc.get("carrot"), NameLoc.get("potatoes")], attr)
|
|
loc_empty_attr = FusedLoc.get([], attr)
|
|
loc_single_attr = FusedLoc.get([NameLoc.get("apple")], attr)
|
|
|
|
try:
|
|
FusedLoc.get([NameLoc.get("x")])
|
|
except ValueError as e:
|
|
# CHECK: fused strict error: FusedLoc.get would collapse
|
|
print("fused strict error:", str(e)[:35])
|
|
else:
|
|
assert False, "expected ValueError from strict FusedLoc.get"
|
|
|
|
ctx = None
|
|
|
|
assert not isinstance(loc_single, FusedLoc)
|
|
assert isinstance(loc_single, NameLoc)
|
|
# CHECK: fused str: loc("apple")
|
|
print("fused str:", str(loc_single))
|
|
# CHECK: fused repr: NameLoc(loc("apple"))
|
|
print("fused repr:", repr(loc_single))
|
|
|
|
assert isinstance(loc, FusedLoc)
|
|
# CHECK: fused str: loc(fused["apple", "banana"])
|
|
print("fused str:", str(loc))
|
|
# CHECK: fused repr: FusedLoc(loc(fused["apple", "banana"]))
|
|
print("fused repr:", repr(loc))
|
|
# CHECK: fused locations: [NameLoc(loc("apple")), NameLoc(loc("banana"))]
|
|
print("fused locations:", loc.locations)
|
|
# CHECK: fused metadata: None
|
|
print("fused metadata:", loc.metadata)
|
|
|
|
assert isinstance(loc_attr, FusedLoc)
|
|
# CHECK: fused metadata: "sauteed"
|
|
print("fused metadata:", loc_attr.metadata)
|
|
# CHECK: fused str: loc(fused<"sauteed">["carrot", "potatoes"])
|
|
print("fused str:", str(loc_attr))
|
|
# CHECK: fused repr: FusedLoc(loc(fused<"sauteed">["carrot", "potatoes"]))
|
|
print("fused repr:", repr(loc_attr))
|
|
# CHECK: fused locations: [NameLoc(loc("carrot")), NameLoc(loc("potatoes"))]
|
|
print("fused locations:", loc_attr.locations)
|
|
|
|
assert not isinstance(loc_empty, FusedLoc)
|
|
assert isinstance(loc_empty, UnknownLoc)
|
|
# CHECK: fused str: loc(unknown)
|
|
print("fused str:", str(loc_empty))
|
|
# CHECK: fused repr: UnknownLoc(loc(unknown))
|
|
print("fused repr:", repr(loc_empty))
|
|
|
|
assert isinstance(loc_empty_attr, FusedLoc)
|
|
# CHECK: fused str: loc(fused<"sauteed">[unknown])
|
|
print("fused str:", str(loc_empty_attr))
|
|
# CHECK: fused repr: FusedLoc(loc(fused<"sauteed">[unknown]))
|
|
print("fused repr:", repr(loc_empty_attr))
|
|
# CHECK: fused locations: [UnknownLoc(loc(unknown))]
|
|
print("fused locations:", loc_empty_attr.locations)
|
|
|
|
assert isinstance(loc_single_attr, FusedLoc)
|
|
# CHECK: fused str: loc(fused<"sauteed">["apple"])
|
|
print("fused str:", str(loc_single_attr))
|
|
# CHECK: fused repr: FusedLoc(loc(fused<"sauteed">["apple"]))
|
|
print("fused repr:", repr(loc_single_attr))
|
|
# CHECK: fused locations: [NameLoc(loc("apple"))]
|
|
print("fused locations:", loc_single_attr.locations)
|
|
|
|
|
|
run(testFused)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testCast
|
|
def testCast():
|
|
with Context() as ctx:
|
|
unknown = UnknownLoc.get()
|
|
as_unknown = UnknownLoc(unknown)
|
|
assert isinstance(as_unknown, UnknownLoc)
|
|
|
|
try:
|
|
FileLineColLoc(unknown)
|
|
except ValueError as e:
|
|
# CHECK: cast error: Cannot cast location to FileLineColLoc (from loc(unknown))
|
|
print("cast error:", str(e))
|
|
else:
|
|
assert False, "expected ValueError"
|
|
|
|
ctx = None
|
|
|
|
|
|
run(testCast)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testLocationCapsule
|
|
def testLocationCapsule():
|
|
with Context() as ctx:
|
|
loc1 = FileLineColLoc.get("foo.txt", 123, 56)
|
|
# CHECK: mlir.ir.Location._CAPIPtr
|
|
loc_capsule = loc1._CAPIPtr
|
|
print(loc_capsule)
|
|
loc2 = Location._CAPICreate(loc_capsule)
|
|
assert loc2 == loc1
|
|
assert loc2.context is ctx
|
|
|
|
|
|
run(testLocationCapsule)
|