add device<->host transfer mappers

arraycontext/impl/pytato/utils.py

@@ -25,6 +25,8 @@
 
 from typing import TYPE_CHECKING, Any, Dict, Mapping, Set, Tuple
 
+from pyopencl import CommandQueue
+from pyopencl.tools import AllocatorBase
 from pytato.array import (
     AbstractResultWithNamedArrays,
     Array,
@@ -36,7 +38,7 @@
     make_placeholder,
 )
 from pytato.target.loopy import LoopyPyOpenCLTarget
-from pytato.transform import CopyMapper
+from pytato.transform import ArrayOrNames, CopyMapper
 from pytools import UniqueNameGenerator, memoize_method
 
 from arraycontext.impl.pyopencl.taggable_cl_array import Axis as ClAxis
@@ -124,4 +126,60 @@ def get_loopy_target(self) -> "lp.PyOpenCLTarget":
 
 # }}}
 
+# {{{ Transfer mappers
+
+
+class TransferToDeviceMapper(CopyMapper):
+    def __init__(self, queue: CommandQueue, allocator: AllocatorBase = None) -> None:
+        super().__init__()
+        self.queue = queue
+        self.allocator = allocator
+
+    def map_data_wrapper(self, expr: DataWrapper) -> Array:
+        import numpy as np
+
+        import arraycontext.impl.pyopencl.taggable_cl_array as tga
+        if isinstance(expr.data, np.ndarray):
+            data = tga.to_device(self.queue, expr.data, allocator=self.allocator)
+            return DataWrapper(
+                data=data,
+                shape=expr.shape,
+                axes=expr.axes,
+                tags=expr.tags,
+                non_equality_tags=expr.non_equality_tags)
+
+        return super().map_data_wrapper(expr)
+
+
+class TransferToHostMapper(CopyMapper):
+    def __init__(self, queue: CommandQueue, allocator: AllocatorBase = None) -> None:
+        super().__init__()
+        self.queue = queue
+        self.allocator = allocator
+
+    def map_data_wrapper(self, expr: DataWrapper) -> Array:
+        import arraycontext.impl.pyopencl.taggable_cl_array as tga
+        if isinstance(expr.data, tga.TaggableCLArray):
+            data = expr.data.get()
+            return DataWrapper(
+                data=data,
+                shape=expr.shape,
+                axes=expr.axes,
+                tags=expr.tags,
+                non_equality_tags=expr.non_equality_tags)
+
+        return super().map_data_wrapper(expr)
+
+
+def transfer_to_device(expr: ArrayOrNames, queue: CommandQueue,
+                       allocator: AllocatorBase = None) -> ArrayOrNames:
+    return TransferToDeviceMapper(queue, allocator)(expr)
+
+
+def transfer_to_host(expr: ArrayOrNames, queue: CommandQueue,
+                     allocator: AllocatorBase = None) -> ArrayOrNames:
+    return TransferToHostMapper(queue, allocator)(expr)
+
+# }}}
+
 # vim: foldmethod=marker

test/test_pytato_arraycontext.py

@@ -192,6 +192,57 @@ def twice(x):
             assert res == 198
 
 
+def test_transfer(actx_factory):
+    import numpy as np
+
+    import pytato as pt
+    actx = actx_factory()
+
+    # {{{ simple tests
+
+    a = actx.from_numpy(np.array([0, 1, 2, 3]))
+
+    from arraycontext.impl.pyopencl.taggable_cl_array import TaggableCLArray
+    assert isinstance(a.data, TaggableCLArray)
+
+    from arraycontext.impl.pytato.utils import transfer_to_device, transfer_to_host
+
+    ah = transfer_to_host(a, actx.queue, actx.allocator)
+    assert ah != a
+    assert isinstance(ah.data, np.ndarray)
+
+    ahh = transfer_to_host(ah, actx.queue, actx.allocator)
+    assert isinstance(ahh.data, np.ndarray)
+    assert ah != ahh  # copied DataWrappers compare unequal
+    assert ah != a
+
+    ad = transfer_to_device(ah, actx.queue, actx.allocator)
+    assert isinstance(ad.data, TaggableCLArray)
+    assert ad != ah
+    assert ad != a  # copied DataWrappers compare unequal
+    assert np.array_equal(a.data.get(), ad.data.get())
+
+    add = transfer_to_device(ad, actx.queue, actx.allocator)
+    assert add != ad  # copied DataWrappers compare unequal
+
+    # }}}
+
+    # {{{ test with DictOfNamedArrays
+
+    dag = pt.make_dict_of_named_arrays({
+        "a_expr": a + 2
+        })
+
+    dagh = transfer_to_host(dag, actx.queue, actx.allocator)
+    assert dagh != dag
+    assert isinstance(dagh["a_expr"].expr.bindings["_in0"].data, np.ndarray)
+
+    daghd = transfer_to_device(dagh, actx.queue, actx.allocator)
+    assert isinstance(daghd["a_expr"].expr.bindings["_in0"].data, TaggableCLArray)
+
+    # }}}
+
+
 if __name__ == "__main__":
     import sys
     if len(sys.argv) > 1: