FazBrowse GitHub Viewer
|
Trending
|
URL:
|
Home
Tools:
[Download Repo ZIP]
[View Raw Code]
[Original HTTPS Page]
modelscript/packages/language/src/codegen/emit_webgpu.ts at main · modelscript/modelscript · GitHub
Uh oh!
There was an error while loading.
Please reload this page
.
modelscript
/
modelscript
Public
Uh oh!
There was an error while loading.
Please reload this page
.
Notifications
You must be signed in to change notification settings
Fork
3
Star
12
Code
Issues
1
Pull requests
5
Discussions
Actions
Projects
Security and quality
0
Insights
Additional navigation options
Code
Issues
Pull requests
Discussions
Actions
Projects
Security and quality
Insights
Expand file tree
Breadcrumbs
modelscript
/
packages
/
language
/
src
/
codegen
/
emit_webgpu.ts
Copy path
More file actions
More file actions
Latest commit
History
History
History
105 lines (87 loc) · 3.98 KB
Breadcrumbs
modelscript
/
packages
/
language
/
src
/
codegen
/
emit_webgpu.ts
Copy path
File metadata and controls
105 lines (87 loc) · 3.98 KB
Raw
Copy raw file
Download raw file
Open symbols panel
Edit and raw actions
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
import
{
LanguageOptions
}
from
"../dsl.js"
;
export
function
generateWebGPUEmitter
(
grammarDef
:
LanguageOptions
)
:
string
{
const
tileSize
=
grammarDef
.
targets
?.
webgpu
?.
tileSize
||
16
;
const
[
wgX
,
wgY
]
=
grammarDef
.
targets
?.
webgpu
?.
workgroupSize
||
[
tileSize
,
tileSize
]
;
return
`// --- Auto-Generated WebGPU Backend (Phase 5) ---
// Cross-platform, browser-native tensor execution via WGSL Compute Shaders
export function emit_wgsl_prelude(): string {
return \`
struct TensorDim {
rows: u32,
cols: u32,
}
@group(0) @binding(0) var<storage, read> dimA: TensorDim;
@group(0) @binding(1) var<storage, read> dimB: TensorDim;
@group(0) @binding(2) var<storage, read> A: array<f32>;
@group(0) @binding(3) var<storage, read> B: array<f32>;
@group(0) @binding(4) var<storage, read_write> C: array<f32>;
\`;
}
export function emit_wgsl_matmul(): string {
// Tiled Matrix Multiplication Shader
return \`
const TILE_SIZE =
${
tileSize
}
u;
var<workgroup> tileA: array<f32,
${
tileSize
*
tileSize
}
>;
var<workgroup> tileB: array<f32,
${
tileSize
*
tileSize
}
>;
@compute @workgroup_size(
${
wgX
}
,
${
wgY
}
)
fn matmul_main(
@builtin(global_invocation_id) global_id: vec3<u32>,
@builtin(local_invocation_id) local_id: vec3<u32>,
@builtin(workgroup_id) group_id: vec3<u32>
) {
let row = global_id.y;
let col = global_id.x;
let K = dimA.cols;
let numTiles = (K + TILE_SIZE - 1u) / TILE_SIZE;
var acc = 0.0;
for (var t = 0u; t < numTiles; t++) {
let tiledColA = t * TILE_SIZE + local_id.x;
let tiledRowB = t * TILE_SIZE + local_id.y;
// Load into shared memory
if (row < dimA.rows && tiledColA < K) {
tileA[local_id.y * TILE_SIZE + local_id.x] = A[row * K + tiledColA];
} else {
tileA[local_id.y * TILE_SIZE + local_id.x] = 0.0;
}
if (tiledRowB < K && col < dimB.cols) {
tileB[local_id.y * TILE_SIZE + local_id.x] = B[tiledRowB * dimB.cols + col];
} else {
tileB[local_id.y * TILE_SIZE + local_id.x] = 0.0;
}
workgroupBarrier();
for (var k = 0u; k < TILE_SIZE; k++) {
acc += tileA[local_id.y * TILE_SIZE + k] * tileB[k * TILE_SIZE + local_id.x];
}
workgroupBarrier();
}
if (row < dimA.rows && col < dimB.cols) {
C[row * dimB.cols + col] = acc;
}
}
\`;
}
export function emit_webgpu_host_spmv(nnz: u32, numRows: u32, numCols: u32): string {
// Generates JS host code to map WASM linear memory into WebGPU Storage Buffers
return \`
async function executeSpMV(device, valBufferOffset, colPtrOffset, rowIdxOffset, xBufferOffset, yBufferOffset) {
// Load WASM memory slices
const vals = new Float32Array(wasmMemory.buffer, valBufferOffset, \${nnz});
const colPtrs = new Uint32Array(wasmMemory.buffer, colPtrOffset, \${numCols} + 1);
const rowIdxs = new Uint32Array(wasmMemory.buffer, rowIdxOffset, \${nnz});
const xVec = new Float32Array(wasmMemory.buffer, xBufferOffset, \${numCols});
// Create GPU buffers
const gpuVals = device.createBuffer({ size: vals.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
const gpuColPtrs = device.createBuffer({ size: colPtrs.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
const gpuRowIdxs = device.createBuffer({ size: rowIdxs.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
const gpuX = device.createBuffer({ size: xVec.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
const gpuY = device.createBuffer({ size: \${numRows} * 4, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC });
// Write to GPU
device.queue.writeBuffer(gpuVals, 0, vals);
device.queue.writeBuffer(gpuColPtrs, 0, colPtrs);
device.queue.writeBuffer(gpuRowIdxs, 0, rowIdxs);
device.queue.writeBuffer(gpuX, 0, xVec);
}
\`;
}
`
;
}
Back
|
FazBrowse Home
|
New Git URL