Skip to content

Commit

Permalink
avoid variables from underscore under common/src/*.hpp
Browse files Browse the repository at this point in the history
  • Loading branch information
Yuuichi Asahi committed Oct 24, 2024
1 parent d620636 commit f0040e7
Show file tree
Hide file tree
Showing 5 changed files with 192 additions and 178 deletions.
42 changes: 21 additions & 21 deletions common/src/KokkosFFT_Helpers.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,26 +13,26 @@
namespace KokkosFFT {
namespace Impl {
template <typename ViewType, std::size_t DIM = 1>
auto get_shift(const ViewType& inout, axis_type<DIM> _axes, int direction = 1) {
auto get_shift(const ViewType& inout, axis_type<DIM> axes, int direction = 1) {
// Convert the input axes to be in the range of [0, rank-1]
std::vector<int> axes;
std::vector<int> non_negative_axes;
for (std::size_t i = 0; i < DIM; i++) {
int axis = KokkosFFT::Impl::convert_negative_axis(inout, _axes.at(i));
axes.push_back(axis);
int axis = KokkosFFT::Impl::convert_negative_axis(inout, axes.at(i));
non_negative_axes.push_back(axis);
}

// Assert if the elements are overlapped
constexpr int rank = ViewType::rank();
KOKKOSFFT_THROW_IF(KokkosFFT::Impl::has_duplicate_values(axes),
KOKKOSFFT_THROW_IF(KokkosFFT::Impl::has_duplicate_values(non_negative_axes),
"Axes overlap");
KOKKOSFFT_THROW_IF(
KokkosFFT::Impl::is_out_of_range_value_included(axes, rank),
KokkosFFT::Impl::is_out_of_range_value_included(non_negative_axes, rank),
"Axes include an out-of-range index."
"Axes must be in the range of [-rank, rank-1].");

axis_type<rank> shift = {0};
for (int i = 0; i < static_cast<int>(DIM); i++) {
int axis = axes.at(i);
int axis = non_negative_axes.at(i);
shift.at(axis) = inout.extent(axis) / 2 * direction;
}
return shift;
Expand All @@ -48,9 +48,9 @@ void roll(const ExecutionSpace& exec_space, ViewType& inout, axis_type<1> shift,
ViewType tmp("tmp", n0);
int len = (n0 - 1) / 2 + 1;

auto [_shift0, _shift1, _shift2] =
auto [s0, s1, s2] =
KokkosFFT::Impl::convert_negative_shift(inout, shift.at(0), 0);
int shift0 = _shift0, shift1 = _shift1, shift2 = _shift2;
int shift0 = s0, shift1 = s1, shift2 = s2;

// shift2 == 0 means shift
if (shift2 == 0) {
Expand Down Expand Up @@ -96,11 +96,11 @@ void roll(const ExecutionSpace& exec_space, ViewType& inout, axis_type<2> shift,
for (int i = 0; static_cast<std::size_t>(i) < DIM1; i++) {
int axis = axes.at(i);

auto [_shift0, _shift1, _shift2] =
auto [s0, s1, s2] =
KokkosFFT::Impl::convert_negative_shift(inout, shift.at(axis), axis);
shift0.at(axis) = _shift0;
shift1.at(axis) = _shift1;
shift2.at(axis) = _shift2;
shift0.at(axis) = s0;
shift1.at(axis) = s1;
shift2.at(axis) = s2;
}

int shift_00 = shift0.at(0), shift_10 = shift0.at(1);
Expand Down Expand Up @@ -230,13 +230,13 @@ void fftshift(const ExecutionSpace& exec_space, ViewType& inout,
"fftshift: View rank must be larger than or equal to 1");

if (axes) {
axis_type<1> _axes{axes.value()};
KokkosFFT::Impl::fftshift_impl(exec_space, inout, _axes);
axis_type<1> tmp_axes{axes.value()};
KokkosFFT::Impl::fftshift_impl(exec_space, inout, tmp_axes);
} else {
constexpr std::size_t rank = ViewType::rank();
constexpr int start = -static_cast<int>(rank);
auto _axes = KokkosFFT::Impl::index_sequence<int, rank, start>();
KokkosFFT::Impl::fftshift_impl(exec_space, inout, _axes);
auto tmp_axes = KokkosFFT::Impl::index_sequence<int, rank, start>();
KokkosFFT::Impl::fftshift_impl(exec_space, inout, tmp_axes);
}
}

Expand Down Expand Up @@ -279,13 +279,13 @@ void ifftshift(const ExecutionSpace& exec_space, ViewType& inout,
static_assert(ViewType::rank() >= 1,
"ifftshift: View rank must be larger than or equal to 1");
if (axes) {
axis_type<1> _axes{axes.value()};
KokkosFFT::Impl::ifftshift_impl(exec_space, inout, _axes);
axis_type<1> tmp_axes{axes.value()};
KokkosFFT::Impl::ifftshift_impl(exec_space, inout, tmp_axes);
} else {
constexpr std::size_t rank = ViewType::rank();
constexpr int start = -static_cast<int>(rank);
auto _axes = KokkosFFT::Impl::index_sequence<int, rank, start>();
KokkosFFT::Impl::ifftshift_impl(exec_space, inout, _axes);
auto tmp_axes = KokkosFFT::Impl::index_sequence<int, rank, start>();
KokkosFFT::Impl::ifftshift_impl(exec_space, inout, tmp_axes);
}
}

Expand Down
49 changes: 27 additions & 22 deletions common/src/KokkosFFT_layouts.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,20 +45,20 @@ auto get_extents(const InViewType& in, const OutViewType& out,

// Get extents for the inner most axes in LayoutRight
// If we allow the FFT on the layoutLeft, this part should be modified
std::vector<int> _in_extents, _out_extents, _fft_extents;
std::vector<int> in_extents_full, out_extents_full, fft_extents_full;
for (std::size_t i = 0; i < rank; i++) {
auto _idx = map.at(i);
auto in_extent = modified_in_shape.at(_idx);
auto out_extent = out.extent(_idx);
_in_extents.push_back(in_extent);
_out_extents.push_back(out_extent);
auto idx = map.at(i);
auto in_extent = modified_in_shape.at(idx);
auto out_extent = out.extent(idx);
in_extents_full.push_back(in_extent);
out_extents_full.push_back(out_extent);

// The extent for transform is always equal to the extent
// of the extent of real type (R2C or C2R)
// For C2C, the in and out extents are the same.
// In the end, we can just use the largest extent among in and out extents.
auto fft_extent = std::max(in_extent, out_extent);
_fft_extents.push_back(fft_extent);
fft_extents_full.push_back(fft_extent);
}

static_assert(!(is_real_v<in_value_type> && is_real_v<out_value_type>),
Expand All @@ -67,11 +67,12 @@ auto get_extents(const InViewType& in, const OutViewType& out,
if constexpr (is_real_v<in_value_type>) {
// Then R2C
if (is_inplace) {
_in_extents.at(inner_most_axis) = _out_extents.at(inner_most_axis) * 2;
in_extents_full.at(inner_most_axis) =
out_extents_full.at(inner_most_axis) * 2;
} else {
KOKKOSFFT_THROW_IF(
_out_extents.at(inner_most_axis) !=
_in_extents.at(inner_most_axis) / 2 + 1,
out_extents_full.at(inner_most_axis) !=
in_extents_full.at(inner_most_axis) / 2 + 1,
"For R2C, the 'output extent' of transform must be equal to "
"'input extent'/2 + 1");
}
Expand All @@ -80,30 +81,34 @@ auto get_extents(const InViewType& in, const OutViewType& out,
if constexpr (is_real_v<out_value_type>) {
// Then C2R
if (is_inplace) {
_out_extents.at(inner_most_axis) = _in_extents.at(inner_most_axis) * 2;
out_extents_full.at(inner_most_axis) =
in_extents_full.at(inner_most_axis) * 2;
} else {
KOKKOSFFT_THROW_IF(
_in_extents.at(inner_most_axis) !=
_out_extents.at(inner_most_axis) / 2 + 1,
in_extents_full.at(inner_most_axis) !=
out_extents_full.at(inner_most_axis) / 2 + 1,
"For C2R, the 'input extent' of transform must be equal to "
"'output extent' / 2 + 1");
}
}

if constexpr (std::is_same_v<array_layout_type, Kokkos::LayoutLeft>) {
std::reverse(_in_extents.begin(), _in_extents.end());
std::reverse(_out_extents.begin(), _out_extents.end());
std::reverse(_fft_extents.begin(), _fft_extents.end());
std::reverse(in_extents_full.begin(), in_extents_full.end());
std::reverse(out_extents_full.begin(), out_extents_full.end());
std::reverse(fft_extents_full.begin(), fft_extents_full.end());
}

// Define subvectors starting from last - DIM
// Dimensions relevant to FFTs
std::vector<int> in_extents(_in_extents.end() - DIM, _in_extents.end());
std::vector<int> out_extents(_out_extents.end() - DIM, _out_extents.end());
std::vector<int> fft_extents(_fft_extents.end() - DIM, _fft_extents.end());

int total_fft_size = std::accumulate(_fft_extents.begin(), _fft_extents.end(),
1, std::multiplies<>());
std::vector<int> in_extents(in_extents_full.end() - DIM,
in_extents_full.end());
std::vector<int> out_extents(out_extents_full.end() - DIM,
out_extents_full.end());
std::vector<int> fft_extents(fft_extents_full.end() - DIM,
fft_extents_full.end());

int total_fft_size = std::accumulate(
fft_extents_full.begin(), fft_extents_full.end(), 1, std::multiplies<>());
int fft_size = std::accumulate(fft_extents.begin(), fft_extents.end(), 1,
std::multiplies<>());
int howmany = total_fft_size / fft_size;
Expand Down
106 changes: 53 additions & 53 deletions common/src/KokkosFFT_padding.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -94,10 +94,10 @@ auto is_crop_or_pad_needed(const ViewType& view,
template <typename ExecutionSpace, typename InViewType, typename OutViewType>
void crop_or_pad_impl(const ExecutionSpace& exec_space, const InViewType& in,
OutViewType& out, shape_type<1> s) {
auto _n0 = s.at(0);
out = OutViewType("out", _n0);
auto s0 = s.at(0);
out = OutViewType("out", s0);

auto n0 = std::min(_n0, in.extent(0));
auto n0 = std::min(s0, in.extent(0));

Kokkos::parallel_for(
"KokkosFFT::crop_or_pad",
Expand All @@ -111,11 +111,11 @@ void crop_or_pad_impl(const ExecutionSpace& exec_space, const InViewType& in,
OutViewType& out, shape_type<2> s) {
constexpr std::size_t DIM = 2;

auto [_n0, _n1] = s;
out = OutViewType("out", _n0, _n1);
auto [s0, s1] = s;
out = OutViewType("out", s0, s1);

int n0 = std::min(_n0, in.extent(0));
int n1 = std::min(_n1, in.extent(1));
int n0 = std::min(s0, in.extent(0));
int n1 = std::min(s1, in.extent(1));

using range_type = Kokkos::MDRangePolicy<
ExecutionSpace,
Expand All @@ -138,12 +138,12 @@ void crop_or_pad_impl(const ExecutionSpace& exec_space, const InViewType& in,
OutViewType& out, shape_type<3> s) {
constexpr std::size_t DIM = 3;

auto [_n0, _n1, _n2] = s;
out = OutViewType("out", _n0, _n1, _n2);
auto [s0, s1, s2] = s;
out = OutViewType("out", s0, s1, s2);

int n0 = std::min(_n0, in.extent(0));
int n1 = std::min(_n1, in.extent(1));
int n2 = std::min(_n2, in.extent(2));
int n0 = std::min(s0, in.extent(0));
int n1 = std::min(s1, in.extent(1));
int n2 = std::min(s2, in.extent(2));

using range_type = Kokkos::MDRangePolicy<
ExecutionSpace,
Expand All @@ -167,13 +167,13 @@ void crop_or_pad_impl(const ExecutionSpace& exec_space, const InViewType& in,
OutViewType& out, shape_type<4> s) {
constexpr std::size_t DIM = 4;

auto [_n0, _n1, _n2, _n3] = s;
out = OutViewType("out", _n0, _n1, _n2, _n3);
auto [s0, s1, s2, s3] = s;
out = OutViewType("out", s0, s1, s2, s3);

int n0 = std::min(_n0, in.extent(0));
int n1 = std::min(_n1, in.extent(1));
int n2 = std::min(_n2, in.extent(2));
int n3 = std::min(_n3, in.extent(3));
int n0 = std::min(s0, in.extent(0));
int n1 = std::min(s1, in.extent(1));
int n2 = std::min(s2, in.extent(2));
int n3 = std::min(s3, in.extent(3));

using range_type = Kokkos::MDRangePolicy<
ExecutionSpace,
Expand All @@ -198,14 +198,14 @@ void crop_or_pad_impl(const ExecutionSpace& exec_space, const InViewType& in,
OutViewType& out, shape_type<5> s) {
constexpr std::size_t DIM = 5;

auto [_n0, _n1, _n2, _n3, _n4] = s;
out = OutViewType("out", _n0, _n1, _n2, _n3, _n4);
auto [s0, s1, s2, s3, s4] = s;
out = OutViewType("out", s0, s1, s2, s3, s4);

int n0 = std::min(_n0, in.extent(0));
int n1 = std::min(_n1, in.extent(1));
int n2 = std::min(_n2, in.extent(2));
int n3 = std::min(_n3, in.extent(3));
int n4 = std::min(_n4, in.extent(4));
int n0 = std::min(s0, in.extent(0));
int n1 = std::min(s1, in.extent(1));
int n2 = std::min(s2, in.extent(2));
int n3 = std::min(s3, in.extent(3));
int n4 = std::min(s4, in.extent(4));

using range_type = Kokkos::MDRangePolicy<
ExecutionSpace,
Expand All @@ -230,15 +230,15 @@ void crop_or_pad_impl(const ExecutionSpace& exec_space, const InViewType& in,
OutViewType& out, shape_type<6> s) {
constexpr std::size_t DIM = 6;

auto [_n0, _n1, _n2, _n3, _n4, _n5] = s;
out = OutViewType("out", _n0, _n1, _n2, _n3, _n4, _n5);
auto [s0, s1, s2, s3, s4, s5] = s;
out = OutViewType("out", s0, s1, s2, s3, s4, s5);

int n0 = std::min(_n0, in.extent(0));
int n1 = std::min(_n1, in.extent(1));
int n2 = std::min(_n2, in.extent(2));
int n3 = std::min(_n3, in.extent(3));
int n4 = std::min(_n4, in.extent(4));
int n5 = std::min(_n5, in.extent(5));
int n0 = std::min(s0, in.extent(0));
int n1 = std::min(s1, in.extent(1));
int n2 = std::min(s2, in.extent(2));
int n3 = std::min(s3, in.extent(3));
int n4 = std::min(s4, in.extent(4));
int n5 = std::min(s5, in.extent(5));

using range_type = Kokkos::MDRangePolicy<
ExecutionSpace,
Expand All @@ -264,16 +264,16 @@ void crop_or_pad_impl(const ExecutionSpace& exec_space, const InViewType& in,
OutViewType& out, shape_type<7> s) {
constexpr std::size_t DIM = 6;

auto [_n0, _n1, _n2, _n3, _n4, _n5, _n6] = s;
out = OutViewType("out", _n0, _n1, _n2, _n3, _n4, _n5, _n6);
auto [s0, s1, s2, s3, s4, s5, s6] = s;
out = OutViewType("out", s0, s1, s2, s3, s4, s5, s6);

int n0 = std::min(_n0, in.extent(0));
int n1 = std::min(_n1, in.extent(1));
int n2 = std::min(_n2, in.extent(2));
int n3 = std::min(_n3, in.extent(3));
int n4 = std::min(_n4, in.extent(4));
int n5 = std::min(_n5, in.extent(5));
int n6 = std::min(_n6, in.extent(6));
int n0 = std::min(s0, in.extent(0));
int n1 = std::min(s1, in.extent(1));
int n2 = std::min(s2, in.extent(2));
int n3 = std::min(s3, in.extent(3));
int n4 = std::min(s4, in.extent(4));
int n5 = std::min(s5, in.extent(5));
int n6 = std::min(s6, in.extent(6));

using range_type = Kokkos::MDRangePolicy<
ExecutionSpace,
Expand Down Expand Up @@ -301,17 +301,17 @@ void crop_or_pad_impl(const ExecutionSpace& exec_space, const InViewType& in,
OutViewType& out, shape_type<8> s) {
constexpr std::size_t DIM = 6;

auto [_n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7] = s;
out = OutViewType("out", _n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7);

int n0 = std::min(_n0, in.extent(0));
int n1 = std::min(_n1, in.extent(1));
int n2 = std::min(_n2, in.extent(2));
int n3 = std::min(_n3, in.extent(3));
int n4 = std::min(_n4, in.extent(4));
int n5 = std::min(_n5, in.extent(5));
int n6 = std::min(_n6, in.extent(6));
int n7 = std::min(_n7, in.extent(7));
auto [s0, s1, s2, s3, s4, s5, s6, s7] = s;
out = OutViewType("out", s0, s1, s2, s3, s4, s5, s6, s7);

int n0 = std::min(s0, in.extent(0));
int n1 = std::min(s1, in.extent(1));
int n2 = std::min(s2, in.extent(2));
int n3 = std::min(s3, in.extent(3));
int n4 = std::min(s4, in.extent(4));
int n5 = std::min(s5, in.extent(5));
int n6 = std::min(s6, in.extent(6));
int n7 = std::min(s7, in.extent(7));

using range_type = Kokkos::MDRangePolicy<
ExecutionSpace,
Expand Down
Loading

0 comments on commit f0040e7

Please sign in to comment.