[SYCL] Add tangle/opportunistic algorithms (#9220)

Enables the following functions to be used with tangle_group and
opportunistic_group arguments:

- group_barrier
- group_broadcast
- any_of_group
- all_of_group
- none_of_group
- reduce_over_group
- exclusive_scan_over_group
- inclusive_scan_over_group

---

A few quick notes to reviewers:

1) This implementation leverages the fact that it is undefined behavior
to use a tangle group or opportunistic group in control flow that does
not match the control flow at the point of construction to avoid using a
mask for most operations. I _think_ it is safe to call the NonUniform
intrinsics directly, because they are already control-flow-aware.

2) In a few places, I've deliberately duplicated the implementation
across tangle group and opportunistic group even though they're the
same. I've done this primarily in an attempt to simplify @JackAKirk's
efforts to merge in his CUDA implementation, because I expect that there
may be some cases where the CUDA implementations of these groups _do_
diverge. If this turns out not to be true, we can tidy things up
afterwards.

3) In general, tangle and opportunistic group are not the same thing.
But I expect their behavior will be identical on all of the SPIR-V
implementations that we're targeting.

---------

Signed-off-by: John Pennycook <[email protected]>

Author: Pennycook

sycl/include/sycl/detail/spirv.hpp

@@ -27,6 +27,8 @@ struct sub_group;
 namespace experimental {
 template <typename ParentGroup> class ballot_group;
 template <size_t PartitionSize, typename ParentGroup> class fixed_size_group;
+template <typename ParentGroup> class tangle_group;
+class opportunistic_group;
 } // namespace experimental
 } // namespace oneapi
 } // namespace ext
@@ -72,6 +74,16 @@ struct group_scope<sycl::ext::oneapi::experimental::fixed_size_group<
   static constexpr __spv::Scope::Flag value = group_scope<ParentGroup>::value;
 };
 
+template <typename ParentGroup>
+struct group_scope<sycl::ext::oneapi::experimental::tangle_group<ParentGroup>> {
+  static constexpr __spv::Scope::Flag value = group_scope<ParentGroup>::value;
+};
+
+template <>
+struct group_scope<::sycl::ext::oneapi::experimental::opportunistic_group> {
+  static constexpr __spv::Scope::Flag value = __spv::Scope::Flag::Subgroup;
+};
+
 // Generic shuffles and broadcasts may require multiple calls to
 // intrinsics, and should use the fewest broadcasts possible
 // - Loop over chunks until remaining bytes < chunk size
@@ -135,6 +147,16 @@ bool GroupAll(
       static_cast<uint32_t>(__spv::GroupOperation::ClusteredReduce),
       static_cast<uint32_t>(pred), PartitionSize);
 }
+template <typename ParentGroup>
+bool GroupAll(ext::oneapi::experimental::tangle_group<ParentGroup>, bool pred) {
+  return __spirv_GroupNonUniformAll(group_scope<ParentGroup>::value, pred);
+}
+template <typename Group>
+bool GroupAll(const ext::oneapi::experimental::opportunistic_group &,
+              bool pred) {
+  return __spirv_GroupNonUniformAll(
+      group_scope<ext::oneapi::experimental::opportunistic_group>::value, pred);
+}
 
 template <typename Group> bool GroupAny(Group, bool pred) {
   return __spirv_GroupAny(group_scope<Group>::value, pred);
@@ -161,6 +183,15 @@ bool GroupAny(
       static_cast<uint32_t>(__spv::GroupOperation::ClusteredReduce),
       static_cast<uint32_t>(pred), PartitionSize);
 }
+template <typename ParentGroup>
+bool GroupAny(ext::oneapi::experimental::tangle_group<ParentGroup>, bool pred) {
+  return __spirv_GroupNonUniformAny(group_scope<ParentGroup>::value, pred);
+}
+bool GroupAny(const ext::oneapi::experimental::opportunistic_group &,
+              bool pred) {
+  return __spirv_GroupNonUniformAny(
+      group_scope<ext::oneapi::experimental::opportunistic_group>::value, pred);
+}
 
 // Native broadcasts map directly to a SPIR-V GroupBroadcast intrinsic
 // FIXME: Do not special-case for half or vec once all backends support all data
@@ -281,6 +312,45 @@ EnableIfNativeBroadcast<T, IdT> GroupBroadcast(
   return __spirv_GroupNonUniformShuffle(group_scope<ParentGroup>::value, OCLX,
                                         OCLId);
 }
