better "nbl_glsl_ext_FFT_getPaddedData"

Author: Erfan-Ahmadi

include/nbl/builtin/glsl/ext/FFT/fft.glsl

@@ -66,12 +66,20 @@ vec2 nbl_glsl_ext_FFT_getData(in uvec3 coordinate, in uint channel);
 void nbl_glsl_ext_FFT_setData(in uvec3 coordinate, in uint channel, in vec2 complex_value);
 #endif
 
+#ifndef _NBL_GLSL_EXT_FFT_GET_PADDED_DATA_DECLARED_
+#define _NBL_GLSL_EXT_FFT_GET_PADDED_DATA_DECLARED_
+vec2 nbl_glsl_ext_FFT_getPaddedData(in uvec3 coordinate, in uint channel);
+#endif
+
 #ifndef _NBL_GLSL_EXT_FFT_GET_DATA_DEFINED_
 #error "You need to define `nbl_glsl_ext_FFT_getData` and mark `_NBL_GLSL_EXT_FFT_GET_DATA_DEFINED_`!"
 #endif
 #ifndef _NBL_GLSL_EXT_FFT_SET_DATA_DEFINED_
 #error "You need to define `nbl_glsl_ext_FFT_setData` and mark `_NBL_GLSL_EXT_FFT_SET_DATA_DEFINED_`!"
 #endif
+#ifndef _NBL_GLSL_EXT_FFT_GET_PADDED_DATA_DEFINED_
+#error "You need to define `nbl_glsl_ext_FFT_getPaddedData` and mark `_NBL_GLSL_EXT_FFT_GET_PADDED_DATA_DEFINED_`!"
+#endif
 
 // Count Leading Zeroes (naive?)
 uint nbl_glsl_ext_FFT_clz(in uint x) 
@@ -127,61 +135,6 @@ uint nbl_glsl_ext_FFT_getDimLength(uvec3 dimension)
     return dimension[pc.direction];
 }
 
-vec2 nbl_glsl_ext_FFT_getPaddedData(in uvec3 coordinate, in uint channel) {
-    uint min_x = 0u;
-    uint max_x = pc.dimension.x + min_x - 1u;
-
-    uint min_y = 0u;
-    uint max_y = pc.dimension.y + min_y - 1u;
-
-    uint min_z = 0u;
-    uint max_z = pc.dimension.z + min_z - 1u;
-
-
-    uvec3 actual_coord = uvec3(0u, 0u, 0u);
-
-    if(_NBL_GLSL_EXT_FFT_CLAMP_TO_EDGE_ == pc.padding_type) {
-        if (coordinate.x < min_x) {
-            actual_coord.x = 0u;
-        } else if(coordinate.x > max_x) {
-            actual_coord.x = pc.dimension.x - 1u;
-        } else {
-            actual_coord.x = coordinate.x - min_x;
-        }
-        
-        if (coordinate.y < min_y) {
-            actual_coord.y = 0u;
-        } else if (coordinate.y > max_y) {
-            actual_coord.y = pc.dimension.y - 1u;
-        } else {
-            actual_coord.y = coordinate.y - min_y;
-        }
-        
-        if (coordinate.z < min_z) {
-            actual_coord.z = 0u;
-        } else if (coordinate.z > max_z) {
-            actual_coord.z = pc.dimension.z - 1u;
-        } else {
-            actual_coord.z = coordinate.z - min_z;
-        }
-        
-    } else if (_NBL_GLSL_EXT_FFT_FILL_WITH_ZERO_ == pc.padding_type) {
-
-        if ( coordinate.x < min_x || coordinate.x > max_x ||
-             coordinate.y < min_y || coordinate.y > max_y ||
-             coordinate.z < min_z || coordinate.z > max_z ) {
-            return vec2(0, 0);
-        }
-        
-        actual_coord.x = coordinate.x - min_x;
-        actual_coord.y = coordinate.y - min_y;
-        actual_coord.z = coordinate.z - min_z;
-
-    }
-    
-    return nbl_glsl_ext_FFT_getData(actual_coord, channel);
-}
-
 void nbl_glsl_ext_FFT()
 {
     // Virtual Threads Calculation

include/nbl/ext/FFT/FFT.h

@@ -106,7 +106,7 @@ class FFT : public core::TotalInterface
 			return (paddedInputDimensions.width * paddedInputDimensions.height * paddedInputDimensions.depth * numChannels) * (sizeof(float) * 2);
 		}
 
