Skip to content

Add a class for interfacing with pyccelized kernels #21

Description

@max-models

Something based on this:

class Pyccelkernel:
    def __init__(self, kernel: Callable[..., Any], use_cupy: bool = False) -> None:
        self._kernel = kernel
        self._use_cupy = use_cupy
        if "cupy" in xp.__name__ or "cupy" in xp.ndarray.__module__:
            self._use_cupy = True

    @staticmethod
    def _convert_to_numpy(value: Any, converted_arrays: list[tuple[Any, np.ndarray]]) -> Any:
        if isinstance(value, xp.ndarray):
            value_np = xp.asnumpy(value)
            converted_arrays.append((value, value_np))
            return value_np

        if isinstance(value, tuple):
            return tuple(Pyccelkernel._convert_to_numpy(item, converted_arrays) for item in value)

        if isinstance(value, list):
            return [Pyccelkernel._convert_to_numpy(item, converted_arrays) for item in value]

        if hasattr(value, "__dict__") and value.__class__.__module__.startswith(("struphy.", "feectools.")):
            value_np = copy.copy(value)
            for name, attr in vars(value).items():
                setattr(value_np, name, Pyccelkernel._convert_to_numpy(attr, converted_arrays))
            return value_np

        return value

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        if self.use_cupy:
            # Convert all args from CuPy to NumPy
            converted_args = []
            args_np = [self._convert_to_numpy(x, converted_args) for x in args]

            # Convert all kwargs from CuPy to NumPy
            converted_kwargs = []
            kwargs_np = {k: self._convert_to_numpy(v, converted_kwargs) for k, v in kwargs.items()}

            # Call kernel
            result = self._kernel(*args_np, **kwargs_np)

            # Copy in-place kernel updates back to CuPy arrays.
            for x, x_np in converted_args:
                x[...] = xp.asarray(x_np)
            for v, v_np in converted_kwargs:
                v[...] = xp.asarray(v_np)

            # Convert NumPy arrays back to CuPy
            if result is None:
                return None
            if isinstance(result, tuple):
                return tuple(xp.asarray(r) if isinstance(r, np.ndarray) else r for r in result)
            if isinstance(result, np.ndarray):
                return xp.asarray(result)
            return result

        else:
            return self._kernel(*args, **kwargs)

    @property
    def name(self) -> str:
        return self.kernel.__name__

    @property
    def kernel(self) -> Callable[..., Any]:
        return self._kernel

    @property
    def use_cupy(self) -> bool:
        return self._use_cupy

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions