GEMMStore Playground

Device: Detecting...WebGPU: Checking...

Upload .js

Kernel SpecificationJavaScript

// GEMMStore Playground - Edit this spec and click "Run Benchmark"

window.CustomKernel = {
    name: "CustomKernel",
    description: "A simple matrix multiplication kernel - edit me!",

inputArgs: [
        {
            name: "matrixA",
            storage_type: "storage",
            type: "float[]",
            fill: "randf32",
            sizes: [1024*1024],
            io_type: "input",
            order: 0
        },
        {
            name: "matrixB",
            storage_type: "storage",
            type: "float[]",
            fill: "randf32",
            sizes: [1024*1024],
            io_type: "input",
            order: 1
        },
        {
            name: "matrixC",
            storage_type: "storage",
            type: "float[]",
            fill: "don't fill",
            sizes: [1024*1024],
            io_type: "output",
            order: 2
        },
        {
            name: "matrixSize",
            storage_type: "uniform",
            type: "int",
            values: [1024]
        }
    ],

tuningArgs: [
        {
            name: "workgroupSizeX",
            type: "fixed",
            values: ["16"],
            description: "Workgroup size in X dimension"
        },
        {
            name: "workgroupSizeY",
            type: "fixed",
            values: ["16"],
            description: "Workgroup size in Y dimension"
        }
    ],

wgsl_shader: `
@compute @workgroup_size({{workgroupSizeX}}, {{workgroupSizeY}})
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let row = global_id.y;
    let col = global_id.x;

if (row >= matrixSize || col >= matrixSize) {
        return;
    }

{{matrixC}}[row * matrixSize + col] = sum;
}`,

requiredFeatures: [],
    measureArgs: ["ms", "GFLOPs"],

js_reference: async function matrixMultiplyGPU(A, B, n) {
        const adapter = await navigator.gpu.requestAdapter();
        const device = await adapter.requestDevice();

device.queue.writeBuffer(bufferA, 0, A);
        device.queue.writeBuffer(bufferB, 0, B);
        device.queue.writeBuffer(uniformBuffer, 0, new Uint32Array([n]));

const computeShader = `
            @group(0) @binding(0) var<storage, read> A: array<f32>;
            @group(0) @binding(1) var<storage, read> B: array<f32>;
            @group(0) @binding(2) var<storage, read_write> C: array<f32>;
            @group(0) @binding(3) var<uniform> n: u32;

@compute @workgroup_size(16, 16)
            fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
                let row = global_id.y;
                let col = global_id.x;
                if (row >= n || col >= n) { return; }
                var sum = 0.0;
                for (var k = 0u; k < n; k++) {
                    sum += A[row * n + k] * B[k * n + col];
                }
                C[row * n + col] = sum;
            }
        `;

const pipeline = device.createComputePipeline({
            layout: 'auto',
            compute: { module: device.createShaderModule({ code: computeShader }), entryPoint: 'main' },
        });

const bindGroup = device.createBindGroup({
            layout: pipeline.getBindGroupLayout(0),
            entries: [
                { binding: 0, resource: { buffer: bufferA } },
                { binding: 1, resource: { buffer: bufferB } },
                { binding: 2, resource: { buffer: bufferC } },
                { binding: 3, resource: { buffer: uniformBuffer } },
            ],
        });

const encoder = device.createCommandEncoder();
        const pass = encoder.beginComputePass();
        pass.setPipeline(pipeline);
        pass.setBindGroup(0, bindGroup);
        pass.dispatchWorkgroups(Math.ceil(n / 16), Math.ceil(n / 16));
        pass.end();
        encoder.copyBufferToBuffer(bufferC, 0, stagingBuffer, 0, n * n * 4);
        device.queue.submit([encoder.finish()]);

await stagingBuffer.mapAsync(GPUMapMode.READ);
        const result = new Float32Array(stagingBuffer.getMappedRange());
        const copy = result.slice();
        stagingBuffer.unmap();
        return copy;
    },

valid_Args: function(A, B, C, n, workgroupSizeX, workgroupSizeY) {
        var first = A.length === n*n && B.length === n*n && C.length === n*n;
        var second = parseInt(workgroupSizeX) * parseInt(workgroupSizeY) <= 256;
        return first && second;
    },

num_Workgroups: function(A, B, C, n, workgroupSizeX, workgroupSizeY) {
        var numWorkgroupsX = Math.ceil(n / parseInt(workgroupSizeX));
        var numWorkgroupsY = Math.ceil(n / parseInt(workgroupSizeY));
        return [numWorkgroupsX, numWorkgroupsY];
    }
};

Benchmark Results

No results yet. Edit the spec and click "Run Benchmark".

Debug Log

Logs will appear here when you run a benchmark.