Skip to content

Converted som kernels to return values instead of in-place mutation #287

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
max-models wants to merge 4 commits into from
Closed

Converted som kernels to return values instead of in-place mutation #287

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
0b0d34a
Converted som kernels to return values instead of in-place mutation
max-models Jul 23, 2026
cacbbc8
[preview docs]
max-models Jul 23, 2026
a5649e9
Merge branch 'devel' into 134-convert-all-pyccelized-kernels-to-funct…
max-models Jul 23, 2026
876c3a6
Merge branch 'devel' into 134-convert-all-pyccelized-kernels-to-funct…
max-models Jul 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 14 additions & 10 deletions src/struphy/feec/banded_to_stencil_kernels.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,27 +1,29 @@
import numpy as np
from numpy import mod, shape
from pyccel.decorators import pure


@pure
def band_to_stencil_1d(arr: "float[:, :]", out: "float[:, :]"):
def band_to_stencil_1d(arr: "float[:, :]", p: "int") -> "float[:, :]":
"""Converts the 2d banded matrix arr, of shape (n, m) and with 2*p + 1 < m bands centered around the diagonal,
into the array out, of shape (n, 2*p + 1), which corresponds to the StencilMatrix format.
into an array of shape (n, 2*p + 1), which corresponds to the StencilMatrix format.
"""
s = shape(arr)
p = shape(out)[1] // 2
out = np.zeros((s[0], 2 * p + 1), dtype=float)

for i in range(s[0]):
for j in range(2 * p + 1):
jj = mod(i - p + j, s[1])
out[i, j] = arr[i, jj]

return out


@pure
def band_to_stencil_2d(arr: "float[:, :, :, :]", out: "float[:, :, :, :]"):
def band_to_stencil_2d(arr: "float[:, :, :, :]", p1: "int", p2: "int") -> "float[:, :, :, :]":
"""Converts a 4d banded matrix to StencilMatrix format (see band_to_stencil_1d)."""
s = shape(arr)
p1 = shape(out)[2] // 2
p2 = shape(out)[3] // 2
out = np.zeros((s[0], s[1], 2 * p1 + 1, 2 * p2 + 1), dtype=float)

for i1 in range(s[0]):
for j1 in range(2 * p1 + 1):
Expand All @@ -32,14 +34,14 @@ def band_to_stencil_2d(arr: "float[:, :, :, :]", out: "float[:, :, :, :]"):
jj2 = mod(i2 - p2 + j2, s[3])
out[i1, i2, j1, j2] = arr[i1, i2, jj1, jj2]

return out


@pure
def band_to_stencil_3d(arr: "float[:, :, :, :, :, :]", out: "float[:, :, :, :, :, :]"):
def band_to_stencil_3d(arr: "float[:, :, :, :, :, :]", p1: "int", p2: "int", p3: "int") -> "float[:, :, :, :, :, :]":
"""Converts a 6d banded matrix to StencilMatrix format (see band_to_stencil_1d)."""
s = shape(arr)
p1 = shape(out)[3] // 2
p2 = shape(out)[4] // 2
p3 = shape(out)[5] // 2
out = np.zeros((s[0], s[1], s[2], 2 * p1 + 1, 2 * p2 + 1, 2 * p3 + 1), dtype=float)

