First commit.
This commit is contained in:
33
shaders/neural/forward_dense_layer1.wgsl
Normal file
33
shaders/neural/forward_dense_layer1.wgsl
Normal file
@@ -0,0 +1,33 @@
|
||||
@group(0) @binding(0) var<storage, read> inputImages : array<f32>;
|
||||
@group(0) @binding(1) var<storage, read_write> weightLayer1 : array<f32>;
|
||||
@group(0) @binding(2) var<storage, read_write> biasLayer1 : array<f32>;
|
||||
@group(0) @binding(3) var<storage, read_write> hiddenActivations : array<f32>;
|
||||
@group(0) @binding(4) var<uniform> layer1Sizes : vec3<u32>;
|
||||
// x = inputDim, y = hiddenDim, z = batchSize
|
||||
|
||||
@compute @workgroup_size(64)
|
||||
fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
|
||||
let index = global_id.x;
|
||||
|
||||
let inputDim = layer1Sizes.x;
|
||||
let hiddenDim = layer1Sizes.y;
|
||||
let batchSize = layer1Sizes.z;
|
||||
|
||||
let totalElements = batchSize * hiddenDim;
|
||||
if (index >= totalElements) { return; }
|
||||
|
||||
let sampleIndex = index / hiddenDim;
|
||||
let hiddenIndex = index % hiddenDim;
|
||||
|
||||
var sum : f32 = biasLayer1[hiddenIndex];
|
||||
|
||||
let baseInputIndex = sampleIndex * inputDim;
|
||||
let baseWeightIndex = hiddenIndex * inputDim;
|
||||
|
||||
for (var i : u32 = 0u; i < inputDim; i++) {
|
||||
sum += inputImages[baseInputIndex + i] *
|
||||
weightLayer1[baseWeightIndex + i];
|
||||
}
|
||||
|
||||
hiddenActivations[index] = max(sum, 0.0);
|
||||
}
|
||||
Reference in New Issue
Block a user