Skip to content

Commit 21cd75b

Browse files
bartgolbartgol
committed
EAMxx: make cld_fraction py test work with both numpy and cupy
1 parent 73feae3 commit 21cd75b

File tree

5 files changed

+112
-21
lines changed

5 files changed

+112
-21
lines changed

‎components/eamxx/src/physics/cld_fraction/cld_fraction_cupy.py‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
import cupy as cp
2+
import numpy as np
3+
4+
# Any initialization step can be done here
5+
# This method is called during CldFraction::initialize_impl
6+
def init ():
7+
pass
8+
9+
#########################################################
10+
def get_cu_array(np_arr):
11+
#########################################################
12+
shape = np_arr.shape
13+
dtype = np_arr.dtype
14+
ptr = np_arr.__array_interface__['data'][0]
15+
strides = np_arr.strides
16+
17+
# The exact size here does not really matter, as we are just creating an
18+
# unmanaged mem block, of which we then simply grab the start address.
19+
# Still, use the correct size for code clarity
20+
size = shape[0]*strides[0]
21+
mem = cp.cuda.UnownedMemory(ptr=ptr,owner=None,size=size)
22+
memptr = cp.cuda.MemoryPointer(mem, 0)
23+
24+
return cp.ndarray(shape=shape,dtype=dtype,memptr=memptr,strides=strides)
25+
26+
#########################################################
27+
def main (ice_threshold, ice_4out_threshold,
28+
qi, liq_cld_frac,
29+
ice_cld_frac, tot_cld_frac,
30+
ice_cld_frac_4out, tot_cld_frac_4out):
31+
#########################################################
32+
33+
cu_qi = get_cu_array(qi)
34+
cu_liq_cld_frac = get_cu_array(liq_cld_frac)
35+
cu_ice_cld_frac = get_cu_array(ice_cld_frac)
36+
cu_tot_cld_frac = get_cu_array(tot_cld_frac)
37+
cu_ice_cld_frac_4out = get_cu_array(ice_cld_frac_4out)
38+
cu_tot_cld_frac_4out = get_cu_array(tot_cld_frac_4out)
39+
40+
cu_ice_cld_frac[:] = 0
41+
cu_ice_cld_frac_4out[:] = 0
42+
cu_ice_cld_frac[cu_qi > ice_threshold] = 1
43+
cu_ice_cld_frac_4out[cu_qi > ice_4out_threshold] = 1
44+
45+
cp.maximum(cu_ice_cld_frac,cu_liq_cld_frac, out=cu_tot_cld_frac)
46+
cp.maximum(cu_ice_cld_frac_4out,cu_liq_cld_frac,out=cu_tot_cld_frac_4out)

‎components/eamxx/src/physics/cld_fraction/cld_fraction.py‎ renamed to ‎components/eamxx/src/physics/cld_fraction/cld_fraction_numpy.py‎

File renamed without changes.

‎components/eamxx/src/physics/cld_fraction/eamxx_cld_fraction_process_interface.cpp‎

Lines changed: 30 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -95,17 +95,30 @@ void CldFraction::run_impl (const double /* dt */)
9595
auto tot_cld_frac_4out = get_field_out("cldfrac_tot_for_analysis");
9696
#ifdef EAMXX_HAS_PYTHON
9797
if (has_py_module()) {
98-
// For now, we run Python code only on CPU
99-
const auto& py_qi = get_py_field_host("qi");
100-
const auto& py_liq_cld_frac = get_py_field_host("cldfrac_liq");
101-
const auto& py_ice_cld_frac = get_py_field_host("cldfrac_ice");
102-
const auto& py_tot_cld_frac = get_py_field_host("cldfrac_tot");
103-
const auto& py_ice_cld_frac_4out = get_py_field_host("cldfrac_ice_for_analysis");
104-
const auto& py_tot_cld_frac_4out = get_py_field_host("cldfrac_tot_for_analysis");
105-
106-
// Sync input to host
107-
liq_cld_frac.sync_to_host();
108-
qi.sync_to_host();
98+
pybind11::array py_qi,
99+
py_liq_cld_frac,
100+
py_ice_cld_frac,
101+
py_tot_cld_frac,
102+
py_ice_cld_frac_4out,
103+
py_tot_cld_frac_4out;
104+
105+
if (m_params.get<std::string>("py_backend")=="device") {
106+
py_qi = get_py_field_dev("qi");
107+
py_liq_cld_frac = get_py_field_dev("cldfrac_liq");
108+
py_ice_cld_frac = get_py_field_dev("cldfrac_ice");
109+
py_tot_cld_frac = get_py_field_dev("cldfrac_tot");
110+
py_ice_cld_frac_4out = get_py_field_dev("cldfrac_ice_for_analysis");
111+
py_tot_cld_frac_4out = get_py_field_dev("cldfrac_tot_for_analysis");
112+
} else {
113+
qi.sync_to_host();
114+
liq_cld_frac.sync_to_host();
115+
py_qi = get_py_field_host("qi");
116+
py_liq_cld_frac = get_py_field_host("cldfrac_liq");
117+
py_ice_cld_frac = get_py_field_host("cldfrac_ice");
118+
py_tot_cld_frac = get_py_field_host("cldfrac_tot");
119+
py_ice_cld_frac_4out = get_py_field_host("cldfrac_ice_for_analysis");
120+
py_tot_cld_frac_4out = get_py_field_host("cldfrac_tot_for_analysis");
121+
}
109122

110123
double ice_threshold = m_params.get<double>("ice_cloud_threshold");
111124
double ice_4out_threshold = m_params.get<double>("ice_cloud_for_analysis_threshold");
@@ -116,11 +129,12 @@ void CldFraction::run_impl (const double /* dt */)
116129
py_ice_cld_frac,py_tot_cld_frac,
117130
py_ice_cld_frac_4out,py_tot_cld_frac_4out);
118131

