From 2c0164d1f30f02a84c00bb135d28e063f5438364 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Wed, 25 Sep 2024 15:49:58 -0500 Subject: [PATCH 01/13] implement TransferTo{Host,Device}Mapper --- arraycontext/impl/pytato/utils.py | 60 ++++++++++++++++++++++++++++++- test/test_pytato_arraycontext.py | 51 ++++++++++++++++++++++++++ 2 files changed, 110 insertions(+), 1 deletion(-) diff --git a/arraycontext/impl/pytato/utils.py b/arraycontext/impl/pytato/utils.py index a5582d18..a4316727 100644 --- a/arraycontext/impl/pytato/utils.py +++ b/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 diff --git a/test/test_pytato_arraycontext.py b/test/test_pytato_arraycontext.py index 7922f383..397a78a9 100644 --- a/test/test_pytato_arraycontext.py +++ b/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: From 216ac2dcb4047cc9c446faa863646d0a6b3b2465 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Wed, 25 Sep 2024 16:37:56 -0500 Subject: [PATCH 02/13] do not support transfer to same type --- arraycontext/impl/pytato/utils.py | 40 ++++++++++++++++--------------- test/test_pytato_arraycontext.py | 10 ++++---- 2 files changed, 25 insertions(+), 25 deletions(-) diff --git a/arraycontext/impl/pytato/utils.py b/arraycontext/impl/pytato/utils.py index a4316727..4073c555 100644 --- a/arraycontext/impl/pytato/utils.py +++ b/arraycontext/impl/pytato/utils.py @@ -139,16 +139,17 @@ 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) + if not isinstance(expr.data, np.ndarray): + raise ValueError("TransferToDeviceMapper: tried to transfer data that " + "is already on the device") - return super().map_data_wrapper(expr) + 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) class TransferToHostMapper(CopyMapper): @@ -159,16 +160,17 @@ def __init__(self, queue: CommandQueue, allocator: AllocatorBase = None) -> None 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) + if not isinstance(expr.data, tga.TaggableCLArray): + raise ValueError("TransferToHostMapper: tried to transfer data that " + "is already on the host") + + data = expr.data.get() + 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, queue: CommandQueue, diff --git a/test/test_pytato_arraycontext.py b/test/test_pytato_arraycontext.py index 397a78a9..c9891a06 100644 --- a/test/test_pytato_arraycontext.py +++ b/test/test_pytato_arraycontext.py @@ -211,10 +211,8 @@ def test_transfer(actx_factory): 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 + with pytest.raises(ValueError): + _ahh = transfer_to_host(ah, actx.queue, actx.allocator) ad = transfer_to_device(ah, actx.queue, actx.allocator) assert isinstance(ad.data, TaggableCLArray) @@ -222,8 +220,8 @@ def test_transfer(actx_factory): 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 + with pytest.raises(ValueError): + _add = transfer_to_device(ad, actx.queue, actx.allocator) # }}} From e32ae2010f4f5002965dbee0f0217a183342d8ce Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Wed, 25 Sep 2024 17:14:50 -0500 Subject: [PATCH 03/13] use {to,from}_numpy --- arraycontext/impl/pytato/utils.py | 28 +++++++++++----------------- test/test_pytato_arraycontext.py | 18 +++++++++++------- 2 files changed, 22 insertions(+), 24 deletions(-) diff --git a/arraycontext/impl/pytato/utils.py b/arraycontext/impl/pytato/utils.py index 4073c555..830a3640 100644 --- a/arraycontext/impl/pytato/utils.py +++ b/arraycontext/impl/pytato/utils.py @@ -25,8 +25,6 @@ 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, @@ -41,6 +39,7 @@ 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 @@ -130,20 +129,18 @@ def get_loopy_target(self) -> "lp.PyOpenCLTarget": class TransferToDeviceMapper(CopyMapper): - def __init__(self, queue: CommandQueue, allocator: AllocatorBase = None) -> None: + def __init__(self, actx: ArrayContext) -> None: super().__init__() - self.queue = queue - self.allocator = allocator + 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, np.ndarray): raise ValueError("TransferToDeviceMapper: tried to transfer data that " "is already on the device") - data = tga.to_device(self.queue, expr.data, allocator=self.allocator) + data = self.actx.from_numpy(expr.data).data return DataWrapper( data=data, shape=expr.shape, @@ -153,10 +150,9 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array: class TransferToHostMapper(CopyMapper): - def __init__(self, queue: CommandQueue, allocator: AllocatorBase = None) -> None: + def __init__(self, actx: ArrayContext) -> None: super().__init__() - self.queue = queue - self.allocator = allocator + self.actx = actx def map_data_wrapper(self, expr: DataWrapper) -> Array: import arraycontext.impl.pyopencl.taggable_cl_array as tga @@ -164,7 +160,7 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array: raise ValueError("TransferToHostMapper: tried to transfer data that " "is already on the host") - data = expr.data.get() + data = self.actx.to_numpy(expr.data) return DataWrapper( data=data, shape=expr.shape, @@ -173,14 +169,12 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array: non_equality_tags=expr.non_equality_tags) -def transfer_to_device(expr: ArrayOrNames, queue: CommandQueue, - allocator: AllocatorBase = None) -> ArrayOrNames: - return TransferToDeviceMapper(queue, allocator)(expr) +def transfer_to_device(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames: + return TransferToDeviceMapper(actx)(expr) -def transfer_to_host(expr: ArrayOrNames, queue: CommandQueue, - allocator: AllocatorBase = None) -> ArrayOrNames: - return TransferToHostMapper(queue, allocator)(expr) +def transfer_to_host(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames: + return TransferToHostMapper(actx)(expr) # }}} diff --git a/test/test_pytato_arraycontext.py b/test/test_pytato_arraycontext.py index c9891a06..0ba47b5b 100644 --- a/test/test_pytato_arraycontext.py +++ b/test/test_pytato_arraycontext.py @@ -200,28 +200,32 @@ def test_transfer(actx_factory): # {{{ simple tests - a = actx.from_numpy(np.array([0, 1, 2, 3])) + 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.queue, actx.allocator) + 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.queue, actx.allocator) + _ahh = transfer_to_host(ah, actx) - ad = transfer_to_device(ah, actx.queue, actx.allocator) + 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.queue, actx.allocator) + _add = transfer_to_device(ad, actx) # }}} @@ -231,11 +235,11 @@ def test_transfer(actx_factory): "a_expr": a + 2 }) - dagh = transfer_to_host(dag, actx.queue, actx.allocator) + 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.queue, actx.allocator) + daghd = transfer_to_device(dagh, actx) assert isinstance(daghd["a_expr"].expr.bindings["_in0"].data, TaggableCLArray) # }}} From c61c7b0acf4b8739e966a9524c8e532c03e406c0 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Wed, 25 Sep 2024 17:25:30 -0500 Subject: [PATCH 04/13] lint fixes --- arraycontext/impl/pytato/utils.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/arraycontext/impl/pytato/utils.py b/arraycontext/impl/pytato/utils.py index 830a3640..921aa19a 100644 --- a/arraycontext/impl/pytato/utils.py +++ b/arraycontext/impl/pytato/utils.py @@ -140,9 +140,10 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array: raise ValueError("TransferToDeviceMapper: tried to transfer data that " "is already on the device") - data = self.actx.from_numpy(expr.data).data + new_dw = self.actx.from_numpy(expr.data) + assert isinstance(new_dw, DataWrapper) return DataWrapper( - data=data, + data=new_dw.data, shape=expr.shape, axes=expr.axes, tags=expr.tags, @@ -155,12 +156,15 @@ def __init__(self, actx: ArrayContext) -> None: 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, From b98a5ef242768ceb4416a68a7adbfd8fbfcb6a40 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 11 Oct 2024 09:15:50 -0500 Subject: [PATCH 05/13] add comment --- arraycontext/impl/pytato/utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/arraycontext/impl/pytato/utils.py b/arraycontext/impl/pytato/utils.py index 921aa19a..1eadb1e1 100644 --- a/arraycontext/impl/pytato/utils.py +++ b/arraycontext/impl/pytato/utils.py @@ -140,6 +140,8 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array: raise ValueError("TransferToDeviceMapper: tried to transfer data that " "is already on the device") + # Ideally, this code should just return self.actx.from_numpy(expr.data).tagged(expr.tags), + # but there seems to be no way to transfer the non_equality_tags in that case. new_dw = self.actx.from_numpy(expr.data) assert isinstance(new_dw, DataWrapper) return DataWrapper( From 25a786e94ad34a2c74067328e99e776f32cf8a93 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 11 Oct 2024 09:17:02 -0500 Subject: [PATCH 06/13] ruff --- arraycontext/impl/pytato/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/arraycontext/impl/pytato/utils.py b/arraycontext/impl/pytato/utils.py index 1eadb1e1..87095c97 100644 --- a/arraycontext/impl/pytato/utils.py +++ b/arraycontext/impl/pytato/utils.py @@ -140,7 +140,8 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array: raise ValueError("TransferToDeviceMapper: tried to transfer data that " "is already on the device") - # Ideally, this code should just return self.actx.from_numpy(expr.data).tagged(expr.tags), + # Ideally, this code should just do + # return self.actx.from_numpy(expr.data).tagged(expr.tags), # but there seems to be no way to transfer the non_equality_tags in that case. new_dw = self.actx.from_numpy(expr.data) assert isinstance(new_dw, DataWrapper) From 0f2777cba7c7779b909c6b1ae44dcc710bd05fbc Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 11 Oct 2024 17:43:07 -0500 Subject: [PATCH 07/13] disable spurious pylint warning --- arraycontext/impl/pytato/utils.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/arraycontext/impl/pytato/utils.py b/arraycontext/impl/pytato/utils.py index 87095c97..26e65522 100644 --- a/arraycontext/impl/pytato/utils.py +++ b/arraycontext/impl/pytato/utils.py @@ -145,6 +145,9 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array: # but there seems to be no way to transfer the non_equality_tags in that case. new_dw = self.actx.from_numpy(expr.data) assert isinstance(new_dw, DataWrapper) + + # https://github.com/pylint-dev/pylint/issues/3893 + # pylint: disable=unexpected-keyword-arg return DataWrapper( data=new_dw.data, shape=expr.shape, @@ -168,6 +171,9 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array: data = self.actx.to_numpy(expr.data) assert isinstance(data, np.ndarray) + + # https://github.com/pylint-dev/pylint/issues/3893 + # pylint: disable=unexpected-keyword-arg return DataWrapper( data=data, shape=expr.shape, From b20d3592efef0b8f28388585c2a9971e0572ad39 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 11 Oct 2024 17:57:31 -0500 Subject: [PATCH 08/13] add to docs --- arraycontext/impl/pytato/__init__.py | 6 ++++++ arraycontext/impl/pytato/utils.py | 15 ++++++++++++--- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/arraycontext/impl/pytato/__init__.py b/arraycontext/impl/pytato/__init__.py index 099738a9..5e3cb749 100644 --- a/arraycontext/impl/pytato/__init__.py +++ b/arraycontext/impl/pytato/__init__.py @@ -20,6 +20,12 @@ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ .. automodule:: arraycontext.impl.pytato.compile + + +Utils +^^^^^ + +.. automodule:: arraycontext.impl.pytato.utils """ __copyright__ = """ Copyright (C) 2020-1 University of Illinois Board of Trustees diff --git a/arraycontext/impl/pytato/utils.py b/arraycontext/impl/pytato/utils.py index 26e65522..163893c8 100644 --- a/arraycontext/impl/pytato/utils.py +++ b/arraycontext/impl/pytato/utils.py @@ -22,6 +22,11 @@ THE SOFTWARE. """ +__doc__ = """ +.. autofunction:: transfer_to_device +.. autofunction:: transfer_to_host +""" + from typing import TYPE_CHECKING, Any, Dict, Mapping, Set, Tuple @@ -129,6 +134,7 @@ def get_loopy_target(self) -> "lp.PyOpenCLTarget": class TransferToDeviceMapper(CopyMapper): + """A mapper to transfer all :class:`DataWrapper` instances to the CL device.""" def __init__(self, actx: ArrayContext) -> None: super().__init__() self.actx = actx @@ -157,6 +163,7 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array: class TransferToHostMapper(CopyMapper): + """A mapper to transfer all :class:`DataWrapper` instances to the host.""" def __init__(self, actx: ArrayContext) -> None: super().__init__() self.actx = actx @@ -169,13 +176,13 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array: 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) + np_data = self.actx.to_numpy(expr.data) + assert isinstance(np_data, np.ndarray) # https://github.com/pylint-dev/pylint/issues/3893 # pylint: disable=unexpected-keyword-arg return DataWrapper( - data=data, + data=np_data, shape=expr.shape, axes=expr.axes, tags=expr.tags, @@ -183,10 +190,12 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array: def transfer_to_device(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames: + """Transfer all :class:`DataWrapper` instances in *expr* to the CL device.""" return TransferToDeviceMapper(actx)(expr) def transfer_to_host(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames: + """Transfer all :class:`DataWrapper` instances in *expr* to the host.""" return TransferToHostMapper(actx)(expr) # }}} From 3d22465f1a9d877bb5f076a0b3beab698741c3ac Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 11 Oct 2024 18:00:56 -0500 Subject: [PATCH 09/13] more doc fixes --- arraycontext/impl/numpy/__init__.py | 4 ++-- arraycontext/impl/pytato/__init__.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/arraycontext/impl/numpy/__init__.py b/arraycontext/impl/numpy/__init__.py index f8ba95e3..c2f884a6 100644 --- a/arraycontext/impl/numpy/__init__.py +++ b/arraycontext/impl/numpy/__init__.py @@ -1,10 +1,10 @@ from __future__ import annotations -""" +__doc__ = """ .. currentmodule:: arraycontext -A mod :`numpy`-based array context. +A :mod:`numpy`-based array context. .. autoclass:: NumpyArrayContext """ diff --git a/arraycontext/impl/pytato/__init__.py b/arraycontext/impl/pytato/__init__.py index 5e3cb749..742258d0 100644 --- a/arraycontext/impl/pytato/__init__.py +++ b/arraycontext/impl/pytato/__init__.py @@ -10,7 +10,7 @@ :class:`~arraycontext.PytatoPyOpenCLArrayContext` uses :mod:`pyopencl` to JIT-compile and execute the array expressions. -Following :mod:`pytato`-based array context are provided: +The following :mod:`pytato`-based array context are provided: .. autoclass:: PytatoPyOpenCLArrayContext .. autoclass:: PytatoJAXArrayContext From 6ec56e689db3c37c31d521511086ac8076719977 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 11 Oct 2024 18:06:44 -0500 Subject: [PATCH 10/13] datawrapper doc --- arraycontext/impl/pytato/utils.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/arraycontext/impl/pytato/utils.py b/arraycontext/impl/pytato/utils.py index 163893c8..5f08d5fb 100644 --- a/arraycontext/impl/pytato/utils.py +++ b/arraycontext/impl/pytato/utils.py @@ -134,7 +134,8 @@ def get_loopy_target(self) -> "lp.PyOpenCLTarget": class TransferToDeviceMapper(CopyMapper): - """A mapper to transfer all :class:`DataWrapper` instances to the CL device.""" + """A mapper to transfer all :class:`~pytato.array.DataWrapper` instances to + the CL device.""" def __init__(self, actx: ArrayContext) -> None: super().__init__() self.actx = actx @@ -163,7 +164,8 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array: class TransferToHostMapper(CopyMapper): - """A mapper to transfer all :class:`DataWrapper` instances to the host.""" + """A mapper to transfer all :class:`~pytato.array.DataWrapper` instances to + the host.""" def __init__(self, actx: ArrayContext) -> None: super().__init__() self.actx = actx @@ -190,12 +192,14 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array: def transfer_to_device(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames: - """Transfer all :class:`DataWrapper` instances in *expr* to the CL device.""" + """Transfer all :class:`~pytato.array.DataWrapper` instances in *expr* to + the CL device.""" return TransferToDeviceMapper(actx)(expr) def transfer_to_host(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames: - """Transfer all :class:`DataWrapper` instances in *expr* to the host.""" + """Transfer all :class:`~pytato.array.DataWrapper` instances in *expr* to + the host.""" return TransferToHostMapper(actx)(expr) # }}} From 54666360068986ec29c61d67e1bfc0cc2281140a Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 11 Oct 2024 18:14:30 -0500 Subject: [PATCH 11/13] more doc fixes --- arraycontext/impl/numpy/__init__.py | 2 -- arraycontext/impl/pytato/__init__.py | 8 ++++---- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/arraycontext/impl/numpy/__init__.py b/arraycontext/impl/numpy/__init__.py index c2f884a6..7e0acc52 100644 --- a/arraycontext/impl/numpy/__init__.py +++ b/arraycontext/impl/numpy/__init__.py @@ -63,8 +63,6 @@ class NumpyNonObjectArray(metaclass=NumpyNonObjectArrayMetaclass): class NumpyArrayContext(ArrayContext): """ A :class:`ArrayContext` that uses :class:`numpy.ndarray` to represent arrays. - - .. automethod:: __init__ """ _loopy_transform_cache: dict[lp.TranslationUnit, lp.ExecutorBase] diff --git a/arraycontext/impl/pytato/__init__.py b/arraycontext/impl/pytato/__init__.py index 742258d0..82ce84e2 100644 --- a/arraycontext/impl/pytato/__init__.py +++ b/arraycontext/impl/pytato/__init__.py @@ -4,13 +4,13 @@ __doc__ = """ .. currentmodule:: arraycontext -A :mod:`pytato`-based array context defers the evaluation of an array until its +A :mod:`pytato`-based array context defers the evaluation of an array until it is frozen. The execution contexts for the evaluations are specific to an -:class:`~arraycontext.ArrayContext` type. For ex. +:class:`~arraycontext.ArrayContext` type. For example, :class:`~arraycontext.PytatoPyOpenCLArrayContext` uses :mod:`pyopencl` to JIT-compile and execute the array expressions. -The following :mod:`pytato`-based array context are provided: +The following :mod:`pytato`-based array contexts are provided: .. autoclass:: PytatoPyOpenCLArrayContext .. autoclass:: PytatoJAXArrayContext @@ -242,7 +242,7 @@ def get_target(self): class PytatoPyOpenCLArrayContext(_BasePytatoArrayContext): """ - A :class:`ArrayContext` that uses :mod:`pytato` data types to represent + An :class:`ArrayContext` that uses :mod:`pytato` data types to represent the arrays targeting OpenCL for offloading operations. .. attribute:: queue From 55bc57ee13d2d268f3ab86d1ded9fa467d49f3cf Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Thu, 14 Nov 2024 08:30:36 -0600 Subject: [PATCH 12/13] Improve mapper names --- arraycontext/impl/pytato/utils.py | 46 ++++++++++++++++++------------- test/test_pytato_arraycontext.py | 14 +++++----- 2 files changed, 34 insertions(+), 26 deletions(-) diff --git a/arraycontext/impl/pytato/utils.py b/arraycontext/impl/pytato/utils.py index 07c24019..5b059992 100644 --- a/arraycontext/impl/pytato/utils.py +++ b/arraycontext/impl/pytato/utils.py @@ -1,6 +1,6 @@ __doc__ = """ -.. autofunction:: transfer_to_device -.. autofunction:: transfer_to_host +.. autofunction:: transfer_from_numpy +.. autofunction:: transfer_to_numpy """ @@ -133,12 +133,14 @@ def get_loopy_target(self) -> "lp.PyOpenCLTarget": # }}} -# {{{ Transfer mappers +# {{{ Transfer mappers -class TransferToDeviceMapper(CopyMapper): - """A mapper to transfer all :class:`~pytato.array.DataWrapper` instances to - the CL device.""" +class TransferFromNumpyMapper(CopyMapper): + """A mapper to transfer arrays contained in :class:`~pytato.array.DataWrapper` + instances to be device arrays, using + :meth:`~arraycontext.ArrayContext.from_numpy`. + """ def __init__(self, actx: ArrayContext) -> None: super().__init__() self.actx = actx @@ -147,7 +149,7 @@ 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 " + raise ValueError("TransferFromNumpyMapper: tried to transfer data that " "is already on the device") # Ideally, this code should just do @@ -166,9 +168,11 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array: non_equality_tags=expr.non_equality_tags) -class TransferToHostMapper(CopyMapper): - """A mapper to transfer all :class:`~pytato.array.DataWrapper` instances to - the host.""" +class TransferToNumpyMapper(CopyMapper): + """A mapper to transfer arrays contained in :class:`~pytato.array.DataWrapper` + instances to be :class:`numpy.ndarray` instances, using + :meth:`~arraycontext.ArrayContext.to_numpy`. + """ def __init__(self, actx: ArrayContext) -> None: super().__init__() self.actx = actx @@ -178,7 +182,7 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array: import arraycontext.impl.pyopencl.taggable_cl_array as tga if not isinstance(expr.data, tga.TaggableCLArray): - raise ValueError("TransferToHostMapper: tried to transfer data that " + raise ValueError("TransferToNumpyMapper: tried to transfer data that " "is already on the host") np_data = self.actx.to_numpy(expr.data) @@ -194,16 +198,20 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array: non_equality_tags=expr.non_equality_tags) -def transfer_to_device(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames: - """Transfer all :class:`~pytato.array.DataWrapper` instances in *expr* to - the CL device.""" - return TransferToDeviceMapper(actx)(expr) +def transfer_from_numpy(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames: + """Transfer arrays contained in :class:`~pytato.array.DataWrapper` + instances to be device arrays, using + :meth:`~arraycontext.ArrayContext.from_numpy`. + """ + return TransferFromNumpyMapper(actx)(expr) -def transfer_to_host(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames: - """Transfer all :class:`~pytato.array.DataWrapper` instances in *expr* to - the host.""" - return TransferToHostMapper(actx)(expr) +def transfer_to_numpy(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames: + """Transfer arrays contained in :class:`~pytato.array.DataWrapper` + instances to be :class:`numpy.ndarray` instances, using + :meth:`~arraycontext.ArrayContext.to_numpy`. + """ + return TransferToNumpyMapper(actx)(expr) # }}} diff --git a/test/test_pytato_arraycontext.py b/test/test_pytato_arraycontext.py index 0ba47b5b..a14df50f 100644 --- a/test/test_pytato_arraycontext.py +++ b/test/test_pytato_arraycontext.py @@ -205,18 +205,18 @@ def test_transfer(actx_factory): 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 + from arraycontext.impl.pytato.utils import transfer_from_numpy, transfer_to_numpy - ah = transfer_to_host(a, actx) + ah = transfer_to_numpy(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) + _ahh = transfer_to_numpy(ah, actx) - ad = transfer_to_device(ah, actx) + ad = transfer_from_numpy(ah, actx) assert isinstance(ad.data, TaggableCLArray) assert ad != ah assert ad != a # copied DataWrappers compare unequal @@ -225,7 +225,7 @@ def test_transfer(actx_factory): assert np.array_equal(a.data.get(), ad.data.get()) with pytest.raises(ValueError): - _add = transfer_to_device(ad, actx) + _add = transfer_from_numpy(ad, actx) # }}} @@ -235,11 +235,11 @@ def test_transfer(actx_factory): "a_expr": a + 2 }) - dagh = transfer_to_host(dag, actx) + dagh = transfer_to_numpy(dag, actx) assert dagh != dag assert isinstance(dagh["a_expr"].expr.bindings["_in0"].data, np.ndarray) - daghd = transfer_to_device(dagh, actx) + daghd = transfer_from_numpy(dagh, actx) assert isinstance(daghd["a_expr"].expr.bindings["_in0"].data, TaggableCLArray) # }}} From 734193222f2cc1b87a7f9b496e4dd79c44f7e442 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Thu, 14 Nov 2024 08:30:47 -0600 Subject: [PATCH 13/13] Bring back numpy actx constructor docs --- arraycontext/impl/numpy/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/arraycontext/impl/numpy/__init__.py b/arraycontext/impl/numpy/__init__.py index 7e0acc52..c2f884a6 100644 --- a/arraycontext/impl/numpy/__init__.py +++ b/arraycontext/impl/numpy/__init__.py @@ -63,6 +63,8 @@ class NumpyNonObjectArray(metaclass=NumpyNonObjectArrayMetaclass): class NumpyArrayContext(ArrayContext): """ A :class:`ArrayContext` that uses :class:`numpy.ndarray` to represent arrays. + + .. automethod:: __init__ """ _loopy_transform_cache: dict[lp.TranslationUnit, lp.ExecutorBase]