Skip to content

Commit 8d25a23

Browse files
ax3lax3l
committed
Helpers: to_numpy/cupy
1 parent 184ce75 commit 8d25a23

File tree

16 files changed

+554
-23
lines changed

16 files changed

+554
-23
lines changed

‎MANIFEST.in‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,10 @@ recursive-include cmake *
66
recursive-include src *
77
recursive-include tests *
88

9+
# avoid accidentially copying compiled Python files
10+
global-exclude */__pycache__/*
11+
global-exclude *.pyc
12+
913
# see .gitignore
1014
prune cmake-build*
1115
prune .spack-env*

‎src/Base/PODVector.cpp‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
#include "pyAMReX.H"
77

88
#include <AMReX_PODVector.H>
9+
#include <AMReX_GpuContainers.H>
910

1011
#include <sstream>
1112

@@ -72,6 +73,16 @@ void make_PODVector(py::module &m, std::string typestr, std::string allocstr)
7273
.def("resize", py::overload_cast<std::size_t, const T&>(&PODVector_type::resize))
7374
.def("reserve", &PODVector_type::reserve)
7475
.def("shrink_to_fit", &PODVector_type::shrink_to_fit)
76+
.def("to_host", [](PODVector_type const & pv) {
77+
PODVector<T, std::allocator<T>> h_data(pv.size());
78+
//py::array_t<T> h_data(pv.size());
79+
amrex::Gpu::copy(amrex::Gpu::deviceToHost,
80+
pv.begin(), pv.end(),
81+
h_data.begin()
82+
//h_data.ptr()
83+
);
84+
return h_data;
85+
})
7586

7687
// front
7788
// back

‎src/Particle/ArrayOfStructs.cpp‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,18 @@ void make_ArrayOfStructs(py::module &m, std::string allocstr)
135135
.def("test_sizes", [](){ })
136136
.def("__setitem__", [](AOSType &aos, int const v, const ParticleType& p){ aos[v] = p; })
137137
.def("__getitem__", [](AOSType &aos, int const v){ return aos[v]; }, py::return_value_policy::reference)
138+
139+
.def("to_host", [](AOSType const & aos) {
140+
ArrayOfStructs<T_ParticleType, std::allocator> h_data;
141+
h_data.resize(aos.size());
142+
//py::array_t<T_ParticleType> h_data(aos.size());
143+
amrex::Gpu::copy(amrex::Gpu::deviceToHost,
144+
aos.begin(), aos.end(),
145+
h_data.begin()
146+
//h_data.ptr()
147+
);
148+
return h_data;
149+
})
138150
;
139151
}
140152

‎src/amrex/Array4.py‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
"""
2+
This file is part of pyAMReX
3+
4+
Copyright 2023 AMReX community
5+
Authors: Axel Huebl
6+
License: BSD-3-Clause-LBNL
7+
"""
8+
9+
10+
def array4_to_numpy(self, copy=False, order="F"):
11+
"""
12+
Provide a Numpy view into an Array4.
13+
14+
Note on the order of indices:
15+
By default, this is as in AMReX in Fortran contiguous order, indexing as
16+
x,y,z. This has performance implications for use in external libraries such
17+
as cupy.
18+
The order="C" option will index as z,y,x and perform better with cupy.
19+
https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074
20+
21+
Parameters
22+
----------
23+
self : amrex.Array4_*
24+
An Array4 class in pyAMReX
25+
copy : bool, optional
26+
Copy the data if true, otherwise create a view (default).
27+
order : string, optional
28+
F order (default) or C. C is faster with external libraries.
29+
30+
Returns
31+
-------
32+
np.array
33+
A numpy n-dimensional array.
34+
"""
35+
import numpy as np
36+
37+
if order == "F":
38+
return np.array(self, copy=copy).T
39+
elif order == "C":
40+
return np.array(self, copy=copy)
41+
else:
42+
raise ValueError("The order argument must be F or C.")
43+
44+
45+
def array4_to_cupy(self, copy=False, order="F"):
46+
"""
47+
Provide a Cupy view into an Array4.
48+
49+
Note on the order of indices:
50+
By default, this is as in AMReX in Fortran contiguous order, indexing as
51+
x,y,z. This has performance implications for use in external libraries such
52+
as cupy.
53+
The order="C" option will index as z,y,x and perform better with cupy.
54+
https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074
55+
56+
Parameters
57+
----------
58+
self : amrex.Array4_*
59+
An Array4 class in pyAMReX
60+
copy : bool, optional
61+
Copy the data if true, otherwise create a view (default).
62+
order : string, optional
63+
F order (default) or C. C is faster with external libraries.
64+
65+
Returns
66+
-------
67+
cupy.array
68+
A cupy n-dimensional array.
69+
70+
Raises
71+
------
72+
ImportError
73+
Raises an exception if cupy is not installed
74+
"""
75+
import cupy as cp
76+
77+
if order == "F":
78+
return cp.array(self, copy=copy).T
79+
elif order == "C":
80+
return cp.array(self, copy=copy)
81+
else:
82+
raise ValueError("The order argument must be F or C.")
83+
84+
85+
def register_Array4_extension(amr):
86+
"""Array4 helper methods"""
87+
import inspect
88+
import sys
89+
90+
# register member functions for every Array4_* type
91+
for _, Array4_type in inspect.getmembers(
92+
sys.modules[amr.__name__],
93+
lambda member: inspect.isclass(member)
94+
and member.__module__ == amr.__name__
95+
and member.__name__.startswith("Array4_"),
96+
):
97+
Array4_type.to_numpy = array4_to_numpy
98+
Array4_type.to_cupy = array4_to_cupy

‎src/amrex/ArrayOfStructs.py‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
"""
2+
This file is part of pyAMReX
3+
4+
Copyright 2023 AMReX community
5+
Authors: Axel Huebl
6+
License: BSD-3-Clause-LBNL
7+
"""
8+
from collections import namedtuple
9+
10+
11+
def aos_to_numpy(self, copy=False):
12+
"""
13+
Provide Numpy views into a ArrayOfStructs.
14+
15+
Parameters
16+
----------
17+
self : amrex.ArrayOfStructs_*
18+
An ArrayOfStructs class in pyAMReX
19+
copy : bool, optional
20+
Copy the data if true, otherwise create a view (default).
21+
22+
Returns
23+
-------
24+
namedtuple
25+
A tuple with real and int components that are each lists
26+
of 1D numpy arrays.
27+
"""
28+
import numpy as np
29+
30+
if self.size() == 0:
31+
raise ValueError("AoS is empty.")
32+
33+
if copy:
34+
# This supports a device-to-host copy.
35+
#
36+
# todo: validate of the to_host() returned object
37+
# lifetime is always managed correctly by
38+
# Python's GC - otherwise copy twice via copy=True
39+
return np.array(self.to_host(), copy=False)
40+
else:
41+
return np.array(self, copy=False)
42+
43+
44+
def aos_to_cupy(self, copy=False):
45+
"""
46+
Provide Cupy views into a ArrayOfStructs.
47+
48+
Parameters
49+
----------
50+
self : amrex.ArrayOfStructs_*
51+
An ArrayOfStructs class in pyAMReX
52+
copy : bool, optional
53+
Copy the data if true, otherwise create a view (default).
54+
55+
Returns
56+
-------
57+
namedtuple
58+
A tuple with real and int components that are each lists
59+
of 1D numpy arrays.
60+
61+
Raises
62+
------
63+
ImportError
64+
Raises an exception if cupy is not installed
65+
"""
66+
import cupy as cp
67+
68+
SoA_cp = namedtuple(type(self).__name__ + "_cp", ["real", "int"])
69+
70+
soa_view = SoA_cp([], [])
71+
72+
if self.size() == 0:
73+
raise ValueError("AoS is empty.")
74+
75+
return cp.array(self, copy=copy)
76+
77+
78+
def register_AoS_extension(amr):
79+
"""ArrayOfStructs helper methods"""
80+
import inspect
81+
import sys
82+
83+
# register member functions for every ArrayOfStructs_* type
84+
for _, AoS_type in inspect.getmembers(
85+
sys.modules[amr.__name__],
86+
lambda member: inspect.isclass(member)
87+
and member.__module__ == amr.__name__
88+
and member.__name__.startswith("ArrayOfStructs_"),
89+
):
90+
AoS_type.to_numpy = aos_to_numpy
91+
AoS_type.to_cupy = aos_to_cupy

‎src/amrex/MultiFab.py‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
"""
2+
This file is part of pyAMReX
3+
4+
Copyright 2023 AMReX community
5+
Authors: Axel Huebl
6+
License: BSD-3-Clause-LBNL
7+
"""
8+
9+
10+
def mf_to_numpy(self, copy=False, order="F"):
11+
"""
12+
Provide a Numpy view into a MultiFab.
13+
14+
Note on the order of indices:
15+
By default, this is as in AMReX in Fortran contiguous order, indexing as
16+
x,y,z. This has performance implications for use in external libraries such
17+
as cupy.
18+
The order="C" option will index as z,y,x and perform better with cupy.
19+
https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074
20+
21+
Parameters
22+
----------
23+
self : amrex.MultiFab
24+
A MultiFab class in pyAMReX
25+
copy : bool, optional
26+
Copy the data if true, otherwise create a view (default).
27+
order : string, optional
28+
F order (default) or C. C is faster with external libraries.
29+
30+
Returns
31+
-------
32+
list of np.array
33+
A list of numpy n-dimensional arrays, for each local block in the
34+
MultiFab.
35+
"""
36+
views = []
37+
for mfi in self:
38+
views.append(self.array(mfi).to_numpy(copy, order))
39+
40+
return views
41+
42+
43+
def mf_to_cupy(self, copy=False, order="F"):
44+
"""
45+
Provide a Cupy view into a MultiFab.
46+
47+
Note on the order of indices:
48+
By default, this is as in AMReX in Fortran contiguous order, indexing as
49+
x,y,z. This has performance implications for use in external libraries such
50+
as cupy.
51+
The order="C" option will index as z,y,x and perform better with cupy.
52+
https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074
53+
54+
Parameters
55+
----------
56+
self : amrex.MultiFab
57+
A MultiFab class in pyAMReX
58+
copy : bool, optional
59+
Copy the data if true, otherwise create a view (default).
60+
order : string, optional
61+
F order (default) or C. C is faster with external libraries.
62+
63+
Returns
64+
-------
65+
list of cupy.array
66+
A list of cupy n-dimensional arrays, for each local block in the
67+
MultiFab.
68+
69+
Raises
70+
------
71+
ImportError
72+
Raises an exception if cupy is not installed
73+
"""
74+
views = []
75+
for mfi in self:
76+
views.append(self.array(mfi).to_cupy(copy, order))
77+
78+
return views
79+
80+
81+
def register_MultiFab_extension(amr):
82+
"""MultiFab helper methods"""
83+
import inspect
84+
import sys
85+
86+
# register member functions for every MultiFab* type
87+
for _, MultiFab_type in inspect.getmembers(
88+
sys.modules[amr.__name__],
89+
lambda member: inspect.isclass(member)
90+
and member.__module__ == amr.__name__
91+
and member.__name__.startswith("MultiFab"),
92+
):
93+
MultiFab_type.to_numpy = mf_to_numpy
94+
MultiFab_type.to_cupy = mf_to_cupy

0 commit comments

Comments
 (0)