-
Notifications
You must be signed in to change notification settings - Fork 421
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #7074 from crtrott/space-aware-accessor
SpaceAwareAccessor
- Loading branch information
Showing
6 changed files
with
409 additions
and
10 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,161 @@ | ||
| //@HEADER | ||
| // ************************************************************************ | ||
| // | ||
| // Kokkos v. 4.0 | ||
| // Copyright (2022) National Technology & Engineering | ||
| // Solutions of Sandia, LLC (NTESS). | ||
| // | ||
| // Under the terms of Contract DE-NA0003525 with NTESS, | ||
| // the U.S. Government retains certain rights in this software. | ||
| // | ||
| // Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions. | ||
| // See https://kokkos.org/LICENSE for license information. | ||
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | ||
| // | ||
| //@HEADER | ||
|
|
||
| #ifndef KOKKOS_IMPL_PUBLIC_INCLUDE | ||
| static_assert(false, | ||
| "Including non-public Kokkos header files is not allowed."); | ||
| #endif | ||
|
|
||
| #ifndef KOKKOS_MDSPAN_ACCESSOR_HPP | ||
| #define KOKKOS_MDSPAN_ACCESSOR_HPP | ||
|
|
||
| #include <Kokkos_Macros.hpp> | ||
| #include <Kokkos_Concepts.hpp> | ||
| #include <Kokkos_Core_fwd.hpp> | ||
|
|
||
| namespace Kokkos { | ||
|
|
||
| // For now use the accessors in Impl namespace, as an | ||
| // implementation detail for rebasing View on mdspan | ||
| namespace Impl { | ||
|
|
||
| template <class MemorySpace, class NestedAccessor> | ||
| struct SpaceAwareAccessor { | ||
| using element_type = typename NestedAccessor::element_type; | ||
| using reference = typename NestedAccessor::reference; | ||
| using data_handle_type = typename NestedAccessor::data_handle_type; | ||
|
|
||
| using offset_policy = | ||
| SpaceAwareAccessor<MemorySpace, typename NestedAccessor::offset_policy>; | ||
|
|
||
| using memory_space = MemorySpace; | ||
|
|
||
| static_assert(is_memory_space_v<memory_space>); | ||
|
|
||
| KOKKOS_DEFAULTED_FUNCTION | ||
| constexpr SpaceAwareAccessor() = default; | ||
|
|
||
| template <class OtherNestedAccessorType, | ||
| std::enable_if_t<std::is_constructible_v<NestedAccessor, | ||
| OtherNestedAccessorType>, | ||
| int> = 0> | ||
| KOKKOS_FUNCTION constexpr SpaceAwareAccessor( | ||
| const SpaceAwareAccessor<MemorySpace, OtherNestedAccessorType>& | ||
| other) noexcept | ||
| : nested_acc(other.nested_acc) {} | ||
|
|
||
| template <class OtherNestedAccessorType, | ||
| std::enable_if_t<std::is_constructible_v<NestedAccessor, | ||
| OtherNestedAccessorType>, | ||
| int> = 0> | ||
| KOKKOS_FUNCTION constexpr SpaceAwareAccessor( | ||
| const SpaceAwareAccessor<AnonymousSpace, OtherNestedAccessorType>& | ||
| other) noexcept | ||
| : nested_acc(other.nested_acc) {} | ||
|
|
||
| KOKKOS_FUNCTION | ||
| SpaceAwareAccessor(const NestedAccessor& acc) : nested_acc(acc) {} | ||
|
|
||
| KOKKOS_FUNCTION | ||
| explicit operator NestedAccessor() const { return nested_acc; } | ||
|
|
||
| KOKKOS_FUNCTION | ||
| constexpr reference access(data_handle_type p, size_t i) const noexcept { | ||
| Kokkos::Impl::runtime_check_memory_access_violation<memory_space>( | ||
| "Kokkos::SpaceAwareAccessor ERROR: attempt to access inaccessible " | ||
| "memory space"); | ||
| return nested_acc.access(p, i); | ||
| } | ||
|
|
||
| KOKKOS_FUNCTION | ||
| constexpr data_handle_type offset(data_handle_type p, size_t i) const | ||
| noexcept { | ||
| return nested_acc.offset(p, i); | ||
| } | ||
|
|
||
| private: | ||
| // We either compile with our custom mdspan impl | ||
| // in which case we discover inside it whether no_unique_address | ||
| // works, or we use C++23 in which case it better be available | ||
| #ifdef _MDSPAN_NO_UNIQUE_ADDRESS | ||
| _MDSPAN_NO_UNIQUE_ADDRESS | ||
| #else | ||
| [[no_unique_address]] | ||
| #endif | ||
| NestedAccessor nested_acc; | ||
| template <class, class> | ||
| friend struct SpaceAwareAccessor; | ||
| }; | ||
|
|
||
| template <class NestedAccessor> | ||
| struct SpaceAwareAccessor<AnonymousSpace, NestedAccessor> { | ||
| using element_type = typename NestedAccessor::element_type; | ||
| using reference = typename NestedAccessor::reference; | ||
| using data_handle_type = typename NestedAccessor::data_handle_type; | ||
|
|
||
| using offset_policy = | ||
| SpaceAwareAccessor<AnonymousSpace, | ||
| typename NestedAccessor::offset_policy>; | ||
|
|
||
| using memory_space = AnonymousSpace; | ||
|
|
||
| KOKKOS_DEFAULTED_FUNCTION | ||
| constexpr SpaceAwareAccessor() = default; | ||
|
|
||
| template <class OtherMemorySpace, class OtherNestedAccessorType, | ||
| std::enable_if_t<std::is_constructible_v<NestedAccessor, | ||
| OtherNestedAccessorType>, | ||
| int> = 0> | ||
| KOKKOS_FUNCTION constexpr SpaceAwareAccessor( | ||
| const SpaceAwareAccessor<OtherMemorySpace, OtherNestedAccessorType>& | ||
| other) noexcept | ||
| : nested_acc(other.nested_acc) {} | ||
|
|
||
| KOKKOS_FUNCTION | ||
| SpaceAwareAccessor(const NestedAccessor& acc) : nested_acc(acc) {} | ||
|
|
||
| KOKKOS_FUNCTION | ||
| explicit operator NestedAccessor() const { return nested_acc; } | ||
|
|
||
| KOKKOS_FUNCTION | ||
| constexpr reference access(data_handle_type p, size_t i) const noexcept { | ||
| return nested_acc.access(p, i); | ||
| } | ||
|
|
||
| KOKKOS_FUNCTION | ||
| constexpr data_handle_type offset(data_handle_type p, size_t i) const | ||
| noexcept { | ||
| return nested_acc.offset(p, i); | ||
| } | ||
|
|
||
| private: | ||
| // We either compile with our custom mdspan impl | ||
| // in which case we discover inside it whether no_unique_address | ||
| // works, or we use C++23 in which case it better be available | ||
| #ifdef _MDSPAN_NO_UNIQUE_ADDRESS | ||
| _MDSPAN_NO_UNIQUE_ADDRESS | ||
| #else | ||
| [[no_unique_address]] | ||
| #endif | ||
| NestedAccessor nested_acc; | ||
| template <class, class> | ||
| friend struct SpaceAwareAccessor; | ||
| }; | ||
|
|
||
| } // namespace Impl | ||
| } // namespace Kokkos | ||
|
|
||
| #endif |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,98 @@ | ||
| //@HEADER | ||
| // ************************************************************************ | ||
| // | ||
| // Kokkos v. 4.0 | ||
| // Copyright (2022) National Technology & Engineering | ||
| // Solutions of Sandia, LLC (NTESS). | ||
| // | ||
| // Under the terms of Contract DE-NA0003525 with NTESS, | ||
| // the U.S. Government retains certain rights in this software. | ||
| // | ||
| // Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions. | ||
| // See https://kokkos.org/LICENSE for license information. | ||
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | ||
| // | ||
| //@HEADER | ||
|
|
||
| #include <Kokkos_Core.hpp> | ||
| #include <type_traits> | ||
|
|
||
| #include <gtest/gtest.h> | ||
|
|
||
| template <class T, class ExecutionSpace> | ||
| void test_space_aware_accessor() { | ||
| using memory_space_t = typename ExecutionSpace::memory_space; | ||
| using value_type = std::remove_const_t<T>; | ||
| Kokkos::View<value_type*, ExecutionSpace> v("V", 100); | ||
|
|
||
| Kokkos::parallel_for( | ||
| Kokkos::RangePolicy<ExecutionSpace>(0, v.extent(0)), | ||
| KOKKOS_LAMBDA(int i) { v(i) = i; }); | ||
|
|
||
| int errors; | ||
| using acc_t = Kokkos::Impl::SpaceAwareAccessor<memory_space_t, | ||
| Kokkos::default_accessor<T>>; | ||
| acc_t acc{}; | ||
| typename acc_t::data_handle_type ptr = v.data(); | ||
|
|
||
| Kokkos::parallel_reduce( | ||
| Kokkos::RangePolicy<ExecutionSpace>(0, v.extent(0)), | ||
| KOKKOS_LAMBDA(int i, int& error) { | ||
| if (acc.access(ptr, i) != ptr[i]) error++; | ||
| static_assert(std::is_same_v<typename acc_t::element_type, T>); | ||
| static_assert(std::is_same_v<typename acc_t::reference, T&>); | ||
| static_assert(std::is_same_v<typename acc_t::data_handle_type, T*>); | ||
| static_assert(std::is_same_v<typename acc_t::offset_policy, acc_t>); | ||
| static_assert( | ||
| std::is_same_v<typename acc_t::memory_space, memory_space_t>); | ||
| #if defined(KOKKOS_ENABLE_CXX20) || defined(KOKKOS_ENABLE_CXX23) || \ | ||
| defined(KOKKOS_ENABLE_CXX26) | ||
| static_assert(std::is_empty_v<acc_t>); | ||
| #endif | ||
| }, | ||
| errors); | ||
| ASSERT_EQ(errors, 0); | ||
| } | ||
|
|
||
| void test_space_aware_accessor_conversion() { | ||
| using ExecutionSpace = TEST_EXECSPACE; | ||
| using memory_space_t = typename ExecutionSpace::memory_space; | ||
| using T = float; | ||
| using acc_t = Kokkos::Impl::SpaceAwareAccessor<memory_space_t, | ||
| Kokkos::default_accessor<T>>; | ||
| using const_acc_t = | ||
| Kokkos::Impl::SpaceAwareAccessor<memory_space_t, | ||
| Kokkos::default_accessor<const T>>; | ||
| using int_acc_t = | ||
| Kokkos::Impl::SpaceAwareAccessor<memory_space_t, | ||
| Kokkos::default_accessor<int>>; | ||
| using host_acc_t = | ||
| Kokkos::Impl::SpaceAwareAccessor<Kokkos::HostSpace, | ||
| Kokkos::default_accessor<T>>; | ||
| using anon_acc_t = | ||
| Kokkos::Impl::SpaceAwareAccessor<Kokkos::AnonymousSpace, | ||
| Kokkos::default_accessor<T>>; | ||
|
|
||
| Kokkos::parallel_for( | ||
| Kokkos::RangePolicy<ExecutionSpace>(0, 1), KOKKOS_LAMBDA(int) { | ||
| static_assert(std::is_constructible_v<const_acc_t, acc_t>); | ||
| static_assert(std::is_convertible_v<acc_t, const_acc_t>); | ||
| static_assert(!std::is_constructible_v<acc_t, const_acc_t>); | ||
| static_assert(!std::is_constructible_v<acc_t, int_acc_t>); | ||
| static_assert(std::is_constructible_v<acc_t, host_acc_t> == | ||
| std::is_same_v<memory_space_t, Kokkos::HostSpace>); | ||
| static_assert(std::is_constructible_v<anon_acc_t, acc_t>); | ||
| static_assert(std::is_constructible_v<acc_t, anon_acc_t>); | ||
| static_assert(std::is_convertible_v<anon_acc_t, acc_t>); | ||
| static_assert(std::is_convertible_v<acc_t, anon_acc_t>); | ||
| }); | ||
| } | ||
|
|
||
| TEST(TEST_CATEGORY, space_aware_accessor) { | ||
| using ExecutionSpace = TEST_EXECSPACE; | ||
| test_space_aware_accessor<int, ExecutionSpace>(); | ||
| test_space_aware_accessor<double, ExecutionSpace>(); | ||
| test_space_aware_accessor<const int, ExecutionSpace>(); | ||
| test_space_aware_accessor<const double, ExecutionSpace>(); | ||
| test_space_aware_accessor_conversion(); | ||
| } |
Oops, something went wrong.