Compiler Explorer

Source code

#include <https://raw.githubusercontent.com/kokkos/mdspan/single-header/mdspan.hpp>
#include <bit>
#include <cassert>
#include <cmath>
#if defined(_MSC_VER)
#  include <cstdlib> // MSVC's _aligned_malloc
#endif
#include <exception>
#include <functional>
#include <memory>
#include <numeric>
#include <type_traits>

#define TEMPLATED_CONVERSION_TO_DEFAULT_ACCESSOR 1

namespace stdex = std::experimental;

// P2389 (voted into C++ at June 2024 STL plenary)
namespace std {
template<size_t Rank, class IndexType = size_t>
using dims = dextents<IndexType, Rank>;

template<size_t ByteAlignment, class ElementType>
bool is_sufficiently_aligned(ElementType* p)
{
  return bit_cast<uintptr_t>(p) % ByteAlignment == 0;
}

template<class ElementType, size_t ByteAlignment>
class aligned_accessor {
public:
  static constexpr size_t byte_alignment = ByteAlignment;

static_assert(has_single_bit(byte_alignment),
    "byte_alignment must be a power of two.");
  static_assert(byte_alignment >= alignof(ElementType),
    "Insufficient byte alignment for ElementType");

using offset_policy = stdex::default_accessor<ElementType>;
  using element_type = ElementType;
  using reference = ElementType&;
  using data_handle_type = ElementType*;

constexpr aligned_accessor() noexcept = default;

template<
    class OtherElementType,
    size_t OtherByteAlignment>
  requires(is_convertible_v<
    OtherElementType(*)[], element_type(*)[]> && 
    gcd(OtherByteAlignment, byte_alignment) == byte_alignment
  )
  constexpr aligned_accessor(
    aligned_accessor<OtherElementType, OtherByteAlignment>)
      noexcept
  {}

template<class OtherElementType>
  requires(is_convertible_v<
    OtherElementType(*)[], element_type(*)[]>)
  constexpr explicit aligned_accessor(
    stdex::default_accessor<OtherElementType>) noexcept
  {}
 
#if defined(TEMPLATED_CONVERSION_TO_DEFAULT_ACCESSOR)
  template<class OtherElementType>
    requires(is_convertible_v<
      element_type(*)[],
      OtherElementType(*)[]
    >)
  constexpr
    operator stdex::default_accessor<OtherElementType>() const noexcept
#else
  constexpr
    operator stdex::default_accessor<element_type>() const noexcept
#endif
  {
    return {};
  }

constexpr reference
    access(data_handle_type p, size_t i) const noexcept
  {
    return assume_aligned<byte_alignment>(p)[i];
  }

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

} // namespace std

namespace { // (anonymous)

template<size_t byte_alignment>
using aligned_mdspan =
  std::mdspan<float, std::dims<1, int>, std::layout_right,
    std::aligned_accessor<float, byte_alignment>>;

// Interfaces that require 32-byte alignment,
// because they want to do 8-wide SIMD of float.
void
vectorized_axpby(aligned_mdspan<32> y,
  float alpha, aligned_mdspan<32> x, float beta)
{
  assert(x.extent(0) == y.extent(0));
  for (int k = 0; k < x.extent(0); ++k) {
    y[k] = beta * y[k] + alpha * x[k]; 
  }
}

// 1-norm of the vector y
float vectorized_norm(aligned_mdspan<32> y)
{
  float one_norm = 0.0f;
  for (int k = 0; k < y.extent(0); ++k) {
    one_norm += std::fabs(y[k]); 
  }
  return one_norm;
}

// Interfaces that require 16-byte alignment,
// because they want to do 4-wide SIMD of float.
void fill_x(aligned_mdspan<16> x) {
  for (int k = 0; k < x.extent(0); ++k) {
    x[k] = static_cast<float>(k + 2);
  }  
}
void fill_y(aligned_mdspan<16> y) {
  for (int k = 0; k < y.extent(0); ++k) {
    y[k] = static_cast<float>(k - 1);
  }  
}

// Helper functions for making overaligned array allocations.

template<class ElementType>
struct delete_raw {
  void operator()(ElementType* p) const {
    std::free(p);
  }
};

template<class ElementType>
using allocation =
  std::unique_ptr<ElementType[], delete_raw<ElementType>>;

template<class ElementType, std::size_t byte_alignment>
allocation<ElementType> allocate_raw(const std::size_t num_elements)
{
  const std::size_t num_bytes = num_elements * sizeof(ElementType);
  float* ptr = reinterpret_cast<float*>(
#if defined(_MSC_VER)
    _aligned_malloc(byte_alignment, num_bytes)
#else
    std::aligned_alloc(byte_alignment, num_bytes)
#endif
  );
  return {ptr, delete_raw<ElementType>{}};
}

float user_function(size_t num_elements, float alpha, float beta)
{
  constexpr size_t max_byte_alignment = 32;
  auto x_alloc = allocate_raw<float, max_byte_alignment>(num_elements);
  auto y_alloc = allocate_raw<float, max_byte_alignment>(num_elements);

aligned_mdspan<max_byte_alignment> x(x_alloc.get(), num_elements);
  aligned_mdspan<max_byte_alignment> y(y_alloc.get(), num_elements);

// Implicit conversion from 32-byte aligned to 16-byte aligned
  fill_x(x);
  fill_y(y);

// No conversion: interfaces expect 32-byte aligned and get it
  vectorized_axpby(y, alpha, x, beta);
  return vectorized_norm(y);
}

} // namespace (anonymous)

namespace test_conversion_to_default_accessor {

template<class ElementType>
void take_default_accessor_generic(stdex::default_accessor<ElementType>) {}

template<class ElementType>
  requires(std::is_const_v<ElementType>)
void take_default_accessor_generic_const(stdex::default_accessor<ElementType>) {}

void take_default_accessor(stdex::default_accessor<float>) {}

void take_default_accessor_const(stdex::default_accessor<const float>) {}

void test() {
  // Test new templated conversion operator to default_accessor.
  {
    std::aligned_accessor<float, 32> aligned_acc_f_nc;
    [[maybe_unused]] stdex::default_accessor<float> acc_f_nc{ aligned_acc_f_nc };
    [[maybe_unused]] stdex::default_accessor<float> acc_f_nc_2 = aligned_acc_f_nc;
#if defined(TEMPLATED_CONVERSION_TO_DEFAULT_ACCESSOR)
    [[maybe_unused]] stdex::default_accessor<const float> acc_f_c{ aligned_acc_f_nc };
    [[maybe_unused]] stdex::default_accessor<const float> acc_f_c_2 = aligned_acc_f_nc;
#endif

// CTAD didn't work before anyway.
    //[[maybe_unused]] stdex::default_accessor acc_f{ aligned_acc_f_nc };

take_default_accessor(aligned_acc_f_nc);
#if defined(TEMPLATED_CONVERSION_TO_DEFAULT_ACCESSOR)
    take_default_accessor_const(aligned_acc_f_nc);
#endif

// Doesn't work either way.
    //take_default_accessor_generic(aligned_acc_f_nc);
  }
  {
    std::aligned_accessor<const float, 32> aligned_acc_f_c;
    [[maybe_unused]] stdex::default_accessor<const float> acc_f_c{ aligned_acc_f_c };
    [[maybe_unused]] stdex::default_accessor<const float> acc_f_c_2 = aligned_acc_f_c;

// CTAD didn't work before anyway.
    //[[maybe_unused]] stdex::default_accessor acc_f{ aligned_acc_f_c };

take_default_accessor_const(aligned_acc_f_c);

// Neither of these work either way.
    //take_default_accessor_generic(aligned_acc_f_c);
    //take_default_accessor_generic_const(aligned_acc_f_c);
  }
}
}

int main(int argc, char* argv[])
{
  float result = user_function(10, 1.0f, -1.0f);
  assert(result == 30.0f); // 3 + 3 + ... + 3 = 30
  test_conversion_to_default_accessor::test();

return 0;
}