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

Implement fft along axis #5

Merged
merged 5 commits into from
Dec 5, 2023
+1,861 −533
Merged

Implement fft along axis #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
1c0fb59
Main CMake was missing
Dec 5, 2023
d9ec528
Introduce utils for internal transpose
Dec 5, 2023
d674aa1
Reconstruct transpose helpers
Dec 5, 2023
3a243f9
use explicit fftw execution functions for each case to use pointers f…
Dec 5, 2023
9658108
Internal transpose capability when needed
Dec 5, 2023
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
22 changes: 22 additions & 0 deletions CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
cmake_minimum_required(VERSION 3.23)
project(kokkos-fft LANGUAGES CXX)

# Add cmake helpers for FFTW
list(INSERT CMAKE_MODULE_PATH 0 "${CMAKE_SOURCE_DIR}/cmake")

find_package(Kokkos CONFIG)
if(NOT kokkos_FOUND)
add_subdirectory(tpls/kokkos)
endif()

# Googletest
include(CTest)
if(BUILD_TESTING)
find_package(GTest CONFIG)
if(NOT GTest_FOUND)
add_subdirectory(tpls/googletest)
endif()
endif()

add_subdirectory(common)
add_subdirectory(fft)
20 changes: 13 additions & 7 deletions common/src/KokkosFFT_layouts.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ namespace KokkosFFT {
using array_layout_type = typename InViewType::array_layout;

// index map after transpose over axis
auto map = get_map_axes(in, _axes);
auto [map, map_inv] = get_map_axes(in, _axes);

constexpr std::size_t rank = InViewType::rank;
int inner_most_axis = std::is_same_v<array_layout_type, typename Kokkos::LayoutLeft> ? 0 : rank - 1;
Expand All @@ -48,7 +48,7 @@ namespace KokkosFFT {
throw std::runtime_error("If the input type is real, the output type should be complex");
}
}

if(std::is_floating_point<out_value_type>::value) {
// Then C2R
if(is_complex<in_value_type>::value) {
Expand Down Expand Up @@ -77,8 +77,7 @@ namespace KokkosFFT {
using array_layout_type = typename InViewType::array_layout;

// index map after transpose over axis
// [TO DO] Implement this later
// auto map = get_map_axes(in, _axes);
auto [map, map_inv] = get_map_axes(in, _axes);

static_assert(InViewType::rank() >= DIM,
"KokkosFFT::get_map_axes: Rank of View must be larger thane or equal to the Rank of FFT axes.");
Expand All @@ -93,17 +92,24 @@ namespace KokkosFFT {
// Get extents for the inner most axes in LayoutRight
// If we allow the FFT on the layoutLeft, this part should be modified
for(std::size_t i=0; i<rank; i++) {
_in_extents.push_back(in.extent(i));
_out_extents.push_back(out.extent(i));
auto _idx = map.at(i);
_in_extents.push_back(in.extent(_idx));
_out_extents.push_back(out.extent(_idx));

// 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(i), out.extent(i));
auto fft_extent = std::max(in.extent(_idx), out.extent(_idx));
_fft_extents.push_back(fft_extent);
}

if(std::is_same<array_layout_type, Kokkos::LayoutLeft>::value) {
std::reverse(_in_extents.begin(), _in_extents.end());
std::reverse(_out_extents.begin(), _out_extents.end());
std::reverse(_fft_extents.begin(), _fft_extents.end());
}

// Define subvectors starting from last - DIM
// Dimensions relevant to FFTs
std::vector<int> in_extents(_in_extents.end()-DIM, _in_extents.end());
Expand Down
Loading