More-than-2-per-thread FFT implemented, not yet tested

Author: Fletterio

include/nbl/builtin/hlsl/memory_accessor.hlsl

@@ -257,6 +257,36 @@ struct MemoryAdaptor<BaseAccessor, 0>
     }
 };
 
+// ---------------------------------------------- Offset Accessor ----------------------------------------------------
+
+template<class BaseAccessor, class AccessorType, uint32_t Offset>
+struct OffsetAccessor
+{
+    BaseAccessor accessor;
+
+    void set(uint32_t idx, NBL_REF_ARG(AccessorType) x) {accessor.set(idx + Offset, x);}
+
+    AccessorType get(uint32_t idx) {return accessor.get(idx + Offset);}
+
+    // TODO: figure out the `enable_if` syntax for this
+    void workgroupExecutionAndMemoryBarrier() {accessor.workgroupExecutionAndMemoryBarrier();}
+};
+
+// Dynamic offset version
+template<class BaseAccessor, class AccessorType>
+struct DynamicOffsetAccessor
+{
+    BaseAccessor accessor;
+    uint32_t offset;
+
+    void set(uint32_t idx, NBL_REF_ARG(AccessorType) x) {accessor.set(idx + offset, x);}
+
+    AccessorType get(uint32_t idx) {return accessor.get(idx + offset);}
+
+    // TODO: figure out the `enable_if` syntax for this
+    void workgroupExecutionAndMemoryBarrier() {accessor.workgroupExecutionAndMemoryBarrier();}
+};
+
 }
 }
 

include/nbl/builtin/hlsl/mpl.hlsl

@@ -63,6 +63,11 @@ struct rotr
     static const T value = (S >= 0) ? ((X >> r) | (X << (N - r))) : (X << (-r)) | (X >> (N - (-r)));
 };
 
+template<uint64_t N>
+struct is_pot : bool_constant< (N > 0 && !(N & (N - 1))) > {};
+
+template<uint64_t N>
+NBL_CONSTEXPR_STATIC_INLINE bool is_pot_v = is_pot<N>::value;
 
 }
 }

include/nbl/builtin/hlsl/type_traits.hlsl

@@ -526,6 +526,9 @@ using is_unbounded_array = std::is_unbounded_array<T>;
 template<class T>
 using is_scalar = std::is_scalar<T>;
 
+template<class T>
+NBL_CONSTEXPR_STATIC_INLINE bool is_scalar_v = is_scalar<T>::value;
+
 template<class T>
 struct is_signed : impl::base_type_forwarder<std::is_signed, T> {};
 
@@ -535,6 +538,9 @@ struct is_unsigned : impl::base_type_forwarder<std::is_unsigned, T> {};
 template<class T>
 struct is_integral : impl::base_type_forwarder<std::is_integral, T> {};
 
+template<class T>
+NBL_CONSTEXPR_STATIC_INLINE bool is_integral_v = is_integral<T>::value;
+
 template<class T>
 struct is_floating_point : impl::base_type_forwarder<std::is_floating_point, T> {};
 
@@ -583,6 +589,9 @@ using extent = std::extent<T, I>;
 template<bool B, class T = void>
 using enable_if = std::enable_if<B, T>;
 
+template<bool B, class T = void>
+using enable_if_t = typename enable_if<B, T>::type;
+
 template<class T>
 using alignment_of = std::alignment_of<T>;
 

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

@@ -5,6 +5,8 @@
 #include "nbl/builtin/hlsl/workgroup/basic.hlsl"
 #include "nbl/builtin/hlsl/glsl_compat/core.hlsl"
 #include "nbl/builtin/hlsl/workgroup/shuffle.hlsl"
