Skip to content

Commit fbdd98f

Browse files
ashermancinelliashermancinelli
authored
[mlir][python] Add pythonic interface for GPUFuncOp (#163596)
The func dialect provides a more pythonic interface for constructing operations, but the gpu dialect does not; this is the first PR to provide the same conveniences for the gpu dialect, starting with the gpu.func op.
1 parent b71515c commit fbdd98f

File tree

2 files changed

+239
-0
lines changed

2 files changed

+239
-0
lines changed

‎mlir/python/mlir/dialects/gpu/__init__.py‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,5 +3,151 @@
33
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
44

55
from .._gpu_ops_gen import *
6+
from .._gpu_ops_gen import _Dialect
67
from .._gpu_enum_gen import *
78
from ..._mlir_libs._mlirDialectsGPU import *
9+
from typing import Callable, Sequence, Union, Optional, List
10+
11+
try:
12+
from ...ir import (
13+
FunctionType,
14+
TypeAttr,
15+
StringAttr,
16+
UnitAttr,
17+
Block,
18+
InsertionPoint,
19+
ArrayAttr,
20+
Type,
21+
DictAttr,
22+
Attribute,
23+
DenseI32ArrayAttr,
24+
)
25+
from .._ods_common import (
26+
get_default_loc_context as _get_default_loc_context,
27+
_cext as _ods_cext,
28+
)
29+
except ImportError as e:
30+
raise RuntimeError("Error loading imports from extension module") from e
31+
32+
33+
@_ods_cext.register_operation(_Dialect, replace=True)
34+
class GPUFuncOp(GPUFuncOp):
35+
__doc__ = GPUFuncOp.__doc__
36+
37+
KERNEL_ATTR_NAME = "gpu.kernel"
38+
KNOWN_BLOCK_SIZE_ATTR_NAME = "known_block_size"
39+
KNOWN_GRID_SIZE_ATTR_NAME = "known_grid_size"
40+
41+
FUNCTION_TYPE_ATTR_NAME = "function_type"
42+
SYM_NAME_ATTR_NAME = "sym_name"
43+
ARGUMENT_ATTR_NAME = "arg_attrs"
44+
RESULT_ATTR_NAME = "res_attrs"
45+
46+
def __init__(
47+
self,
48+
function_type: Union[FunctionType, TypeAttr],
49+
sym_name: Optional[Union[str, StringAttr]] = None,
50+
kernel: Optional[bool] = None,
51+
workgroup_attrib_attrs: Optional[Sequence[dict]] = None,
52+
private_attrib_attrs: Optional[Sequence[dict]] = None,
53+
known_block_size: Optional[Union[Sequence[int], DenseI32ArrayAttr]] = None,
54+
known_grid_size: Optional[Union[Sequence[int], DenseI32ArrayAttr]] = None,
55+
loc=None,
56+
ip=None,
57+
body_builder: Optional[Callable[[GPUFuncOp], None]] = None,
58+
):
59+
"""
60+
Create a GPUFuncOp with the provided `function_type`, `sym_name`,
61+
`kernel`, `workgroup_attrib_attrs`, `private_attrib_attrs`, `known_block_size`,
62+
`known_grid_size`, and `body_builder`.
63+
- `function_type` is a FunctionType or a TypeAttr.
64+
- `sym_name` is a string or a StringAttr representing the function name.
65+
- `kernel` is a boolean representing whether the function is a kernel.
66+
- `workgroup_attrib_attrs` is an optional list of dictionaries.
67+
- `private_attrib_attrs` is an optional list of dictionaries.
68+
- `known_block_size` is an optional list of integers or a DenseI32ArrayAttr representing the known block size.
69+
- `known_grid_size` is an optional list of integers or a DenseI32ArrayAttr representing the known grid size.
70+
- `body_builder` is an optional callback. When provided, a new entry block
71+
is created and the callback is invoked with the new op as argument within
72+
an InsertionPoint context already set for the block. The callback is
73+
expected to insert a terminator in the block.
74+
"""
75+
function_type = (
76+
TypeAttr.get(function_type)
77+
if not isinstance(function_type, TypeAttr)
78+
else function_type
79+
)
80+
super().__init__(
81+
function_type,
82+
workgroup_attrib_attrs=workgroup_attrib_attrs,
83+
private_attrib_attrs=private_attrib_attrs,
84+
loc=loc,
85+
ip=ip,
86+
)
87+
88+
if isinstance(sym_name, str):
89+
self.attributes[self.SYM_NAME_ATTR_NAME] = StringAttr.get(sym_name)
90+
elif isinstance(sym_name, StringAttr):
91+
self.attributes[self.SYM_NAME_ATTR_NAME] = sym_name
92+
else:
93+
raise ValueError("sym_name must be a string or a StringAttr")
94+
95+
if kernel:
96+
self.attributes[self.KERNEL_ATTR_NAME] = UnitAttr.get()
97+
98+
if known_block_size is not None:
99+
if isinstance(known_block_size, Sequence):
100+
block_size = DenseI32ArrayAttr.get(known_block_size)
101+
self.attributes[self.KNOWN_BLOCK_SIZE_ATTR_NAME] = block_size
102+
elif isinstance(known_block_size, DenseI32ArrayAttr):
103+
self.attributes[self.KNOWN_BLOCK_SIZE_ATTR_NAME] = known_block_size
104+
else:
105+
raise ValueError(
106+
"known_block_size must be a list of integers or a DenseI32ArrayAttr"
107+
)
108+
109+
if known_grid_size is not None:
110+
if isinstance(known_grid_size, Sequence):
111+
grid_size = DenseI32ArrayAttr.get(known_grid_size)
112+
self.attributes[self.KNOWN_GRID_SIZE_ATTR_NAME] = grid_size
113+
elif isinstance(known_grid_size, DenseI32ArrayAttr):
114+
self.attributes[self.KNOWN_GRID_SIZE_ATTR_NAME] = known_grid_size
115+
else:
116+
raise ValueError(
117+
"known_grid_size must be a list of integers or a DenseI32ArrayAttr"
118+
)
119+
120+
if body_builder is not None:
121+
with InsertionPoint(self.add_entry_block()):
122+
body_builder(self)
123+
124+
@property
125+
def name(self) -> StringAttr:
126+
return StringAttr(self.attributes[self.SYM_NAME_ATTR_NAME])
127+
128+
@property
129+
def is_kernel(self) -> bool:
130+
return self.KERNEL_ATTR_NAME in self.attributes
131+
132+
def add_entry_block(self) -> Block:
133+
if len(self.body.blocks) > 0:
134+
raise RuntimeError(f"Entry block already exists for {self.name.value}")
135+
136+
function_type = self.function_type.value
137+
return self.body.blocks.append(
138+
*function_type.inputs,
139+
arg_locs=[self.location for _ in function_type.inputs],
140+
)
141+
142+
@property
143+
def entry_block(self) -> Block:
144+
if len(self.body.blocks) == 0:
145+
raise RuntimeError(
146+
f"Entry block does not exist for {self.name.value}."
147+
+ " Do you need to call the add_entry_block() method on this GPUFuncOp?"
148+
)
149+
return self.body.blocks[0]
150+
151+
@property
152+
def arguments(self) -> Sequence[Type]:
153+
return self.function_type.value.inputs

‎mlir/test/python/dialects/gpu/dialect.py‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# RUN: %PYTHON %s | FileCheck %s
22

33
from mlir.ir import *
4+
import mlir.ir as ir
45
import mlir.dialects.gpu as gpu
56
import mlir.dialects.gpu.passes
67
from mlir.passmanager import *
@@ -64,3 +65,95 @@ def testObjectAttr():
6465
# CHECK: #gpu.object<#nvvm.target, kernels = <[#gpu.kernel_metadata<"kernel", () -> ()>]>, "BC\C0\DE5\14\00\00\05\00\00\00b\0C0$MY\BEf">
6566
print(o)
6667
assert o.kernels == kernelTable
68+
69+
70+
# CHECK-LABEL: testGPUFuncOp
71+
@run
72+
def testGPUFuncOp():
73+
assert gpu.GPUFuncOp.__doc__ is not None
74+
module = Module.create()
75+
with InsertionPoint(module.body):
76+
gpu_module_name = StringAttr.get("gpu_module")
77+
gpumodule = gpu.GPUModuleOp(gpu_module_name)
78+
block = gpumodule.bodyRegion.blocks.append()
79+
80+
def builder(func: gpu.GPUFuncOp) -> None:
81+
gpu.GlobalIdOp(gpu.Dimension.x)
82+
gpu.ReturnOp([])
83+
84+
with InsertionPoint(block):
85+
name = StringAttr.get("kernel0")
86+
func_type = ir.FunctionType.get(inputs=[], results=[])
87+
type_attr = TypeAttr.get(func_type)
88+
func = gpu.GPUFuncOp(type_attr, name)
89+
func.attributes["sym_name"] = name
90+
func.attributes["gpu.kernel"] = UnitAttr.get()
91+
92+
try:
93+
func.entry_block
94+
assert False, "Expected RuntimeError"
95+
except RuntimeError as e:
96+
assert (
97+
str(e)
98+
== "Entry block does not exist for kernel0. Do you need to call the add_entry_block() method on this GPUFuncOp?"
99+
)
100+
101+
block = func.add_entry_block()
102+
with InsertionPoint(block):
103+
builder(func)
104+
105+
try:
106+
func.add_entry_block()
107+
assert False, "Expected RuntimeError"
108+
except RuntimeError as e:
109+
assert str(e) == "Entry block already exists for kernel0"
110+
111+
func = gpu.GPUFuncOp(
112+
func_type,
113+
sym_name="kernel1",
114+
kernel=True,
115+
body_builder=builder,
116+
known_block_size=[1, 2, 3],
117+
known_grid_size=DenseI32ArrayAttr.get([4, 5, 6]),
118+
)
119+
120+
assert func.name.value == "kernel1"
121+
assert func.function_type.value == func_type
122+
assert func.arg_attrs == None
123+
assert func.res_attrs == None
124+
assert func.arguments == []
125+
assert func.entry_block == func.body.blocks[0]
126+
assert func.is_kernel
127+
assert func.known_block_size == DenseI32ArrayAttr.get(
128+
[1, 2, 3]
129+
), func.known_block_size
130+
assert func.known_grid_size == DenseI32ArrayAttr.get(
131+
[4, 5, 6]
132+
), func.known_grid_size
133+
134+
func = gpu.GPUFuncOp(
135+
func_type,
136+
sym_name="non_kernel_func",
137+
body_builder=builder,
138+
)
139+
assert not func.is_kernel
140+
assert func.known_block_size is None
141+
assert func.known_grid_size is None
142+
143+
print(module)
144+
145+
# CHECK: gpu.module @gpu_module
146+
# CHECK: gpu.func @kernel0() kernel {
147+
# CHECK: %[[VAL_0:.*]] = gpu.global_id x
148+
# CHECK: gpu.return
149+
# CHECK: }
150+
# CHECK: gpu.func @kernel1() kernel attributes
151+
# CHECK-SAME: known_block_size = array<i32: 1, 2, 3>
152+
# CHECK-SAME: known_grid_size = array<i32: 4, 5, 6>
153+
# CHECK: %[[VAL_0:.*]] = gpu.global_id x
154+
# CHECK: gpu.return
155+
# CHECK: }
156+
# CHECK: gpu.func @non_kernel_func() {
157+
# CHECK: %[[VAL_0:.*]] = gpu.global_id x
158+
# CHECK: gpu.return
159+
# CHECK: }

0 commit comments

Comments
 (0)