[SYCL] Allow alignment property to be used for group load/store (#16882)

It makes possible to provide alignment<value> property to the load/store
operations indicating the known alignment of the pointer. It will allow
to avoid expensive dynamic alignment checks.

Author: againull

sycl/doc/extensions/experimental/sycl_ext_oneapi_group_load_store.asciidoc

@@ -101,11 +101,13 @@ in the group.
  and default constructible.
 * `Properties` is an instance of `sycl::ext::oneapi::experimental::properties`
 
+_Mandates_: If `Properties` contains the `alignment` property, `InputIteratorT` must be a pointer.
+
 _Effects_: Loads single element from `in_iter` to `out` by using the `g` group
 object to identify memory location as `in_iter` + `g.get_local_linear_id()`.
 
-Properties may provide xref:optimization_properties[assertions] which can
-enable better optimizations.
+Properties may provide xref:optimization_properties[assertions] or the `alignment` property
+which can enable better optimizations.
 
 ==== `sycl::vec` Overload
 
@@ -132,6 +134,8 @@ in the group.
  and default constructible.
 * `Properties` is an instance of `sycl::ext::oneapi::experimental::properties`
 
+_Mandates_: If `Properties` contains the `alignment` property, `InputIteratorT` must be a pointer.
+
 _Effects_: Loads `N` elements from `in_iter` to `out`
 using the `g` group object.
 Properties may specify xref:data_placement[data placement].
@@ -140,8 +144,9 @@ Default data placement is a blocked one:
 in striped case:
 `out[i]` = `in_iter[g.get_local_linear_id() + g.get_local_linear_range() * i];`
 for `i` between `0` and `N`.
-Properties may also provide xref:optimization_properties[assertions] which can
-enable better optimizations.
+Properties may also provide xref:optimization_properties[assertions] or the `alignment` property
+which can enable better optimizations.
+
 
 ==== Fixed-size Array Overload
 
@@ -169,6 +174,8 @@ work-group or sub-group.
  and default constructible.
 * `Properties` is an instance of `sycl::ext::oneapi::experimental::properties`
 
+_Mandates_: If `Properties` contains the `alignment` property, `InputIteratorT` must be a pointer.
+
 _Effects_: Loads `ElementsPerWorkItem` elements from `in_iter` to `out`
 using the `g` group object.
 Properties may specify xref:data_placement[data placement].
@@ -177,8 +184,9 @@ Default placement is a blocked one:
 in striped case:
 `out[i]` = `in_iter[g.get_local_linear_id() + g.get_local_linear_range() * i];`
 for `i` between `0` and `ElementsPerWorkItem`.
-Properties may also provide xref:optimization_properties[assertions] which can
-enable better optimizations.
+Properties may also provide xref:optimization_properties[assertions] or the `alignment` property
+which can enable better optimizations.
+
 
 
 === Store API
@@ -209,11 +217,13 @@ in the group.
  and default constructible.
 * `Properties` is an instance of `sycl::ext::oneapi::experimental::properties`
 
+_Mandates_: If `Properties` contains the `alignment` property, `OutputIteratorT` must be a pointer.
+
 _Effects_: Stores single element `in` to `out_iter` by using the `g` group
 object to identify memory location as `out_iter` + `g.get_local_linear_id()`
 
-Properties may provide xref:optimization_properties[assertions] which can
-enable better optimizations.
+Properties may provide xref:optimization_properties[assertions] or the `alignment` property
+which can enable better optimizations.
 
 
 ==== `sycl::vec` Overload
@@ -241,6 +251,8 @@ in the group.
  and default constructible.
 * `Properties` is an instance of `sycl::ext::oneapi::experimental::properties`
 
+_Mandates_: If `Properties` contains the `alignment` property, `OutputIteratorT` must be a pointer.
+
 _Effects_: Stores `N` elements from `in` vec to `out_iter`
 using the `g` group object.
 Properties may specify xref:data_placement[data placement].
@@ -249,8 +261,8 @@ Default placement is a blocked one:
 in striped case:
 `out_iter[g.get_local_linear_id() + g.get_local_linear_range() * i]` = `in[i];`
 for `i` between `0` and `N`.
-Properties may also provide xref:optimization_properties[assertions] which can
-enable better optimizations.
+Properties may also provide xref:optimization_properties[assertions] or the `alignment` property
+which can enable better optimizations.
 
 
 ==== Fixed-size Array Overload
@@ -280,6 +292,8 @@ work-group or sub-group.
  and default constructible.
 * `Properties` is an instance of `sycl::ext::oneapi::experimental::properties`
 
+_Mandates_: If `Properties` contains the `alignment` property, `OutputIteratorT` must be a pointer.
+
 _Effects_: Stores `ElementsPerWorkItem` elements from `in` span to `out_iter`
 using the `g` group object.
 
@@ -289,8 +303,9 @@ Default placement is a blocked one:
 in striped case:
 `out_iter[g.get_local_linear_id() + g.get_local_linear_range() * i]` = `in[i];`
 for `i` between `0` and `ItemsPerWorkItem`.
-Properties may also provide xref:optimization_properties[assertions] which can
-enable better optimizations.
+Properties may also provide xref:optimization_properties[assertions] or the `alignment` property
+which can enable better optimizations.
+
 
 === Data Placement
 
@@ -442,6 +457,23 @@ so the implementation can rely on `get_max_local_range()` range size:
 
 If partition is uneven the behavior is undefined.
 
+== Alignment
+
+If `InputIteratorT`/`OutputIteratorT` is a pointer then the following property can be used
+to provide an alignment of the pointer. It can allow to avoid dynamic alignment check.
+
+```c++
+namespace sycl::ext::oneapi::experimental {
+struct alignment_key {
+  template <int K>
+  using value_t = property_value<alignment_key, std::integral_constant<int, K>>;
+};
+
+template<int K>
+inline constexpr alignment_key::value_t<K> alignment;
+} // namespace sycl::ext::oneapi::experimental
+```
+
 == Usage Example
 
 Example shows the simplest case without local memory usage of blocked load
@@ -458,8 +490,8 @@ constexpr std::size_t block_count = 2;
 constexpr std::size_t size = block_count * block_size * items_per_thread;
 
 sycl::queue q;
-T* input = sycl::malloc_device<T>(size, q);
-T* output = sycl::malloc_device<T>(size, q);
+T* input = sycl::aligned_alloc_device<T>(16, size, q);
+T* output = sycl::aligned_alloc_device<T>(16, size, q);
 
 q.submit([&](sycl::handler& cgh) {
     cgh.parallel_for(
@@ -472,7 +504,7 @@ q.submit([&](sycl::handler& cgh) {
             auto offset = g.get_group_id(0) * g.get_local_range(0) *
                                        items_per_thread;
 
-            auto props = sycl_exp::properties{sycl_exp::contiguous_memory};
+            auto props = sycl_exp::properties{sycl_exp::contiguous_memory, sycl_exp::alignment<16>};
 
             sycl_exp::group_load(g, input + offset, sycl::span{ data }, props);
 

sycl/include/sycl/ext/oneapi/experimental/group_load_store.hpp

@@ -10,6 +10,7 @@
 
 #pragma once
 
+#include <sycl/ext/oneapi/experimental/annotated_ptr/annotated_ptr.hpp>
 #include <sycl/ext/oneapi/properties/properties.hpp>
 #include <sycl/group_barrier.hpp>
 #include <sycl/sycl_span.hpp>
@@ -255,25 +256,29 @@ constexpr auto get_block_op_ptr(IteratorT iter,
   }
 }
 
-template <int RequiredAlign, typename IteratorType>
-bool is_aligned(IteratorType iter) {
+template <int RequiredAlign, typename IteratorType, typename Properties>
+bool is_aligned(IteratorType iter, [[maybe_unused]] Properties props) {
   using value_type = remove_decoration_t<
       typename std::iterator_traits<IteratorType>::value_type>;
+
+  if constexpr (Properties::template has_property<alignment_key>()) {
+    if (Properties::template get_property<alignment_key>().value >=
+        RequiredAlign)
+      return true;
+  }
+
   return alignof(value_type) >= RequiredAlign ||
          reinterpret_cast<uintptr_t>(&*iter) % RequiredAlign == 0;
 }
 
-} // namespace detail
-
-// Load API span overload.
 template <typename Group, typename InputIteratorT, typename OutputT,
           std::size_t ElementsPerWorkItem,
           typename Properties = decltype(properties())>
 std::enable_if_t<detail::verify_load_types<InputIteratorT, OutputT> &&
                  detail::is_generic_group_v<Group> &&
                  is_property_list_v<Properties>>
-group_load(Group g, InputIteratorT in_ptr,
-           span<OutputT, ElementsPerWorkItem> out, Properties props = {}) {
+group_load_impl(Group g, InputIteratorT in_ptr,
+                span<OutputT, ElementsPerWorkItem> out, Properties props = {}) {
   constexpr bool blocked = detail::isBlocked(props);
   using use_naive =
       detail::merged_properties_t<Properties,
@@ -286,7 +291,7 @@ group_load(Group g, InputIteratorT in_ptr,
     group_barrier(g);
     return;
   } else if constexpr (!std::is_same_v<Group, sycl::sub_group>) {
-    return group_load(g, in_ptr, out, use_naive{});
+    return group_load_impl(g, in_ptr, out, use_naive{});
   } else {
     auto ptr = detail::get_block_op_ptr<ElementsPerWorkItem>(in_ptr, props);
     static constexpr auto deduced_address_space =
@@ -297,12 +302,12 @@ group_load(Group g, InputIteratorT in_ptr,
                     access::address_space::generic_space) {
         if (auto local_ptr = detail::dynamic_address_cast<
                 access::address_space::local_space>(ptr)) {
-          return group_load(g, local_ptr, out, props);
+          return group_load_impl(g, local_ptr, out, props);
         } else if (auto global_ptr = detail::dynamic_address_cast<
                        access::address_space::global_space>(ptr)) {
-          return group_load(g, global_ptr, out, props);
+          return group_load_impl(g, global_ptr, out, props);
         } else {
-          return group_load(g, in_ptr, out, use_naive{});
+          return group_load_impl(g, in_ptr, out, use_naive{});
         }
       } else {
         using value_type = remove_decoration_t<
@@ -314,8 +319,8 @@ group_load(Group g, InputIteratorT in_ptr,
         constexpr int ReqAlign =
             detail::RequiredAlignment<detail::operation_type::load,
                                       deduced_address_space>::value;
-        if (!detail::is_aligned<ReqAlign>(in_ptr))
-          return group_load(g, in_ptr, out, use_naive{});
+        if (!detail::is_aligned<ReqAlign>(in_ptr, props))
+          return group_load_impl(g, in_ptr, out, use_naive{});
 
         // We know the pointer is aligned and the address space is known. Do the
         // optimized load.
@@ -353,20 +358,21 @@ group_load(Group g, InputIteratorT in_ptr,
         }
       }
     } else {
-      return group_load(g, in_ptr, out, use_naive{});
+      return group_load_impl(g, in_ptr, out, use_naive{});
     }
+
+    return;
   }
 }
 
-// Store API span overload.
 template <typename Group, typename InputT, std::size_t ElementsPerWorkItem,
           typename OutputIteratorT,
           typename Properties = decltype(properties())>
 std::enable_if_t<detail::verify_store_types<InputT, OutputIteratorT> &&
                  detail::is_generic_group_v<Group> &&
                  is_property_list_v<Properties>>
-group_store(Group g, const span<InputT, ElementsPerWorkItem> in,
-            OutputIteratorT out_ptr, Properties props = {}) {
+group_store_impl(Group g, const span<InputT, ElementsPerWorkItem> in,
+                 OutputIteratorT out_ptr, Properties props = {}) {
   constexpr bool blocked = detail::isBlocked(props);
   using use_naive =
       detail::merged_properties_t<Properties,
@@ -379,7 +385,7 @@ group_store(Group g, const span<InputT, ElementsPerWorkItem> in,
     group_barrier(g);
     return;
   } else if constexpr (!std::is_same_v<Group, sycl::sub_group>) {
-    return group_store(g, in, out_ptr, use_naive{});
+    return group_store_impl(g, in, out_ptr, use_naive{});
   } else {
     auto ptr = detail::get_block_op_ptr<ElementsPerWorkItem>(out_ptr, props);
 
@@ -390,12 +396,12 @@ group_store(Group g, const span<InputT, ElementsPerWorkItem> in,
                     access::address_space::generic_space) {
         if (auto local_ptr = detail::dynamic_address_cast<
                 access::address_space::local_space>(ptr)) {
-          return group_store(g, in, local_ptr, props);
+          return group_store_impl(g, in, local_ptr, props);
         } else if (auto global_ptr = detail::dynamic_address_cast<
                        access::address_space::global_space>(ptr)) {
-          return group_store(g, in, global_ptr, props);
+          return group_store_impl(g, in, global_ptr, props);
         } else {
-          return group_store(g, in, out_ptr, use_naive{});
+          return group_store_impl(g, in, out_ptr, use_naive{});
         }
       } else {
         using block_info = typename detail::BlockTypeInfo<
@@ -406,8 +412,8 @@ group_store(Group g, const span<InputT, ElementsPerWorkItem> in,
         constexpr int ReqAlign =
             detail::RequiredAlignment<detail::operation_type::store,
                                       deduced_address_space>::value;
-        if (!detail::is_aligned<ReqAlign>(out_ptr))
-          return group_store(g, in, out_ptr, use_naive{});
+        if (!detail::is_aligned<ReqAlign>(out_ptr, props))
+          return group_store_impl(g, in, out_ptr, use_naive{});
 
         std::remove_const_t<remove_decoration_t<
             typename std::iterator_traits<OutputIteratorT>::value_type>>
@@ -424,10 +430,41 @@ group_store(Group g, const span<InputT, ElementsPerWorkItem> in,
                                         sycl::bit_cast<block_op_type>(values));
       }
     } else {
-      return group_store(g, in, out_ptr, use_naive{});
+      return group_store_impl(g, in, out_ptr, use_naive{});
     }
   }
 }
+} // namespace detail
+
+// Load API span overload.
+template <typename Group, typename InputIteratorT, typename OutputT,
+          std::size_t ElementsPerWorkItem,
+          typename Properties = decltype(properties())>
+std::enable_if_t<detail::verify_load_types<InputIteratorT, OutputT> &&
+                 detail::is_generic_group_v<Group> &&
+                 is_property_list_v<Properties>>
+group_load(Group g, InputIteratorT in_ptr,
+           span<OutputT, ElementsPerWorkItem> out, Properties props = {}) {
+  static_assert(std::is_pointer_v<InputIteratorT> ||
+                    !Properties::template has_property<alignment_key>(),
+                "group_load requires a pointer if alignment property is used");
+  detail::group_load_impl(g, in_ptr, out, props);
+}
+
+// Store API span overload.
+template <typename Group, typename InputT, std::size_t ElementsPerWorkItem,
+          typename OutputIteratorT,
+          typename Properties = decltype(properties())>
+std::enable_if_t<detail::verify_store_types<InputT, OutputIteratorT> &&
+                 detail::is_generic_group_v<Group> &&
+                 is_property_list_v<Properties>>
+group_store(Group g, const span<InputT, ElementsPerWorkItem> in,
+            OutputIteratorT out_ptr, Properties props = {}) {
+  static_assert(std::is_pointer_v<OutputIteratorT> ||
+                    !Properties::template has_property<alignment_key>(),
+                "group_store requires a pointer if alignment property is used");
+  detail::group_store_impl(g, in, out_ptr, props);
+}
 
 // Load API scalar.
 template <typename Group, typename InputIteratorT, typename OutputT,
@@ -437,7 +474,10 @@ std::enable_if_t<detail::verify_load_types<InputIteratorT, OutputT> &&
                  is_property_list_v<Properties>>
 group_load(Group g, InputIteratorT in_ptr, OutputT &out,
            Properties properties = {}) {
-  group_load(g, in_ptr, span<OutputT, 1>(&out, 1), properties);
+  static_assert(std::is_pointer_v<InputIteratorT> ||
+                    !Properties::template has_property<alignment_key>(),
+                "group_load requires a pointer if alignment property is used");
+  detail::group_load_impl(g, in_ptr, span<OutputT, 1>(&out, 1), properties);
 }
 
 // Store API scalar.
@@ -448,7 +488,11 @@ std::enable_if_t<detail::verify_store_types<InputT, OutputIteratorT> &&
                  is_property_list_v<Properties>>
 group_store(Group g, const InputT &in, OutputIteratorT out_ptr,
             Properties properties = {}) {
-  group_store(g, span<const InputT, 1>(&in, 1), out_ptr, properties);
+  static_assert(std::is_pointer_v<OutputIteratorT> ||
+                    !Properties::template has_property<alignment_key>(),
+                "group_store requires a pointer if alignment property is used");
+  detail::group_store_impl(g, span<const InputT, 1>(&in, 1), out_ptr,
+                           properties);
 }
 
 // Load API sycl::vec overload.
@@ -459,7 +503,10 @@ std::enable_if_t<detail::verify_load_types<InputIteratorT, OutputT> &&
                  is_property_list_v<Properties>>
 group_load(Group g, InputIteratorT in_ptr, sycl::vec<OutputT, N> &out,
            Properties properties = {}) {
-  group_load(g, in_ptr, span<OutputT, N>(&out[0], N), properties);
+  static_assert(std::is_pointer_v<InputIteratorT> ||
+                    !Properties::template has_property<alignment_key>(),
+                "group_load requires a pointer if alignment property is used");
+  detail::group_load_impl(g, in_ptr, span<OutputT, N>(&out[0], N), properties);
 }
 
 // Store API sycl::vec overload.
@@ -470,7 +517,11 @@ std::enable_if_t<detail::verify_store_types<InputT, OutputIteratorT> &&
                  is_property_list_v<Properties>>
 group_store(Group g, const sycl::vec<InputT, N> &in, OutputIteratorT out_ptr,
             Properties properties = {}) {
-  group_store(g, span<const InputT, N>(&in[0], N), out_ptr, properties);
+  static_assert(std::is_pointer_v<OutputIteratorT> ||
+                    !Properties::template has_property<alignment_key>(),
+                "group_store requires a pointer if alignment property is used");
+  detail::group_store_impl(g, span<const InputT, N>(&in[0], N), out_ptr,
+                           properties);
 }
 
 #else