add device<->host transfer mappers

arraycontext/impl/pytato/utils.py

@@ -36,9 +36,10 @@
     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 import ArrayContext
 from arraycontext.impl.pyopencl.taggable_cl_array import Axis as ClAxis
 
 
@@ -124,4 +125,61 @@ def get_loopy_target(self) -> "lp.PyOpenCLTarget":
 
 # }}}
 
+# {{{ Transfer mappers
+
+
+class TransferToDeviceMapper(CopyMapper):
+    def __init__(self, actx: ArrayContext) -> None:
+        super().__init__()
+        self.actx = actx
+
+    def map_data_wrapper(self, expr: DataWrapper) -> Array:
+        import numpy as np
+
+        if not isinstance(expr.data, np.ndarray):
+            raise ValueError("TransferToDeviceMapper: tried to transfer data that "
+                             "is already on the device")
+
+        new_dw = self.actx.from_numpy(expr.data)
+        assert isinstance(new_dw, DataWrapper)
+        return DataWrapper(
+            data=new_dw.data,
+            shape=expr.shape,
+            axes=expr.axes,
+            tags=expr.tags,
+            non_equality_tags=expr.non_equality_tags)
+
+
+class TransferToHostMapper(CopyMapper):
+    def __init__(self, actx: ArrayContext) -> None:
+        super().__init__()
+        self.actx = actx
+
+    def map_data_wrapper(self, expr: DataWrapper) -> Array:
+        import numpy as np
+
+        import arraycontext.impl.pyopencl.taggable_cl_array as tga
+        if not isinstance(expr.data, tga.TaggableCLArray):
+            raise ValueError("TransferToHostMapper: tried to transfer data that "
+                             "is already on the host")
+
+        data = self.actx.to_numpy(expr.data)
+        assert isinstance(data, np.ndarray)
+        return DataWrapper(
+            data=data,
+            shape=expr.shape,
+            axes=expr.axes,
+            tags=expr.tags,
+            non_equality_tags=expr.non_equality_tags)
+
+
+def transfer_to_device(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames:
+    return TransferToDeviceMapper(actx)(expr)
+
+
+def transfer_to_host(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames:
+    return TransferToHostMapper(actx)(expr)
+
+# }}}
+
 # vim: foldmethod=marker

test/test_pytato_arraycontext.py

@@ -192,6 +192,59 @@ 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])).tagged(FooTag())
+
+    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)
+    assert ah != a
+    assert a.tags == ah.tags
+    assert a.non_equality_tags == ah.non_equality_tags
+    assert isinstance(ah.data, np.ndarray)
+
+    with pytest.raises(ValueError):
+        _ahh = transfer_to_host(ah, actx)
+
+    ad = transfer_to_device(ah, actx)
+    assert isinstance(ad.data, TaggableCLArray)
+    assert ad != ah
+    assert ad != a  # copied DataWrappers compare unequal
+    assert ad.tags == ah.tags
+    assert ad.non_equality_tags == ah.non_equality_tags
+    assert np.array_equal(a.data.get(), ad.data.get())
+
+    with pytest.raises(ValueError):
+        _add = transfer_to_device(ad, actx)
+
+    # }}}
+
+    # {{{ test with DictOfNamedArrays
+
+    dag = pt.make_dict_of_named_arrays({
+        "a_expr": a + 2
+        })
+
+    dagh = transfer_to_host(dag, actx)
+    assert dagh != dag
+    assert isinstance(dagh["a_expr"].expr.bindings["_in0"].data, np.ndarray)
+
+    daghd = transfer_to_device(dagh, actx)
+    assert isinstance(daghd["a_expr"].expr.bindings["_in0"].data, TaggableCLArray)
+
+    # }}}
+
+
 if __name__ == "__main__":
     import sys
     if len(sys.argv) > 1: