Skip to content

Debug UniformGrid class capture in lambda #320

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

Draft
Sichao25 wants to merge 1 commit into
base: develop
Choose a base branch
from
+47 −14
Draft

Debug UniformGrid class capture in lambda #320

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
14 changes: 10 additions & 4 deletions src/pcms/field/evaluator/uniform_grid.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,13 +106,16 @@ class UniformGridPointEvaluator : public PointEvaluator<Real, LayoutPolicy>
}

// Order-1: multilinear interpolation
// Create local copy to avoid capturing reference in device lambdas in some
// compilers
auto grid_copy = grid_;
Kokkos::View<LO**, DeviceMemorySpace> cell_dim_indices("cell_dim_indices",
num_points, Dim);
Kokkos::parallel_for(
"compute_cell_dim_indices",
Kokkos::RangePolicy<DeviceMemorySpace::execution_space>(0, num_points),
KOKKOS_CLASS_LAMBDA(const LO i) {
auto dim_idx = grid_.GetDimensionedIndex(cell_indices(i));
auto dim_idx = grid_copy.GetDimensionedIndex(cell_indices(i));
for (unsigned d = 0; d < Dim; ++d) {
cell_dim_indices(i, d) = dim_idx[d];
}
Expand All @@ -132,7 +135,7 @@ class UniformGridPointEvaluator : public PointEvaluator<Real, LayoutPolicy>
"compute_parametric_coords",
Kokkos::RangePolicy<DeviceMemorySpace::execution_space>(0, num_points),
KOKKOS_CLASS_LAMBDA(const LO i) {
auto cell_bbox = grid_.GetCellBBOX(cell_indices(i));
auto cell_bbox = grid_copy.GetCellBBOX(cell_indices(i));
for (unsigned d = 0; d < Dim; ++d) {
Real coord = coordinates(i, d);
Real cell_min = cell_bbox.center[d] - cell_bbox.half_width[d];
Expand Down Expand Up @@ -256,6 +259,9 @@ class UniformGridEvaluatorFactory : public FieldEvaluatorFactory<Real>
"is_out_of_bounds", num_points);

// Parallel reduction to compute cell indices and count out-of-bounds points
// Create local copy to avoid capturing reference in device lambdas in some
// compilers
auto grid_copy = grid_;
size_t num_out_of_bounds = 0;
Kokkos::parallel_reduce(
"localize_points_on_device",
Expand All @@ -266,13 +272,13 @@ class UniformGridEvaluatorFactory : public FieldEvaluatorFactory<Real>
point[d] = coordinates_device(i, d);
}

bool out_of_bounds = !grid_.IsPointInBounds(point);
bool out_of_bounds = !grid_copy.IsPointInBounds(point);
is_out_of_bounds_device(i) = out_of_bounds;
if (out_of_bounds) {
local_count += 1;
}

cell_indices_device(i) = grid_.ClosestCellID(point);
cell_indices_device(i) = grid_copy.ClosestCellID(point);
},
num_out_of_bounds);

Expand Down
47 changes: 37 additions & 10 deletions src/pcms/field/evaluator/uniform_grid_spline.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,23 @@
namespace pcms
{

namespace detail
{

struct InitCoordsFunctor
{
Kokkos::View<Real*, DeviceMemorySpace> coords;
Real bot_left;
Real delta;

KOKKOS_FUNCTION void operator()(LO i) const
{
coords(i) = bot_left + delta * i;
}
};

} // namespace detail

template <typename LayoutPolicy =
detail::default_layout_for_memory_space_t<DeviceMemorySpace>>
class UniformGridSplinePointEvaluator2D
Expand All @@ -34,12 +51,20 @@ class UniformGridSplinePointEvaluator2D
{
Real dx = grid_.edge_length[0] / grid_.divisions[0];
Real dy = grid_.edge_length[1] / grid_.divisions[1];
for (LO ix = 0; ix <= grid_.divisions[0]; ++ix) {
x_coords_(ix) = grid_.bot_left[0] + dx * ix;
}
for (LO iy = 0; iy <= grid_.divisions[1]; ++iy) {
y_coords_(iy) = grid_.bot_left[1] + dy * iy;
}
Real bot_left_x = grid_.bot_left[0];
Real bot_left_y = grid_.bot_left[1];
LO nx = grid_.divisions[0];
LO ny = grid_.divisions[1];

Kokkos::parallel_for(
"init_x_coords",
Kokkos::RangePolicy<DeviceMemorySpace::execution_space>(0, nx + 1),
detail::InitCoordsFunctor{x_coords_, bot_left_x, dx});

Kokkos::parallel_for(
"init_y_coords",
Kokkos::RangePolicy<DeviceMemorySpace::execution_space>(0, ny + 1),
detail::InitCoordsFunctor{y_coords_, bot_left_y, dy});
}

void Evaluate(
Expand Down Expand Up @@ -149,8 +174,8 @@ class UniformGridSplinePointEvaluator2D
const UniformGrid<2>& grid_;
UniformGridFieldLocalizationHint<2> hint_;
Real fill_value_;
Kokkos::View<Real*, HostMemorySpace> x_coords_;
Kokkos::View<Real*, HostMemorySpace> y_coords_;
Kokkos::View<Real*, DeviceMemorySpace> x_coords_;
Kokkos::View<Real*, DeviceMemorySpace> y_coords_;
};

class UniformGridSplineEvaluatorFactory2D : public FieldEvaluatorFactory<Real>
Expand Down Expand Up @@ -219,6 +244,8 @@ class UniformGridSplineEvaluatorFactory2D : public FieldEvaluatorFactory<Real>

// Parallel reduction to compute cell indices and count out-of-bounds points
size_t num_out_of_bounds = 0;
auto grid_copy =
grid_; // Create local copy to avoid capturing reference in some compilers
Kokkos::parallel_reduce(
"localize_points_on_device_spline",
Kokkos::RangePolicy<DeviceMemorySpace::execution_space>(0, num_points),
Expand All @@ -228,13 +255,13 @@ class UniformGridSplineEvaluatorFactory2D : public FieldEvaluatorFactory<Real>
point[d] = coordinates_d(i, d);
}

bool out_of_bounds = !grid_.IsPointInBounds(point);
bool out_of_bounds = !grid_copy.IsPointInBounds(point);
is_out_of_bounds_device(i) = out_of_bounds;
if (out_of_bounds) {
local_count += 1;
}

cell_indices_device(i) = grid_.ClosestCellID(point);
cell_indices_device(i) = grid_copy.ClosestCellID(point);
},
num_out_of_bounds);

Expand Down
Loading