Skip to content

add aligned accessor as part of the library API #446

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

Open
nmm0 wants to merge 12 commits into
base: stable
Choose a base branch
from
Open

add aligned accessor as part of the library API #446

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
9515404
Port `aligned_accessor` from the Raft fork of Kokkos to mainline.
vitor1001 Jun 18, 2025
8673523
Apply suggestions from code review
vitor1001 Jun 19, 2025
19d5184
#412: add tests for aligned accessor
nmm0 Jan 28, 2026
4753b39
#412: fix some compiler warnings
nmm0 Jan 28, 2026
d5ff70b
#412: update aligned accessor example to use library aligned accessor
nmm0 Jan 28, 2026
80b86b3
#412: add missing include for <cstdlib> to obtain aligned_alloc in te…
nmm0 Jan 28, 2026
abf6d07
#412: use MSVCs _aligned_Malloc
nmm0 Jan 29, 2026
5ef9114
#412: remove accidental rebase artifact
nmm0 Feb 3, 2026
f827aad
#412: add test with structure type whose size is different from align…
nmm0 Feb 3, 2026
0a01127
#412: address reviewer comments
nmm0 Feb 4, 2026
2cfc8a5
make aligned accessor example more in line with standard c++
nmm0 Jun 29, 2026
ab5f549
#412: fix compilation issues on C++14-20
nmm0 Jul 7, 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
8 changes: 8 additions & 0 deletions examples/aligned_accessor/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1 +1,9 @@
mdspan_add_example(aligned_accessor)

if(MDSPAN_ENABLE_OPENMP)
if(OpenMP_CXX_FOUND)
target_link_libraries(aligned_accessor OpenMP::OpenMP_CXX)
target_compile_definitions(aligned_accessor MDSPAN_ENABLE_OPENMP)
message(STATUS "Building aligned accessor example with OpenMP")
endif()
endif()
161 changes: 19 additions & 142 deletions examples/aligned_accessor/aligned_accessor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,15 @@
#include <memory>
#include <type_traits>
#include <string> // stoi
#ifdef MDSPAN_ENABLE_OPENMP
#include <omp.h>
#endif

// mfh 2022/08/08: This is based on my comment on RAPIDS RAFT issue 725:
// https://github.com/rapidsai/raft/pull/725#discussion_r937991701

