5. pystencils

import numpy as np
import pystencils as ps
import sympy as sp
from pystencils import create_kernel
from pystencils import plot as plt

sp.init_printing()

input_arr = np.random.rand(1024, 1024)
output_arr = np.zeros_like(input_arr)

def numpy_kernel():
    output_arr[1:-1, 1:-1] = (
        input_arr[2:, 1:-1]
        + input_arr[:-2, 1:-1]
        + input_arr[1:-1, 2:]
        + input_arr[1:-1, :-2]
    )

%%timeit
numpy_kernel()

1.77 ms ± 12.6 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

src, dst = ps.fields(src=input_arr, dst=output_arr)

symbolic_description = ps.Assignment(
    dst[0, 0], (src[1, 0] + src[-1, 0] + src[0, 1] + src[0, -1]) / 4
)

symbolic_description

plt.figure(figsize=(3, 3))
ps.stencil.plot_expression(symbolic_description.rhs)

kernel = ps.create_kernel(symbolic_description).compile()

def pystencils_kernel():
    kernel(src=input_arr, dst=output_arr)

%%timeit
pystencils_kernel()

408 μs ± 188 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

my_field = ps.fields("f(3) : double[2D]")

field_access = my_field[1, 0](1)

field_access

isinstance(field_access, sp.Symbol)

True

img_field = ps.fields("img(4): [2D]")

w1, w2 = sp.symbols("w_1 w_2")
color = 2
sobel_x = (
    -w2 * img_field[-1, 0](color)
    - w1 * img_field[-1, -1](color)
    - w1 * img_field[-1, +1](color)
    + w2 * img_field[+1, 0](color)
    + w1 * img_field[+1, -1](color)
    - w1 * img_field[+1, +1](color)
) ** 2

sobel_x

sobel_x = sobel_x.subs(w1, 0.5)

sobel_x

dst_field = ps.fields("dst: [2D]")
update_rule = ps.Assignment(dst_field[0, 0], sobel_x)

update_rule

ast = create_kernel(update_rule, cpu_openmp=False)
compiled_kernel = ast.compile()

src_arr = np.zeros([20, 30])
dst_arr = np.zeros_like(src_arr)

dst, src = ps.fields(dst=dst_arr, src=src_arr)

grad_x, grad_y = sp.symbols("grad_x, grad_y")

symbolic_description = [
    ps.Assignment(grad_x, (src[1, 0] - src[-1, 0]) / 2),
    ps.Assignment(grad_y, (src[0, 1] - src[0, -1]) / 2),
    ps.Assignment(dst[0, 0], grad_x + grad_y),
]
kernel = ps.create_kernel(symbolic_description)

symbolic_description

@ps.kernel
def symbolic_description_using_function():
    grad_x @= (src[1, 0] - src[-1, 0]) / 2
    grad_y @= (src[0, 1] - src[0, -1]) / 2
    dst[0, 0] @= grad_x + grad_y


symbolic_description_using_function

sp.Piecewise((1.0, src[0, 1] > 0), (0.0, True)) + src[1, 0]

@ps.kernel
def kernel_with_piecewise():
    grad_x @= (src[1, 0] - src[-1, 0]) / 2 if src[-1, 0] > 0 else 0.0


kernel_with_piecewise

@ps.kernel
def somewhat_longer_dummy_kernel(s):
    s.a @= src[0, 1] + src[-1, 0]
    s.b @= 2 * src[1, 0] + src[0, -1]
    s.c @= src[0, 1] + 2 * src[1, 0] + src[-1, 0] + src[0, -1] - src[0, 0]
    dst[0, 0] @= s.a + s.b + s.c


ac = ps.AssignmentCollection(
    main_assignments=somewhat_longer_dummy_kernel[-1:],
    subexpressions=somewhat_longer_dummy_kernel[:-1],
)
ac

Subexpressions:

$$a \leftarrow_{} {src}_{(0,1)} + {src}_{(-1,0)}$$

$$b \leftarrow_{} 2 {src}_{(1,0)} + {src}_{(0,-1)}$$

$$c \leftarrow_{} - {src}_{(0,0)} + 2 {src}_{(1,0)} + {src}_{(0,1)} + {src}_{(0,-1)} + {src}_{(-1,0)}$$

