[Mlir-commits] [mlir] [mlir][python] Add pythonic interface for GPUFuncOp (PR #163596)
Maksim Levental
llvmlistbot at llvm.org
Wed Oct 15 15:06:12 PDT 2025
================
@@ -3,5 +3,121 @@
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from .._gpu_ops_gen import *
+from .._gpu_ops_gen import _Dialect
from .._gpu_enum_gen import *
from ..._mlir_libs._mlirDialectsGPU import *
+from typing import Callable, Sequence, Union, Optional
+
+try:
+ from ...ir import (
+ FunctionType,
+ TypeAttr,
+ StringAttr,
+ UnitAttr,
+ Block,
+ InsertionPoint,
+ ArrayAttr,
+ Type,
+ DictAttr,
+ Attribute,
+ )
+ from .._ods_common import (
+ get_default_loc_context as _get_default_loc_context,
+ _cext as _ods_cext,
+ )
+except ImportError as e:
+ raise RuntimeError("Error loading imports from extension module") from e
+
+
+FUNCTION_TYPE_ATTRIBUTE_NAME = "function_type"
+KERNEL_ATTRIBUTE_NAME = "gpu.kernel"
+SYM_NAME_ATTRIBUTE_NAME = "sym_name"
+ARGUMENT_ATTRIBUTE_NAME = "arg_attrs"
+RESULT_ATTRIBUTE_NAME = "res_attrs"
+
+
+ at _ods_cext.register_operation(_Dialect, replace=True)
+class GPUFuncOp(GPUFuncOp):
----------------
makslevental wrote:
yea i guess it should just work:
```python
if workgroup_attrib_attrs is not None: attributes["workgroup_attrib_attrs"] = (workgroup_attrib_attrs if (
isinstance(workgroup_attrib_attrs, _ods_ir.Attribute) or
not _ods_ir.AttrBuilder.contains('DictArrayAttr')) else
_ods_ir.AttrBuilder.get('DictArrayAttr')(workgroup_attrib_attrs, context=_ods_context))
```
so you can support these with just like `Optional[Sequence[dict]] = None`. we have an `AttrBuilder` already for `DictArrayAttr` that'll do the right thing: https://github.com/llvm/llvm-project/blob/aed73d2afeecace1e38f9a85ec94d0aa05e86bda/mlir/python/mlir/ir.py#L213-L215
https://github.com/llvm/llvm-project/pull/163596
More information about the Mlir-commits
mailing list