1// One explicit time step (symplectic Euler) of the von Karman plate.
2//
3// State layout (structure of arrays, n = nx * ny nodes):
4// pos: [u | v | w | wLo] wLo carries the rounding error of the w update
5// vel: [du | dv | dw]
6// Each 16x16 workgroup loads its tile plus a 2-node halo of u, v, w, dw into
7// shared memory once, so the stencils below never touch global memory.
9@group(0) @binding(0) var<uniform> P: Params;
10@group(0) @binding(1) var<uniform> SP: StepParams;
11@group(0) @binding(2) var<storage, read> posIn: array<f32>;
12@group(0) @binding(3) var<storage, read> velIn: array<f32>;
13@group(0) @binding(4) var<storage, read_write> posOut: array<f32>;
14@group(0) @binding(5) var<storage, read_write> velOut: array<f32>;
15@group(0) @binding(6) var<storage, read_write> hist: array<f32>; // [slot * n + node] acceleration of w
17const TILE = 16;
18const HALO = 2;
19const TW = 20; // TILE + 2 * HALO
20var<workgroup> tw: array<f32, 400>;
21// Membrane force contributions of the 17x17 cells around the tile, 4 nodes x 3 components each.
22const CT = 17;
23var<workgroup> cellF: array<f32, 3468>;
25// Tile origin in global node coordinates (set once per invocation).
26var<private> ox: i32;
27var<private> oy: i32;
28var<private> ihx2: f32;
29var<private> ihy2: f32;
30var<private> ihxy: f32;
32fn id(i: i32, j: i32) -> u32 { return u32(j) * P.nx + u32(i); }
33fn lidx(i: i32, j: i32) -> u32 { return u32(j - oy) * u32(TW) + u32(i - ox); }
34// Global indices; the tile covers [ox, ox + TW) x [oy, oy + TW).
35fn W(i: i32, j: i32) -> f32 { return tw[lidx(i, j)]; }
36fn U(i: i32, j: i32) -> f32 { return posIn[id(i, j)]; }
37fn V(i: i32, j: i32) -> f32 { return posIn[P.n + id(i, j)]; }
38fn DW(i: i32, j: i32) -> f32 { return velIn[2u * P.n + id(i, j)]; }
40// Weighted bending "moments" at nodes; zero where the second difference is undefined.
41fn Ma(i: i32, j: i32) -> f32 {
42 let nx = i32(P.nx); let ny = i32(P.ny);
43 if (i < 1 || i > nx - 2 || j < 0 || j > ny - 1) { return 0.0; }
44 let a = (W(i + 1, j) - 2.0 * W(i, j) + W(i - 1, j)) * ihx2;
45 var b = 0.0;
46 if (j >= 1 && j <= ny - 2) { b = (W(i, j + 1) - 2.0 * W(i, j) + W(i, j - 1)) * ihy2; }
47 return trapWeight(i, j) * (a + P.nu * b);
48}
49fn Mb(i: i32, j: i32) -> f32 {
50 let nx = i32(P.nx); let ny = i32(P.ny);
51 if (j < 1 || j > ny - 2 || i < 0 || i > nx - 1) { return 0.0; }
52 let b = (W(i, j + 1) - 2.0 * W(i, j) + W(i, j - 1)) * ihy2;
53 var a = 0.0;
54 if (i >= 1 && i <= nx - 2) { a = (W(i + 1, j) - 2.0 * W(i, j) + W(i - 1, j)) * ihx2; }
55 return trapWeight(i, j) * (b + P.nu * a);
56}
57// Mixed second difference at a cell center; zero outside the plate.
58fn Cc(ci: i32, cj: i32) -> f32 {
59 if (ci < 0 || ci > i32(P.nx) - 2 || cj < 0 || cj > i32(P.ny) - 2) { return 0.0; }
60 return (W(ci + 1, cj + 1) - W(ci + 1, cj) - W(ci, cj + 1) + W(ci, cj)) * ihxy;
61}
63// Membrane forces of the bilinear cell (ci, cj) on its four nodes (2x2 Gauss),
64// written to cellF at local cell index lc.
65fn cellForces(ci: i32, cj: i32, lc: u32) {
66 let base = lc * 12u;
67 if (ci < 0 || ci > i32(P.nx) - 2 || cj < 0 || cj > i32(P.ny) - 2) {
68 for (var q = 0u; q < 12u; q++) { cellF[base + q] = 0.0; }
69 return;
70 }
71 let Un = vec4f(U(ci, cj), U(ci + 1, cj), U(ci + 1, cj + 1), U(ci, cj + 1));
72 let Vn = vec4f(V(ci, cj), V(ci + 1, cj), V(ci + 1, cj + 1), V(ci, cj + 1));
73 let Wn = vec4f(W(ci, cj), W(ci + 1, cj), W(ci + 1, cj + 1), W(ci, cj + 1));
74 let q = 0.25 * P.hx * P.hy;
75 let ihx = 1.0 / P.hx; let ihy = 1.0 / P.hy;
76 let g3 = 0.57735026918962576;
77 var Fu = vec4f(0.0); var Fv = vec4f(0.0); var Fw = vec4f(0.0);
78 for (var gp = 0; gp < 4; gp++) {
79 var xi = -g3; if ((gp & 1) != 0) { xi = g3; }
80 var eta = -g3; if ((gp & 2) != 0) { eta = g3; }
81 let dNdx = vec4f(-(1.0 - eta), 1.0 - eta, 1.0 + eta, -(1.0 + eta)) * (0.5 * ihx);
82 let dNdy = vec4f(-(1.0 - xi), -(1.0 + xi), 1.0 + xi, 1.0 - xi) * (0.5 * ihy);
83 let ux = dot(dNdx, Un); let uy = dot(dNdy, Un);
84 let vx = dot(dNdx, Vn); let vy = dot(dNdy, Vn);
85 let wx = dot(dNdx, Wn); let wy = dot(dNdy, Wn);
86 let exx = ux + 0.5 * wx * wx;
87 let eyy = vy + 0.5 * wy * wy;
88 let gxy = uy + vx + wx * wy;
89 let Nxx = P.C * (exx + P.nu * eyy);
90 let Nyy = P.C * (eyy + P.nu * exx);
91 let Nxy = P.G * gxy;
92 Fu -= q * (Nxx * dNdx + Nxy * dNdy);
93 Fv -= q * (Nyy * dNdy + Nxy * dNdx);
94 Fw -= q * ((Nxx * wx + Nxy * wy) * dNdx + (Nxy * wx + Nyy * wy) * dNdy);
95 }
96 for (var a = 0u; a < 4u; a++) {
97 cellF[base + a * 3u] = Fu[a];
98 cellF[base + a * 3u + 1u] = Fv[a];
99 cellF[base + a * 3u + 2u] = Fw[a];
100 }
101}
103// Stored membrane force on local node `a` of cell (ci, cj).
104fn cellForce(ci: i32, cj: i32, a: u32) -> vec3f {
105 let lc = u32(cj - oy - 1) * u32(CT) + u32(ci - ox - 1);
106 let base = lc * 12u + a * 3u;
107 return vec3f(cellF[base], cellF[base + 1u], cellF[base + 2u]);
108}
110@compute @workgroup_size(16, 16)
111fn main(@builtin(global_invocation_id) gid: vec3u, @builtin(local_invocation_id) lid: vec3u, @builtin(workgroup_id) wid: vec3u) {
112 let n = P.n;
113 ox = i32(wid.x) * TILE - HALO;
114 oy = i32(wid.y) * TILE - HALO;
115 ihx2 = 1.0 / (P.hx * P.hx); ihy2 = 1.0 / (P.hy * P.hy); ihxy = 1.0 / (P.hx * P.hy);
116 // Cooperative load of the w tile (400 entries, 256 threads).
117 let lt = lid.y * 16u + lid.x;
118 for (var e = lt; e < 400u; e += 256u) {
119 let gi = ox + i32(e % u32(TW)); let gj = oy + i32(e / u32(TW));
120 var ww = 0.0;
121 if (gi >= 0 && gi < i32(P.nx) && gj >= 0 && gj < i32(P.ny)) { ww = posIn[2u * n + id(gi, gj)]; }
122 tw[e] = ww;
123 }
124 workgroupBarrier();
125 // Membrane forces of the 17x17 cells adjacent to the tile's nodes.
126 for (var e = lt; e < 289u; e += 256u) {
127 let lci = i32(e % u32(CT)); let lcj = i32(e / u32(CT));
128 cellForces(ox + 1 + lci, oy + 1 + lcj, e);
129 }
130 workgroupBarrier();
132 let i = i32(gid.x); let j = i32(gid.y);
133 if (i >= i32(P.nx) || j >= i32(P.ny)) { return; }
134 let k = id(i, j);
137 let bend = (Ma(i - 1, j) - 2.0 * Ma(i, j) + Ma(i + 1, j)) * ihx2
138 + (Mb(i, j - 1) - 2.0 * Mb(i, j) + Mb(i, j + 1)) * ihy2
139 + 2.0 * (1.0 - P.nu) * (Cc(i, j) - Cc(i - 1, j) - Cc(i, j - 1) + Cc(i - 1, j - 1)) * ihxy;
140 var F = vec3f(0.0, 0.0, -P.D * P.hx * P.hy * bend);
141 F += cellForce(i, j, 0u);
142 F += cellForce(i - 1, j, 1u);
143 F += cellForce(i - 1, j - 1, 2u);
144 F += cellForce(i, j - 1, 3u);
145 let mass = P.rhoh * P.hx * P.hy * trapWeight(i, j);
146 var acc = F * (1.0 / mass);
148 // Loss: frequency-independent on all components, frequency-dependent (graph Laplacian) on w.
149 let vel = vec3f(velIn[k], velIn[n + k], velIn[2u * n + k]);
150 var lap = 0.0;
151 if (i > 0) { lap += (DW(i - 1, j) - vel.z) * ihx2; }
152 if (i < i32(P.nx) - 1) { lap += (DW(i + 1, j) - vel.z) * ihx2; }
153 if (j > 0) { lap += (DW(i, j - 1) - vel.z) * ihy2; }
154 if (j < i32(P.ny) - 1) { lap += (DW(i, j + 1) - vel.z) * ihy2; }
155 acc -= 2.0 * SP.sigma0 * vel;
156 acc.z += 2.0 * SP.sigma1 * lap;
158 if (SP.audioSlot >= 0) { hist[u32(SP.audioSlot) * n + k] = acc.z; }
160 var vNew = vel + P.dt * acc;
161 let pOld = vec3f(posIn[k], posIn[n + k], posIn[2u * n + k]);
162 var pNew = pOld + P.dt * vNew;
163 // Compensated update of w: the increment dt * v can be far below the
164 // float32 resolution of w, so the rounding error is carried separately.
165 let y = P.dt * vNew.z + posIn[3u * n + k];
166 let t = pOld.z + y;
167 var lo = (pOld.z - t) + y;
168 pNew.z = t;
169 if ((SP.flags & 1u) != 0u) {
170 let dxp = nodeX(i) - P.pushX; let dyp = nodeY(j) - P.pushY;
171 if (dxp * dxp + dyp * dyp <= P.pushR2) { pNew.z = SP.pushW; vNew.z = SP.pushV; lo = 0.0; }
172 }
173 posOut[k] = pNew.x; posOut[n + k] = pNew.y; posOut[2u * n + k] = pNew.z; posOut[3u * n + k] = lo;
174 velOut[k] = vNew.x; velOut[n + k] = vNew.y; velOut[2u * n + k] = vNew.z;
175}