Main Assignments:

$${dst}_{(0,0)} \leftarrow_{} a + b + c$$

ac.operation_count

{'adds': 8,
 'muls': 2,
 'divs': 0,
 'sqrts': 0,
 'fast_sqrts': 0,
 'fast_inv_sqrts': 0,
 'fast_div': 0}

opt_ac = ps.simp.subexpression_substitution_in_existing_subexpressions(ac)
opt_ac

Subexpressions:

$$a \leftarrow_{} {src}_{(0,1)} + {src}_{(-1,0)}$$

$$b \leftarrow_{} 2 {src}_{(1,0)} + {src}_{(0,-1)}$$

$$c \leftarrow_{} - {src}_{(0,0)} + a + b$$

Main Assignments:

$${dst}_{(0,0)} \leftarrow_{} a + b + c$$

opt_ac.operation_count

{'adds': 6,
 'muls': 1,
 'divs': 0,
 'sqrts': 0,
 'fast_sqrts': 0,
 'fast_inv_sqrts': 0,
 'fast_div': 0}

kernel = ps.create_kernel(ps.Assignment(dst[0, 0], src[2, 0] + src[-1, 0]))
ps.show_code(kernel)

FUNC_PREFIX void kernel(double * RESTRICT  _data_dst, double * RESTRICT const _data_src)
{
   for (int64_t ctr_0 = 2; ctr_0 < 18; ctr_0 += 1)
   {
      for (int64_t ctr_1 = 2; ctr_1 < 28; ctr_1 += 1)
      {
         _data_dst[30*ctr_0 + ctr_1] = _data_src[30*ctr_0 + ctr_1 + 60] + _data_src[30*ctr_0 + ctr_1 - 30];
      }
   }
}

gl_spec = [
    (0, 2),  # 0 ghost layers at the left, 2 at the right border
    (1, 0),
]  # 1 ghost layer at the lower y, one at the upper y coordinate
kernel = ps.create_kernel(
    ps.Assignment(dst[0, 0], src[2, 0] + src[-1, 0]), ghost_layers=gl_spec
)
ps.show_code(kernel)

FUNC_PREFIX void kernel(double * RESTRICT  _data_dst, double * RESTRICT const _data_src)
{
   for (int64_t ctr_0 = 0; ctr_0 < 18; ctr_0 += 1)
   {
      for (int64_t ctr_1 = 1; ctr_1 < 30; ctr_1 += 1)
      {
         _data_dst[30*ctr_0 + ctr_1] = _data_src[30*ctr_0 + ctr_1 + 60] + _data_src[30*ctr_0 + ctr_1 - 30];
      }
   }
}

invalid_description = [
    ps.Assignment(dst[1, 0], src[1, 0] + src[-1, 0]),
    ps.Assignment(dst[0, 0], src[1, 0] - src[-1, 0]),
]
try:
    invalid_kernel = ps.create_kernel(invalid_description)
    assert False, "Should never be executed"
except ValueError as e:
    print(e)

Field dst is written at two different locations

valid_kernel = ps.create_kernel(ps.Assignment(src[0, 0], 2 * src[0, 0] + 42))

v = ps.fields("v(2): double[2D]")
valid_kernel = ps.create_kernel(
    [
        ps.Assignment(v[0, 0](1), 2 * v[0, 0](1) + 42),
        ps.Assignment(v[0, 1](0), 2 * v[1, 0](0) + 42),
    ]
)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[38], line 2
      1 v = ps.fields("v(2): double[2D]")
----> 2 valid_kernel = ps.create_kernel(
      3     [
      4         ps.Assignment(v[0, 0](1), 2 * v[0, 0](1) + 42),
      5         ps.Assignment(v[0, 1](0), 2 * v[1, 0](0) + 42),
      6     ]
      7 )

File /usr/lib/python3.12/site-packages/pystencils/kernelcreation.py:87, in create_kernel(assignments, config, **kwargs)
     85     return create_indexed_kernel(assignments, config=config)
     86 else:
---> 87     return create_domain_kernel(assignments, config=config)

