More fft utils

include/nbl/builtin/hlsl/bitreverse.hlsl

@@ -0,0 +1,47 @@
+#ifndef _NBL_BUILTIN_HLSL_BITREVERSE_INCLUDED_
+#define _NBL_BUILTIN_HLSL_BITREVERSE_INCLUDED_
+
+
+#include <nbl/builtin/hlsl/cpp_compat.hlsl>
+
+namespace nbl
+{
+namespace hlsl
+{
+
+template<typename T, uint16_t Bits NBL_FUNC_REQUIRES(is_unsigned_v<T>&& Bits <= sizeof(T) * 8)
+/**
+* @brief Takes the binary representation of `value` as a string of `Bits` bits and returns a value of the same type resulting from reversing the string
+*
+* @tparam T Type of the value to operate on.
+* @tparam Bits The length of the string of bits used to represent `value`.
+*
+* @param [in] value The value to bitreverse.
+*/
+T bitReverseAs(T value)
+{
+	return bitReverse<T>(value) >> promote<T, scalar_type_t<T> >(scalar_type_t <T>(sizeof(T) * 8 - Bits));
+}
+
+template<typename T NBL_FUNC_REQUIRES(is_unsigned_v<T>)
+/**
+* @brief Takes the binary representation of `value` and returns a value of the same type resulting from reversing the string of bits as if it was `bits` long.
+* Keep in mind `bits` cannot exceed `8 * sizeof(T)`.
+*
+* @tparam T type of the value to operate on.
+*
+* @param [in] value The value to bitreverse.
+* @param [in] bits The length of the string of bits used to represent `value`.
+*/
+T bitReverseAs(T value, uint16_t bits)
+{
+	return bitReverse<T>(value) >> promote<T, scalar_type_t<T> >(scalar_type_t <T>(sizeof(T) * 8 - bits));
+}
+
+
+}
+}
+
+
+
+#endif

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

@@ -14,34 +14,63 @@ namespace hlsl
 namespace fft
 {
 
-// template parameter N controls the number of dimensions of the input
-// template parameter M controls the number of dimensions to pad up to PoT
-// "axes" indicates which dimensions to pad up to PoT
-template <uint16_t N, uint16_t M NBL_FUNC_REQUIRES(M <= N)
-inline vector<uint64_t, 3> padDimensions(NBL_CONST_REF_ARG(vector<uint32_t, N>) dimensions, NBL_CONST_REF_ARG(vector<uint16_t, M>) axes, bool realFFT = false)
+
+template <uint16_t N NBL_FUNC_REQUIRES(N > 0 && N <= 4)
+/**
+* @brief Returns the size of the full FFT computed, in terms of number of complex elements.
+*
+* @tparam N Number of dimensions of the signal to perform FFT on.
+*
+* @param [in] dimensions Size of the signal.
+* @param [in] realFFT Indicates whether the signal is real. False by default.
+* @param [in] firstAxis Indicates which axis the FFT is performed on first. Only relevant for real-valued signals. Must be less than N. 0 by default.
+*/
+inline vector<uint64_t, N> padDimensions(NBL_CONST_REF_ARG(vector<uint32_t, N>) dimensions, bool realFFT = false, uint16_t firstAxis = 0u)
 {
     vector<uint32_t, N> newDimensions = dimensions;
-    uint16_t axisCount = 0;
-    for (uint16_t i = 0u; i < M; i++)
+    for (uint16_t i = 0u; i < N; i++)
     {
         newDimensions[i] = hlsl::roundUpToPoT(newDimensions[i]);
-        if (realFFT && !axisCount++)
-            newDimensions[i] /= 2;
     }
+    if (realFFT)
+        newDimensions[firstAxis] /= 2;
     return newDimensions;
 }
 
-// template parameter N controls the number of dimensions of the input
-// template parameter M controls the number of dimensions we run an FFT along AND store the result
-// "axes" indicates which dimensions we run an FFT along AND store the result
-template <uint16_t N, uint16_t M NBL_FUNC_REQUIRES(M <= N)
-inline uint64_t getOutputBufferSize(NBL_CONST_REF_ARG(vector<uint32_t, N>) inputDimensions, uint32_t numChannels, NBL_CONST_REF_ARG(vector<uint16_t, M>) axes, bool realFFT = false, bool halfFloats = false)
+template <uint16_t N NBL_FUNC_REQUIRES(N > 0 && N <= 4)
+/**
+* @brief Returns the size required by a buffer to hold the result of the FFT of a signal after a certain pass.
+*
+* @tparam N Number of dimensions of the signal to perform FFT on.
+*
+* @param [in] numChannels Number of channels of the signal.
+* @param [in] inputDimensions Size of the signal.
+* @param [in] passIx Which pass the size is being computed for.
+* @param [in] axisPassOrder Order of the axis in which the FFT is computed in. Default is xyzw.
+* @param [in] realFFT True if the signal is real. False by default.
+* @param [in] halfFloats True if using half-precision floats. False by default.
+*/
+inline uint64_t getOutputBufferSize(
+    uint32_t numChannels,
+    NBL_CONST_REF_ARG(vector<uint32_t, N>) inputDimensions,
+    uint16_t passIx,
+    NBL_CONST_REF_ARG(vector<uint16_t, N>) axisPassOrder = _static_cast<vector<uint16_t, N> >(uint16_t4(0, 1, 2, 3)),
+    bool realFFT = false,
+    bool halfFloats = false
+)
 {
-    const vector<uint64_t, 3> paddedDims = padDimensions<N, M>(inputDimensions, axes);
-    const uint64_t numberOfComplexElements = paddedDims[0] * paddedDims[1] * paddedDims[2] * uint64_t(numChannels);
+    const vector<uint32_t, N> paddedDimensions = padDimensions<N>(inputDimensions, realFFT, axisPassOrder[0]);
+    vector<bool, N> axesDone = promote<vector<bool, N>, bool>(false);
+    for (uint16_t i = 0; i <= passIx; i++)
+        axesDone[axisPassOrder[i]] = true;
+    const vector<uint32_t, N> passOutputDimension = lerp(inputDimensions, paddedDimensions, axesDone);
+    uint64_t numberOfComplexElements = uint64_t(numChannels);
+    for (uint16_t i = 0; i < N; i++)
+        numberOfComplexElements *= uint64_t(passOutputDimension[i]);
     return numberOfComplexElements * (halfFloats ? sizeof(complex_t<float16_t>) : sizeof(complex_t<float32_t>));
 }
 
+
 // 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<bool inverse, typename Scalar>
@@ -95,13 +124,6 @@ void unpack(NBL_REF_ARG(complex_t<Scalar>) lo, NBL_REF_ARG(complex_t<Scalar>) hi
     lo = x;
 }
 
-// Bit-reverses T as a binary string of length given by Bits
-template<typename T, uint16_t Bits NBL_FUNC_REQUIRES(is_integral_v<T> && Bits <= sizeof(T) * 8)
-T bitReverseAs(T value)
-{
-    return hlsl::bitReverse<uint32_t>(value) >> (sizeof(T) * 8 - Bits);
-}
-
 }
 }
 }

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

@@ -7,11 +7,7 @@
 #include "nbl/builtin/hlsl/cpp_compat/basic.h"
 #include "nbl/builtin/hlsl/spirv_intrinsics/core.hlsl"
 #include "nbl/builtin/hlsl/type_traits.hlsl"
-<<<<<<< HEAD
-#include "nbl/builtin/hlsl/bit.hlsl"
-=======
 #include "nbl/builtin/hlsl/spirv_intrinsics/glsl.std.450.hlsl"
->>>>>>> master
 
 namespace nbl 
 {

include/nbl/builtin/hlsl/math/intutil.hlsl

@@ -58,10 +58,9 @@ NBL_CONSTEXPR_FORCED_INLINE_FUNC Integer align(Integer alignment, Integer size,
     return address = nextAlignedAddr;
 }
 
+// ------------------------------------- CPP ONLY ----------------------------------------------------------
 #ifndef __HLSL_VERSION
 
-// Have to wait for the HLSL patch for `is_enum`. Would also have to figure out how to do it without initializer lists for HLSL use. 
-
 //! Get bitmask from variadic arguments passed. 
 /*
     For example if you were to create bitmask for vertex attributes