Skip to content

Commit 7b8a160

Browse files
committed
[tuner] use python binding to get indexing maps for root op
Signed-off-by: Bangtian Liu <[email protected]>
1 parent 104ab43 commit 7b8a160

File tree

1 file changed

+2
-9
lines changed

1 file changed

+2
-9
lines changed

tuner/tuner/dispatch_parser.py

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -61,15 +61,8 @@ def get_problem_size(self) -> ProblemSize:
6161
contraction_dims = linalg.infer_contraction_dimensions(root_op)
6262
assert contraction_dims, "no contraction dimensions"
6363

64-
# TODO(Bangtian): Expose Python bindings for getting indexing maps.
65-
indexing_maps_attr = None
66-
for attr in root_op.opview.attributes:
67-
if attr.name == "indexing_maps" and isinstance(attr.attr, ir.ArrayAttr):
68-
indexing_maps_attr = attr.attr
69-
break
70-
71-
assert indexing_maps_attr, "indexing_maps attribute not found"
72-
maps = [attr.value for attr in indexing_maps_attr]
64+
res_maps = linalg.get_indexing_maps(root_op)
65+
maps = [map_attr.value for map_attr in res_maps]
7366
lhs_dims = get_map_result_dim_positions(maps[0])
7467
rhs_dims = get_map_result_dim_positions(maps[1])
7568
res_dims = get_map_result_dim_positions(maps[2])

0 commit comments

Comments
 (0)