Skip to content

[FRONTEND] Add get_int_attr to Operation for out-of-tree walk #8892

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
peterbell10 merged 2 commits into from
Dec 4, 2025
Merged

[FRONTEND] Add get_int_attr to Operation for out-of-tree walk #8892

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0ff7a35
[FRONTEND] Add get_int_attr to Operation for out-of-tree walk
jjvraw Dec 3, 2025
2bca18a
Merge branch 'main' into op-get-int-attr
peterbell10 Dec 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions python/src/ir.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -635,6 +635,13 @@ void init_triton_ir(py::module &&m) {
return py::none();
return py::str(ret.getValue().str());
})
.def("get_int_attr",
[](Operation &self, const std::string &name) -> py::object {
auto ret = self.getAttrOfType<IntegerAttr>(name);
if (!ret)
return py::none();
return py::int_(ret.getInt());
})
.def("get_bool_attr",
[](Operation &self, const std::string &name) -> py::object {
auto ret = self.getAttrOfType<BoolAttr>(name);
Expand Down
10 changes: 9 additions & 1 deletion python/test/unit/runtime/test_bindings.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
import torch
import math

_BLOCK_SIZE = 16


@triton.jit
def add_helper(x, y):
Expand Down Expand Up @@ -50,14 +52,20 @@ def walk_fn(op):
op.get_str_attr("sym_name")
if name == "tt.call":
op.get_flat_symbol_ref_attr("callee")
if name == "tt.make_range":
assert 0 == op.get_int_attr("start")
assert _BLOCK_SIZE == op.get_int_attr("end")
if name == "arith.constant":
val = op.get_int_attr("value")
assert val is None or isinstance(val, int)

kernel = add_kernel
args = [
torch.empty((32, 32), device=device), # in_ptr0
torch.empty((32, 32), device=device), # in_ptr1
1024, # n_elements
torch.empty((32, 32), device=device), # out_ptr
16, # BLOCK_SIZE
_BLOCK_SIZE, # BLOCK_SIZE
]
target = triton.runtime.driver.active.get_current_target()
backend = triton.compiler.compiler.make_backend(target)
Expand Down
Loading