rework specializations for native, emulated funcs

Author: keptsecret

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

@@ -19,7 +19,7 @@ namespace hlsl
 namespace subgroup2
 {
 
-template<typename Config, class BinOp, int32_t _ItemsPerInvocation=1, class device_capabilities=void, bool OverrideUseNativeInstrinsics=true NBL_PRIMARY_REQUIRES(is_configuration_v<Config>)
+template<typename Config, class BinOp, int32_t _ItemsPerInvocation=1, class device_capabilities=void NBL_PRIMARY_REQUIRES(is_configuration_v<Config>)
 struct ArithmeticParams
 {
     using config_t = Config;
@@ -30,15 +30,15 @@ struct ArithmeticParams
     NBL_CONSTEXPR_STATIC_INLINE int32_t ItemsPerInvocation = _ItemsPerInvocation;
     // if OverrideUseNativeInstrinsics is true, tries to use native spirv intrinsics
     // if OverrideUseNativeInstrinsics is false, will always use emulated versions
-    NBL_CONSTEXPR_STATIC_INLINE bool UseNativeIntrinsics = device_capabilities_traits<device_capabilities>::shaderSubgroupArithmetic && OverrideUseNativeInstrinsics /*&& /*some heuristic for when its faster*/;
+    NBL_CONSTEXPR_STATIC_INLINE bool UseNativeIntrinsics = device_capabilities_traits<device_capabilities>::shaderSubgroupArithmetic /*&& /*some heuristic for when its faster*/;
 };
 
 template<typename Params>
-struct reduction : impl::reduction<Params,Params::ItemsPerInvocation,Params::UseNativeIntrinsics> {};
+struct reduction : impl::reduction<Params,typename Params::binop_t,Params::ItemsPerInvocation,Params::UseNativeIntrinsics> {};
 template<typename Params>
-struct inclusive_scan : impl::inclusive_scan<Params,Params::ItemsPerInvocation,Params::UseNativeIntrinsics> {};
+struct inclusive_scan : impl::inclusive_scan<Params,typename Params::binop_t,Params::ItemsPerInvocation,Params::UseNativeIntrinsics> {};
 template<typename Params>
-struct exclusive_scan : impl::exclusive_scan<Params,Params::ItemsPerInvocation,Params::UseNativeIntrinsics> {};
+struct exclusive_scan : impl::exclusive_scan<Params,typename Params::binop_t,Params::ItemsPerInvocation,Params::UseNativeIntrinsics> {};
 
 }
 }

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

@@ -4,14 +4,15 @@
 #ifndef _NBL_BUILTIN_HLSL_SUBGROUP2_ARITHMETIC_PORTABILITY_IMPL_INCLUDED_
 #define _NBL_BUILTIN_HLSL_SUBGROUP2_ARITHMETIC_PORTABILITY_IMPL_INCLUDED_
 
-// #include "nbl/builtin/hlsl/glsl_compat/subgroup_shuffle.hlsl"
-// #include "nbl/builtin/hlsl/glsl_compat/subgroup_arithmetic.hlsl"
+#include "nbl/builtin/hlsl/glsl_compat/subgroup_shuffle.hlsl"
+#include "nbl/builtin/hlsl/glsl_compat/subgroup_arithmetic.hlsl"
 
-// #include "nbl/builtin/hlsl/subgroup/ballot.hlsl"
+#include "nbl/builtin/hlsl/subgroup/ballot.hlsl"
+#include "nbl/builtin/hlsl/subgroup2/ballot.hlsl"
 
-// #include "nbl/builtin/hlsl/functional.hlsl"
+#include "nbl/builtin/hlsl/functional.hlsl"
 
-#include "nbl/builtin/hlsl/subgroup/arithmetic_portability_impl.hlsl"
+// #include "nbl/builtin/hlsl/subgroup/arithmetic_portability_impl.hlsl"
 
 namespace nbl
 {
@@ -23,12 +24,14 @@ namespace subgroup2
 namespace impl
 {
 
-template<class Params, uint32_t ItemsPerInvocation, bool native>
+// BinOp needed to specialize native
+template<class Params, class BinOp, uint32_t ItemsPerInvocation, bool native>
 struct inclusive_scan
 {
     using type_t = typename Params::type_t;
     using scalar_t = typename Params::scalar_t;
     using binop_t = typename Params::binop_t;
+    // assert binop_t == BinOp
     using exclusive_scan_op_t = subgroup::impl::exclusive_scan<binop_t, native>;
 
     // NBL_CONSTEXPR_STATIC_INLINE uint32_t ItemsPerInvocation = vector_traits<T>::Dimension;
@@ -52,13 +55,13 @@ struct inclusive_scan
     }
 };
 
-template<class Params, uint32_t ItemsPerInvocation, bool native>
+template<class Params, class BinOp, uint32_t ItemsPerInvocation, bool native>
 struct exclusive_scan
 {
     using type_t = typename Params::type_t;
     using scalar_t = typename Params::scalar_t;
     using binop_t = typename Params::binop_t;
-    using inclusive_scan_op_t = subgroup2::impl::inclusive_scan<Params, ItemsPerInvocation, native>;
+    using inclusive_scan_op_t = subgroup2::impl::inclusive_scan<Params, binop_t, ItemsPerInvocation, native>;
 
     // NBL_CONSTEXPR_STATIC_INLINE uint32_t ItemsPerInvocation = vector_traits<T>::Dimension;
 
@@ -78,7 +81,7 @@ struct exclusive_scan
     }
 };
 
