From aa384a5b1461e930619d178aecf81542ec422573 Mon Sep 17 00:00:00 2001
From: Frederik Hennig <frederik.hennig@fau.de>
Date: Tue, 18 Feb 2025 15:59:17 +0100
Subject: [PATCH] fix rank check for linear3d indexing
---
src/pystencils/codegen/__init__.py | 2 ++
src/pystencils/codegen/gpu_indexing.py | 14 ++++++++------
tests/kernelcreation/test_gpu.py | 12 +++++++++++-
3 files changed, 21 insertions(+), 7 deletions(-)
diff --git a/src/pystencils/codegen/__init__.py b/src/pystencils/codegen/__init__.py
index d06c18382..1b2cd2ffb 100644
--- a/src/pystencils/codegen/__init__.py
+++ b/src/pystencils/codegen/__init__.py
@@ -7,6 +7,7 @@ from .parameters import Parameter
from .kernel import Kernel, GpuKernel
from .driver import create_kernel, get_driver
from .functions import Lambda
+from .errors import CodegenError
__all__ = [
"Target",
@@ -18,4 +19,5 @@ __all__ = [
"Lambda",
"create_kernel",
"get_driver",
+ "CodegenError",
]
diff --git a/src/pystencils/codegen/gpu_indexing.py b/src/pystencils/codegen/gpu_indexing.py
index afd2958c1..2d22ec624 100644
--- a/src/pystencils/codegen/gpu_indexing.py
+++ b/src/pystencils/codegen/gpu_indexing.py
@@ -131,7 +131,7 @@ class ManualLaunchConfiguration(GpuLaunchConfiguration):
class DynamicBlockSizeLaunchConfiguration(GpuLaunchConfiguration):
"""GPU launch configuration that permits the user to set a block size and dynamically computes the grid size.
-
+
The actual launch grid size is computed from the user-defined ``user_block_size`` and the number of work items
in the kernel's iteration space as follows.
For each dimension :math:`c \\in \\{ x, y, z \\}`,
@@ -201,7 +201,7 @@ class DynamicBlockSizeLaunchConfiguration(GpuLaunchConfiguration):
class GpuIndexing:
"""Factory for GPU indexing objects required during code generation.
-
+
This class acts as a helper class for the code generation driver.
It produces both the `ThreadMapping` required by the backend,
as well as factories for the launch configuration required later by the runtime system.
@@ -259,6 +259,12 @@ class GpuIndexing:
work_items_expr = self._get_work_items()
rank = len(work_items_expr)
+ if rank > 3:
+ raise CodegenError(
+ "Cannot create a launch grid configuration using the Linear3D indexing scheme"
+ f" for a {rank}-dimensional kernel."
+ )
+
num_work_items = cast(
_Dim3Lambda,
tuple(Lambda.from_expression(self._ctx, wit) for wit in work_items_expr),
@@ -328,10 +334,6 @@ class GpuIndexing:
match ispace:
case FullIterationSpace():
dimensions = ispace.dimensions_in_loop_order()[::-1]
- if len(dimensions) > 3:
- raise NotImplementedError(
- f"Cannot create a GPU threads range for an {len(dimensions)}-dimensional iteration space"
- )
from ..backend.ast.analysis import collect_undefined_symbols as collect
diff --git a/tests/kernelcreation/test_gpu.py b/tests/kernelcreation/test_gpu.py
index 75239c9b1..10b37e610 100644
--- a/tests/kernelcreation/test_gpu.py
+++ b/tests/kernelcreation/test_gpu.py
@@ -11,7 +11,6 @@ from pystencils import (
CreateKernelConfig,
create_kernel,
Target,
- assignment_from_stencil,
)
from pystencils.slicing import (
@@ -77,6 +76,17 @@ def test_indexing_options(
cp.testing.assert_allclose(dst_arr, expected)
+def test_invalid_indexing_schemes():
+ src, dst = fields("src, dst: [4D]")
+ asm = Assignment(src.center(0), dst.center(0))
+
+ cfg = CreateKernelConfig(target=Target.CUDA)
+ cfg.gpu.indexing_scheme = "linear3d"
+
+ with pytest.raises(Exception):
+ create_kernel(asm, cfg)
+
+
def test_averaging_kernel():
size = (40, 55)
src_arr = np.random.rand(*size)
--
GitLab