+#include "nbl/builtin/hlsl/mpl.hlsl"
+#include "nbl/builtin/hlsl/memory_accessor.hlsl"
 
 namespace nbl 
 {
@@ -159,21 +161,96 @@ struct FFT<2,true, Scalar, device_capabilities>
 
 /*
 
-// then define 4,8,16 in terms of calling the FFT<2> and doing the special radix steps before/after
-template<uint16_t K, bool Inverse, class device_capabilities>
-struct FFT
+// Forward FFT
+template<uint32_t K, typename Scalar, class device_capabilities>
+struct FFT<K,false,device_capabilities>
 {
-    template<typename Accessor, typename ShaderMemoryAccessor>
-    static void __call(NBL_REF_ARG(Accessor) accessor, NBL_REF_ARG(ShaderMemoryAccessor) sharedmemAccessor)
+    template<typename Accessor, typename SharedMemoryAccessor>
+    static enable_if_t<mpl::is_pot_v<K>, 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++)
+        {
+            [unroll(K/2)]
+            for (uint32_t virtualThread = 0; virtualThread < virtualThreadCount; virtualThread++)
+            {
+                const uint32_t virtualThreadID = virtualThread * _NBL_HLSL_WORKGROUP_SIZE_ + SubgroupContiguousIndex();
+
+                const uint32_t lsb = virtualThread & (stride - 1);
+                const uint32_t loIx = ((virtualThread ^ lsb) << 1) | lsb;
+                const uint32_t hiIx = loIx | stride;
+                
+                complex_t<Scalar> lo = accessor.get(loIx * _NBL_HLSL_WORKGROUP_SIZE_);
+                complex_t<Scalar> hi = accessor.get(hiIx * _NBL_HLSL_WORKGROUP_SIZE_);
+                
+                fft::DIF<Scalar>::radix2(fft::twiddle<false,Scalar>(virtualThreadID & (stride - 1), stride),lo,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
+        OffsetAccessor < Accessor, complex_t<Scalar> > offsetAccessor;
+        [unroll(K/2)]
+        for (uint32_t k = 0; k < K; k += 2)
+        {
+            if (k)
+            sharedmemAccessor.executionAndMemoryBarrier();
+            offsetAccessor.offset = _NBL_HLSL_WORKGROUP_SIZE_*k;
+            FFT<2,false, Scalar, device_capabilities>::template __call(offsetAccessor,sharedmemAccessor);
+        }
+        accessor = offsetAccessor.accessor;
+    }
+};
+
+// Inverse FFT
+template<uint32_t K, typename Scalar, class device_capabilities>
+struct FFT<K,true,device_capabilities>
+{
+    template<typename Accessor, typename SharedMemoryAccessor>
+    static enable_if_t<mpl::is_pot_v<K>, void> __call(NBL_REF_ARG(Accessor) accessor, NBL_REF_ARG(SharedMemoryAccessor) sharedmemAccessor)
     {
-        if (!Inverse)
+        // do K/2 small workgroup FFTs
+        OffsetAccessor < Accessor, complex_t<Scalar> > offsetAccessor;
+        [unroll(K/2)]
+        for (uint32_t k = 0; k < K; k += 2)
         {
-           ... special steps ...
+            if (k)
+            sharedmemAccessor.executionAndMemoryBarrier();
+            offsetAccessor.offset = _NBL_HLSL_WORKGROUP_SIZE_*k;
+            FFT<2,true, Scalar, device_capabilities>::template __call(offsetAccessor,sharedmemAccessor);
         }
-        FFT<2,Inverse,device_capabilities>::template __call<Accessor,SharedMemoryAccessor>(access,sharedMemAccessor);
-        if (Inverse)
+        accessor = offsetAccessor.accessor;
+        
+        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++)
         {
-           ... special steps ...
+            [unroll(K/2)]
+            for (uint32_t virtualThread = 0; virtualThread < virtualThreadCount; virtualThread++)
+            {
+                const uint32_t lsb = virtualThread & (stride - 1);
+                const uint32_t loIx = ((virtualThread ^ lsb) << 1) | lsb;
+                const uint32_t hiIx = loIx | stride;
+                
+                complex_t<Scalar> lo = accessor.get(loIx * _NBL_HLSL_WORKGROUP_SIZE_);
+                complex_t<Scalar> hi = accessor.get(hiIx * _NBL_HLSL_WORKGROUP_SIZE_);
+                
+                fft::DIF<Scalar>::radix2(fft::twiddle<true,Scalar>(virtualThreadID & (stride - 1), stride),lo,hi);
+                
+                accessor.set(loIx, lo);
+                accessor.set(hiIx, hi);
+            }
+            accessor.memoryBarrier(); // no execution barrier just making sure writes propagate to accessor
+            stride <<= 1;
         }
     }
 };