Reese Levine
238856ec8f
* Basic JIT compilation for mul_mat, get_rows, and scale (#17) * scale jit working * preliminary working jit for getrows and mulmat, needs refining * simplified mul_mat preprocessing switch statement * get_rows fixes, mul_mat refinement * formatted + last edits * removed some extraneous prints * fixed get_rows, fixed workgroup dispatch in mul_mat. no gibberish * small fix * some changes, working * get_rows and mul_mat jit fixed and working * Update formatting * formatting * Add header --------- Co-authored-by: Neha Abbas <nehaabbas@ReeseLevines-MacBook-Pro.local> Co-authored-by: Reese Levine <reeselevine1@gmail.com> * Start work on all-encompassing shader library * refactor argmax, set_rows * Refactor all but flashattention, mat mul * flashattention and matrix multiplication moved to new format * clean up preprocessing * Formatting * remove duplicate constants * Split large shaders into multiple static strings --------- Co-authored-by: neha-ha <137219201+neha-ha@users.noreply.github.com>
64 lines
1.4 KiB
WebGPU Shading Language
64 lines
1.4 KiB
WebGPU Shading Language
#ifdef INPLACE
|
|
@group(0) @binding(1)
|
|
var<uniform> params: Params;
|
|
|
|
fn store_scale(val: f32, offset: u32) {
|
|
src[offset] = val;
|
|
}
|
|
#else
|
|
@group(0) @binding(1)
|
|
var<storage, read_write> dst: array<f32>;
|
|
|
|
@group(0) @binding(2)
|
|
var<uniform> params: Params;
|
|
|
|
fn store_scale(val: f32, offset: u32) {
|
|
dst[offset] = val;
|
|
}
|
|
#endif
|
|
|
|
struct Params {
|
|
offset_src: u32,
|
|
offset_dst: u32,
|
|
|
|
// Strides (in elements)
|
|
stride_src1: u32,
|
|
stride_src2: u32,
|
|
stride_src3: u32,
|
|
|
|
stride_dst1: u32,
|
|
stride_dst2: u32,
|
|
stride_dst3: u32,
|
|
|
|
ne: u32,
|
|
ne0: u32,
|
|
ne1: u32,
|
|
ne2: u32,
|
|
|
|
scale: f32,
|
|
bias: f32
|
|
};
|
|
|
|
@group(0) @binding(0)
|
|
var<storage, read_write> src: array<f32>;
|
|
|
|
@compute @workgroup_size(WG_SIZE)
|
|
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
if (gid.x >= params.ne) {
|
|
return;
|
|
}
|
|
|
|
var i = gid.x;
|
|
let i3 = i / (params.ne2 * params.ne1 * params.ne0);
|
|
i = i % (params.ne2 * params.ne1 * params.ne0);
|
|
let i2 = i / (params.ne1 * params.ne0);
|
|
i = i % (params.ne1 * params.ne0);
|
|
let i1 = i / params.ne0;
|
|
let i0 = i % params.ne0;
|
|
|
|
let i_src = params.offset_src + i3 * params.stride_src3 + i2 * params.stride_src2 + i1 * params.stride_src1 + i0;
|
|
let i_dst = params.offset_dst + i3 * params.stride_dst3 + i2 * params.stride_dst2 + i1 * params.stride_dst1 + i0;
|
|
|
|
store_scale(src[i_src] * params.scale + params.bias, i_dst);
|
|
}
|