Account non-contiguous arys in PyOpenCLActx.to_numpy.

Author: kaushikcfd

arraycontext/impl/pyopencl/__init__.py

@@ -37,6 +37,8 @@
 import numpy as np
 from typing_extensions import Self, override
 
+from pytools import memoize_method
+
 from arraycontext.container.traversal import (
     rec_map_array_container,
     rec_map_container,
@@ -62,7 +64,7 @@
 if TYPE_CHECKING:
     from collections.abc import Callable, Mapping
 
-    from numpy.typing import NDArray
+    from numpy.typing import DTypeLike, NDArray
 
     import loopy as lp
     import pyopencl as cl
@@ -263,12 +265,74 @@ def to_numpy(self, array: Array) -> np.ndarray:
     def to_numpy(self, array: ContainerOrScalarT) -> ContainerOrScalarT:
         ...
 
+    @memoize_method
+    def _get_to_numpy_noncontiguous_copy_kernel(
+        self, dtype: DTypeLike, ndim: int
+    ) -> lp.TranslationUnit:
+        """
+        Returns a translation unit containing a loopy kernel that:
+
+        - Accepts a PyOpenCL array ``inp`` with per-axis strides exposed as
+          ``s0, s1, ..., s{ndim-1}``.
+        - Produces a contiguous, row-major (C-order) output array ``out`` of
+          the same shape, with elements copied from the corresponding
+          coordinates in ``inp``.
+        """
+        import loopy as lp
+
+        from arraycontext import make_loopy_program
+
+        inames = tuple(f"i{iaxis}" for iaxis in range(ndim))
+        shape_names = tuple(f"n{iaxis}" for iaxis in range(ndim))
+        domain = (
+            "{ ["
+            + ", ".join(inames)
+            + "] : "
+            + " and ".join(
+                f"0 <= {iname} < {shape_name}"
+                for iname, shape_name in zip(inames, shape_names, strict=True)
+            )
+            + " }"
+        )
+
+        indices = ", ".join(inames)
+        t_unit = make_loopy_program(
+            [domain],
+            [f"out[{indices}] = inp[{indices}]"],
+            kernel_data=[
+                lp.GlobalArg("out", dtype=dtype, shape=lp.auto),
+                lp.GlobalArg(
+                    "inp",
+                    dtype=dtype,
+                    strides=tuple(f"s{i}" for i in range(ndim)),
+                    offset=lp.auto,
+                    shape=lp.auto,
+                ),
+                lp.ValueArg(",".join([f"s{i}" for i in range(ndim)]), dtype=np.int64),
+                lp.ValueArg(",".join([f"n{i}" for i in range(ndim)]), dtype=np.int64),
+            ],
+            name=f"to_numpy_noncontiguous_copy_{ndim}d",
+        )
+        return t_unit
+
     @override
     def to_numpy(self,
                  array: ArrayOrContainerOrScalar
                  ) -> NumpyOrContainerOrScalar:
         def _to_numpy(ary):
-            return ary.get(queue=self.queue)
+            if ary.flags.forc:
+                # pyopencl supports host transfers only for contiguous arrays.
+                return ary.get(queue=self.queue)
+
+            result = self.call_loopy(
+                self._get_to_numpy_noncontiguous_copy_kernel(ary.dtype, ary.ndim),
+                inp=ary,
+                **{
+                    f"s{i}": stride // ary.dtype.itemsize
+                    for i, stride in enumerate(ary.strides)
+                },
+            )["out"]
+            return result.get(queue=self.queue)
 
         return with_array_context(
             self._rec_map_container(_to_numpy, array),

test/test_arraycontext.py

@@ -1659,6 +1659,25 @@ def test_linspace(actx_factory: ArrayContextFactory, args, kwargs):
     assert np.allclose(actx_linspace, np_linspace)
 
 
+# {{{ test_to_numpy_transpose
+
+def test_to_numpy_transpose(actx_factory: ArrayContextFactory):
+    # fails prior to <https://github.com/inducer/arraycontext/pull/357> for
+    # pyopencl actx -- cl_array.Array.transpose generates non-contiguous
+    # arrays requiring non-trivial logic for to host copies.
+    actx = actx_factory()
+    rng = np.random.default_rng()
+    np_ary = rng.random((256, 256, 256))
+    ary = actx.from_numpy(np_ary)
+    axis_perm = (0, 2, 1)
+
+    np.testing.assert_allclose(
+        actx.to_numpy(actx.np.transpose(ary, axis_perm)),
+        np.transpose(np_ary, axis_perm))
+
+# }}}
+
+
 if __name__ == "__main__":
     import sys
     if len(sys.argv) > 1: