Skip to content

Batched: allowing dynrankview in get_extent get_stride #3036

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

Merged
lucbv merged 8 commits into from
Apr 20, 2026
Merged

Batched: allowing dynrankview in get_extent get_stride #3036

Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
8e41e93
Batched: allowing dynrankview in get_extent get_stride
lucbv Apr 16, 2026
7de6178
Update batched/KokkosBatched_Util.hpp
lucbv Apr 17, 2026
9467162
Update batched/KokkosBatched_Util.hpp
lucbv Apr 17, 2026
8054162
Update batched/KokkosBatched_Util.hpp
lucbv Apr 17, 2026
2373c33
Update batched/KokkosBatched_Util.hpp
lucbv Apr 17, 2026
4ef0fa5
Re-adding the DynRankView branch to check rank
lucbv Apr 17, 2026
a64faf5
Fix if constexpr static_assert statement
lucbv Apr 17, 2026
511fcd1
clang-format
lucbv Apr 17, 2026
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
30 changes: 24 additions & 6 deletions batched/KokkosBatched_Util.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,9 @@
#include <complex>
#include <iostream>

#include <Kokkos_Assert.hpp>
#include "Kokkos_Complex.hpp"
#include "Kokkos_DynRankView.hpp"

#include "KokkosKernels_config.h"
#include "KokkosKernels_Macros.hpp"
Expand Down Expand Up @@ -677,9 +679,17 @@ KOKKOS_INLINE_FUNCTION void fma_bounds_check(ViewType v, SizeType m, SizeType n,
namespace Impl {
template <typename ViewType>
KOKKOS_INLINE_FUNCTION int get_extent_int(const ViewType &v, const int r) {
static_assert(Kokkos::is_view_v<ViewType>, "KokkosBatched: ViewType is not a Kokkos::View.");
constexpr std::size_t V_rank = ViewType::rank();
static_assert(V_rank <= 2, "KokkosBatched: ViewType must have rank 0, 1 or 2.");
// Check for view and dynrankview
if constexpr (Kokkos::is_view_v<ViewType>) {
static_assert(ViewType::rank() <= 2, "KokkosBatched: ViewType must have rank 0, 1 or 2.");
} else if constexpr (Kokkos::is_dyn_rank_view_v<ViewType>) {
KOKKOS_EXPECTS((v.rank() <= 2));
} else {
Comment thread
lucbv marked this conversation as resolved.
static_assert(Kokkos::is_view_v<ViewType> || Kokkos::is_dyn_rank_view_v<ViewType>,
"KokkosBatched: ViewType must be a Kokkos::View or a Kokkos::DynRankView");
}

const std::size_t V_rank = v.rank();

if (r == 0) {
int V_extent_0 = V_rank < 1 ? 1 : v.extent_int(0);
Expand All @@ -694,9 +704,17 @@ KOKKOS_INLINE_FUNCTION int get_extent_int(const ViewType &v, const int r) {

template <typename ViewType>
KOKKOS_INLINE_FUNCTION std::size_t get_stride(const ViewType &v, const int r) {
static_assert(Kokkos::is_view_v<ViewType>, "KokkosBatched: ViewType is not a Kokkos::View.");
constexpr std::size_t V_rank = ViewType::rank();
static_assert(V_rank <= 2, "KokkosBatched: ViewType must have rank 0, 1 or 2.");
// Check for view and dynrankview
if constexpr (Kokkos::is_view_v<ViewType>) {
static_assert(ViewType::rank() <= 2, "KokkosBatched: ViewType must have rank 0, 1 or 2.");
} else if constexpr (Kokkos::is_dyn_rank_view_v<ViewType>) {
KOKKOS_EXPECTS((v.rank() <= 2));
} else {
static_assert(Kokkos::is_view_v<ViewType> || Kokkos::is_dyn_rank_view_v<ViewType>,
"KokkosBatched: ViewType must be a Kokkos::View or a Kokkos::DynRankView");
}
Comment thread
lucbv marked this conversation as resolved.

const std::size_t V_rank = v.rank();

if (r == 0) {
std::size_t V_stride_0 = V_rank < 1 ? 1 : v.stride(0);
Expand Down
Loading