some indexing fixes for 3-level reduce/scan

Author: keptsecret

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

@@ -51,22 +51,23 @@ struct ArithmeticConfiguration
     NBL_CONSTEXPR_STATIC_INLINE uint16_t VirtualWorkgroupSize = uint16_t(0x1u) << virtual_wg_t::value;
     static_assert(VirtualWorkgroupSize<=WorkgroupSize*SubgroupSize);
 
-    NBL_CONSTEXPR_STATIC_INLINE uint16_t __SubgroupsPerVirtualWorkgroupLog2 = mpl::max_v<uint16_t, WorkgroupSizeLog2-SubgroupSizeLog2, SubgroupSizeLog2>;
-    NBL_CONSTEXPR_STATIC_INLINE uint16_t __SubgroupsPerVirtualWorkgroup = uint16_t(0x1u) << __SubgroupsPerVirtualWorkgroupLog2;
-
     using items_per_invoc_t = impl::items_per_invocation<virtual_wg_t, _ItemsPerInvocation>;
     // NBL_CONSTEXPR_STATIC_INLINE uint32_t2 ItemsPerInvocation;    TODO? doesn't allow inline definitions for uint32_t2 for some reason, uint32_t[2] as well ; declaring out of line results in not constant expression
     NBL_CONSTEXPR_STATIC_INLINE uint16_t ItemsPerInvocation_0 = items_per_invoc_t::value0;
     NBL_CONSTEXPR_STATIC_INLINE uint16_t ItemsPerInvocation_1 = items_per_invoc_t::value1;
     NBL_CONSTEXPR_STATIC_INLINE uint16_t ItemsPerInvocation_2 = items_per_invoc_t::value2;
     static_assert(ItemsPerInvocation_1<=4, "3 level scan would have been needed with this config!");
 
+    NBL_CONSTEXPR_STATIC_INLINE uint16_t __ItemsPerVirtualWorkgroupLog2 = mpl::max_v<uint16_t, WorkgroupSizeLog2-SubgroupSizeLog2, SubgroupSizeLog2>;
+    NBL_CONSTEXPR_STATIC_INLINE uint16_t __ItemsPerVirtualWorkgroup = uint16_t(0x1u) << __ItemsPerVirtualWorkgroupLog2;
+    NBL_CONSTEXPR_STATIC_INLINE uint16_t __SubgroupsPerVirtualWorkgroup = __ItemsPerVirtualWorkgroup / ItemsPerInvocation_1;
+
     NBL_CONSTEXPR_STATIC_INLINE uint32_t SharedScratchElementCount = conditional_value<LevelCount==1,uint16_t,
         0,
         conditional_value<LevelCount==3,uint16_t,
-            SubgroupSize*ItemsPerInvocation_2,
-            0
-            >::value + SubgroupSize*ItemsPerInvocation_1
+            SubgroupSize*ItemsPerInvocation_2+__ItemsPerVirtualWorkgroup,
+            SubgroupSize*ItemsPerInvocation_1
+            >::value
         >::value;
 
     static bool electLast()

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

@@ -245,7 +245,7 @@ struct reduce<Config, BinOp, 3, device_capabilities>
 
         const uint32_t invocationIndex = workgroup::SubgroupContiguousIndex();
         // level 1 scan
-        const uint32_t lv1_smem_size = Config::SubgroupsSize*Config::ItemsPerInvocation_1;
+        const uint32_t lv1_smem_size = Config::__ItemsPerVirtualWorkgroup;
         subgroup2::reduction<params_lv1_t> reduction1;
         if (glsl::gl_SubgroupID() < Config::SubgroupSize*Config::ItemsPerInvocation_2)
         {
@@ -303,8 +303,8 @@ struct scan<Config, BinOp, Exclusive, 3, device_capabilities>
 
         const uint32_t invocationIndex = workgroup::SubgroupContiguousIndex();
         // level 1 scan
-        const uint32_t lv1_smem_size = Config::SubgroupsSize*Config::ItemsPerInvocation_1;
-        const uint32_t lv1_num_invoc = Config::SubgroupsSize*Config::ItemsPerInvocation_2;
+        const uint32_t lv1_smem_size = Config::__ItemsPerVirtualWorkgroup;
+        const uint32_t lv1_num_invoc = Config::SubgroupSize*Config::ItemsPerInvocation_2;
         subgroup2::exclusive_scan<params_lv1_t> exclusiveScan1;
         if (glsl::gl_SubgroupID() < lv1_num_invoc)
         {