Skip to content

Commit ff0f1dd

Browse files
authored
[mlir:python] Small optimization to get_op_result_or_results. (#123866)
* We can call .results without figuring out whether we have an Operation or an OpView, and that's likely the common case anyway. * If we have one or more results, we can return them directly, with no need for a call to get_op_result_or_value. We're guaranteed that .results returns a PyOpResultList, so we have either an OpResult or sequence of OpResults, just as the API expects. This saves a few 100ms during IR construction in an LLM JAX benchmark.
1 parent 7986e0c commit ff0f1dd

File tree

1 file changed

+10
-9
lines changed

1 file changed

+10
-9
lines changed

mlir/python/mlir/dialects/_ods_common.py

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -133,15 +133,16 @@ def get_op_results_or_values(
133133
def get_op_result_or_op_results(
134134
op: _Union[_cext.ir.OpView, _cext.ir.Operation],
135135
) -> _Union[_cext.ir.Operation, _cext.ir.OpResult, _Sequence[_cext.ir.OpResult]]:
136-
if isinstance(op, _cext.ir.OpView):
137-
op = op.operation
138-
return (
139-
list(get_op_results_or_values(op))
140-
if len(op.results) > 1
141-
else get_op_result_or_value(op)
142-
if len(op.results) > 0
143-
else op
144-
)
136+
results = op.results
137+
num_results = len(results)
138+
if num_results == 1:
139+
return results[0]
140+
elif num_results > 1:
141+
return results
142+
elif isinstance(op, _cext.ir.OpView):
143+
return op.operation
144+
else:
145+
return op
145146

146147
ResultValueTypeTuple = _cext.ir.Operation, _cext.ir.OpView, _cext.ir.Value
147148
ResultValueT = _Union[ResultValueTypeTuple]

0 commit comments

Comments
 (0)