Skip to content
Snippets Groups Projects
Commit aa384a5b authored by Frederik Hennig's avatar Frederik Hennig
Browse files

fix rank check for linear3d indexing

parent 28859ef6
No related branches found
No related tags found
1 merge request!449GPU Indexing Schemes and Launch Configurations
Pipeline #74225 passed
...@@ -7,6 +7,7 @@ from .parameters import Parameter ...@@ -7,6 +7,7 @@ from .parameters import Parameter
from .kernel import Kernel, GpuKernel from .kernel import Kernel, GpuKernel
from .driver import create_kernel, get_driver from .driver import create_kernel, get_driver
from .functions import Lambda from .functions import Lambda
from .errors import CodegenError
__all__ = [ __all__ = [
"Target", "Target",
...@@ -18,4 +19,5 @@ __all__ = [ ...@@ -18,4 +19,5 @@ __all__ = [
"Lambda", "Lambda",
"create_kernel", "create_kernel",
"get_driver", "get_driver",
"CodegenError",
] ]
...@@ -131,7 +131,7 @@ class ManualLaunchConfiguration(GpuLaunchConfiguration): ...@@ -131,7 +131,7 @@ class ManualLaunchConfiguration(GpuLaunchConfiguration):
class DynamicBlockSizeLaunchConfiguration(GpuLaunchConfiguration): class DynamicBlockSizeLaunchConfiguration(GpuLaunchConfiguration):
"""GPU launch configuration that permits the user to set a block size and dynamically computes the grid size. """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 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. in the kernel's iteration space as follows.
For each dimension :math:`c \\in \\{ x, y, z \\}`, For each dimension :math:`c \\in \\{ x, y, z \\}`,
...@@ -201,7 +201,7 @@ class DynamicBlockSizeLaunchConfiguration(GpuLaunchConfiguration): ...@@ -201,7 +201,7 @@ class DynamicBlockSizeLaunchConfiguration(GpuLaunchConfiguration):
class GpuIndexing: class GpuIndexing:
"""Factory for GPU indexing objects required during code generation. """Factory for GPU indexing objects required during code generation.
This class acts as a helper class for the code generation driver. This class acts as a helper class for the code generation driver.
It produces both the `ThreadMapping` required by the backend, It produces both the `ThreadMapping` required by the backend,
as well as factories for the launch configuration required later by the runtime system. as well as factories for the launch configuration required later by the runtime system.
...@@ -259,6 +259,12 @@ class GpuIndexing: ...@@ -259,6 +259,12 @@ class GpuIndexing:
work_items_expr = self._get_work_items() work_items_expr = self._get_work_items()
rank = len(work_items_expr) 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( num_work_items = cast(
_Dim3Lambda, _Dim3Lambda,
tuple(Lambda.from_expression(self._ctx, wit) for wit in work_items_expr), tuple(Lambda.from_expression(self._ctx, wit) for wit in work_items_expr),
...@@ -328,10 +334,6 @@ class GpuIndexing: ...@@ -328,10 +334,6 @@ class GpuIndexing:
match ispace: match ispace:
case FullIterationSpace(): case FullIterationSpace():
dimensions = ispace.dimensions_in_loop_order()[::-1] 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 from ..backend.ast.analysis import collect_undefined_symbols as collect
......
...@@ -11,7 +11,6 @@ from pystencils import ( ...@@ -11,7 +11,6 @@ from pystencils import (
CreateKernelConfig, CreateKernelConfig,
create_kernel, create_kernel,
Target, Target,
assignment_from_stencil,
) )
from pystencils.slicing import ( from pystencils.slicing import (
...@@ -77,6 +76,17 @@ def test_indexing_options( ...@@ -77,6 +76,17 @@ def test_indexing_options(
cp.testing.assert_allclose(dst_arr, expected) 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(): def test_averaging_kernel():
size = (40, 55) size = (40, 55)
src_arr = np.random.rand(*size) src_arr = np.random.rand(*size)
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment