Merge remote-tracking branch 'remotes/origin/sync-subgroup-fix'

Author: devsh

examples_tests

@@ -15,14 +15,19 @@ namespace hlsl
 {
 namespace spirv
 {
+[[vk::ext_capability(spv::CapabilityGroupNonUniformBallot)]]
 [[vk::ext_builtin_input(spv::BuiltInSubgroupEqMask)]]
 static const uint32_t4 BuiltInSubgroupEqMask;
+[[vk::ext_capability(spv::CapabilityGroupNonUniformBallot)]]
 [[vk::ext_builtin_input(spv::BuiltInSubgroupGeMask)]]
 static const uint32_t4 BuiltInSubgroupGeMask;
+[[vk::ext_capability(spv::CapabilityGroupNonUniformBallot)]]
 [[vk::ext_builtin_input(spv::BuiltInSubgroupGtMask)]]
 static const uint32_t4 BuiltInSubgroupGtMask;
+[[vk::ext_capability(spv::CapabilityGroupNonUniformBallot)]]
 [[vk::ext_builtin_input(spv::BuiltInSubgroupLeMask)]]
 static const uint32_t4 BuiltInSubgroupLeMask;
+[[vk::ext_capability(spv::CapabilityGroupNonUniformBallot)]]
 [[vk::ext_builtin_input(spv::BuiltInSubgroupLtMask)]]
 static const uint32_t4 BuiltInSubgroupLtMask;
 

include/nbl/builtin/hlsl/spirv_intrinsics/subgroup_ballot.hlsl

@@ -0,0 +1,64 @@
+// Copyright (C) 2025 - DevSH Graphics Programming Sp. z O.O.
+// This file is part of the "Nabla Engine".
+// For conditions of distribution and use, see copyright notice in nabla.h
+#ifndef _NBL_BUILTIN_HLSL_SUBGROUP2_ARITHMETIC_PARAMS_INCLUDED_
+#define _NBL_BUILTIN_HLSL_SUBGROUP2_ARITHMETIC_PARAMS_INCLUDED_
+
+#ifdef __HLSL_VERSION
+#include "nbl/builtin/hlsl/device_capabilities_traits.hlsl"
+#include "nbl/builtin/hlsl/subgroup2/ballot.hlsl"
+#endif
+#include "nbl/builtin/hlsl/concepts.hlsl"
+
+namespace nbl
+{
+namespace hlsl
+{
+namespace subgroup2
+{
+
+#ifdef __HLSL_VERSION
+template<typename Config, class BinOp, int32_t _ItemsPerInvocation=1, class device_capabilities=void NBL_PRIMARY_REQUIRES(is_configuration_v<Config> && is_scalar_v<typename BinOp::type_t>)
+struct ArithmeticParams
+{
+    using config_t = Config;
+    using binop_t = BinOp;
+    using scalar_t = typename BinOp::type_t;
+    using type_t = vector<scalar_t, _ItemsPerInvocation>;
+    using device_traits = device_capabilities_traits<device_capabilities>;
+
+    NBL_CONSTEXPR_STATIC_INLINE int32_t ItemsPerInvocation = _ItemsPerInvocation;
+    NBL_CONSTEXPR_STATIC_INLINE bool UseNativeIntrinsics = device_capabilities_traits<device_capabilities>::shaderSubgroupArithmetic /*&& /*some heuristic for when its faster*/;
+    // TODO add a IHV enum to device_capabilities_traits to check !is_nvidia
+};
+#endif
+
+#ifndef __HLSL_VERSION
+#include <sstream>
+#include <string>
+struct SArithmeticParams
+{
+    void init(const uint16_t _SubgroupSizeLog2, const uint16_t _ItemsPerInvocation)
+    {
+        SubgroupSizeLog2 = _SubgroupSizeLog2;
+        ItemsPerInvocation = _ItemsPerInvocation;
+    }
+
+    // alias should provide Binop and device_capabilities template parameters
+    std::string getParamTemplateStructString()
+    {
+        std::ostringstream os;
+        os << "nbl::hlsl::subgroup2::ArithmeticParams<nbl::hlsl::subgroup2::Configuration<" << SubgroupSizeLog2 << ">, Binop," << ItemsPerInvocation << ", device_capabilities>;";
+        return os.str();
+    }
+
+    uint32_t SubgroupSizeLog2;
+    uint32_t ItemsPerInvocation;
+};
+#endif
+
+}
+}
+}
+
+#endif

include/nbl/builtin/hlsl/subgroup2/arithmetic_params.hlsl

@@ -4,13 +4,8 @@
 #ifndef _NBL_BUILTIN_HLSL_SUBGROUP2_ARITHMETIC_PORTABILITY_INCLUDED_
 #define _NBL_BUILTIN_HLSL_SUBGROUP2_ARITHMETIC_PORTABILITY_INCLUDED_
 
-
 #include "nbl/builtin/hlsl/device_capabilities_traits.hlsl"
-
-#include "nbl/builtin/hlsl/subgroup2/ballot.hlsl"
 #include "nbl/builtin/hlsl/subgroup2/arithmetic_portability_impl.hlsl"
-#include "nbl/builtin/hlsl/concepts.hlsl"
-
 
 namespace nbl
 {
@@ -19,20 +14,6 @@ namespace hlsl
 namespace subgroup2
 {
 
-template<typename Config, class BinOp, int32_t _ItemsPerInvocation=1, class device_capabilities=void NBL_PRIMARY_REQUIRES(is_configuration_v<Config> && is_scalar_v<typename BinOp::type_t>)
-struct ArithmeticParams
-{
-    using config_t = Config;
-    using binop_t = BinOp;
-    using scalar_t = typename BinOp::type_t;
-    using type_t = vector<scalar_t, _ItemsPerInvocation>;
-    using device_traits = device_capabilities_traits<device_capabilities>;
-
-    NBL_CONSTEXPR_STATIC_INLINE int32_t ItemsPerInvocation = _ItemsPerInvocation;
-    NBL_CONSTEXPR_STATIC_INLINE bool UseNativeIntrinsics = device_capabilities_traits<device_capabilities>::shaderSubgroupArithmetic /*&& /*some heuristic for when its faster*/;
-    // TODO add a IHV enum to device_capabilities_traits to check !is_nvidia
-};
-
 template<typename Params>
 struct reduction : impl::reduction<Params,typename Params::binop_t,Params::ItemsPerInvocation,Params::UseNativeIntrinsics> {};
 template<typename Params>

include/nbl/builtin/hlsl/subgroup2/arithmetic_portability.hlsl

@@ -158,6 +158,10 @@ struct inclusive_scan<Params, BinOp, 1, false>
 
     static scalar_t __call(scalar_t value)
     {
+        // sync up each subgroup invocation so it runs in lockstep
+        // not ideal because might not write to shared memory but a storage class is needed
+        spirv::memoryBarrier(spv::ScopeSubgroup, spv::MemorySemanticsWorkgroupMemoryMask | spv::MemorySemanticsAcquireMask);
+
         binop_t op;
         const uint32_t subgroupInvocation = glsl::gl_SubgroupInvocationID();
 
@@ -185,6 +189,10 @@ struct exclusive_scan<Params, BinOp, 1, false>
 
     scalar_t operator()(scalar_t value)
     {
+        // sync up each subgroup invocation so it runs in lockstep
+        // not ideal because might not write to shared memory but a storage class is needed
+        spirv::memoryBarrier(spv::ScopeSubgroup, spv::MemorySemanticsWorkgroupMemoryMask | spv::MemorySemanticsAcquireMask);
+
         scalar_t left = hlsl::mix(binop_t::identity, glsl::subgroupShuffleUp<scalar_t>(value,1), bool(glsl::gl_SubgroupInvocationID()));
         return inclusive_scan<Params, BinOp, 1, false>::__call(left);
     }
@@ -203,8 +211,11 @@ struct reduction<Params, BinOp, 1, false>
 
     scalar_t operator()(scalar_t value)
     {
-        binop_t op;
+        // sync up each subgroup invocation so it runs in lockstep
+        // not ideal because might not write to shared memory but a storage class is needed
+        spirv::memoryBarrier(spv::ScopeSubgroup, spv::MemorySemanticsWorkgroupMemoryMask | spv::MemorySemanticsAcquireMask);
 
+        binop_t op;
         const uint32_t SubgroupSizeLog2 = config_t::SizeLog2;
         [unroll]
         for (uint32_t i = 0; i < integral_constant<uint32_t,SubgroupSizeLog2>::value; i++)

include/nbl/builtin/hlsl/subgroup2/arithmetic_portability_impl.hlsl

@@ -181,6 +181,8 @@ struct SItemsPerInvoc
 };
 }
 
+#include <sstream>
+#include <string>
 struct SArithmeticConfiguration
 {
     void init(const uint16_t _WorkgroupSizeLog2, const uint16_t _SubgroupSizeLog2, const uint16_t _ItemsPerInvocation)
@@ -203,6 +205,13 @@ struct SArithmeticConfiguration
         #undef DEFINE_ASSIGN
     }
 
+    std::string getConfigTemplateStructString()
+    {
+        std::ostringstream os;
+        os << "nbl::hlsl::workgroup2::ArithmeticConfiguration<" << WorkgroupSizeLog2 << "," << SubgroupSizeLog2 << "," << ItemsPerInvocation_0 << ">;";
+        return os.str();
+    }
+
     #define DEFINE_ASSIGN(TYPE,ID,...) TYPE ID;
     #include "impl/arithmetic_config_def.hlsl"
     #undef DEFINE_ASSIGN

include/nbl/builtin/hlsl/workgroup2/arithmetic_config.hlsl

@@ -0,0 +1,20 @@
+// Copyright (C) 2023 - DevSH Graphics Programming Sp. z O.O.
+// This file is part of the "Nabla Engine".
+// For conditions of distribution and use, see copyright notice in nabla.h
+#ifndef _NBL_BUILTIN_HLSL_WORKGROUP2_BASIC_INCLUDED_
+#define _NBL_BUILTIN_HLSL_WORKGROUP2_BASIC_INCLUDED_
+
+#include "nbl/builtin/hlsl/workgroup/basic.hlsl"
+
+namespace nbl
+{
+namespace hlsl
+{
+namespace workgroup2
+{
+// empty
+}
+}
+}
+
+#endif

include/nbl/builtin/hlsl/workgroup2/basic.hlsl

@@ -7,6 +7,7 @@
 #include "nbl/builtin/hlsl/workgroup/broadcast.hlsl"
 #include "nbl/builtin/hlsl/glsl_compat/subgroup_basic.hlsl"
 #include "nbl/builtin/hlsl/subgroup2/ballot.hlsl"
+#include "nbl/builtin/hlsl/subgroup2/arithmetic_params.hlsl"
 #include "nbl/builtin/hlsl/subgroup2/arithmetic_portability.hlsl"
 #include "nbl/builtin/hlsl/mpl.hlsl"
 #include "nbl/builtin/hlsl/workgroup2/arithmetic_config.hlsl"

include/nbl/builtin/hlsl/workgroup2/shared_scan.hlsl

@@ -332,6 +332,7 @@ LIST_BUILTIN_RESOURCE(NBL_RESOURCES_TO_EMBED "hlsl/subgroup/arithmetic_portabili
 LIST_BUILTIN_RESOURCE(NBL_RESOURCES_TO_EMBED "hlsl/subgroup/fft.hlsl")
 #subgroup2
 LIST_BUILTIN_RESOURCE(NBL_RESOURCES_TO_EMBED "hlsl/subgroup2/ballot.hlsl")
+LIST_BUILTIN_RESOURCE(NBL_RESOURCES_TO_EMBED "hlsl/subgroup2/arithmetic_params.hlsl")
 LIST_BUILTIN_RESOURCE(NBL_RESOURCES_TO_EMBED "hlsl/subgroup2/arithmetic_portability.hlsl")
 LIST_BUILTIN_RESOURCE(NBL_RESOURCES_TO_EMBED "hlsl/subgroup2/arithmetic_portability_impl.hlsl")
 #shared header between C++ and HLSL
@@ -346,6 +347,7 @@ LIST_BUILTIN_RESOURCE(NBL_RESOURCES_TO_EMBED "hlsl/workgroup/scratch_size.hlsl")
 LIST_BUILTIN_RESOURCE(NBL_RESOURCES_TO_EMBED "hlsl/workgroup/shared_scan.hlsl")
 LIST_BUILTIN_RESOURCE(NBL_RESOURCES_TO_EMBED "hlsl/workgroup/shuffle.hlsl")
 #workgroup2
+LIST_BUILTIN_RESOURCE(NBL_RESOURCES_TO_EMBED "hlsl/workgroup2/basic.hlsl")
 LIST_BUILTIN_RESOURCE(NBL_RESOURCES_TO_EMBED "hlsl/workgroup2/arithmetic_config.hlsl")
 LIST_BUILTIN_RESOURCE(NBL_RESOURCES_TO_EMBED "hlsl/workgroup2/impl/virtual_wg_size_def.hlsl")
 LIST_BUILTIN_RESOURCE(NBL_RESOURCES_TO_EMBED "hlsl/workgroup2/impl/items_per_invoc_def.hlsl")