From 6a6e57f08bec708036387c78919f7d4e028d86d4 Mon Sep 17 00:00:00 2001
From: zy69guqi <richard.angersbach@fau.de>
Date: Thu, 20 Mar 2025 18:19:03 +0100
Subject: [PATCH] Fix wrong rank being used for obtaining default block sizes
---
src/pystencils/codegen/gpu_indexing.py | 4 +++-
1 file changed, 3 insertions(+), 1 deletion(-)
diff --git a/src/pystencils/codegen/gpu_indexing.py b/src/pystencils/codegen/gpu_indexing.py
index 43b612bd7..f5606da02 100644
--- a/src/pystencils/codegen/gpu_indexing.py
+++ b/src/pystencils/codegen/gpu_indexing.py
@@ -260,6 +260,7 @@ class DynamicBlockSizeLaunchConfiguration(GpuLaunchConfiguration):
def __init__(
self,
+ rank: int,
num_work_items: _Dim3Lambda,
hw_props: HardwareProperties,
assume_warp_aligned_block_size: bool,
@@ -270,7 +271,7 @@ class DynamicBlockSizeLaunchConfiguration(GpuLaunchConfiguration):
self._assume_warp_aligned_block_size = assume_warp_aligned_block_size
- default_bs = GpuLaunchConfiguration.get_default_block_size(len(num_work_items))
+ default_bs = GpuLaunchConfiguration.get_default_block_size(rank)
self._default_block_size = default_bs
self._init_block_size: dim3 = default_bs
self._compute_block_size: (
@@ -598,6 +599,7 @@ class GpuIndexing:
def factory():
return DynamicBlockSizeLaunchConfiguration(
+ rank,
num_work_items,
self._hw_props,
self._assume_warp_aligned_block_size,
--
GitLab