| 40 | } |
| 41 | |
| 42 | getUserCode(): string { |
| 43 | const opString = this.component === 'real' ? |
| 44 | 'return real * expR - imag * expI;' : |
| 45 | 'return real * expI + imag * expR;'; |
| 46 | const userCode = ` |
| 47 | fn unaryOpComplex(real: f32, expR: f32, imag: f32, expI: f32) -> f32 { |
| 48 | ${opString} |
| 49 | } |
| 50 | |
| 51 | fn mulMatDFT(batch: i32, index: i32) -> f32 { |
| 52 | let indexRatio = f32(index) / f32(uniforms.realShape[1]); |
| 53 | let exponentMultiplierTimesIndexRatio = |
| 54 | uniforms.exponentMultiplier * indexRatio; |
| 55 | |
| 56 | var result = 0.0; |
| 57 | |
| 58 | for (var i = 0; i < uniforms.realShape[1]; i = i + 1) { |
| 59 | // x = (-2|2 * PI / N) * index * i; |
| 60 | let x = exponentMultiplierTimesIndexRatio * f32(i); |
| 61 | let expR = cos(x); |
| 62 | let expI = sin(x); |
| 63 | let real = getReal(batch, i); |
| 64 | let imag = getImag(batch, i); |
| 65 | |
| 66 | result = result + |
| 67 | unaryOpComplex(real, expR, imag, expI) / uniforms.denominator; |
| 68 | } |
| 69 | |
| 70 | return result; |
| 71 | } |
| 72 | |
| 73 | ${main('index')} { |
| 74 | if (index < uniforms.size) { |
| 75 | let coords = getOutputCoords(); |
| 76 | setOutputAtIndex(index, mulMatDFT(coords[0], coords[1])); |
| 77 | } |
| 78 | } |
| 79 | `; |
| 80 | return userCode; |
| 81 | } |
| 82 | } |