Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

PyOpenCL Kernel Implementations #799

Draft
wants to merge 1 commit into
base: toast3
Choose a base branch
from
Draft

PyOpenCL Kernel Implementations #799

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/toast/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -169,6 +169,7 @@ add_subdirectory(io)
add_subdirectory(accelerator)
add_subdirectory(tests)
add_subdirectory(jax)
add_subdirectory(opencl)
add_subdirectory(ops)
add_subdirectory(templates)
add_subdirectory(scripts)
Expand Down
8 changes: 8 additions & 0 deletions src/toast/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,14 @@
* Values "0", "false", or "no" will disable runtime support for hybrid GPU pipelines.
* Requires TOAST_GPU_OPENMP or TOAST_GPU_JAX to be enabled.

TOAST_OPENCL=<value>
* Values "1", "true", or "yes" will enable runtime support for pyopencl.
* Requires pyopencl to be available / importable.

TOAST_OPENCL_DEFAULT=<value>
* Default OpenCL device type, where supported values are "CPU", "GPU",
and "OCLGRIND".

OMP_NUM_THREADS=<integer>
* Toast uses OpenMP threading in several places and the concurrency is set by the
usual environment variable.
Expand Down
47 changes: 33 additions & 14 deletions src/toast/_libtoast/ops_pixels_healpix.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -598,7 +598,8 @@ void pixels_healpix_nest_inner(
int64_t n_samp,
int64_t idet,
uint8_t mask,
bool use_flags
bool use_flags,
bool compute_submaps
) {
const double zaxis[3] = {0.0, 0.0, 1.0};
int32_t p_indx = pixel_index[idet];
Expand All @@ -618,8 +619,10 @@ void pixels_healpix_nest_inner(
if (use_flags && ((flags[isamp] & mask) != 0)) {
pixels[poff] = -1;
} else {
sub_map = (int64_t)(pixels[poff] / n_pix_submap);
hsub[sub_map] = 1;
if (compute_submaps) {
sub_map = (int64_t)(pixels[poff] / n_pix_submap);
hsub[sub_map] = 1;
}
}

return;
Expand All @@ -639,7 +642,9 @@ void pixels_healpix_ring_inner(
int64_t n_samp,
int64_t idet,
uint8_t mask,
bool use_flags) {
bool use_flags,
bool compute_submaps
) {
const double zaxis[3] = {0.0, 0.0, 1.0};
int32_t p_indx = pixel_index[idet];
int32_t q_indx = quat_index[idet];
Expand All @@ -658,8 +663,10 @@ void pixels_healpix_ring_inner(
if (use_flags && ((flags[isamp] & mask) != 0)) {
pixels[poff] = -1;
} else {
sub_map = (int64_t)(pixels[poff] / n_pix_submap);
hsub[sub_map] = 1;
if (compute_submaps) {
sub_map = (int64_t)(pixels[poff] / n_pix_submap);
hsub[sub_map] = 1;
}
}

return;
Expand Down Expand Up @@ -1163,6 +1170,7 @@ void init_ops_pixels_healpix(py::module & m) {
int64_t n_pix_submap,
int64_t nside,
bool nest,
bool compute_submaps,
bool use_accel
) {
auto & omgr = OmpManager::get();
Expand Down Expand Up @@ -1195,10 +1203,14 @@ void init_ops_pixels_healpix(py::module & m) {
);
int64_t n_view = temp_shape[0];

// Optionally compute the hit submaps
uint8_t * raw_hsub = extract_buffer <uint8_t> (
hit_submaps, "hit_submaps", 1, temp_shape, {-1}
);
int64_t n_submap = temp_shape[0];
if (! compute_submaps) {
raw_hsub = omgr.null_ptr <uint8_t> ();
}

// Optionally use flags
bool use_flags = true;
Expand All @@ -1225,6 +1237,7 @@ void init_ops_pixels_healpix(py::module & m) {
int64_t * dev_pixels = omgr.device_ptr(raw_pixels);
Interval * dev_intervals = omgr.device_ptr(raw_intervals);
uint8_t * dev_flags = omgr.device_ptr(raw_flags);
uint8_t * dev_hsub = omgr.device_ptr(raw_hsub);

// Make sure the lookup table exists on device
size_t utab_bytes = 0x100 * sizeof(int64_t);
Expand Down Expand Up @@ -1258,9 +1271,9 @@ void init_ops_pixels_healpix(py::module & m) {
n_det, \
n_samp, \
shared_flag_mask, \
compute_submaps, \
use_flags \
) \
map(tofrom : raw_hsub[0 : n_submap])
)
{
if (nest) {
# pragma omp target teams distribute parallel for collapse(3) \
Expand All @@ -1269,6 +1282,7 @@ void init_ops_pixels_healpix(py::module & m) {
dev_pixels, \
dev_quats, \
dev_flags, \
dev_hsub, \
dev_intervals, \
dev_utab \
)
Expand All @@ -1293,14 +1307,15 @@ void init_ops_pixels_healpix(py::module & m) {
raw_pixel_index,
dev_quats,
dev_flags,
raw_hsub,
dev_hsub,
dev_pixels,
n_pix_submap,
adjusted_isamp,
n_samp,
idet,
shared_flag_mask,
use_flags
use_flags,
compute_submaps
);
}
}
Expand All @@ -1312,6 +1327,7 @@ void init_ops_pixels_healpix(py::module & m) {
dev_pixels, \
dev_quats, \
dev_flags, \
dev_hsub, \
dev_intervals, \
dev_utab \
)
Expand All @@ -1335,14 +1351,15 @@ void init_ops_pixels_healpix(py::module & m) {
raw_pixel_index,
dev_quats,
dev_flags,
raw_hsub,
dev_hsub,
dev_pixels,
n_pix_submap,
adjusted_isamp,
n_samp,
idet,
shared_flag_mask,
use_flags
use_flags,
compute_submaps
);
}
}
Expand Down Expand Up @@ -1376,7 +1393,8 @@ void init_ops_pixels_healpix(py::module & m) {
n_samp,
idet,
shared_flag_mask,
use_flags
use_flags,
compute_submaps
);
}
}
Expand Down Expand Up @@ -1404,7 +1422,8 @@ void init_ops_pixels_healpix(py::module & m) {
n_samp,
idet,
shared_flag_mask,
use_flags
use_flags,
compute_submaps
);
}
}
Expand Down
2 changes: 2 additions & 0 deletions src/toast/accelerator/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,11 @@
accel_data_update_device,
accel_data_update_host,
accel_enabled,
accel_wait,
accel_get_device,
use_accel_jax,
use_accel_omp,
use_accel_opencl,
use_hybrid_pipelines,
)
from .kernel_registry import ImplementationType, kernel
Loading