+template <typename ParentGroup, typename T, typename IdT>
+EnableIfNativeBroadcast<T, IdT>
+GroupBroadcast(ext::oneapi::experimental::tangle_group<ParentGroup> g, T x,
+               IdT local_id) {
+  // Remap local_id to its original numbering in ParentGroup.
+  auto LocalId = detail::IdToMaskPosition(g, local_id);
+
+  // TODO: Refactor to avoid duplication after design settles.
+  using GroupIdT = typename GroupId<ParentGroup>::type;
+  GroupIdT GroupLocalId = static_cast<GroupIdT>(LocalId);
+  using OCLT = detail::ConvertToOpenCLType_t<T>;
+  using WidenedT = WidenOpenCLTypeTo32_t<OCLT>;
+  using OCLIdT = detail::ConvertToOpenCLType_t<GroupIdT>;
+  WidenedT OCLX = detail::convertDataToType<T, OCLT>(x);
+  OCLIdT OCLId = detail::convertDataToType<GroupIdT, OCLIdT>(GroupLocalId);
+
+  return __spirv_GroupNonUniformBroadcast(group_scope<ParentGroup>::value, OCLX,
+                                          OCLId);
+}
+template <typename T, typename IdT>
+EnableIfNativeBroadcast<T, IdT>
+GroupBroadcast(const ext::oneapi::experimental::opportunistic_group &g, T x,
+               IdT local_id) {
+  // Remap local_id to its original numbering in sub-group
+  auto LocalId = detail::IdToMaskPosition(g, local_id);
+
+  // TODO: Refactor to avoid duplication after design settles.
+  using GroupIdT = typename GroupId<sycl::ext::oneapi::sub_group>::type;
+  GroupIdT GroupLocalId = static_cast<GroupIdT>(LocalId);
+  using OCLT = detail::ConvertToOpenCLType_t<T>;
+  using WidenedT = WidenOpenCLTypeTo32_t<OCLT>;
+  using OCLIdT = detail::ConvertToOpenCLType_t<GroupIdT>;
+  WidenedT OCLX = detail::convertDataToType<T, OCLT>(x);
+  OCLIdT OCLId = detail::convertDataToType<GroupIdT, OCLIdT>(GroupLocalId);
+
+  return __spirv_GroupNonUniformBroadcast(
+      group_scope<ext::oneapi::experimental::opportunistic_group>::value, OCLX,
+      OCLId);
+}
 
 template <typename Group, typename T, typename IdT>
 EnableIfBitcastBroadcast<T, IdT> GroupBroadcast(Group g, T x, IdT local_id) {
@@ -956,6 +1026,18 @@ ControlBarrier(Group, memory_scope FenceScope, memory_order Order) {
 #endif
 }
 
+template <typename Group>
+struct is_tangle_or_opportunistic_group : std::false_type {};
+
+template <typename ParentGroup>
+struct is_tangle_or_opportunistic_group<
+    sycl::ext::oneapi::experimental::tangle_group<ParentGroup>>
+    : std::true_type {};
+
+template <>
+struct is_tangle_or_opportunistic_group<
+    sycl::ext::oneapi::experimental::opportunistic_group> : std::true_type {};
+
 // TODO: Refactor to avoid duplication after design settles
 #define __SYCL_GROUP_COLLECTIVE_OVERLOAD(Instruction)                          \
   template <__spv::GroupOperation Op, typename Group, typename T>              \
@@ -1037,6 +1119,24 @@ ControlBarrier(Group, memory_scope FenceScope, memory_order Order) {
       }                                                                        \
       return tmp;                                                              \
     }                                                                          \
+  }                                                                            \
+  template <__spv::GroupOperation Op, typename Group, typename T>              \
+  inline typename std::enable_if_t<                                            \
+      is_tangle_or_opportunistic_group<Group>::value, T>                       \
+      Group##Instruction(Group, T x) {                                         \
+    using ConvertedT = detail::ConvertToOpenCLType_t<T>;                       \
+                                                                               \
+    using OCLT = std::conditional_t<                                           \
+        std::is_same<ConvertedT, cl_char>() ||                                 \
+            std::is_same<ConvertedT, cl_short>(),                              \
+        cl_int,                                                                \
+        std::conditional_t<std::is_same<ConvertedT, cl_uchar>() ||             \
+                               std::is_same<ConvertedT, cl_ushort>(),          \
+                           cl_uint, ConvertedT>>;                              \
+    OCLT Arg = x;                                                              \
+    OCLT Ret = __spirv_GroupNonUniform##Instruction(                           \
+        group_scope<Group>::value, static_cast<unsigned int>(Op), Arg);        \
+    return Ret;                                                                \
   }
 
 __SYCL_GROUP_COLLECTIVE_OVERLOAD(SMin)

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

