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

diversify Target enum; remove Backend enum

parent 2d17a6d1
No related merge requests found
Pipeline #64089 failed with stages
in 1 minute and 34 seconds
"""Module to generate stencil kernels in C or CUDA using sympy expressions and call them as Python functions""" """Module to generate stencil kernels in C or CUDA using sympy expressions and call them as Python functions"""
from .enums import Backend, Target from .enums import Target
from .defaults import DEFAULTS from .defaults import DEFAULTS
from . import fd from . import fd
from . import stencil as stencil from . import stencil as stencil
...@@ -39,7 +39,6 @@ __all__ = [ ...@@ -39,7 +39,6 @@ __all__ = [
"create_kernel", "create_kernel",
"KernelFunction", "KernelFunction",
"Target", "Target",
"Backend",
"show_code", "show_code",
"to_dot", "to_dot",
"get_code_obj", "get_code_obj",
......
...@@ -15,11 +15,8 @@ from .defaults import DEFAULTS ...@@ -15,11 +15,8 @@ from .defaults import DEFAULTS
class CreateKernelConfig: class CreateKernelConfig:
"""Options for create_kernel.""" """Options for create_kernel."""
target: Target = Target.CPU target: Target = Target.GenericCPU
"""The code generation target. """The code generation target."""
TODO: Enhance `Target` from enum to a larger target spec, e.g. including vectorization architecture, ...
"""
jit: JitBase | None = None jit: JitBase | None = None
"""Just-in-time compiler used to compile and load the kernel for invocation from the current Python environment. """Just-in-time compiler used to compile and load the kernel for invocation from the current Python environment.
...@@ -94,12 +91,10 @@ class CreateKernelConfig: ...@@ -94,12 +91,10 @@ class CreateKernelConfig:
# Infer JIT # Infer JIT
if self.jit is None: if self.jit is None:
match self.target: if self.target.is_cpu():
case Target.CPU: from .backend.jit import LegacyCpuJit
from .backend.jit import LegacyCpuJit self.jit = LegacyCpuJit()
else:
self.jit = LegacyCpuJit() raise NotImplementedError(
case _: f"No default JIT compiler implemented yet for target {self.target}"
raise NotImplementedError( )
f"No default JIT compiler implemented yet for target {self.target}"
)
from enum import Enum, auto from enum import Flag, auto
class Target(Enum): class Target(Flag):
""" """
The Target enumeration represents all possible targets that can be used for the code generation. The Target enumeration represents all possible targets that can be used for the code generation.
""" """
CPU = auto()
""" # ------------------ Component Flags - Do Not Use Directly! -------------------------------------------
Target CPU architecture.
_CPU = auto()
_VECTOR = auto()
_X86 = auto()
_SSE = auto()
_AVX = auto()
_AVX512 = auto()
_ARM = auto()
_NEON = auto()
_SVE = auto()
_GPU = auto()
_CUDA = auto()
_AUTOMATIC = auto()
# ------------------ Actual Targets -------------------------------------------------------------------
CurrentCPU = _CPU | _AUTOMATIC
""" """
GPU = auto() Auto-best CPU target.
`CurrentCPU` causes the code generator to automatically select a CPU target according to CPUs found
on the current machine and runtime environment.
""" """
Target GPU architecture.
GenericCPU = _CPU
"""Generic CPU target.
Generate the kernel for a generic multicore CPU architecture. This opens up all architecture-independent
optimizations including OpenMP, but no vectorization.
""" """
CPU = GenericCPU
"""Alias for backward-compatibility"""
class Backend(Enum): X86_SSE = _CPU | _VECTOR | _X86 | _SSE
""" """x86 architecture with SSE vector extensions."""
The Backend enumeration represents all possible backends that can be used for the code generation.
Backends and targets must be combined with care. For example CPU as a target and CUDA as a backend makes no sense. X86_AVX = _CPU | _VECTOR | _X86 | _AVX
""" """x86 architecture with AVX vector extensions."""
C = auto()
X86_AVX512 = _CPU | _VECTOR | _X86 | _AVX512
"""x86 architecture with AVX512 vector extensions."""
ARM_NEON = _CPU | _VECTOR | _ARM | _NEON
"""ARM architecture with NEON vector extensions"""
ARM_SVE = _CPU | _VECTOR | _ARM | _SVE
"""ARM architecture with SVE vector extensions"""
CurrentGPU = _GPU | _AUTOMATIC
""" """
Use the C Backend of pystencils. Auto-best GPU target.
`CurrentGPU` causes the code generator to automatically select a GPU target according to GPU devices
found on the current machine and runtime environment.
""" """
CUDA = auto()
GenericCUDA = _GPU | _CUDA
""" """
Use the CUDA backend to generate code for NVIDIA GPUs. Generic CUDA GPU target.
Generate a CUDA kernel for a generic Nvidia GPU.
""" """
GPU = GenericCUDA
"""Alias for backward compatibility."""
def is_automatic(self) -> bool:
return Target._AUTOMATIC in self
def is_cpu(self) -> bool:
return Target._CPU in self
def is_vector_cpu(self) -> bool:
return self.is_cpu() and Target._VECTOR in self
def is_gpu(self) -> bool:
return Target._GPU in self
...@@ -72,10 +72,8 @@ def create_kernel( ...@@ -72,10 +72,8 @@ def create_kernel(
kernel_body = typify(kernel_body) kernel_body = typify(kernel_body)
match config.target: match config.target:
case Target.CPU: case Target.GenericCPU:
from .backend.platforms import GenericCpu from .backend.platforms import GenericCpu
# TODO: CPU platform should incorporate instruction set info, OpenMP, etc.
platform = GenericCpu(ctx) platform = GenericCpu(ctx)
case _: case _:
# TODO: CUDA/HIP platform # TODO: CUDA/HIP platform
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment