This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new ac880557 [FEAT][IR] Expose operator attribute names and attribute maps
(#2)
ac880557 is described below
commit ac880557e2b0b8f1feeae20872ef3df5dd5113ea
Author: Linzhang Li <[email protected]>
AuthorDate: Tue Oct 20 16:54:45 2026 -0400
[FEAT][IR] Expose operator attribute names and attribute maps (#2)
### Summary
This PR exposes the operator registry's attributes to Python and to
other FFI frontends:
| Python API | Global function | Result |
|---|---|---|
| `Op.list_attr_names()` | `ir.ListOpAttrNames()` | Names of all
registered attributes, in unspecified order |
| `Op.get_attr_map(name)` | `ir.GetOpAttrColumn(name)` | An `OpAttrMap`
for `name`, or `None` if no operator registered it |
`OpAttrMap` mirrors C++ `OpAttrMap`:
```python
m = Op.get_attr_map("TCallEffectKind")
op in m # C++ count(op)
m[op] # value, KeyError if op has none
m.get(expr, None) # value or default; non-Op expressions yield default
```
### Motivation
C++ code reads operator attributes through `Op::GetAttrMap<T>(name)`,
usually held in a
`static` for the process lifetime; a lookup then reads entry `op->index`
of the registry's
storage directly. `GetAttrMap` is a template and `Op::GetAttrColumn` is
private, so code outside
C++ cannot get the same view. Its only option is `op.get_attr(name)`,
one FFI call per operator
and attribute.
Frontends written in other languages, such as Rust ports of TIRx passes,
read attributes like
`TCallEffectKind`, `TVectorizable`, `FTileOpDispatch` and
`<target>.FLowerIntrinsic` on hot
paths. They can now fetch an attribute map once, cache it, and look
operators up without an FFI
call, just as C++ `OpAttrMap` does. In Python, `Op.get_attr_map` gives
the same map-level API
instead of handing out the raw storage.
### Semantics
- The map reads the registry's own storage, not a copy. Later `set_attr`
calls, overrides,
`reset_attr` removals, and operators registered later are all visible
through a cached map.
- A missing attribute returns `None` instead of raising, so callers need
no separate
`has_attr` check.
- The map is read-only; register values with `Op.set_attr`.
- Like the rest of the operator registry, maps are not synchronized: do
not read them while
another thread registers operators or attributes.
### Implementation
- `OpRegistry::ListAttrNames()` enumerates the attribute names.
- `ir.ListOpAttrNames` and `ir.GetOpAttrColumn` are registered as
lambdas next to
`ir.ListOpNames` and `ir.RegisterOp` in `src/ir/op.cc`. They go through
the file-local
`OpRegistry` because `Op::GetAttrColumn` is private.
`ir.GetOpAttrColumn` returns the live
`ffi.List` that backs the map (entry `op.index` holds the value of `op`,
unset entries are
`None`); it is the building block for the Python and Rust maps, not a
user-facing API.
- `python/tvm_tirx/ir/op.py` adds `Op.list_attr_names`,
`Op.get_attr_map` and `OpAttrMap`,
which is exported from `tvm_tirx.ir` and listed in the API reference.
### Tests
`tests/python/ir/test_attrs.py`:
- `test_op_attr_map_names_and_values`: the attribute names include known
attributes, the map
agrees with `Op.get_attr` for every registered operator, and a missing
attribute returns
`None`.
- `test_op_attr_map_lookup_without_value`: `in`, `[]` and `get` for an
operator without a
value and for a non-Op expression.
- `test_op_attr_map_is_live`: a cached map observes an operator
registered later, an override
and a removal.
```bash
uv run --no-sync pytest -q tests/python/ir # 56 passed
```
Signed-off-by: yuchuan <[email protected]>
---
docs/reference/python/ir.rst | 1 +
python/tvm_tirx/ir/__init__.py | 2 +-
python/tvm_tirx/ir/op.py | 81 ++++++++++++++++++++++++++++++++++++++++++
src/ir/op.cc | 13 +++++++
tests/python/ir/test_attrs.py | 62 ++++++++++++++++++++++++++++++++
5 files changed, 158 insertions(+), 1 deletion(-)
diff --git a/docs/reference/python/ir.rst b/docs/reference/python/ir.rst
index b8f0895f..28215a22 100644
--- a/docs/reference/python/ir.rst
+++ b/docs/reference/python/ir.rst
@@ -94,6 +94,7 @@ Expr
Var
BaseFunc
Op
+ OpAttrMap
Stmt
----
diff --git a/python/tvm_tirx/ir/__init__.py b/python/tvm_tirx/ir/__init__.py
index a05391c4..938d0199 100644
--- a/python/tvm_tirx/ir/__init__.py
+++ b/python/tvm_tirx/ir/__init__.py
@@ -68,7 +68,7 @@ from . import prim
from .function import BaseFunc, CallingConv
from .global_info import GlobalInfo
from .module import IRModule
-from .op import Op, register_op_attr
+from .op import Op, OpAttrMap, register_op_attr
from .stmt import (
Stmt,
SeqStmt,
diff --git a/python/tvm_tirx/ir/op.py b/python/tvm_tirx/ir/op.py
index b3e8b552..aeb2405c 100644
--- a/python/tvm_tirx/ir/op.py
+++ b/python/tvm_tirx/ir/op.py
@@ -69,6 +69,38 @@ class Op(Expr):
"""
return _ffi_api.ListOpNames()
+ @staticmethod
+ def list_attr_names():
+ """List registered attribute column names in unspecified order.
+
+ Returns
+ -------
+ list[str]
+ The attribute column names, such as ``"TCallEffectKind"``.
+ """
+ return _ffi_api.ListOpAttrNames()
+
+ @staticmethod
+ def get_attr_map(attr_name):
+ """Get the map from operators to their values of an attribute.
+
+ This is the Python counterpart of C++ ``Op::GetAttrMap``. The map
reads the
+ registry's storage directly, so later registrations, overrides and
removals
+ are visible through it, and it can be cached.
+
+ Parameters
+ ----------
+ attr_name : str
+ Attribute name, such as ``"TCallEffectKind"``.
+
+ Returns
+ -------
+ OpAttrMap or None
+ The attribute map, or None when no operator has registered the
attribute.
+ """
+ column = _ffi_api.GetOpAttrColumn(attr_name)
+ return None if column is None else OpAttrMap(attr_name, column)
+
def set_attr(self, attr_name, value, override=False):
"""Set an operator attribute.
@@ -139,6 +171,55 @@ class Op(Expr):
self._set_signature(arg_names, arg_docs, ty_arg_names, ty_arg_docs,
value_tail, type_tail)
+class OpAttrMap:
+ """The values of one operator attribute, keyed by operator.
+
+ Mirrors C++ ``OpAttrMap``. Obtain one with :py:meth:`Op.get_attr_map`.
Lookups
+ see later registrations, overrides and removals; register values with
+ :py:meth:`Op.set_attr`. Like the registry itself, the map is not
synchronized
+ with concurrent registration.
+ """
+
+ def __init__(self, attr_name, column):
+ self._attr_name = attr_name
+ self._column = column
+
+ @property
+ def attr_name(self):
+ """The attribute name."""
+ return self._attr_name
+
+ def __contains__(self, op):
+ if not isinstance(op, Op):
+ return False
+ index = op.index
+ return index < len(self._column) and self._column[index] is not None
+
+ def __getitem__(self, op):
+ if op not in self:
+ name = op.name if isinstance(op, Op) else op
+ raise KeyError(f"Attribute {self._attr_name} has not been
registered for {name}")
+ return self._column[op.index]
+
+ def get(self, expr, default=None):
+ """Get the attribute of an operator with a fallback.
+
+ Parameters
+ ----------
+ expr : Expr
+ The operator. Any other expression yields ``default``.
+
+ default : object
+ The value to return when ``expr`` is not an Op or has no value.
+
+ Returns
+ -------
+ object
+ The registered value or ``default``.
+ """
+ return self[expr] if expr in self else default
+
+
def register_op_attr(op_name, attr_key, value=None, override=False):
"""Register an operator property of an operator by name.
diff --git a/src/ir/op.cc b/src/ir/op.cc
index f2533bfd..3e9df183 100644
--- a/src/ir/op.cc
+++ b/src/ir/op.cc
@@ -64,6 +64,12 @@ class OpRegistry {
bool HasAttrMap(const ffi::String& name) const { return attrs_.count(name); }
+ ffi::Array<ffi::String> ListAttrNames() const {
+ ffi::Array<ffi::String> names;
+ for (const auto& [name, column] : attrs_) names.push_back(name);
+ return names;
+ }
+
ffi::List<ffi::Any> GetAttrColumn(const ffi::String& name) const {
auto column = attrs_.Get(name);
TVM_FFI_CHECK(column.has_value(), InternalError)
@@ -321,6 +327,13 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef()
.def("ir.GetOp", &Op::Get)
.def("ir.ListOpNames", &Op::ListNames)
+ .def("ir.ListOpAttrNames", []() { return
OpRegistry::Global()->ListAttrNames(); })
+ .def("ir.GetOpAttrColumn",
+ [](ffi::String name) -> ffi::Optional<ffi::List<ffi::Any>> {
+ OpRegistry* registry = OpRegistry::Global();
+ if (!registry->HasAttrMap(name)) return std::nullopt;
+ return registry->GetAttrColumn(name);
+ })
.def("ir.RegisterOp",
[](ffi::String name, ffi::String doc) {
TVM_FFI_CHECK(!OpRegistry::Global()->Contains(name),
AttributeError)
diff --git a/tests/python/ir/test_attrs.py b/tests/python/ir/test_attrs.py
index c39a9c4e..4d27e46a 100644
--- a/tests/python/ir/test_attrs.py
+++ b/tests/python/ir/test_attrs.py
@@ -129,6 +129,68 @@ def
test_op_set_signature_checks_counts_without_types_or_attrs():
assert call.attrs is not None
+def test_op_attr_map_names_and_values():
+ names = tvm_tirx.ir.Op.list_attr_names()
+ assert {"TCallEffectKind", "TGlobalSymbol"} <= set(names)
+ for name in names:
+ assert tvm_tirx.ir.Op.get("prim.exp").has_attr(name)
+
+ attr_map = tvm_tirx.ir.Op.get_attr_map("TCallEffectKind")
+ assert attr_map.attr_name == "TCallEffectKind"
+ for op_name in tvm_tirx.ir.Op.list_op_names():
+ op = tvm_tirx.ir.Op.get(op_name)
+ value = op.get_attr("TCallEffectKind")
+ assert (op in attr_map) == (value is not None), op_name
+ assert attr_map.get(op) == value, op_name
+ if value is not None:
+ assert attr_map[op] == value, op_name
+
+ assert tvm_tirx.ir.Op.get_attr_map("test.MissingAttrMap") is None
+
+
+def test_op_attr_map_lookup_without_value():
+ key = "test.AttrMapLookup"
+ tvm_tirx.ir._ffi_api.RegisterOp("test.attr_map_lookup.set", "")
+ tvm_tirx.ir._ffi_api.RegisterOp("test.attr_map_lookup.unset", "")
+ tvm_tirx.ir.Op.get("test.attr_map_lookup.set").set_attr(key, 1)
+ unset = tvm_tirx.ir.Op.get("test.attr_map_lookup.unset")
+
+ attr_map = tvm_tirx.ir.Op.get_attr_map(key)
+ assert unset not in attr_map
+ assert attr_map.get(unset) is None
+ assert attr_map.get(unset, 5) == 5
+ with pytest.raises(KeyError, match="test.attr_map_lookup.unset"):
+ attr_map[unset]
+
+ # A non-Op expression has no attributes, as in C++ OpAttrMap::get(Expr,
default).
+ var = tvm_tirx.ir.Var("x", tvm_tirx.ir.PrimType("int32"))
+ assert var not in attr_map
+ assert attr_map.get(var, 5) == 5
+
+
+def test_op_attr_map_is_live():
+ key = "test.LiveAttrMap"
+ tvm_tirx.ir._ffi_api.RegisterOp("test.live_attr_map.first", "")
+ first = tvm_tirx.ir.Op.get("test.live_attr_map.first")
+ assert key not in tvm_tirx.ir.Op.list_attr_names()
+ first.set_attr(key, 7)
+ assert key in tvm_tirx.ir.Op.list_attr_names()
+
+ attr_map = tvm_tirx.ir.Op.get_attr_map(key)
+ assert attr_map[first] == 7
+
+ # A cached map sees an Op registered later, an override and a removal.
+ tvm_tirx.ir._ffi_api.RegisterOp("test.live_attr_map.second", "")
+ second = tvm_tirx.ir.Op.get("test.live_attr_map.second")
+ assert second not in attr_map
+ second.set_attr(key, 8)
+ assert attr_map[second] == 8
+ first.set_attr(key, 9, override=True)
+ assert attr_map[first] == 9
+ first.reset_attr(key)
+ assert first not in attr_map
+
+
def test_generated_op_api_constructs_dictionary_attrs():
name = "test.python_dict_attrs_api"
tvm_tirx.ir._ffi_api.RegisterOp(name, "Dictionary attribute API test")
