Skip to content

Draft: Generalise usage of Structs for nested array access

Markus Holzer requested to merge holzer/pystencils:struct into master

In this, MR Structs are introduced in a more general form than they are used in the index kernel. The structs here can hold data and pointers to fields. This makes it possible to iterate over a struct and extract field pointers in each loop iteration. The extracted fields are then updated in the normal loop nest.

The idea can be illustrated in a small example:

import numpy as np
import pystencils as ps
from pystencils.typing import BasicType, FieldPointerSymbol, PointerType
from pystencils.struct import Struct

dtype = BasicType(np.float64)
f = ps.fields(f'f(1): double[3d]')
g = ps.fields(f'g(1): double[3d]')

struct_src = Struct("src")
struct_src.add_member(PointerType(dtype, const=False, restrict=False, double_pointer=True))
struct_dst = Struct("dst")
struct_dst.add_member(PointerType(dtype, const=False, restrict=False, double_pointer=True))

update_rule = [ps.Assignment(FieldPointerSymbol("f", dtype, const=True), struct_src[0]),
               ps.Assignment(FieldPointerSymbol("g", dtype, const=False), struct_dst[0]),
ast = ps.create_kernel(update_rule)

This produces the following C-Code:

FUNC_PREFIX void kernel(double **  _data_dst, double **  _data_src, int64_t const _size_dst, int64_t const _size_f_0, int64_t const _size_f_1, int64_t const _size_f_2, int64_t const _stride_f_0, int64_t const _stride_f_1, int64_t const _stride_f_2, int64_t const _stride_g_0, int64_t const _stride_g_1, int64_t const _stride_g_2)
   for (int64_t ctr_0 = 0; ctr_0 < _size_dst; ctr_0 += 1)
      double * RESTRICT _data_f = _data_src[ctr_0];
      double * RESTRICT  _data_g = _data_dst[ctr_0];
      for (int64_t ctr_1 = 0; ctr_1 < _size_f_0; ctr_1 += 1)
         for (int64_t ctr_2 = 0; ctr_2 < _size_f_1; ctr_2 += 1)
            for (int64_t ctr_3 = 0; ctr_3 < _size_f_2; ctr_3 += 1)
               _data_g[_stride_g_0*ctr_1 + _stride_g_1*ctr_2 + _stride_g_2*ctr_3] = _data_f[_stride_f_0*ctr_1 + _stride_f_1*ctr_2 + _stride_f_2*ctr_3];

Thus the struct is used as a container for an arbitrary number of subarrays that are all updated at once. Since the struct only holds a single pointer per Element in the above example we can represent it as a double pointer **

Edited by Markus Holzer

Merge request reports