Merge pull request #698 from Devsh-Graphics-Programming/nahim_complex_fft_workgroup

FFT Workgroup Ops

Author: devshgraphicsprogramming

examples_tests

@@ -12,10 +12,20 @@ namespace nbl
 {
 namespace hlsl
 {
+
+// TODO: make a common `nbl/builtin/hlsl/__ref.hlsl`
+// TODO: also refactor `bda::__base_ref` into just `__ref` and make it a typedef
+template<uint32_t StorageClass, typename T>
+using __spv_ptr_t = spirv::pointer_t<StorageClass,T>;
+
+template<uint32_t StorageClass, typename T>
+[[vk::ext_instruction(spv::OpCopyObject)]]
+__spv_ptr_t<StorageClass,T> addrof([[vk::ext_reference]] T v);
+
 namespace bda
 {
 template<typename T>
-using __spv_ptr_t = spirv::pointer_t<spv::StorageClassPhysicalStorageBuffer, T>;
+using __spv_ptr_t = spirv::pointer_t<spv::StorageClassPhysicalStorageBuffer,T>;
 
 template<typename T>
 struct __ptr;

include/nbl/builtin/hlsl/bda/__ref.hlsl

@@ -32,9 +32,8 @@ namespace hlsl
 {
 
 template<class T, class U>
-T bit_cast(U val)
+enable_if_t<sizeof(T)==sizeof(U),T> bit_cast(U val)
 {
-    static_assert(sizeof(T)==sizeof(U));
     return spirv::bitcast<T,U>(val);
 }
 
@@ -92,7 +91,7 @@ uint16_t clz(uint64_t N)
 template<>
 uint16_t clz<1>(uint64_t N) { return uint16_t(1u-N&1); }
 
-}
+} //namespace impl
 
 template<typename T>
 uint16_t countl_zero(T n)

include/nbl/builtin/hlsl/bit.hlsl

@@ -191,6 +191,10 @@ const static complex_t< SCALAR > multiplies< complex_t< SCALAR > >::identity = {
 template<> \
 const static complex_t< SCALAR > divides< complex_t< SCALAR > >::identity = { promote< SCALAR , uint32_t>(1),  promote< SCALAR , uint32_t>(0)};
 
+COMPLEX_ARITHMETIC_IDENTITIES(float16_t)
+COMPLEX_ARITHMETIC_IDENTITIES(float16_t2)
+COMPLEX_ARITHMETIC_IDENTITIES(float16_t3)
+COMPLEX_ARITHMETIC_IDENTITIES(float16_t4)  
 COMPLEX_ARITHMETIC_IDENTITIES(float32_t)
 COMPLEX_ARITHMETIC_IDENTITIES(float32_t2)
 COMPLEX_ARITHMETIC_IDENTITIES(float32_t3)
@@ -287,6 +291,10 @@ COMPLEX_COMPOUND_ASSIGN_IDENTITY(minus, SCALAR) \
 COMPLEX_COMPOUND_ASSIGN_IDENTITY(multiplies, SCALAR) \
 COMPLEX_COMPOUND_ASSIGN_IDENTITY(divides, SCALAR)
 
+COMPLEX_COMPOUND_ASSIGN_IDENTITIES(float16_t)
+COMPLEX_COMPOUND_ASSIGN_IDENTITIES(float16_t2)
+COMPLEX_COMPOUND_ASSIGN_IDENTITIES(float16_t3)
+COMPLEX_COMPOUND_ASSIGN_IDENTITIES(float16_t4)
 COMPLEX_COMPOUND_ASSIGN_IDENTITIES(float32_t)
 COMPLEX_COMPOUND_ASSIGN_IDENTITIES(float32_t2)
 COMPLEX_COMPOUND_ASSIGN_IDENTITIES(float32_t3)

include/nbl/builtin/hlsl/complex.hlsl

@@ -14,11 +14,11 @@ namespace fft
 
 // Computes the kth element in the group of N roots of unity
 // Notice 0 <= k < N/2, rotating counterclockwise in the forward (DIF) transform and clockwise in the inverse (DIT)
-template<typename Scalar, bool inverse>
-complex_t<Scalar> twiddle(uint32_t k, uint32_t N)
+template<bool inverse, typename Scalar>
+complex_t<Scalar> twiddle(uint32_t k, uint32_t halfN)
 {
     complex_t<Scalar> retVal;
-    const Scalar kthRootAngleRadians = 2.f * numbers::pi<Scalar> * Scalar(k) / Scalar(N);
+    const Scalar kthRootAngleRadians = numbers::pi<Scalar> * Scalar(k) / Scalar(halfN);
     retVal.real( cos(kthRootAngleRadians) );
     if (! inverse)
         retVal.imag( sin(kthRootAngleRadians) );
@@ -27,7 +27,7 @@ complex_t<Scalar> twiddle(uint32_t k, uint32_t N)
     return retVal;                         
 }
 
-template<typename Scalar, bool inverse> 
+template<bool inverse, typename Scalar> 
 struct DIX 
 { 
     static void radix2(NBL_CONST_REF_ARG(complex_t<Scalar>) twiddle, NBL_REF_ARG(complex_t<Scalar>) lo, NBL_REF_ARG(complex_t<Scalar>) hi)
@@ -49,10 +49,10 @@ struct DIX
 };
 
 template<typename Scalar>
-using DIT = DIX<Scalar, true>;
+using DIT = DIX<true, Scalar>;
 
 template<typename Scalar>
-using DIF = DIX<Scalar, false>;
+using DIF = DIX<false, Scalar>;
 }
 }
 }

include/nbl/builtin/hlsl/fft/common.hlsl

@@ -199,7 +199,7 @@ struct bitfieldExtract<T, false, true>
     } 
 };
 
-}
+} //namespace impl
 
 template<typename T>
 T bitfieldExtract( T val, uint32_t offsetBits, uint32_t numBits )

include/nbl/builtin/hlsl/glsl_compat/core.hlsl

@@ -5,6 +5,7 @@
 #define _NBL_BUILTIN_HLSL_MEMORY_ACCESSOR_INCLUDED_
 
 #include "nbl/builtin/hlsl/glsl_compat/core.hlsl"
+#include "nbl/builtin/hlsl/member_test_macros.hlsl"
 
 // weird namespace placing, see the comment where the macro is defined
 GENERATE_METHOD_TESTER(atomicExchange)
@@ -33,104 +34,189 @@ struct pair
     second_type second;
 };
 
+namespace accessor_adaptors
+{
+namespace impl
+{
+// only base class to use integral_constant because we need to use void to indicate a dynamic value and all values are valid
+template<typename IndexType, typename Offset>
+struct OffsetBase
+{
+    NBL_CONSTEXPR IndexType offset = Offset::value;
+};
+template<typename IndexType>
+struct OffsetBase<IndexType,void>
+{
+    IndexType offset;
+};
+
+template<typename IndexType, uint64_t ElementStride, uint64_t SubElementStride, typename Offset>
+struct StructureOfArraysStrides
+{
+    NBL_CONSTEXPR IndexType elementStride = ElementStride;
+    NBL_CONSTEXPR IndexType subElementStride = SubElementStride;
+
+    //static_assert(elementStride>0 && subElementStride>0);
+};
+template<typename IndexType, typename Offset>
+struct StructureOfArraysStrides<IndexType,0,0,Offset> : OffsetBase<IndexType,Offset>
+{
+    IndexType elementStride;
+    IndexType subElementStride;
+};
+#if 0 // don't seem to be able to specialize one at a time
+template<typename IndexType, uint64_t ElementStride, typename Offset>
+struct StructureOfArraysStrides<IndexType,ElementStride,0,Offset> : OffsetBase<IndexType,Offset>
+{
+    NBL_CONSTEXPR IndexType elementStride = ElementStride;
+    IndexType subElementStride;
+};
+template<typename IndexType, uint64_t SubElementStride, typename Offset>
+struct StructureOfArraysStrides<IndexType,0,SubElementStride,Offset> : OffsetBase<IndexType,Offset>
+{
+    IndexType elementStride;
+    NBL_CONSTEXPR IndexType subElementStride = SubElementStride;
+};
+#endif
+
+
+template<typename IndexType, uint64_t ElementStride, uint64_t SubElementStride, typename Offset>
+struct StructureOfArraysBase : StructureOfArraysStrides<IndexType,ElementStride,SubElementStride,Offset>
+{
+    IndexType getIx(const IndexType ix, const IndexType el)
+    {
+        using base_t = StructureOfArraysStrides<IndexType,ElementStride,SubElementStride,Offset>;
+        return base_t::elementStride*ix+base_t::subElementStride*el+OffsetBase<IndexType,Offset>::offset;
+    }
+};
+
+// maybe we should have our own std::array
+template<typename T, uint64_t count>
+struct array
+{
+    T data[count];
+};
+}
 
-// TODO: find some cool way to SFINAE the default into `_NBL_HLSL_WORKGROUP_SIZE_` if defined, and something like 1 otherwise
-template<class BaseAccessor, typename AccessType, typename IndexType=uint32_t, typename Strides=pair<integral_constant<IndexType,1>,integral_constant<IndexType,_NBL_HLSL_WORKGROUP_SIZE_> > >
-struct MemoryAdaptor // TODO: rename to something nicer like StructureOfArrays and add a `namespace accessor_adaptors`
+// TODO: some CRTP thing to forward through atomics and barriers
+
+// If you want static strides pass `Stride=pair<integral_constant<IndexType,ElementStride>,integral_constant<IndexType,SubElementStride> >`
+template<class BaseAccessor, typename AccessType, typename IndexType=uint32_t, uint64_t ElementStride=0, uint64_t SubElementStride=0, typename _Offset=integral_constant<IndexType,0> >
+struct StructureOfArrays : impl::StructureOfArraysBase<IndexType,ElementStride,SubElementStride,_Offset>
 {
+    using base_t = impl::StructureOfArraysBase<IndexType,ElementStride,SubElementStride,_Offset>;
     // Question: should the `BaseAccessor` let us know what this is?
     using access_t = AccessType;
     using index_t = IndexType;
-    NBL_CONSTEXPR index_t ElementStride = Strides::first_type::value;
-    NBL_CONSTEXPR index_t SubElementStride = Strides::second_type::value;
 
     BaseAccessor accessor;
-    
-    access_t get(const index_t ix)
-    { 
-        access_t retVal;
-        get<access_t>(ix,retVal);
-        return retVal; 
-    }
 
     // Question: shall we go back to requiring a `access_t get(index_t)` on the `BaseAccessor`, then we could `enable_if` check the return type (via `has_method_get`) matches and we won't get Nasty HLSL copy-in copy-out conversions
     template<typename T>
     enable_if_t<sizeof(T)%sizeof(access_t)==0,void> get(const index_t ix, NBL_REF_ARG(T) value)
-    { 
+    {
         NBL_CONSTEXPR uint64_t SubElementCount = sizeof(T)/sizeof(access_t);
-        access_t aux[SubElementCount];
-        for (uint64_t i=0; i<SubElementCount; i++)
-            accessor.get(ix*ElementStride+i*SubElementStride,aux[i]);
-        value = bit_cast<T,access_t[SubElementCount]>(aux);
+        // `vector` for now, we'll use `array` later
+        vector<access_t,SubElementCount> aux;
+        for (index_t i=0; i<SubElementCount; i++)
+            accessor.get(base_t::getIx(ix,i),aux[i]);
+        value = bit_cast<T,vector<access_t,SubElementCount> >(aux);
     }
 
     template<typename T>
     enable_if_t<sizeof(T)%sizeof(access_t)==0,void> set(const index_t ix, NBL_CONST_REF_ARG(T) value)
     { 
         NBL_CONSTEXPR uint64_t SubElementCount = sizeof(T)/sizeof(access_t);
-        access_t aux[SubElementCount] = bit_cast<access_t[SubElementCount],T>(value);
-        for (uint64_t i=0; i<SubElementCount; i++)
-            accessor.set(ix*ElementStride+i*SubElementStride,aux[i]);
+        // `vector` for now, we'll use `array` later
+        vector<access_t,SubElementCount> aux;
+        aux = bit_cast<vector<access_t,SubElementCount>,T>(value);
+        for (index_t i=0; i<SubElementCount; i++)
+            accessor.set(base_t::getIx(ix,i),aux[i]);
+
     }
-    
+
     template<typename T, typename S=BaseAccessor>
     enable_if_t<
-        sizeof(T)==sizeof(access_t) && is_same_v<S,BaseAccessor> && is_same_v<has_method_atomicExchange<S,index_t,access_t>::return_type,access_t>,void
+        sizeof(T)==sizeof(access_t) && is_same_v<S,BaseAccessor> && is_same_v<typename has_method_atomicExchange<S,index_t,access_t>::return_type,access_t>,void
     > atomicExchange(const index_t ix, const T value, NBL_REF_ARG(T) orig)
     {
-       orig = bit_cast<T,access_t>(accessor.atomicExchange(ix,bit_cast<access_t,T>(value)));
+       orig = bit_cast<T,access_t>(accessor.atomicExchange(getIx(ix),bit_cast<access_t,T>(value)));
     }
     template<typename T, typename S=BaseAccessor>
     enable_if_t<
-        sizeof(T)==sizeof(access_t) && is_same_v<S,BaseAccessor> && is_same_v<has_method_atomicCompSwap<S,index_t,access_t,access_t>::return_type,access_t>,void
+        sizeof(T)==sizeof(access_t) && is_same_v<S,BaseAccessor> && is_same_v<typename has_method_atomicCompSwap<S,index_t,access_t,access_t>::return_type,access_t>,void
     > atomicCompSwap(const index_t ix, const T value, const T comp, NBL_REF_ARG(T) orig)
     {
-       orig = bit_cast<T,access_t>(accessor.atomicCompSwap(ix,bit_cast<access_t,T>(comp),bit_cast<access_t,T>(value)));
+       orig = bit_cast<T,access_t>(accessor.atomicCompSwap(getIx(ix),bit_cast<access_t,T>(comp),bit_cast<access_t,T>(value)));
     }
 
     template<typename T, typename S=BaseAccessor>
     enable_if_t<
-        sizeof(T)==sizeof(access_t) && is_same_v<S,BaseAccessor> && is_same_v<has_method_atomicAnd<S,index_t,access_t>::return_type,access_t>,void
+        sizeof(T)==sizeof(access_t) && is_same_v<S,BaseAccessor> && is_same_v<typename has_method_atomicAnd<S,index_t,access_t>::return_type,access_t>,void
     > atomicAnd(const index_t ix, const T value, NBL_REF_ARG(T) orig)
     {
-       orig = bit_cast<T,access_t>(accessor.atomicAnd(ix,bit_cast<access_t,T>(value)));
+       orig = bit_cast<T,access_t>(accessor.atomicAnd(getIx(ix),bit_cast<access_t,T>(value)));
     }
     template<typename T, typename S=BaseAccessor>
     enable_if_t<
-        sizeof(T)==sizeof(access_t) && is_same_v<S,BaseAccessor> && is_same_v<has_method_atomicOr<S,index_t,access_t>::return_type,access_t>,void
+        sizeof(T)==sizeof(access_t) && is_same_v<S,BaseAccessor> && is_same_v<typename has_method_atomicOr<S,index_t,access_t>::return_type,access_t>,void
     > atomicOr(const index_t ix, const T value, NBL_REF_ARG(T) orig)
     {
-       orig = bit_cast<T,access_t>(accessor.atomicOr(ix,bit_cast<access_t,T>(value)));
+       orig = bit_cast<T,access_t>(accessor.atomicOr(getIx(ix),bit_cast<access_t,T>(value)));
     }
     template<typename T, typename S=BaseAccessor>
     enable_if_t<
-        sizeof(T)==sizeof(access_t) && is_same_v<S,BaseAccessor> && is_same_v<has_method_atomicXor<S,index_t,access_t>::return_type,access_t>,void
+        sizeof(T)==sizeof(access_t) && is_same_v<S,BaseAccessor> && is_same_v<typename has_method_atomicXor<S,index_t,access_t>::return_type,access_t>,void
     > atomicXor(const index_t ix, const T value, NBL_REF_ARG(T) orig)
     {
-       orig = bit_cast<T,access_t>(accessor.atomicXor(ix,bit_cast<access_t,T>(value)));
+       orig = bit_cast<T,access_t>(accessor.atomicXor(getIx(ix),bit_cast<access_t,T>(value)));
     }
 
     // This has the upside of never calling a `(uint32_t)(uint32_t,uint32_t)` overload of `atomicAdd` because it checks the return type!
     // If someone makes a `(float)(uint32_t,uint32_t)` they will break this detection code, but oh well.
     template<typename T>
-    enable_if_t<is_same_v<has_method_atomicAdd<BaseAccessor,index_t,T>::return_type,T>,void> atomicAdd(const index_t ix, const T value, NBL_REF_ARG(T) orig)
+    enable_if_t<is_same_v<typename has_method_atomicAdd<BaseAccessor,index_t,T>::return_type,T>,void> atomicAdd(const index_t ix, const T value, NBL_REF_ARG(T) orig)
     {
-       orig = accessor.atomicAdd(ix,value);
+       orig = accessor.atomicAdd(getIx(ix),value);
     }
     template<typename T>
-    enable_if_t<is_same_v<has_method_atomicMin<BaseAccessor,index_t,T>::return_type,T>,void> atomicMin(const index_t ix, const T value, NBL_REF_ARG(T) orig)
+    enable_if_t<is_same_v<typename has_method_atomicMin<BaseAccessor,index_t,T>::return_type,T>,void> atomicMin(const index_t ix, const T value, NBL_REF_ARG(T) orig)
     {
-       orig = accessor.atomicMin(ix,value);
+       orig = accessor.atomicMin(getIx(ix),value);
     }
     template<typename T>
-    enable_if_t<is_same_v<has_method_atomicMax<BaseAccessor,index_t,T>::return_type,T>,void> atomicMax(const index_t ix, const T value, NBL_REF_ARG(T) orig)
+    enable_if_t<is_same_v<typename has_method_atomicMax<BaseAccessor,index_t,T>::return_type,T>,void> atomicMax(const index_t ix, const T value, NBL_REF_ARG(T) orig)
+    {
+        orig = accessor.atomicMax(getIx(ix),value);
+    }
+    
+    template<typename S=BaseAccessor>
+    enable_if_t<
+        is_same_v<S,BaseAccessor> && is_same_v<typename has_method_workgroupExecutionAndMemoryBarrier<S>::return_type,void>,void
+    > workgroupExecutionAndMemoryBarrier()
     {
-        orig = accessor.atomicMax(ix,value);
+        accessor.workgroupExecutionAndMemoryBarrier();
     }
+};
+
+// ---------------------------------------------- Offset Accessor ----------------------------------------------------
+
+template<class BaseAccessor, typename IndexType=uint32_t, typename _Offset=void>
+struct Offset : impl::OffsetBase<IndexType,_Offset>
+{
+    using base_t = impl::OffsetBase<IndexType,_Offset>;
+
+    BaseAccessor accessor;
+
+    template <typename T>
+    void set(uint32_t idx, T value) {accessor.set(idx+base_t::offset,value); }
+
+    template <typename T> 
+    void get(uint32_t idx, NBL_REF_ARG(T) value) {accessor.get(idx+base_t::offset,value);}
 
     template<typename S=BaseAccessor>
     enable_if_t<
-        is_same_v<S,BaseAccessor> && is_same_v<has_method_workgroupExecutionAndMemoryBarrier<S>::return_type,void>,void
+        is_same_v<S,BaseAccessor> && is_same_v<typename has_method_workgroupExecutionAndMemoryBarrier<S>::return_type,void>,void
     > workgroupExecutionAndMemoryBarrier()
     {
         accessor.workgroupExecutionAndMemoryBarrier();
@@ -139,5 +225,5 @@ struct MemoryAdaptor // TODO: rename to something nicer like StructureOfArrays a
 
 }
 }
-
+}
 #endif

include/nbl/builtin/hlsl/memory_accessor.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/mpl.hlsl

@@ -54,6 +54,11 @@ using pointer_t = vk::SpirvOpaqueType<spv::OpTypePointer,vk::Literal<vk::integra
 
 //! General Operations
 
+// The holy operation that makes addrof possible
+template<uint32_t StorageClass, typename T>
+[[vk::ext_instruction(spv::OpCopyObject)]]
+pointer_t<StorageClass,T> copyObject([[vk::ext_reference]] T v);
+
 // Here's the thing with atomics, it's not only the data type that dictates whether you can do an atomic or not.
 // It's the storage class that has the most effect (shared vs storage vs image) and we can't check that easily
 template<typename T> // integers operate on 2s complement so same op for signed and unsigned