Small math refactor, doing away with workgroup size log 2 requirement

Author: Fletterio

include/nbl/builtin/hlsl/workgroup/fft.hlsl

@@ -16,6 +16,7 @@ namespace workgroup
 {
 namespace fft
 {
+
 // ---------------------------------- Utils -----------------------------------------------
 
 template<typename SharedMemoryAccessor, typename Scalar>
@@ -166,33 +167,25 @@ struct FFT<K, false, Scalar, device_capabilities>
     template<typename Accessor, typename SharedMemoryAccessor>
     static enable_if_t< (mpl::is_pot_v<K> && K > 2), void > __call(NBL_REF_ARG(Accessor) accessor, NBL_REF_ARG(SharedMemoryAccessor) sharedmemAccessor)
     {
-        static const uint32_t virtualThreadCount = K >> 1;
-        static const uint16_t passes = mpl::log2<K>::value - 1;
-        uint32_t stride = K >> 1;
-        //[unroll(passes)]
-        for (uint16_t pass = 0; pass < passes; pass++)
+        for (uint32_t stride = (K >> 1) * _NBL_HLSL_WORKGROUP_SIZE_; stride > _NBL_HLSL_WORKGROUP_SIZE_; stride >>= 1)
         {
             //[unroll(K/2)]
-            for (uint32_t virtualThread = 0; virtualThread < virtualThreadCount; virtualThread++)
+            for (uint32_t virtualThreadID = SubgroupContiguousIndex(); virtualThreadID < (K >> 1) * _NBL_HLSL_WORKGROUP_SIZE_; virtualThreadID += _NBL_HLSL_WORKGROUP_SIZE_)
             {
-                const uint32_t virtualThreadID = (virtualThread << _NBL_HLSL_WORKGROUP_SIZE_LOG_2_) | SubgroupContiguousIndex();
-
-                const uint32_t lsb = virtualThread & (stride - 1);
-                const uint32_t loIx = ((virtualThread ^ lsb) << 1) | lsb;
+                const uint32_t loIx = ((virtualThreadID & (~(stride - 1))) << 1) | (virtualThreadID & (stride - 1));
                 const uint32_t hiIx = loIx | stride;
 
-                complex_t<Scalar> lo = accessor.get((loIx << _NBL_HLSL_WORKGROUP_SIZE_LOG_2_) | SubgroupContiguousIndex());
-                complex_t<Scalar> hi = accessor.get((hiIx << _NBL_HLSL_WORKGROUP_SIZE_LOG_2_) | SubgroupContiguousIndex());
+                complex_t<Scalar> lo = accessor.get(loIx);
+                complex_t<Scalar> hi = accessor.get(hiIx);
 
-                hlsl::fft::DIF<Scalar>::radix2(hlsl::fft::twiddle<false,Scalar>(virtualThreadID & ((stride << _NBL_HLSL_WORKGROUP_SIZE_LOG_2_) - 1), stride << _NBL_HLSL_WORKGROUP_SIZE_LOG_2_),lo,hi);
+                hlsl::fft::DIF<Scalar>::radix2(hlsl::fft::twiddle<false,Scalar>(virtualThreadID & (stride - 1), stride),lo,hi);
 
-                accessor.set((loIx << _NBL_HLSL_WORKGROUP_SIZE_LOG_2_) | SubgroupContiguousIndex(), lo);
-                accessor.set((hiIx << _NBL_HLSL_WORKGROUP_SIZE_LOG_2_) | SubgroupContiguousIndex(), hi);
+                accessor.set(loIx, lo);
+                accessor.set(hiIx, hi);
             }
             accessor.memoryBarrier(); // no execution barrier just making sure writes propagate to accessor
-            stride >>= 1;
         }
-        
+
         // do K/2 small workgroup FFTs
         DynamicOffsetAccessor < Accessor, complex_t<Scalar> > offsetAccessor;
         //[unroll(K/2)]
@@ -225,41 +218,34 @@ struct FFT<K, true, Scalar, device_capabilities>
             FFT<2,true, Scalar, device_capabilities>::template __call(offsetAccessor,sharedmemAccessor);
         }
         accessor = offsetAccessor.accessor;
-        
-        static const uint32_t virtualThreadCount = K >> 1;
-        static const uint16_t passes = mpl::log2<K>::value - 1;
-        uint32_t stride = 2;
-        //[unroll(passes)]
-        for (uint16_t pass = 0; pass < passes; pass++)
+      
+        for (uint32_t stride = _NBL_HLSL_WORKGROUP_SIZE_ << 1; stride <= (K >> 1) * _NBL_HLSL_WORKGROUP_SIZE_; stride <<= 1)
         {
             //[unroll(K/2)]
-            for (uint32_t virtualThread = 0; virtualThread < virtualThreadCount; virtualThread++)
+            for (uint32_t virtualThreadID = SubgroupContiguousIndex(); virtualThreadID < (K >> 1) * _NBL_HLSL_WORKGROUP_SIZE_; virtualThreadID += _NBL_HLSL_WORKGROUP_SIZE_)
             {
-                const uint32_t virtualThreadID = (virtualThread << _NBL_HLSL_WORKGROUP_SIZE_LOG_2_) | SubgroupContiguousIndex();
-
-                const uint32_t lsb = virtualThread & (stride - 1);
-                const uint32_t loIx = ((virtualThread ^ lsb) << 1) | lsb;
+                const uint32_t loIx = ((virtualThreadID & (~(stride - 1))) << 1) | (virtualThreadID & (stride - 1));
                 const uint32_t hiIx = loIx | stride;
+                
+                complex_t<Scalar> lo = accessor.get(loIx);
+                complex_t<Scalar> hi = accessor.get(hiIx);
 
-                complex_t<Scalar> lo = accessor.get((loIx << _NBL_HLSL_WORKGROUP_SIZE_LOG_2_) | SubgroupContiguousIndex());
-                complex_t<Scalar> hi = accessor.get((hiIx << _NBL_HLSL_WORKGROUP_SIZE_LOG_2_) | SubgroupContiguousIndex());
-
-                hlsl::fft::DIT<Scalar>::radix2(hlsl::fft::twiddle<true,Scalar>(virtualThreadID & ((stride << _NBL_HLSL_WORKGROUP_SIZE_LOG_2_) - 1), stride << _NBL_HLSL_WORKGROUP_SIZE_LOG_2_),lo,hi);
+                hlsl::fft::DIT<Scalar>::radix2(hlsl::fft::twiddle<true,Scalar>(virtualThreadID & (stride - 1), stride), lo,hi);
 
                 // Divide by special factor at the end
-                if (passes - 1 == pass)
+                if ( (K >> 1) * _NBL_HLSL_WORKGROUP_SIZE_ == stride)
                 {
                     divides_assign< complex_t<Scalar> > divAss;
-                    divAss(lo, virtualThreadCount);
-                    divAss(hi, virtualThreadCount);  
+                    divAss(lo, K >> 1);
+                    divAss(hi, K >> 1);  
                 }
 
-                accessor.set((loIx << _NBL_HLSL_WORKGROUP_SIZE_LOG_2_) | SubgroupContiguousIndex(), lo);
-                accessor.set((hiIx << _NBL_HLSL_WORKGROUP_SIZE_LOG_2_) | SubgroupContiguousIndex(), hi);
+                accessor.set(loIx, lo);
+                accessor.set(hiIx, hi);
             }
             accessor.memoryBarrier(); // no execution barrier just making sure writes propagate to accessor
-            stride <<= 1;
         }
+        
     }
 };