Repository navigation
Expand file tree
/
Copy pathkernel.wgsl
More file actions
123 lines (101 loc) · 5.19 KB
/
Copy pathkernel.wgsl
File metadata and controls
123 lines (101 loc) · 5.19 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
/*
bindings are blueprints to what is in memory...
Those @bindings like will create a structure in memory, for instance:
obs: the size of the arrays must be known at "compile time", so this blueprint is tricky to make (have to replace placeholders in the source string).
Remember to just place the correct type and in the correct order that you create the buffers in the python
file.
In CUDA, we have variables with special names that have pre-defined functions or assignments, here we don't... To identify
variables as certain functions or special "link" in the GPU, we use @xxx decorators. For instance, @builtin(global_invocation_id) will mark a variable as the index of the thread (to locate a position in a vector, for instance)
*/
// note: position is 8 bytes (ok), kind is only 4, no ok... use pad to make 8. parameters is 16, ok!
struct Source {
position: vec2<i32>,
kind: i32,
pad:i32,
parameters: vec4<f32>,
};
@group(0) @binding(0) var<storage, read_write> c1x: array<array<f32, [[M]]>, [[N]]>;
@group(0) @binding(1) var<storage, read_write> c1y: array<array<f32, [[M]]>, [[N]]>;
@group(0) @binding(2) var<storage, read_write> c2x: array<array<f32, [[M]]>, [[N]]>;
@group(0) @binding(3) var<storage, read_write> c2y: array<array<f32, [[M]]>, [[N]]>;
@group(0) @binding(4) var<storage, read_write> c3: array<array<f32, [[M]]>, [[N]]>;
@group(0) @binding(5) var<storage, read_write> c4: array<array<f32, [[M]]>, [[N]]>;
@group(0) @binding(6) var<storage, read_write> wire: array<array<f32, [[M]]>, [[N]]>;
@group(0) @binding(7) var<uniform> dx: f32;
@group(0) @binding(8) var<uniform> dy: f32;
@group(0) @binding(9) var<uniform> dt: f32;
@group(0) @binding(10) var<storage, read_write> t: f32;
@group(0) @binding(11) var<storage, read_write> Ezix_1: array<array<f32, [[M]]>, [[N]]>;
@group(0) @binding(12) var<storage, read_write> Eziy_1: array<array<f32, [[M]]>, [[N]]>;
@group(0) @binding(13) var<storage, read_write> Hxi_1: array<array<f32, [[M]]>, [[N]]>;
@group(0) @binding(14) var<storage, read_write> Hyi_1: array<array<f32, [[M]]>, [[N]]>;
@group(0) @binding(15) var<storage, read_write> Ezix_2: array<array<f32, [[M]]>, [[N]]>;
@group(0) @binding(16) var<storage, read_write> Eziy_2: array<array<f32, [[M]]>, [[N]]>;
@group(0) @binding(17) var<storage, read_write> Hxi_2: array<array<f32, [[M]]>, [[N]]>;
@group(0) @binding(18) var<storage, read_write> Hyi_2: array<array<f32, [[M]]>, [[N]]>;
@group(0) @binding(19) var<storage, read_write> sources: array<Source>;
var<private> pi : f32 = 3.141592;
fn generate_source(kind : i32, parameters : vec4<f32>, t : f32) -> f32 {
let p1 = parameters.x;
let p2 = parameters.y;
let p3 = parameters.z;
let p4 = parameters.w;
if (kind == 1) {
return p1*sin(2.0*pi*p2*t);
}
if (kind == 2) {
return p1*exp(-(t-p2)*(t-p2)/(2.0*p3));
}
if (kind == 3) {
return p1*sin(2.0*pi*p2*t)*exp(-(t-p3)*(t-p3)/(2.0*p4));
}
return 0.0;
}
@compute @workgroup_size([[workgoup_size]], [[workgoup_size]])
fn leap_frog1(@builtin(global_invocation_id) index: vec3<u32>) {
let idx_row = index.x;
let idx_col = index.y;
if ( (i32(idx_row)>=[[N]]-1) || (i32(idx_col)>=[[M]]-1) ){
return;
}
let G1x = c1x[idx_row+1][idx_col+1]*Ezix_1[idx_row+1][idx_col+1];
let G1y = c1y[idx_row+1][idx_col+1]*Eziy_1[idx_row+1][idx_col+1];
let G2 = wire[idx_row+1][idx_col+1]*Hyi_1[idx_row+1][idx_col+1];
let G3 = wire[idx_row+0][idx_col+1]*Hyi_1[idx_row+0][idx_col+1];
let G4 = wire[idx_row+1][idx_col+1]*Hxi_1[idx_row+1][idx_col+1];
let G5 = wire[idx_row+1][idx_col+0]*Hxi_1[idx_row+1][idx_col+0];
Ezix_2[idx_row+1][idx_col+1] = G1x + c2x[idx_row+1][idx_col+1]*( - ( G4 - G5 )/dy );
Eziy_2[idx_row+1][idx_col+1] = G1y + c2y[idx_row+1][idx_col+1]*( ( G2 - G3 )/dx );
}
@compute @workgroup_size([[workgoup_size]], [[workgoup_size]])
fn leap_frog2(@builtin(global_invocation_id) index: vec3<u32>) {
let idx_row = index.x;
let idx_col = index.y;
if ( (i32(idx_row)>=[[N]]-1) || (i32(idx_col)>=[[M]]-1) ){
return;
}
let K1 = wire[idx_row+0][idx_col+0]*(Ezix_2[idx_row+0][idx_col+0] + Eziy_2[idx_row+0][idx_col+0]);
let K2 = wire[idx_row+0][idx_col+1]*(Ezix_2[idx_row+0][idx_col+1] + Eziy_2[idx_row+0][idx_col+1]);
let K3 = wire[idx_row+1][idx_col+0]*(Ezix_2[idx_row+1][idx_col+0] + Eziy_2[idx_row+1][idx_col+0]);
let K4 = wire[idx_row+0][idx_col+0]*(Ezix_2[idx_row+0][idx_col+0] + Eziy_2[idx_row+0][idx_col+0]);
Hxi_2[idx_row][idx_col] = Hxi_1[idx_row][idx_col] + c3[idx_row][idx_col]*( ( K1 - K2 )/dy );
Hyi_2[idx_row][idx_col] = Hyi_1[idx_row][idx_col] + c4[idx_row][idx_col]*( ( K3 - K4 )/dx );
}
@compute @workgroup_size(1, 1)
fn update_source(@builtin(global_invocation_id) index: vec3<u32>) {
if ( (i32(index.x)>=1) || (i32(index.y)>=1) ){
return;
}
t += dt;
for (var i : u32 = 0; i<2; i++) {
let x = sources[i].position.x;
let y = sources[i].position.y;
let kind = sources[i].kind;
let parameters = sources[i].parameters;
let v = generate_source(kind, parameters, t);
Ezix_1[x][y] += v/2.0;
Eziy_1[x][y] += v/2.0;
Ezix_2[x][y] += v/2.0;
Eziy_2[x][y] += v/2.0;
}
}