Cold-Air-Networks/mirror-llvm-project
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 7

[mlir][python] Add cluster_size to gpu.launch_func python binding (#177811)

Browse source
This commit is contained in:
Ivan Butygin 2026-01-26 13:21:29 +03:00 committed by GitHub
parent 3ed48305ab
commit 77ae87ac07
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 17 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
mlir/python/mlir/dialects/gpu/__init__.py
View file

@ -186,6 +186,7 @@ class LaunchFuncOp(LaunchFuncOp):
async_dependencies: Optional[List[Value]] = None,
dynamic_shared_memory_size: Optional[Value] = None,
async_object=None,
cluster_size: Optional[Tuple[Any, Any, Any]] = None,
*,
loc=None,
ip=None,
@ -202,6 +203,11 @@ class LaunchFuncOp(LaunchFuncOp):
block_size_x, block_size_y, block_size_z = map(
_convert_literal_to_constant, block_size
)
cluster_size_x, cluster_size_y, cluster_size_z = (
map(_convert_literal_to_constant, cluster_size)
if cluster_size
else (None, None, None)
)
super().__init__(
async_token,
@ -214,6 +220,9 @@ class LaunchFuncOp(LaunchFuncOp):
block_size_y,
block_size_z,
kernel_operands,
clusterSizeX=cluster_size_x,
clusterSizeY=cluster_size_y,
clusterSizeZ=cluster_size_z,
dynamicSharedMemorySize=dynamic_shared_memory_size,
asyncObject=async_object,
loc=loc,
@ -229,6 +238,7 @@ def launch_func(
async_dependencies: Optional[List[Value]] = None,
dynamic_shared_memory_size: Optional[Value] = None,
async_object=None,
cluster_size: Optional[Tuple[Any, Any, Any]] = None,
*,
loc=None,
ip=None,
@ -241,6 +251,7 @@ def launch_func(
async_dependencies=async_dependencies,
dynamic_shared_memory_size=dynamic_shared_memory_size,
async_object=async_object,
cluster_size=cluster_size,
loc=loc,
ip=ip,
)

7
mlir/test/python/dialects/gpu/dialect.py
View file

@ -187,6 +187,7 @@ def testGPULaunchFuncOp():
c1 = arith.constant(T.index(), 1)
grid_sizes = (1, 1, 1)
block_sizes = (1, 1, 1)
cluster_sizes = (1, 1, 1)
token = gpu.wait()
token = gpu.launch_func(
async_dependencies=[token],
@ -194,6 +195,7 @@ def testGPULaunchFuncOp():
grid_size=grid_sizes,
block_size=block_sizes,
kernel_operands=[],
cluster_size=cluster_sizes,
)
gpu.wait(async_dependencies=[token])
func.ReturnOp([])
@ -215,7 +217,10 @@ def testGPULaunchFuncOp():
# CHECK: %[[CONSTANT_4:.*]] = arith.constant 1 : index
# CHECK: %[[CONSTANT_5:.*]] = arith.constant 1 : index
# CHECK: %[[CONSTANT_6:.*]] = arith.constant 1 : index
# CHECK: %[[LAUNCH_FUNC_0:.*]] = gpu.launch_func async {{\[}}%[[WAIT_0]]] @gpu_module::@kernel blocks in (%[[CONSTANT_1]], %[[CONSTANT_2]], %[[CONSTANT_3]]) threads in (%[[CONSTANT_4]], %[[CONSTANT_5]], %[[CONSTANT_6]])
# CHECK: %[[CONSTANT_7:.*]] = arith.constant 1 : index
# CHECK: %[[CONSTANT_8:.*]] = arith.constant 1 : index
# CHECK: %[[CONSTANT_9:.*]] = arith.constant 1 : index
# CHECK: %[[LAUNCH_FUNC_0:.*]] = gpu.launch_func async {{\[}}%[[WAIT_0]]] @gpu_module::@kernel clusters in (%[[CONSTANT_7]], %[[CONSTANT_8]], %[[CONSTANT_9]]) blocks in (%[[CONSTANT_1]], %[[CONSTANT_2]], %[[CONSTANT_3]]) threads in (%[[CONSTANT_4]], %[[CONSTANT_5]], %[[CONSTANT_6]])
# CHECK: %[[WAIT_1:.*]] = gpu.wait async {{\[}}%[[LAUNCH_FUNC_0]]]
# CHECK: return
# CHECK: }