Converted som kernels to return values instead of in-place mutation - #287
Converted som kernels to return values instead of in-place mutation#287max-models wants to merge 4 commits into
Conversation
|
📖 Docs preview: https://struphy-hub.github.io/docs-preview/pr-287/ |
spossann
left a comment
There was a problem hiding this comment.
We do no want to allocate arrays too often, see my comment below. For stencil date size objects we want to pass out.
| """ | ||
| p = args_solve.degree[c] | ||
|
|
||
| f_eval_aux = zeros((n1, n2, n3), dtype=float) |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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])
# ...There was a problem hiding this comment.
Yes something along these lines is needed 👍
There was a problem hiding this comment.
I'm adding this functionality to cunumpy so that it can be used in both struphy and feectools max-models/cunumpy#22
This PR adds returns to some of the pyccelized kernels instead of updating (small) arrays in-place.
Related to #134, but some kernels don't really fit. So I just converted a few of the obvious kernels.
Why do this? Firstly, I think it makes the code a bit easier to follow, because it's not always obvious which arrays are mutated and which are not. Secondly, for the cupy/numpy switch, we still have to do something to convert the arrays to numpy before calling the kernel. Afterwards, some of the arrays may need to be converted back to cupy, but in principle this is only needed for the mutated arrays, which we don't have a clear way of specifying. If we instead always return one (or multiple) arrays, then it's always clear that the returned numpy array should be converted to cupy afterwards.
For some pyccelized kernels, like accumulation kernels, I don't think it makes sense to return new arrays, since they are too big, so only in-place mutation is realistic. Maybe we can update the
Pyccelkernel.__call__method in way where we can specify which arrays should be mutated and which not. I'm not sure of what is best here.