87 lines
2.0 KiB
WebGPU Shading Language
87 lines
2.0 KiB
WebGPU Shading Language
const PI2 : f32 = 6.28318530718;
|
|
|
|
@group(0) @binding(0) var<storage, read> inputReal : array<f32>;
|
|
@group(0) @binding(1) var<storage, read> inputImag : array<f32>;
|
|
@group(0) @binding(2) var<storage, read_write> outputReal : array<f32>;
|
|
@group(0) @binding(3) var<storage, read_write> outputImag : array<f32>;
|
|
@group(0) @binding(4) var<storage, read> params : array<f32>;
|
|
|
|
var<workgroup> wr : array<f32, 64>;
|
|
var<workgroup> wi : array<f32, 64>;
|
|
|
|
fn bitReverse( x : u32, bits : u32 ) -> u32 {
|
|
var n : u32 = x;
|
|
var r : u32 = 0u;
|
|
for ( var i : u32 = 0u; i < bits; i = i + 1u ) {
|
|
r = ( r << 1u ) | ( n & 1u );
|
|
n = n >> 1u;
|
|
}
|
|
return r;
|
|
}
|
|
|
|
@compute @workgroup_size( 64 )
|
|
fn main( @builtin(local_invocation_id) lid : vec3<u32>,
|
|
@builtin(workgroup_id) gid : vec3<u32> ) {
|
|
|
|
let N : u32 = u32( params[ 1u ] );
|
|
|
|
if ( N != 64u ) {
|
|
return;
|
|
}
|
|
|
|
let row : u32 = gid.x;
|
|
let col : u32 = lid.x;
|
|
let idx : u32 = row * N + col;
|
|
|
|
let bits : u32 = 6u;
|
|
let rev : u32 = bitReverse( col, bits );
|
|
let srcIdx : u32 = row * N + rev;
|
|
|
|
wr[ col ] = inputReal[ srcIdx ];
|
|
wi[ col ] = inputImag[ srcIdx ];
|
|
|
|
workgroupBarrier();
|
|
|
|
var step : u32 = 1u;
|
|
while ( step < N ) {
|
|
|
|
let jump : u32 = step << 1u;
|
|
let twiddleAngle : f32 = -PI2 / f32( jump );
|
|
|
|
let pairIndex : u32 = ( col / jump ) * jump + ( col % step );
|
|
let matchIndex : u32 = pairIndex + step;
|
|
|
|
let k : u32 = col % step;
|
|
let angle : f32 = twiddleAngle * f32( k );
|
|
|
|
let c : f32 = cos( angle );
|
|
let s : f32 = sin( angle );
|
|
|
|
let er : f32 = wr[ pairIndex ];
|
|
let ei : f32 = wi[ pairIndex ];
|
|
|
|
let or : f32 = wr[ matchIndex ];
|
|
let oi : f32 = wi[ matchIndex ];
|
|
|
|
let tr : f32 = c * or - s * oi;
|
|
let ti : f32 = s * or + c * oi;
|
|
|
|
if ( col % jump < step ) {
|
|
wr[ pairIndex ] = er + tr;
|
|
wi[ pairIndex ] = ei + ti;
|
|
} else {
|
|
wr[ matchIndex ] = er - tr;
|
|
wi[ matchIndex ] = ei - ti;
|
|
}
|
|
|
|
workgroupBarrier();
|
|
|
|
step = jump;
|
|
}
|
|
|
|
let invScale : f32 = 1.0 / f32( N );
|
|
outputReal[ idx ] = wr[ col ] * invScale;
|
|
outputImag[ idx ] = wi[ col ] * invScale;
|
|
|
|
}
|