-		static core::smart_refctd_ptr<video::IGPUSpecializedShader> createShader(video::IVideoDriver* driver, DataType inputType, uint32_t maxPaddedDimensionSize);
+		static core::smart_refctd_ptr<video::IGPUSpecializedShader> createShader(video::IVideoDriver* driver, DataType inputType, uint32_t maxDimensionSize);
 
 		_NBL_STATIC_INLINE_CONSTEXPR uint32_t MAX_DESCRIPTOR_COUNT = 2u;
 		static inline void updateDescriptorSet(

src/nbl/ext/FFT/FFT.cpp

@@ -98,9 +98,9 @@ core::SRange<const video::IGPUDescriptorSetLayout::SBinding> FFT::getDefaultBind
 	return {bnd, bnd+sizeof(bnd)/sizeof(IGPUDescriptorSetLayout::SBinding)};
 }
 
-core::smart_refctd_ptr<video::IGPUSpecializedShader> FFT::createShader(video::IVideoDriver* driver, DataType inputType, uint32_t maxPaddedDimensionSize)
+core::smart_refctd_ptr<video::IGPUSpecializedShader> FFT::createShader(video::IVideoDriver* driver, DataType inputType, uint32_t maxDimensionSize)
 {
-	assert(core::isPoT(maxPaddedDimensionSize));
+	uint32_t const maxPaddedDimensionSize = core::roundUpToPoT(maxDimensionSize);
 
 	const char* sourceFmt =
 R"===(#version 430 core
@@ -130,6 +130,7 @@ layout(local_size_x=_NBL_GLSL_EXT_FFT_BLOCK_SIZE_X_DEFINED_, local_size_y=_NBL_G
  
 #define _NBL_GLSL_EXT_FFT_GET_DATA_DEFINED_
 #define _NBL_GLSL_EXT_FFT_SET_DATA_DEFINED_
+#define _NBL_GLSL_EXT_FFT_GET_PADDED_DATA_DEFINED_
 #include "nbl/builtin/glsl/ext/FFT/fft.glsl"
 
 // Input Descriptor
@@ -187,7 +188,7 @@ layout(set=_NBL_GLSL_EXT_FFT_OUTPUT_SET_DEFINED_, binding=_NBL_GLSL_EXT_FFT_OUTP
 
 // Get/Set Data Function
 
-vec2 nbl_glsl_ext_FFT_getData(in uvec3 coordinate, in uint channel)
+nbl_glsl_complex nbl_glsl_ext_FFT_getData(in uvec3 coordinate, in uint channel)
 {
 	vec2 retValue = vec2(0, 0);
 #if USE_SSBO_FOR_INPUT > 0
@@ -201,13 +202,27 @@ vec2 nbl_glsl_ext_FFT_getData(in uvec3 coordinate, in uint channel)
 	return retValue;
 }
 
-void nbl_glsl_ext_FFT_setData(in uvec3 coordinate, in uint channel, in vec2 complex_value)
+void nbl_glsl_ext_FFT_setData(in uvec3 coordinate, in uint channel, in nbl_glsl_complex complex_value)
 {
 	uvec3 dimension = pc.padded_dimension;
 	uint index = channel * (dimension.x * dimension.y * dimension.z) + coordinate.z * (dimension.x * dimension.y) + coordinate.y * (dimension.x) + coordinate.x;
 	outData[index].complex_value = complex_value;
 }
 
+nbl_glsl_complex nbl_glsl_ext_FFT_getPaddedData(in uvec3 coordinate, in uint channel) {
+
+	uvec3 max_coord = pc.dimension - uvec3(1u);
+	uvec3 clamped_coord = min(coordinate, max_coord);
+	
+	bool is_out_of_range = any(bvec3(coordinate!=clamped_coord));
+
+	if (_NBL_GLSL_EXT_FFT_FILL_WITH_ZERO_ == pc.padding_type && is_out_of_range) {
+		return nbl_glsl_complex(0, 0);
+	}
+	
+	return nbl_glsl_ext_FFT_getData(clamped_coord, channel);
+}
+
 void main()
 {
 	nbl_glsl_ext_FFT();
@@ -271,9 +286,9 @@ layout(set=0, binding=1) restrict buffer OutBuffer
 
 void main()
 {
-    float power = length(in_data[0].complex_value);
-    vec2 normalized_data = in_data[gl_GlobalInvocationID.x].complex_value / power;
-    out_data[gl_GlobalInvocationID.x].complex_value = normalized_data;
+	float power = length(in_data[0].complex_value);
+	vec2 normalized_data = in_data[gl_GlobalInvocationID.x].complex_value / power;
+	out_data[gl_GlobalInvocationID.x].complex_value = normalized_data;
 }
 )===";