119-
// Sync outputs to dev
120-
ice_cld_frac.sync_to_dev();
121-
tot_cld_frac.sync_to_dev();
122-
ice_cld_frac_4out.sync_to_dev();
123-
tot_cld_frac_4out.sync_to_dev();
132+
if (m_params.get<std::string>("py_backend")=="host") {
133+
ice_cld_frac.sync_to_dev();
134+
tot_cld_frac.sync_to_dev();
135+
ice_cld_frac_4out.sync_to_dev();
136+
tot_cld_frac_4out.sync_to_dev();
137+
}
124138
} else
125139
#endif
126140
{

‎components/eamxx/tests/single-process/cld_fraction/CMakeLists.txt‎

Lines changed: 35 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,14 @@ CreateADUnitTest(cld_fraction_standalone
2020
FIXTURES_SETUP cldfrac_cpp)
2121

2222
if (EAMXX_ENABLE_PYTHON)
23+
include (BuildCprnc)
24+
BuildCprnc()
25+
2326
# Configure yaml files to run directory
2427
set (POSTFIX py)
25-
set (PY_MODULE_NAME "cld_fraction")
28+
set (PY_MODULE_NAME "cld_fraction_numpy")
2629
set (PY_MODULE_PATH ${SCREAM_SOURCE_DIR}/src/physics/cld_fraction)
30+
set (PY_BACKEND "host")
2731
configure_file(${CMAKE_CURRENT_SOURCE_DIR}/input.yaml
2832
${CMAKE_CURRENT_BINARY_DIR}/input_py.yaml)
2933
configure_file(${CMAKE_CURRENT_SOURCE_DIR}/output.yaml
@@ -35,10 +39,7 @@ if (EAMXX_ENABLE_PYTHON)
3539
LABELS cld_fraction physics
3640
FIXTURES_SETUP cldfrac_py)
3741

38-
# Finally, compare output of the two tests
39-
include (BuildCprnc)
40-
BuildCprnc()
41-
42+
# Compare output of py and cpp tests
4243
set (SRC_FILE "cldfrac_standalone_output_cpp.INSTANT.nsteps_x1.np1.${RUN_T0}.nc")
4344
set (TGT_FILE "cldfrac_standalone_output_py.INSTANT.nsteps_x1.np1.${RUN_T0}.nc")
4445
set (TEST_NAME cldfrac_standalone_cpp_vs_py)
@@ -49,6 +50,35 @@ if (EAMXX_ENABLE_PYTHON)
4950
LABELS "cldfrac;infrastructure"
5051
FIXTURES_REQUIRED "cldfrac_py;cldfrac_cpp")
5152

53+
if (Kokkos_ENABLE_CUDA)
54+
# Also run with cupy instead of numpy
55+
set (PY_MODULE_NAME "cld_fraction_cupy")
56+
set (PY_BACKEND "device")
57+
set (POSTFIX "cupy")
58+
configure_file(${CMAKE_CURRENT_SOURCE_DIR}/input.yaml
59+
${CMAKE_CURRENT_BINARY_DIR}/input_cupy.yaml)
60+
configure_file(${CMAKE_CURRENT_SOURCE_DIR}/output.yaml
61+
${CMAKE_CURRENT_BINARY_DIR}/output_cupy.yaml)
62+
63+
# Test the process with cupy impl
64+
CreateUnitTestFromExec(cld_fraction_standalone_cupy cld_fraction_standalone
65+
EXE_ARGS "--args -ifile=input_cupy.yaml"
66+
LABELS cld_fraction physics
67+
FIXTURES_SETUP cldfrac_cupy)
68+
69+
# Compare output of cupy and cpp tests
70+
set (SRC_FILE "cldfrac_standalone_output_cpp.INSTANT.nsteps_x1.np1.${RUN_T0}.nc")
71+
set (TGT_FILE "cldfrac_standalone_output_cupy.INSTANT.nsteps_x1.np1.${RUN_T0}.nc")
72+
set (TEST_NAME cldfrac_standalone_cpp_vs_cupy)
73+
add_test (NAME ${TEST_NAME}
74+
COMMAND cmake -P ${CMAKE_BINARY_DIR}/bin/CprncTest.cmake ${SRC_FILE} ${TGT_FILE}
75+
WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
76+
set_tests_properties(${TEST_NAME} PROPERTIES
77+
LABELS "cldfrac;infrastructure"
78+
FIXTURES_REQUIRED "cldfrac_cupy;cldfrac_cpp")
79+
80+
endif()
81+
5282
# Run an ml emulator for cld-fraction
5383
set (PY_MODULE_NAME "cld_fraction_ml")
5484
set (PY_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR})

‎components/eamxx/tests/single-process/cld_fraction/input.yaml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ eamxx:
1515
ice_cloud_for_analysis_threshold: 1e-5
1616
py_module_name: ${PY_MODULE_NAME}
1717
py_module_path: ${PY_MODULE_PATH}
18+
py_backend: ${PY_BACKEND}
1819

1920
grids_manager:
2021
type: mesh_free

0 commit comments

Comments
 (0)