intel/llvm
1
0
Fork 0
mirror of https://github.com/intel/llvm.git synced 2026-01-23 07:58:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

[MLIR][Python] Add arg_attrs and res_attrs to gpu func (#168475)

Browse Source 
I missed these attributes when I added the wrapper for GPUFuncOp in
fbdd98f74f.
This commit is contained in:
Asher Mancinelli
2025-11-18 07:55:11 -08:00
committed by GitHub
parent 4d093683ce
commit 47d9d735a7
2 changed files with 11 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
mlir/python/mlir/dialects/gpu/__init__.py
View File

@@ -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 @@ class GPUFuncOp(GPUFuncOp):
)
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,

12
mlir/test/python/dialects/gpu/dialect.py
View File

@@ -133,9 +133,10 @@ def testGPUFuncOp():
), 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 testGPUFuncOp():
# 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