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

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

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