-template<class Params, uint32_t ItemsPerInvocation, bool native>
+template<class Params, class BinOp, uint32_t ItemsPerInvocation, bool native>
 struct reduction
 {
     using type_t = typename Params::type_t;
@@ -103,74 +106,98 @@ struct reduction
 
 // specs for N=1 uses subgroup funcs
 // specialize native
-// #define SPECIALIZE(NAME,BINOP,SUBGROUP_OP) template<typename T> struct NAME<BINOP<T>,true> \
-// { \
-//     using type_t = T; \
-//  \
-//     type_t operator()(NBL_CONST_REF_ARG(type_t) v) {return glsl::subgroup##SUBGROUP_OP<type_t>(v);} \
-// }
+#define SPECIALIZE(NAME,BINOP,SUBGROUP_OP) template<class Params, typename T> struct NAME<Params,BINOP<T>,1,true> \
+{ \
+    using type_t = T; \
+ \
+    type_t operator()(NBL_CONST_REF_ARG(type_t) v) {return glsl::subgroup##SUBGROUP_OP<type_t>(v);} \
+}
 
-// #define SPECIALIZE_ALL(BINOP,SUBGROUP_OP) SPECIALIZE(reduction,BINOP,SUBGROUP_OP); \
-//     SPECIALIZE(inclusive_scan,BINOP,Inclusive##SUBGROUP_OP); \
-//     SPECIALIZE(exclusive_scan,BINOP,Exclusive##SUBGROUP_OP);
+#define SPECIALIZE_ALL(BINOP,SUBGROUP_OP) SPECIALIZE(reduction,BINOP,SUBGROUP_OP); \
+    SPECIALIZE(inclusive_scan,BINOP,Inclusive##SUBGROUP_OP); \
+    SPECIALIZE(exclusive_scan,BINOP,Exclusive##SUBGROUP_OP);
 
-// SPECIALIZE_ALL(bit_and,And);
-// SPECIALIZE_ALL(bit_or,Or);
-// SPECIALIZE_ALL(bit_xor,Xor);
+SPECIALIZE_ALL(bit_and,And);
+SPECIALIZE_ALL(bit_or,Or);
+SPECIALIZE_ALL(bit_xor,Xor);
 
-// SPECIALIZE_ALL(plus,Add);
-// SPECIALIZE_ALL(multiplies,Mul);
+SPECIALIZE_ALL(plus,Add);
+SPECIALIZE_ALL(multiplies,Mul);
 
-// SPECIALIZE_ALL(minimum,Min);
-// SPECIALIZE_ALL(maximum,Max);
+SPECIALIZE_ALL(minimum,Min);
+SPECIALIZE_ALL(maximum,Max);
 
-// #undef SPECIALIZE_ALL
-// #undef SPECIALIZE
+#undef SPECIALIZE_ALL
+#undef SPECIALIZE
 
 // specialize portability
-template<class Params, bool native>
-struct inclusive_scan<Params, 1, native>
+template<class Params, class BinOp>
+struct inclusive_scan<Params, BinOp, 1, false>
 {
     using type_t = typename Params::type_t;
     using scalar_t = typename Params::scalar_t;
     using binop_t = typename Params::binop_t;
-    using op_t = subgroup::impl::inclusive_scan<binop_t, native>;
     // assert T == scalar type, binop::type == T
+    using config_t = typename Params::config_t;
 
-    type_t operator()(NBL_CONST_REF_ARG(type_t) value)
+    // affected by https://github.com/microsoft/DirectXShaderCompiler/issues/7006
+    // NBL_CONSTEXPR_STATIC_INLINE uint32_t SubgroupSizeLog2 = config_t::SizeLog2;
+
+    type_t operator()(type_t value)
     {
-        op_t op;
-        return op(value);
+        return __call(value);
+    }
+
+    static type_t __call(type_t value)
+    {
+        binop_t op;
+        const uint32_t subgroupInvocation = glsl::gl_SubgroupInvocationID();
+    
+        type_t rhs = glsl::subgroupShuffleUp<type_t>(value, 1u); // all invocations must execute the shuffle, even if we don't apply the op() to all of them
+        // TODO waiting on mix intrinsic fix from bxdf branch, value = op(value, hlsl::mix(rhs, binop_t::identity, subgroupInvocation < 1u));
+        value = op(value, subgroupInvocation<1u ? binop_t::identity : rhs);
+    
+        const uint32_t SubgroupSizeLog2 = config_t::SizeLog2;
+        [unroll]
+        for (uint32_t i = 1; i < integral_constant<uint32_t,SubgroupSizeLog2>::value; i++)
+        {
+            const uint32_t step = i * 2;
+            rhs = glsl::subgroupShuffleUp<type_t>(value, step);
+            // TODO value = op(value, hlsl::mix(rhs, binop_t::identity, subgroupInvocation < step));
+            value = op(value, subgroupInvocation<step ? binop_t::identity : rhs);
+        }
+        return value;
     }
 };
 
-template<class Params, bool native>
-struct exclusive_scan<Params, 1, native>
+template<class Params, class BinOp>
+struct exclusive_scan<Params, BinOp, 1, false>
 {
     using type_t = typename Params::type_t;
     using scalar_t = typename Params::scalar_t;
     using binop_t = typename Params::binop_t;
-    using op_t = subgroup::impl::exclusive_scan<binop_t, native>;
 
     type_t operator()(NBL_CONST_REF_ARG(type_t) value)
     {
-        op_t op;
-        return op(value);
+        value = inclusive_scan<Params, BinOp, 1, false>::__call(value);
+        // can't risk getting short-circuited, need to store to a var
+        type_t left = glsl::subgroupShuffleUp<type_t>(value,1);
+        // the first invocation doesn't have anything in its left so we set to the binop's identity value for exlusive scan
+        return bool(glsl::gl_SubgroupInvocationID()) ? left:binop_t::identity;
     }
 };
 
-template<class Params, bool native>
-struct reduction<Params, 1, native>
+template<class Params, class BinOp>
+struct reduction<Params, BinOp, 1, false>
 {
     using type_t = typename Params::type_t;
     using scalar_t = typename Params::scalar_t;
     using binop_t = typename Params::binop_t;
-    using op_t = subgroup::impl::reduction<binop_t, native>;
 
     scalar_t operator()(NBL_CONST_REF_ARG(type_t) value)
     {
-        op_t op;
-        return op(value);
+        // take the last subgroup invocation's value for the reduction
+        return subgroup::BroadcastLast<type_t>(inclusive_scan<Params, BinOp, 1, false>::__call(value));
     }
 };