for i1 in range(s[0]):
for j1 in range(2 * p1 + 1):
Expand All @@ -53,3 +55,5 @@ def band_to_stencil_3d(arr: "float[:, :, :, :, :, :]", out: "float[:, :, :, :, :
for j3 in range(2 * p3 + 1):
jj3 = mod(i3 - p3 + j3, s[5])
out[i1, i2, i3, j1, j2, j3] = arr[i1, i2, i3, jj1, jj2, jj3]

return out
126 changes: 101 additions & 25 deletions src/struphy/feec/local_projectors_kernels.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
from struphy.kernel_arguments.local_projectors_args_kernels import LocalProjectorsArguments


def compute_shifts(IoH: "bool[:]", p: "int[:]", B_nbasis: "int[:]", shift: "int[:]"):
def compute_shifts(IoH: "bool[:]", p: "int[:]", B_nbasis: "int[:]") -> "int[:]":
"""This function computes by how much we must shift the indices in case we loop over the evaluation points.

Parameters
Expand All @@ -19,9 +19,13 @@ def compute_shifts(IoH: "bool[:]", p: "int[:]", B_nbasis: "int[:]", shift: "int[
B_nbasis: 1d int array
Array with the number of B-splines in each direction.

shifts : 1d int array
Returns
-------
shift : 1d int array
array of 3 ints, each one denotes the amout by which we must shift the indices to loop around the quasi-points for each spatial direction.
"""
shift = zeros(3, dtype=int)

for i, ioh in enumerate(IoH):
# Histopolation
if ioh:
Expand All @@ -33,6 +37,8 @@ def compute_shifts(IoH: "bool[:]", p: "int[:]", B_nbasis: "int[:]", shift: "int[
else:
shift[i] = -2 * B_nbasis[i]

return shift


def get_local_problem_size(periodic: "bool[:]", p: "int[:]", IoH: "bool[:]"):
"""Determines the number of interpolation or histopolation weights present for a fixed index i.
Expand Down Expand Up @@ -82,9 +88,11 @@ def get_dofs_local_1_form_ec_component_weighted(
basis1: "float[:]",
basis2: "float[:]",
arezeroc: "int[:]",
f_eval_aux: "float[:,:,:]",
n1: int,
n2: int,
n3: int,
c: int,
):
) -> "float[:,:,:]":
"""Kernel for evaluating the degrees of freedom for the c-th component of 1-forms. This function is for local commuting projetors.

Parameters
Expand All @@ -107,14 +115,21 @@ def get_dofs_local_1_form_ec_component_weighted(
arezeroc : 1d int array
Array of zeros or ones. A one means that for this particular set of quadrature points, in the c-th direction, the basis function is not zero for at least one of them.

f_eval_aux : 3d float array
Output array where the evaluated degrees of freedom are stored. It is passed to this function with zeros in each entry.
n1, n2, n3 : int
Shape of the output array of evaluated degrees of freedom.

c : int
This int tell us whichone of the three components of the 1-form vector we are dealing with. It must be 0, 1 or 2.

Returns
-------
f_eval_aux : 3d float array
The evaluated degrees of freedom.
"""
p = args_solve.degree[c]

f_eval_aux = zeros((n1, n2, n3), dtype=float)

@spossann spossann Jul 24, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is possible though that we do not want to allocate a new array every time we call this function. Especially for large arrays, which should be allocated only once at the start. In that case one should pass out to the kernel.

As a rule, we should pass out whenever the array has stencil vector size, 1d, 2d or 3d.

@max-models max-models Jul 24, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree, I think I have to rethink this a bit. There will definitely be a lot of cases where we would like to pass an out.

Maybe we should add an outputs: tuple argument to the PyccelKernel class https://github.com/struphy-hub/struphy/blob/59a288f2a1326dc5b4910fe45c36d0936dce2263/src/struphy/utils/pyccel.py

In that way, we can do:

interpolate = Pyccelkernel(
    some_interpolation_kernel,
    outputs=(5,),
)

interpolate(
    x,
    y,
    z,
    basis,
    coeffs,
    out,      # argument 5
)

And inside the __call__ method of Pyccelkernel, we only convert the outputs arrays back to numpy arrays:

def __call__(self, *args: Any, **kwargs: Any) -> Any:
    if self.use_cupy:
        # ...

        result = self._kernel(*args_np, **kwargs_np)

        # Only copy back mutated arrays
        for i in self._outputs:
            if isinstance(args[i], xp.ndarray):
                args[i][...] = xp.asarray(args_np[i])

        # ...

@spossann spossann Jul 24, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes something along these lines is needed 👍

@max-models max-models Jul 24, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok, I will try that!

@max-models max-models Aug 5, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm adding this functionality to cunumpy so that it can be used in both struphy and feectools max-models/cunumpy#22


for i in range(shape(f_eval_aux)[0]):
if c == 0:
computei = arezeroc[i] != 0
Expand Down Expand Up @@ -165,14 +180,18 @@ def get_dofs_local_1_form_ec_component_weighted(
* args_solve.wts2[args_solve.inv_index_translation2[k], ii]
)

return f_eval_aux


@stack_array("shp")
def get_dofs_local_1_form_ec_component(
args_solve: LocalProjectorsArguments,
f3: "float[:,:,:]",
f_eval_aux: "float[:,:,:]",
n1: int,
n2: int,
n3: int,
c: int,
):
) -> "float[:,:,:]":
"""Kernel for evaluating the degrees of freedom for the c-th component of 1-forms. This function is for local commuting projetors.

Parameters
Expand All @@ -183,13 +202,20 @@ def get_dofs_local_1_form_ec_component(
f3 : 3d float array
Evaluation for the c-th component of the 1-form function over all the interpolation points in e_a and e_b (a!=b!=c), as well as all the Gauss-Legendre quadrature point in e_c.

f_eval_aux : 3d float array
Output array where the evaluated degrees of freedom are stored. It is passed to this function with zeros in each entry.
n1, n2, n3 : int
Shape of the output array of evaluated degrees of freedom.

c : int
This integer determines which of the three components of the 1-form vector we are working on. Must be 0,1 or 2.

Returns
-------
f_eval_aux : 3d float array
The evaluated degrees of freedom.
"""

f_eval_aux = zeros((n1, n2, n3), dtype=float)

shp = zeros(3, dtype=int)
shp[:] = shape(f_eval_aux)

Expand Down Expand Up @@ -220,14 +246,18 @@ def get_dofs_local_1_form_ec_component(
for ii in range(p):
f_eval_aux[i, j, k] += f3[i, j, in_start + ii] * wts[inv_index_translation[k], ii]

return f_eval_aux


@stack_array("shp")
def get_dofs_local_2_form_ec_component(
args_solve: LocalProjectorsArguments,
fc: "float[:,:,:]",
f_eval_aux: "float[:,:,:]",
n1: int,
n2: int,
n3: int,
c: int,
):
) -> "float[:,:,:]":
"""Kernel for evaluating the degrees of freedom for the c-th component of 2-forms. This function is for local commuting projetors.

Parameters
Expand All @@ -239,13 +269,20 @@ def get_dofs_local_2_form_ec_component(
Evaluation for the c-th component of the 2-form function over all the interpolation points in e_c, as well as all the Gauss-Legendre quadrature point in e_a and e_b.
Of the two spatial directions different from e_c, e_a is the one with the smaller index, and e_b is the one with the larger index.

f_eval_aux : 3d float array
Output array where the evaluated degrees of freedom are stored. It is passed to this function with zeros in each entry.
n1, n2, n3 : int
Shape of the output array of evaluated degrees of freedom.

c : int
This integer determines which of the three components of the 2-form vector we are working on. Must be 0, 1 or 2.

Returns
-------
f_eval_aux : 3d float array
The evaluated degrees of freedom.
"""

f_eval_aux = zeros((n1, n2, n3), dtype=float)

shp = zeros(3, dtype=int)
shp[:] = shape(f_eval_aux)

Expand Down Expand Up @@ -283,6 +320,8 @@ def get_dofs_local_2_form_ec_component(
* args_solve.wts1[args_solve.inv_index_translation1[j], kk]
)

return f_eval_aux


def get_dofs_local_2_form_ec_component_weighted(
args_solve: LocalProjectorsArguments,
Expand All @@ -292,9 +331,11 @@ def get_dofs_local_2_form_ec_component_weighted(
basis2: "float[:]",
arezero_a: "int[:]",
arezero_b: "int[:]",
f_eval_aux: "float[:,:,:]",
n1: int,
n2: int,
n3: int,
c: int,
):
) -> "float[:,:,:]":
"""Kernel for evaluating the degrees of freedom for the c-th component of 2-forms. This function is for local commuting projetors.

Parameters
Expand All @@ -321,13 +362,20 @@ def get_dofs_local_2_form_ec_component_weighted(
arezero_b : 1d int array
Array zeros or ones. A one means that for this particular set of quadrature points, in the e_b direction, the basis function is not zero for at least one of them.

f_eval_aux : 3d float array
Output array where the evaluated degrees of freedom are stored. It is passed to this function with zeros in each entry.
n1, n2, n3 : int
Shape of the output array of evaluated degrees of freedom.

c : int
This int tell us whichone of the three components of the 2-form vector we are dealing with. It must be 0, 1 or 2.

Returns
-------
f_eval_aux : 3d float array
The evaluated degrees of freedom.
"""

f_eval_aux = zeros((n1, n2, n3), dtype=float)

for i in range(shape(f_eval_aux)[0]):
if c == 0:
computei = abs(basis0[i]) >= 10.0 ** (-16)
Expand Down Expand Up @@ -390,9 +438,17 @@ def get_dofs_local_2_form_ec_component_weighted(
* args_solve.wts1[args_solve.inv_index_translation1[j], kk]
)

return f_eval_aux


@stack_array("shp")
def get_dofs_local_3_form(args_solve: LocalProjectorsArguments, faux: "float[:,:,:]", f_eval: "float[:,:,:]"):
def get_dofs_local_3_form(
args_solve: LocalProjectorsArguments,
faux: "float[:,:,:]",
n1: int,
n2: int,
n3: int,
) -> "float[:,:,:]":
"""Kernel for evaluating the degrees of freedom for 3-forms. This function is for local commuting projetors.

Parameters
Expand All @@ -403,9 +459,16 @@ def get_dofs_local_3_form(args_solve: LocalProjectorsArguments, faux: "float[:,:
faux : 3d float array
Evaluation for the 3-form function over all the Gauss-Legendre quadrature point in e1, e2 and e3.

f_eval : 3d float array
Output array where the evaluated degrees of freedom are stored. It is passed to this function with zeros in each entry.
n1, n2, n3 : int
Shape of the output array of evaluated degrees of freedom.

Returns
-------
f_eval : 3d float array
The evaluated degrees of freedom.
"""
f_eval = zeros((n1, n2, n3), dtype=float)

shp = zeros(3, dtype=int)
shp[:] = shape(f_eval)

Expand All @@ -429,6 +492,8 @@ def get_dofs_local_3_form(args_solve: LocalProjectorsArguments, faux: "float[:,:
* args_solve.wts2[args_solve.inv_index_translation2[k], kk]
)

return f_eval


def get_dofs_local_3_form_weighted(
args_solve: LocalProjectorsArguments,
Expand All @@ -439,8 +504,10 @@ def get_dofs_local_3_form_weighted(
arezero0: "int[:]",
arezero1: "int[:]",
arezero2: "int[:]",
f_eval: "float[:,:,:]",
):
n1: int,
n2: int,
n3: int,
) -> "float[:,:,:]":
"""Kernel for evaluating the degrees of freedom for 3-forms. This function is for local commuting projetors.

Parameters
Expand Down Expand Up @@ -469,10 +536,17 @@ def get_dofs_local_3_form_weighted(
arezero2 : 1d int array
Array of zeros or ones. A one means that for this particular set of quadrature points, in the third direction, the basis function is not zero for at least one of them.

f_eval : 3d float array
Output array where the evaluated degrees of freedom are stored. It is passed to this function with zeros in each entry.
n1, n2, n3 : int
Shape of the output array of evaluated degrees of freedom.

Returns
-------
f_eval : 3d float array
The evaluated degrees of freedom.
"""

f_eval = zeros((n1, n2, n3), dtype=float)

for i in range(shape(f_eval)[0]):
if arezero0[i] != 0:
for j in range(shape(f_eval)[1]):
Expand All @@ -499,6 +573,8 @@ def get_dofs_local_3_form_weighted(
* args_solve.wts2[args_solve.inv_index_translation2[k], kk]
)

return f_eval


# We need a functions that tell us which of the quasi-interpolation points to take for a any given i
def select_quasi_points(i: int, p: int, Nbasis: int, periodic: bool):
Expand Down
Loading
Loading