Draft: Generalise usage of Structs for nested array access
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]),
ps.Assignment(g.center, f.center)]
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