|
@@ -0,0 +1,57 @@
|
|
|
|
|
+@group(0) @binding(0)
|
|
|
|
|
+var<uniform> u_input: texture_2d<f32>; // readonly image2D
|
|
|
|
|
+
|
|
|
|
|
+@group(0) @binding(1)
|
|
|
|
|
+var<storage, write> u_output: texture_2d<f32>; // writeonly image2D
|
|
|
|
|
+
|
|
|
|
|
+let PI: f32 = 3.14159265358979323846;
|
|
|
|
|
+
|
|
|
|
|
+struct Complex {
|
|
|
|
|
+ x: f32,
|
|
|
|
|
+ y: f32,
|
|
|
|
|
+};
|
|
|
|
|
+
|
|
|
|
|
+fn multiply_complex(a: Complex, b: Complex) -> Complex {
|
|
|
|
|
+ return Complex(
|
|
|
|
|
+ a.x * b.x - a.y * b.y,
|
|
|
|
|
+ a.y * b.x + a.x * b.y
|
|
|
|
|
+ );
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+fn add_complex(a: Complex, b: Complex) -> Complex{
|
|
|
|
|
+ return Complex(
|
|
|
|
|
+ a.x + b.x,
|
|
|
|
|
+ a.y + b.y
|
|
|
|
|
+ )
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+fn butterfly_operation(a: Complex, b: Complex, twiddle: Complex) -> vec4<f32> {
|
|
|
|
|
+ let twiddle_b: Complex = multiply_complex(twiddle, b);
|
|
|
|
|
+ let result: vec4<f32> = vec4<f32>(a.x + twiddle_b.x, a.y + twiddle_b.y, a.x - twiddle_b.x, a.y - twiddle_b.y);
|
|
|
|
|
+ return result;
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+@compute @workgroup_size(256)
|
|
|
|
|
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
|
|
|
|
|
+ let pixel_coord: vec2<u32> = vec2<u32>(global_id.x, global_id.y);
|
|
|
|
|
+
|
|
|
|
|
+ let thread_count: u32 = u32(u_total_count >> 1); // Integer division by 2
|
|
|
|
|
+ let thread_idx: u32 = pixel_coord.x;
|
|
|
|
|
+
|
|
|
|
|
+ let in_idx: u32 = thread_idx & (u_subseq_count - 1u);
|
|
|
|
|
+ let out_idx: u32 = ((thread_idx - in_idx) << 1) + in_idx;
|
|
|
|
|
+
|
|
|
|
|
+ let angle: f32 = -PI * (f32(in_idx) / f32(u_subseq_count));
|
|
|
|
|
+ let twiddle: Complex = Complex(cos(angle), sin(angle));
|
|
|
|
|
+
|
|
|
|
|
+ let a: vec4<f32> = textureSample(u_input, pixel_coord);
|
|
|
|
|
+ let b: vec4<f32> = textureSample(u_input, vec2<u32>(pixel_coord.x + thread_count, pixel_coord.y));
|
|
|
|
|
+
|
|
|
|
|
+ // Transforming two complex sequences independently and simultaneously
|
|
|
|
|
+
|
|
|
|
|
+ let result0: vec4<f32> = butterfly_operation(Complex(a.x, a.y), Complex(b.x, b.y), twiddle);
|
|
|
|
|
+ let result1: vec4<f32> = butterfly_operation(Complex(a.z, a.w), Complex(b.z, b.w), twiddle);
|
|
|
|
|
+
|
|
|
|
|
+ textureStore(u_output, vec2<u32>(out_idx, pixel_coord.y), vec4<f32>(result0.xy, result1.xy));
|
|
|
|
|
+ textureStore(u_output, vec2<u32>(out_idx + u_subseq_count, pixel_coord.y), vec4<f32>(result0.zw, result1.zw));
|
|
|
|
|
+}
|