namespace {
using Kokkos::aligned_accessor;

using test_value_type = float;
constexpr std::size_t min_overalignment_factor = 8;
Expand All @@ -43,129 +47,6 @@ constexpr std::size_t min_byte_alignment = min_overalignment_factor * sizeof(flo
// Some compilers have trouble optimizing loops with unsigned or 64-bit index types.
using index_type = int;


// Prefer std::assume_aligned if available, as it is in the C++ Standard.
// Otherwise, use a compiler-specific equivalent if available.

// NOTE (mfh 2022/08/08) BYTE_ALIGNMENT must be unsigned and a power of 2.
#if defined(__cpp_lib_assume_aligned)
# define MDSPAN_IMPL_ASSUME_ALIGNED( ELEMENT_TYPE, POINTER, BYTE_ALIGNMENT ) (std::assume_aligned< BYTE_ALIGNMENT >( POINTER ))
constexpr char assume_aligned_method[] = "std::assume_aligned";
#elif defined(__ICL)
# define MDSPAN_IMPL_ASSUME_ALIGNED( ELEMENT_TYPE, POINTER, BYTE_ALIGNMENT ) POINTER
constexpr char assume_aligned_method[] = "(none)";
#elif defined(__ICC)
# define MDSPAN_IMPL_ASSUME_ALIGNED( ELEMENT_TYPE, POINTER, BYTE_ALIGNMENT ) POINTER
constexpr char assume_aligned_method[] = "(none)";
#elif defined(__clang__)
# define MDSPAN_IMPL_ASSUME_ALIGNED( ELEMENT_TYPE, POINTER, BYTE_ALIGNMENT ) POINTER
constexpr char assume_aligned_method[] = "(none)";
#elif defined(__GNUC__)
// __builtin_assume_aligned returns void*
# define MDSPAN_IMPL_ASSUME_ALIGNED( ELEMENT_TYPE, POINTER, BYTE_ALIGNMENT ) reinterpret_cast< ELEMENT_TYPE* >(__builtin_assume_aligned( POINTER, BYTE_ALIGNMENT ))
constexpr char assume_aligned_method[] = "__builtin_assume_aligned";
#else
# define MDSPAN_IMPL_ASSUME_ALIGNED( ELEMENT_TYPE, POINTER, BYTE_ALIGNMENT ) POINTER
constexpr char assume_aligned_method[] = "(none)";
#endif

// Some compilers other than Clang or GCC like to define __clang__ or __GNUC__.
// Thus, we order the tests from most to least specific.
#if defined(__ICL)
# define MDSPAN_IMPL_ALIGN_VALUE_ATTRIBUTE( BYTE_ALIGNMENT ) __declspec(align_value( BYTE_ALIGNMENT ))
constexpr char align_attribute_method[] = "__declspec(align_value(BYTE_ALIGNMENT))";
#elif defined(__ICC)
# define MDSPAN_IMPL_ALIGN_VALUE_ATTRIBUTE( BYTE_ALIGNMENT ) __attribute__((align_value( BYTE_ALIGNMENT )))
constexpr char align_attribute_method[] = "__attribute__((align_value(BYTE_ALIGNMENT)))";
#elif defined(__clang__)
# define MDSPAN_IMPL_ALIGN_VALUE_ATTRIBUTE( BYTE_ALIGNMENT ) __attribute__((align_value( BYTE_ALIGNMENT )))
constexpr char align_attribute_method[] = "__attribute__((align_value(BYTE_ALIGNMENT)))";
#else
# define MDSPAN_IMPL_ALIGN_VALUE_ATTRIBUTE( BYTE_ALIGNMENT )
constexpr char align_attribute_method[] = "(none)";
#endif

constexpr bool
is_nonzero_power_of_two(const std::size_t x)
{
// Just checking __cpp_lib_int_pow2 isn't enough for some GCC versions.
// The <bit> header exists, but std::has_single_bit does not.
#if defined(__cpp_lib_int_pow2) && __cplusplus >= 202002L
return std::has_single_bit(x);
#else
return x != 0 && (x & (x - 1)) == 0;
#endif
}

template<class ElementType>
constexpr bool
valid_byte_alignment(const std::size_t byte_alignment)
{
return is_nonzero_power_of_two(byte_alignment) && byte_alignment >= alignof(ElementType);
}

// We define aligned_pointer_t through a struct
// so we can check whether the byte alignment is valid.
// This makes it impossible to use the alias
// with an invalid byte alignment.
template<class ElementType, std::size_t byte_alignment>
struct aligned_pointer {
static_assert(valid_byte_alignment<ElementType>(byte_alignment),
"byte_alignment must be a power of two no less than "
"the minimum required alignment of ElementType.");

#if defined(__ICC)
// x86-64 ICC 2021.5.0 emits warning #3186 ("expected typedef declaration") here.
// No other compiler (including Clang, which has a similar type attribute) has this issue.
# pragma warning push
# pragma warning disable 3186
#endif

using type = ElementType* MDSPAN_IMPL_ALIGN_VALUE_ATTRIBUTE( byte_alignment );

#if defined(__ICC)
# pragma warning pop
#endif
};

template<class ElementType, std::size_t byte_alignment>
using aligned_pointer_t = typename aligned_pointer<ElementType, byte_alignment>::type;

template<class ElementType, std::size_t byte_alignment>
aligned_pointer_t<ElementType, byte_alignment>
bless(ElementType* ptr, std::integral_constant<std::size_t, byte_alignment> /* ba */ )
{
return MDSPAN_IMPL_ASSUME_ALIGNED( ElementType, ptr, byte_alignment );
}

template<class ElementType, std::size_t byte_alignment>
struct aligned_accessor {
using offset_policy = Kokkos::default_accessor<ElementType>;
using element_type = ElementType;
using reference = ElementType&;
using data_handle_type = aligned_pointer_t<ElementType, byte_alignment>;

constexpr aligned_accessor() noexcept = default;

MDSPAN_TEMPLATE_REQUIRES(
class OtherElementType,
std::size_t other_byte_alignment,
/* requires */ (std::is_convertible<OtherElementType(*)[], element_type(*)[]>::value && other_byte_alignment == byte_alignment)
)
constexpr aligned_accessor(aligned_accessor<OtherElementType, other_byte_alignment>) noexcept {}

constexpr reference access(data_handle_type p, size_t i) const noexcept {
// This may declare alignment twice, depending on
// if we have an attribute for marking pointer types.
return MDSPAN_IMPL_ASSUME_ALIGNED( ElementType, p, byte_alignment )[i];
}

constexpr typename offset_policy::data_handle_type
offset(data_handle_type p, size_t i) const noexcept {
return p + i;
}
};

template<class ElementType>
struct delete_raw {
void operator()(ElementType* p) const {
Expand Down Expand Up @@ -245,14 +126,14 @@ class aligned_array_allocation {
num_elements(number_of_elements)
{}

aligned_pointer_t<ElementType, byte_alignment> data() const
ElementType *data() const
{
return MDSPAN_IMPL_ASSUME_ALIGNED( ElementType, pointer, byte_alignment );
return Kokkos::assume_aligned< byte_alignment >(pointer);
}

private:
allocation_t<ElementType> allocation{nullptr, delete_raw<ElementType>{}};
aligned_pointer_t<ElementType, byte_alignment> pointer{nullptr};
ElementType *pointer{nullptr};
std::size_t num_elements{0};
};

Expand Down Expand Up @@ -411,9 +292,9 @@ auto benchmark_add_raw_1d(const std::size_t num_trials,
// Assume that x, y, and z all have the same alignment.
template<class ElementType, std::size_t byte_alignment>
void add_aligned_raw_1d(const index_type n,
aligned_pointer_t<const ElementType, byte_alignment> x,
aligned_pointer_t<const ElementType, byte_alignment> y,
aligned_pointer_t<ElementType, byte_alignment> z)
const ElementType * MDSPAN_ALIGN(byte_alignment) x,
const ElementType * MDSPAN_ALIGN(byte_alignment) y,
ElementType * MDSPAN_ALIGN(byte_alignment) z)
{
for (index_type i = 0; i < n; ++i) {
z[i] = x[i] + y[i];
Expand All @@ -427,19 +308,19 @@ auto benchmark_add_aligned_raw_1d(const std::size_t num_trials,
const ElementType x[],
const ElementType y[],
ElementType z[],
std::integral_constant<std::size_t, byte_alignment> ba)
std::integral_constant<std::size_t, byte_alignment>)
{
TICK();
auto x_blessed = bless(x, ba);
auto y_blessed = bless(y, ba);
auto z_blessed = bless(z, ba);
auto x_blessed = Kokkos::assume_aligned< byte_alignment >(x);
auto y_blessed = Kokkos::assume_aligned< byte_alignment >(y);
auto z_blessed = Kokkos::assume_aligned< byte_alignment >(z);
for (std::size_t trial = 0; trial < num_trials; ++trial) {
add_aligned_raw_1d<ElementType, byte_alignment>(n, x_blessed, y_blessed, z_blessed);
}
return TOCK();
}

#ifdef _OPENMP
#ifdef MDSPAN_ENABLE_OPENMP
template<class ElementType, std::size_t byte_alignment>
void add_omp_simd_aligned_mdspan_1d(aligned_mdspan_1d<const ElementType, byte_alignment> x,
aligned_mdspan_1d<const ElementType, byte_alignment> y,
Expand Down Expand Up @@ -623,7 +504,7 @@ auto benchmark_add_omp_aligned_simd_aligned_raw_1d(
}
return TOCK();
}
#endif // _OPENMP
#endif // MDSPAN_ENABLE_OPENMP

template<class ElementType>
void set_elements_of_arrays(const index_type n,
Expand Down Expand Up @@ -668,7 +549,7 @@ int main(int argc, char* argv[])
benchmark_add_aligned_raw_1d(num_trials, n, x_aligned.data(),
y_aligned.data(), z_aligned.data(),
byte_alignment);
#ifdef _OPENMP
#ifdef MDSPAN_ENABLE_OPENMP
auto omp_simd_aligned_mdspan_result =
benchmark_add_omp_simd_aligned_mdspan_1d(num_trials, n, x_aligned.data(),
y_aligned.data(), z_aligned.data(),
Expand Down Expand Up @@ -699,7 +580,7 @@ int main(int argc, char* argv[])
benchmark_add_raw_1d(num_trials, n, x_unaligned.data(),
y_unaligned.data(), z_unaligned.data());

#ifdef _OPENMP
#ifdef MDSPAN_ENABLE_OPENMP
auto omp_simd_mdspan_result =
benchmark_add_omp_simd_mdspan_1d(num_trials, n, x_unaligned.data(),
y_unaligned.data(), z_unaligned.data());
Expand All @@ -710,17 +591,13 @@ int main(int argc, char* argv[])

cout << "Number of trials: " << num_trials << endl
<< "Number of loop iterations per trial: " << n << endl
<< "Way to declare a pointer value aligned, if any: "
<< assume_aligned_method << endl
<< "Way to declare a pointer type aligned, if any: "
<< align_attribute_method << endl
<< "Total time in seconds for non-OpenMP loops:" << endl
<< " aligned mdspan: " << aligned_mdspan_result << endl
<< " unaligned mdspan: " << mdspan_result << endl
<< " aligned raw: " << aligned_raw_result << endl
<< " unaligned raw: " << raw_result << endl;

#ifdef _OPENMP
#ifdef MDSPAN_ENABLE_OPENMP
cout << "Total time in seconds for OpenMP (omp simd) loops:" << endl
<< " omp_simd_aligned_mdspan: " << omp_simd_aligned_mdspan_result << endl
<< " omp_simd_mdspan: " << omp_simd_mdspan_result << endl
Expand Down
Loading
Loading