diff --git a/mlir/python/mlir/dialects/gpu/__init__.py b/mlir/python/mlir/dialects/gpu/__init__.py index 2fbcbb059f87a..d15643ca700e4 100644 --- a/mlir/python/mlir/dialects/gpu/__init__.py +++ b/mlir/python/mlir/dialects/gpu/__init__.py @@ -49,13 +49,13 @@ class GPUFuncOp(GPUFuncOp): FUNCTION_TYPE_ATTR_NAME = "function_type" SYM_NAME_ATTR_NAME = "sym_name" - ARGUMENT_ATTR_NAME = "arg_attrs" - RESULT_ATTR_NAME = "res_attrs" def __init__( self, function_type: Union[FunctionType, TypeAttr], sym_name: Optional[Union[str, StringAttr]] = None, + arg_attrs: Optional[Sequence[dict]] = None, + res_attrs: Optional[Sequence[dict]] = None, kernel: Optional[bool] = None, workgroup_attrib_attrs: Optional[Sequence[dict]] = None, private_attrib_attrs: Optional[Sequence[dict]] = None, @@ -88,6 +88,8 @@ def __init__( ) super().__init__( function_type, + arg_attrs=arg_attrs, + res_attrs=res_attrs, workgroup_attrib_attrs=workgroup_attrib_attrs, private_attrib_attrs=private_attrib_attrs, loc=loc, diff --git a/mlir/test/python/dialects/gpu/dialect.py b/mlir/test/python/dialects/gpu/dialect.py index 3945c99c41091..1a009b7dfa30d 100644 --- a/mlir/test/python/dialects/gpu/dialect.py +++ b/mlir/test/python/dialects/gpu/dialect.py @@ -133,9 +133,10 @@ def builder(func: gpu.GPUFuncOp) -> None: ), func.known_grid_size func = gpu.GPUFuncOp( - func_type, + ir.FunctionType.get(inputs=[T.index()], results=[]), sym_name="non_kernel_func", body_builder=builder, + arg_attrs=[{"gpu.some_attribute": ir.StringAttr.get("foo")}], ) assert not func.is_kernel assert func.known_block_size is None @@ -154,10 +155,11 @@ def builder(func: gpu.GPUFuncOp) -> None: # CHECK: %[[VAL_0:.*]] = gpu.global_id x # CHECK: gpu.return # CHECK: } - # CHECK: gpu.func @non_kernel_func() { - # CHECK: %[[VAL_0:.*]] = gpu.global_id x - # CHECK: gpu.return - # CHECK: } + # CHECK: gpu.func @non_kernel_func( + # CHECK-SAME: %[[ARG0:.*]]: index {gpu.some_attribute = "foo"}) { + # CHECK: %[[GLOBAL_ID_0:.*]] = gpu.global_id x + # CHECK: gpu.return + # CHECK: } # CHECK-LABEL: testGPULaunchFuncOp