File /usr/lib/python3.12/site-packages/pystencils/kernelcreation.py:132, in create_domain_kernel(assignments, config)
    127 # FUTURE WORK from here we shouldn't NEED sympy
    128 # --- check constrains
    129 check = KernelConstraintsCheck(check_independence_condition=not config.skip_independence_check,
    130                                check_double_write_condition=not config.allow_double_writes)
--> 132 check.visit(assignments)
    134 assignments.bound_fields = check.fields_written
    135 assignments.rhs_fields = check.fields_read

File /usr/lib/python3.12/site-packages/pystencils/kernel_contrains_check.py:49, in KernelConstraintsCheck.visit(self, obj)
     47 def visit(self, obj):
     48     if isinstance(obj, (AssignmentCollection, NodeCollection)):
---> 49         [self.visit(e) for e in obj.all_assignments]
     50     elif isinstance(obj, list) or isinstance(obj, tuple):
     51         [self.visit(e) for e in obj]

File /usr/lib/python3.12/site-packages/pystencils/kernel_contrains_check.py:53, in KernelConstraintsCheck.visit(self, obj)
     51     [self.visit(e) for e in obj]
     52 elif isinstance(obj, (sp.Eq, ast.SympyAssignment, Assignment)):
---> 53     self.process_assignment(obj)
     54 elif isinstance(obj, ast.Conditional):
     55     self.scopes.push()

File /usr/lib/python3.12/site-packages/pystencils/kernel_contrains_check.py:81, in KernelConstraintsCheck.process_assignment(self, assignment)
     78 def process_assignment(self, assignment: Union[sp.Eq, ast.SympyAssignment, Assignment]):
     79     # for checks it is crucial to process rhs before lhs to catch e.g. a = a + 1
     80     self.process_expression(assignment.rhs)
---> 81     self.process_lhs(assignment.lhs)

File /usr/lib/python3.12/site-packages/pystencils/kernel_contrains_check.py:102, in KernelConstraintsCheck.process_lhs(self, lhs)
    100 def process_lhs(self, lhs: Union[Field.Access, TypedSymbol, sp.Symbol]):
    101     assert isinstance(lhs, sp.Symbol)
--> 102     self.update_accesses_lhs(lhs)

File /usr/lib/python3.12/site-packages/pystencils/kernel_contrains_check.py:120, in KernelConstraintsCheck.update_accesses_lhs(self, lhs)
    118         if len(reads) > 1 or lhs.offsets != reads[0]:
    119             if self.check_independence_condition:
--> 120                 raise ValueError(f"Field {lhs.field.name} is written at different location than it was read. "
    121                                  f"This means the resulting kernel would not be thread safe")
    122 elif isinstance(lhs, sp.Symbol):
    123     if self.scopes.is_defined_locally(lhs):

ValueError: Field v is written at different location than it was read. This means the resulting kernel would not be thread safe

@ps.kernel
def not_allowed():
    a, b = sp.symbols("a b")
    a @= src[0, 0]
    b @= a + 3
    a @= src[-1, 0]
    dst[0, 0] @= a + b


try:
    ps.create_kernel(not_allowed)
    assert False
except ValueError as e:
    print(e)

Assignments not in SSA form, multiple assignments to a

@ps.kernel
def not_allowed():
    dst[0, 0] @= src[0, 1] + src[1, 0]
    dst[0, 0] @= 2 * dst[0, 0]


try:
    ps.create_kernel(not_allowed)
    assert False
except ValueError as e:
    print(e)

Field dst is written twice at the same location

tmp_var = sp.Symbol("a")


@ps.kernel
def allowed():
    tmp_var @= src[0, 1] + src[1, 0]
    dst[0, 0] @= 2 * tmp_var


ast = ps.create_kernel(allowed)
ps.show_code(ast)

FUNC_PREFIX void kernel(double * RESTRICT  _data_dst, double * RESTRICT const _data_src)
{
   for (int64_t ctr_0 = 1; ctr_0 < 19; ctr_0 += 1)
   {
      for (int64_t ctr_1 = 1; ctr_1 < 29; ctr_1 += 1)
      {
         const double a = _data_src[30*ctr_0 + ctr_1 + 1] + _data_src[30*ctr_0 + ctr_1 + 30];
         _data_dst[30*ctr_0 + ctr_1] = a*2.0;
      }
   }
}