Skip to content
Snippets Groups Projects

Consolidate codegen and JIT modules.

Merged Frederik Hennig requested to merge fhennig/codegen-module into v2.0-dev
All threads resolved!
5 files
+ 33
80
Compare changes
  • Side-by-side
  • Inline
Files
5
@@ -23,6 +23,7 @@ import numpy as np
import matplotlib.pyplot as plt
```
(guide_kernelcreation)=
# Kernel Creation
Once a kernel's assignments are fully assembled, they need to be passed through pystencils' code
@@ -99,12 +100,12 @@ Pystencils supports code generation for a variety of CPU and GPU hardware.
```{code-cell} ipython3
:tags: [remove-cell]
from pystencils.kernelcreation import DefaultKernelCreationDriver
from pystencils.display_utils import show_ir
# from pystencils.kernelcreation import DefaultKernelCreationDriver
# from pystencils.display_utils import show_ir
def _inspect_ir(kernel, cfg=ps.CreateKernelConfig()):
body = DefaultKernelCreationDriver(ps.CreateKernelConfig()).parse_kernel_body(assignments)
show_ir(body)
# def _inspect_ir(kernel, cfg=ps.CreateKernelConfig()):
# body = DefaultKernelCreationDriver(ps.CreateKernelConfig()).parse_kernel_body(assignments)
# show_ir(body)
```
To produce valid output code, the code generator has to figure out the data types of each
@@ -123,28 +124,42 @@ We can observe this behavior by setting up a kernel including several fields wit
```{code-cell} ipython3
from pystencils.sympyextensions import CastFunc
f = ps.fields("f: float64[2D]")
g = ps.fields("g: float32[2D]")
f = ps.fields("f: float32[2D]")
g = ps.fields("g: float16[2D]")
x, y, z = sp.symbols("x, y, z")
assignments = [
ps.Assignment(z, 42),
ps.Assignment(x, f(0) + z),
ps.Assignment(y, g(0)),
ps.Assignment(f(0), CastFunc(y, "float64") + x)
ps.Assignment(x, 42),
ps.Assignment(y, f(0) + x),
ps.Assignment(z, g(0))
]
cfg = ps.CreateKernelConfig(
default_dtype="float32",
index_dtype="int32"
)
kernel = ps.create_kernel(assignments, cfg)
```
Inspecting the above kernel's intermediate representation, we see that `y` has inherited
its data type from the access to `g`.
The symbol `z`, on the other hand, has been assigned the {any}`default_dtype <CreateKernelConfig.default_dtype>`,
which, unless otherwise specified, is `float64`:
Inspecting the above kernel's output code, we see that `x` has received the `float` type,
which was specified via the {py:data}`default_dtype <CreateKernelConfig.default_dtype>` option.
The symbol `y`, on the other hand, has inherited its data type `half` from the access to the field `g`
on its declaration's right-hand side.
Also, we can observe that the loop counters and symbols related to the field's memory layout
are using the `int32` data type, as specified in {py:data}`index_dtype <CreateKernelConfig.index_dtype>`:
:::{admonition} Developer TODO
Don't show the final C code here, but inspect the codegen IR.
That might make it clearer.
:::
```{code-cell} ipython3
:tags: [remove-input]
_inspect_ir(assignments)
ps.show_code(kernel)
```
```{eval-rst}
Loading