@@ -64,6 +64,8 @@ namespace ext::oneapi::experimental {
 // Forward declarations of non-uniform group types for algorithm definitions
 template <typename ParentGroup> class ballot_group;
 template <size_t PartitionSize, typename ParentGroup> class fixed_size_group;
+template <typename ParentGroup> class tangle_group;
+class opportunistic_group;
 
 } // namespace ext::oneapi::experimental
 

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

@@ -111,13 +111,16 @@ class opportunistic_group {
 #endif
   }
 
-private:
+protected:
   sub_group_mask Mask;
 
-protected:
   opportunistic_group(sub_group_mask m) : Mask(m) {}
 
   friend opportunistic_group this_kernel::get_opportunistic_group();
+
+  friend uint32_t
+  sycl::detail::IdToMaskPosition<opportunistic_group>(opportunistic_group Group,
+                                                      uint32_t Id);
 };
 
 namespace this_kernel {
@@ -144,5 +147,10 @@ template <>
 struct is_user_constructed_group<opportunistic_group> : std::true_type {};
 
 } // namespace ext::oneapi::experimental
+
+template <>
+struct is_group<ext::oneapi::experimental::opportunistic_group>
+    : std::true_type {};
+
 } // __SYCL_INLINE_VER_NAMESPACE(_V1)
 } // namespace sycl

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

@@ -112,13 +112,15 @@ template <typename ParentGroup> class tangle_group {
 #endif
   }
 
-private:
+protected:
   sub_group_mask Mask;
 
-protected:
   tangle_group(sub_group_mask m) : Mask(m) {}
 
   friend tangle_group<ParentGroup> get_tangle_group<ParentGroup>(ParentGroup);
+
+  friend uint32_t sycl::detail::IdToMaskPosition<tangle_group<ParentGroup>>(
+      tangle_group<ParentGroup> Group, uint32_t Id);
 };
 
 template <typename Group>
@@ -149,5 +151,10 @@ template <typename ParentGroup>
 struct is_user_constructed_group<tangle_group<ParentGroup>> : std::true_type {};
 
 } // namespace ext::oneapi::experimental
+
+template <typename ParentGroup>
+struct is_group<ext::oneapi::experimental::tangle_group<ParentGroup>>
+    : std::true_type {};
+
 } // __SYCL_INLINE_VER_NAMESPACE(_V1)
 } // namespace sycl

sycl/test-e2e/NonUniformGroups/opportunistic_group_algorithms.cpp

