diff options
| author | Alex Zinenko <zinenko@google.com> | 2023-05-04 14:13:06 +0000 |
|---|---|---|
| committer | Alex Zinenko <zinenko@google.com> | 2023-05-05 12:42:05 +0000 |
| commit | d9064269442d8d6126021db9230bac97a2ea9e13 (patch) | |
| tree | d9478a9e4331a985329a15e12132edb0365331bf /mlir/python | |
| parent | 3b6bc87520273d7ebc3b796b8493ddb401257038 (diff) | |
[mlir] make transform.loop.outline also return the call handle
Outlining is particularly interesting when the outlined function is
replaced with something else, e.g., a microkernel. It is good to have a
handle to the call in this case.
Reviewed By: springerm
Differential Revision: https://reviews.llvm.org/D149849
Diffstat (limited to 'mlir/python')
| -rw-r--r-- | mlir/python/mlir/dialects/_loop_transform_ops_ext.py | 6 |
1 files changed, 4 insertions, 2 deletions
diff --git a/mlir/python/mlir/dialects/_loop_transform_ops_ext.py b/mlir/python/mlir/dialects/_loop_transform_ops_ext.py index a275ea615378..10079d32fd92 100644 --- a/mlir/python/mlir/dialects/_loop_transform_ops_ext.py +++ b/mlir/python/mlir/dialects/_loop_transform_ops_ext.py @@ -39,7 +39,8 @@ class LoopOutlineOp: def __init__( self, - result_type: Type, + function_type: Type, + call_type: Type, target: Union[Operation, Value], *, func_name: Union[str, StringAttr], @@ -47,7 +48,8 @@ class LoopOutlineOp: loc=None, ): super().__init__( - result_type, + function_type, + call_type, _get_op_result_or_value(target), func_name=(func_name if isinstance(func_name, StringAttr) else StringAttr.get(func_name)), |