[Mlir-commits] [mlir] ff0f1dd - [mlir:python] Small optimization to get_op_result_or_results. (#123866)

[llvmlistbot at llvm.org]

Author: Peter Hawkins
Date: 2025-01-22T05:41:31-08:00
New Revision: ff0f1dd341cb2fdc1281a14413b3aa93bf9a20c7

URL: https://github.com/llvm/llvm-project/commit/ff0f1dd341cb2fdc1281a14413b3aa93bf9a20c7
DIFF: https://github.com/llvm/llvm-project/commit/ff0f1dd341cb2fdc1281a14413b3aa93bf9a20c7.diff

LOG: [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.

Added: 
    

Modified: 
    mlir/python/mlir/dialects/_ods_common.py

Removed: 
    


################################################################################
diff  --git a/mlir/python/mlir/dialects/_ods_common.py b/mlir/python/mlir/dialects/_ods_common.py
index d40d936cdc83d6..5b67ab03d6f494 100644
--- a/mlir/python/mlir/dialects/_ods_common.py
+++ b/mlir/python/mlir/dialects/_ods_common.py
@@ -133,15 +133,16 @@ def get_op_results_or_values(
 def get_op_result_or_op_results(
     op: _Union[_cext.ir.OpView, _cext.ir.Operation],
 ) -> _Union[_cext.ir.Operation, _cext.ir.OpResult, _Sequence[_cext.ir.OpResult]]:
-    if isinstance(op, _cext.ir.OpView):
-        op = op.operation
-    return (
-        list(get_op_results_or_values(op))
-        if len(op.results) > 1
-        else get_op_result_or_value(op)
-        if len(op.results) > 0
-        else op
-    )
+    results = op.results
+    num_results = len(results)
+    if num_results == 1:
+        return results[0]
+    elif num_results > 1:
+        return results
+    elif isinstance(op, _cext.ir.OpView):
+        return op.operation
+    else:
+        return op
 
 ResultValueTypeTuple = _cext.ir.Operation, _cext.ir.OpView, _cext.ir.Value
 ResultValueT = _Union[ResultValueTypeTuple]