@@ -0,0 +1,131 @@
+// RUN: %clangxx -fsycl -fsycl-targets=%sycl_triple %s -o %t.out
+// RUN: %GPU_RUN_PLACEHOLDER %t.out
+//
+// UNSUPPORTED: cpu || cuda || hip
+
+#include <sycl/sycl.hpp>
+#include <vector>
+namespace syclex = sycl::ext::oneapi::experimental;
+
+class TestKernel;
+
+constexpr uint32_t SGSize = 32;
+constexpr uint32_t ArbitraryItem = 5;
+
+int main() {
+  sycl::queue Q;
+
+  auto SGSizes = Q.get_device().get_info<sycl::info::device::sub_group_sizes>();
+  if (std::find(SGSizes.begin(), SGSizes.end(), SGSize) == SGSizes.end()) {
+    std::cout << "Test skipped due to missing support for sub-group size 32."
+              << std::endl;
+    return 0;
+  }
+
+  sycl::buffer<size_t, 1> TmpBuf{sycl::range{SGSize}};
+  sycl::buffer<bool, 1> BarrierBuf{sycl::range{SGSize}};
+  sycl::buffer<bool, 1> BroadcastBuf{sycl::range{SGSize}};
+  sycl::buffer<bool, 1> AnyBuf{sycl::range{SGSize}};
+  sycl::buffer<bool, 1> AllBuf{sycl::range{SGSize}};
+  sycl::buffer<bool, 1> NoneBuf{sycl::range{SGSize}};
+  sycl::buffer<bool, 1> ReduceBuf{sycl::range{SGSize}};
+  sycl::buffer<bool, 1> ExScanBuf{sycl::range{SGSize}};
+  sycl::buffer<bool, 1> IncScanBuf{sycl::range{SGSize}};
+
+  const auto NDR = sycl::nd_range<1>{SGSize, SGSize};
+  Q.submit([&](sycl::handler &CGH) {
+    sycl::accessor TmpAcc{TmpBuf, CGH, sycl::write_only};
+    sycl::accessor BarrierAcc{BarrierBuf, CGH, sycl::write_only};
+    sycl::accessor BroadcastAcc{BroadcastBuf, CGH, sycl::write_only};
+    sycl::accessor AnyAcc{AnyBuf, CGH, sycl::write_only};
+    sycl::accessor AllAcc{AllBuf, CGH, sycl::write_only};
+    sycl::accessor NoneAcc{NoneBuf, CGH, sycl::write_only};
+    sycl::accessor ReduceAcc{ReduceBuf, CGH, sycl::write_only};
+    sycl::accessor ExScanAcc{ExScanBuf, CGH, sycl::write_only};
+    sycl::accessor IncScanAcc{IncScanBuf, CGH, sycl::write_only};
+    const auto KernelFunc =
+        [=](sycl::nd_item<1> item) [[sycl::reqd_sub_group_size(SGSize)]] {
+          auto WI = item.get_global_id();
+          auto SG = item.get_sub_group();
+
+          uint32_t OriginalLID = SG.get_local_linear_id();
+
+          // Given the dynamic nature of opportunistic groups, the simplest
+          // case we can reason about is a single work-item. This isn't a very
+          // robust test, but choosing an arbitrary work-item (i.e. rather
+          // than the leader) should test an implementation's ability to handle
+          // arbitrary group membership.
+          if (OriginalLID == ArbitraryItem) {
+            auto OpportunisticGroup =
+                syclex::this_kernel::get_opportunistic_group();
+
+            // This is trivial, but does test that group_barrier can be called.
+            TmpAcc[WI] = 1;
+            sycl::group_barrier(OpportunisticGroup);
+            size_t Visible = TmpAcc[WI];
+            BarrierAcc[WI] = (Visible == 1);
+
+            // Simple check of group algorithms.
+            uint32_t LID = OpportunisticGroup.get_local_linear_id();
+
+            uint32_t BroadcastResult =
+                sycl::group_broadcast(OpportunisticGroup, OriginalLID, 0);
+            BroadcastAcc[WI] = (BroadcastResult == OriginalLID);
+
+            bool AnyResult = sycl::any_of_group(OpportunisticGroup, (LID == 0));
+            AnyAcc[WI] = (AnyResult == true);
+
+            bool AllResult = sycl::all_of_group(OpportunisticGroup, (LID == 0));
+            AllAcc[WI] = (AllResult == true);
+
+            bool NoneResult =
+                sycl::none_of_group(OpportunisticGroup, (LID != 0));
+            NoneAcc[WI] = (NoneResult == true);
+
+            uint32_t ReduceResult =
+                sycl::reduce_over_group(OpportunisticGroup, 1, sycl::plus<>());
+            ReduceAcc[WI] =
+                (ReduceResult == OpportunisticGroup.get_local_linear_range());
+
+            uint32_t ExScanResult = sycl::exclusive_scan_over_group(
+                OpportunisticGroup, 1, sycl::plus<>());
+            ExScanAcc[WI] = (ExScanResult == LID);
+
+            uint32_t IncScanResult = sycl::inclusive_scan_over_group(
+                OpportunisticGroup, 1, sycl::plus<>());
+            IncScanAcc[WI] = (IncScanResult == LID + 1);
+          } else {
+            BarrierAcc[WI] = false;
+            BroadcastAcc[WI] = false;
+            AnyAcc[WI] = false;
+            AllAcc[WI] = false;
+            NoneAcc[WI] = false;
+            ReduceAcc[WI] = false;
+            ExScanAcc[WI] = false;
+            IncScanAcc[WI] = false;
+          }
+        };
+    CGH.parallel_for<TestKernel>(NDR, KernelFunc);
+  });
+
+  sycl::host_accessor BarrierAcc{BarrierBuf, sycl::read_only};
+  sycl::host_accessor BroadcastAcc{BroadcastBuf, sycl::read_only};
+  sycl::host_accessor AnyAcc{AnyBuf, sycl::read_only};
+  sycl::host_accessor AllAcc{AllBuf, sycl::read_only};
+  sycl::host_accessor NoneAcc{NoneBuf, sycl::read_only};
+  sycl::host_accessor ReduceAcc{ReduceBuf, sycl::read_only};
+  sycl::host_accessor ExScanAcc{ExScanBuf, sycl::read_only};
+  sycl::host_accessor IncScanAcc{IncScanBuf, sycl::read_only};
+  for (uint32_t WI = 0; WI < 32; ++WI) {
+    bool ExpectedResult = (WI == ArbitraryItem);
+    assert(BarrierAcc[WI] == ExpectedResult);
+    assert(BroadcastAcc[WI] == ExpectedResult);
+    assert(AnyAcc[WI] == ExpectedResult);
+    assert(AllAcc[WI] == ExpectedResult);
+    assert(NoneAcc[WI] == ExpectedResult);
+    assert(ReduceAcc[WI] == ExpectedResult);
+    assert(ExScanAcc[WI] == ExpectedResult);
+    assert(IncScanAcc[WI] == ExpectedResult);
+  }
+  return 0;
+}