[Mlir-commits] [mlir] [mlir][python] Add pythonic interface for GPUFuncOp (PR #163596)

[Maksim Levental]

================
@@ -3,5 +3,167 @@
 #  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, List
+
+try:
+    from ...ir import (
+        FunctionType,
+        TypeAttr,
+        StringAttr,
+        UnitAttr,
+        Block,
+        InsertionPoint,
+        ArrayAttr,
+        Type,
+        DictAttr,
+        Attribute,
+        DenseI32ArrayAttr,
+    )
+    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
+
+
+KERNEL_ATTRIBUTE_NAME = "gpu.kernel"
+KNOWN_BLOCK_SIZE_ATTRIBUTE_NAME = "gpu.known_block_size"
+KNOWN_GRID_SIZE_ATTRIBUTE_NAME = "gpu.known_grid_size"
+
+FUNCTION_TYPE_ATTRIBUTE_NAME = "function_type"
+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):
+    __doc__ = GPUFuncOp.__doc__
+
+    def __init__(
+        self,
+        function_type: Union[FunctionType, TypeAttr],
+        sym_name: Optional[Union[str, StringAttr]] = None,
+        kernel: Optional[bool] = None,
+        body_builder: Optional[Callable[[GPUFuncOp], None]] = None,
+        known_block_size: Optional[Union[List[int], DenseI32ArrayAttr]] = None,
+        known_grid_size: Optional[Union[List[int], DenseI32ArrayAttr]] = None,
----------------
makslevental wrote:

nit:
```suggestion
        known_block_size: Optional[Union[Sequence[int], DenseI32ArrayAttr]] = None,
        known_grid_size: Optional[Union[Sequence[int], DenseI32ArrayAttr]] = None,
```

because some people might use tuples

https://github.com/llvm/llvm-project/pull/163596