Optimizing ONNX Inference: WebAssembly SIMD vs. WebGPU Compute

#include <wasm_simd128.h>
#include <cstdint>

// Assumes input, output, and weights are 16-byte aligned for safe, fast loads.
// In a real scenario, you'd have padding or alignment checks.
extern "C" void fused_depthwise_conv_relu(float* input, float* output, const float* weights, const float* bias, int width, int height) {
    const int stride = width;
    const v128_t zero = wasm_f32x4_splat(0.0f);

    // We process 4 output pixels (a v128) at a time when possible.
    for (int y = 1; y < height - 1; ++y) {
        for (int x = 1; x < width - 4; x += 4) {
            int base_idx = y * stride + x;

            // Load 3x3 input window for the first of 4 parallel pixels
            // This part is complex; real libraries have efficient im2col/sliding window helpers
            // For simplicity, we manually load and calculate.
            v128_t w0 = wasm_v128_load(weights + 0);
            v128_t w1 = wasm_v128_load(weights + 3);
            v128_t w2 = wasm_v128_load(weights + 6);

            // Load input vectors. This is the most performance-critical part.
            // Unaligned loads are much slower.
            v128_t i0 = wasm_v128_load(input + (base_idx - stride - 1));
            v128_t i1 = wasm_v128_load(input + (base_idx - 1));
            v128_t i2 = wasm_v128_load(input + (base_idx + stride - 1));
            
            // This is a simplified calculation. A real conv is more involved.
            // The point is to demonstrate vector math.
            v128_t acc = wasm_f32x4_splat(0.0f);
            // Simplified dot product simulation
            acc = wasm_f32x4_add(acc, wasm_f32x4_mul(i0, w0));
            acc = wasm_f32x4_add(acc, wasm_f32x4_mul(i1, w1));
            acc = wasm_f32x4_add(acc, wasm_f32x4_mul(i2, w2));
            
            // Add bias
            v128_t b = wasm_v128_load(bias);
            acc = wasm_f32x4_add(acc, b);

            // Apply ReLU activation
            acc = wasm_f32x4_max(acc, zero);

            // Store the 4-pixel result
            wasm_v128_store(output + base_idx, acc);
        }
        // Handle edge cases (pixels not divisible by 4) with scalar code here...
    }
}

# -O3 for max optimization
# -msimd128 to enable SIMD instruction set
# -s WASM=1 to output WASM
# -s EXPORTED_FUNCTIONS="['_fused_depthwise_conv_relu']" to export our function
# -s ALLOW_MEMORY_GROWTH=1 is important for dynamic tensor sizes
emcc -O3 -msimd128 -s WASM=1 -s EXPORTED_FUNCTIONS="['_fused_depthwise_conv_relu']" -s ALLOW_MEMORY_GROWTH=1 -o fused_op.js fused_op.cpp

import { InferenceSession } from 'onnxruntime-web';
import createFusedModule from './fused_op.js';

async function runOptimizedInference(session, inputTensor) {
    // ... standard ONNX inference up to the point before our custom op

    // Assume 'intermediateTensor' is the output from the previous layer
    const intermediateTensor = /* result from session.run() */;

    const fusedModule = await createFusedModule();
    const wasmMemory = fusedModule.instance.exports.memory;

    // 1. Allocate memory in the WASM heap
    const inputSize = intermediateTensor.data.length;
    const inputPtr = fusedModule._malloc(inputSize * 4); // 4 bytes per float
    const outputPtr = fusedModule._malloc(inputSize * 4); // Assuming same size output
    // ... allocate for weights and bias too

    // 2. Copy data from JS ArrayBuffer to WASM memory
    new Float32Array(wasmMemory.buffer, inputPtr, inputSize).set(intermediateTensor.data);
    // ... copy weights and bias

    // 3. Call the exported WASM function
    fusedModule._fused_depthwise_conv_relu(
        inputPtr, 
        outputPtr, 
        weightsPtr, 
        biasPtr, 
        width, 
        height
    );

    // 4. Copy data back from WASM memory to JS
    const outputData = new Float32Array(inputSize);
    outputData.set(new Float32Array(wasmMemory.buffer, outputPtr, inputSize));

    // 5. Free the allocated WASM memory
    fusedModule._free(inputPtr);
    fusedModule._free(outputPtr);
    // ... free weights and bias

    // 6. Create a new ONNX Tensor and continue the rest of the model graph
    const outputTensor = new Tensor('float32', outputData, intermediateTensor.dims);
    // ... run the rest of the model
}

// Bindings for our data buffers
@group(0) @binding(0) var<storage, read> input_data: array<f32>;
@group(0) @binding(1) var<storage, read_write> output_data: array<f32>;

// Uniforms for metadata like tensor size
struct Metadata {
    elements: u32,
};
@group(0) @binding(2) var<uniform> metadata: Metadata;

// A shared memory space for threads within the same workgroup
var<workgroup> shared_max: f32;

// Workgroup size defined in the pipeline setup (e.g., 256)
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>, @builtin(local_invocation_id) local_id: vec3<u32>) {
    let idx = global_id.x;
    if (idx >= metadata.elements) {
        return;
    }

    // --- Stage 1: Find max value in parallel ---
    // Each thread finds the max in its portion of the data
    // This is a simplified reduction. A real one would be multi-staged.
    if (local_id.x == 0u) {
        shared_max = -3.4e38; // f32.min
    }
    workgroupBarrier();

    // In a real implementation, you'd do a parallel reduction within the workgroup
    // For this example, we'll do a simpler (less efficient) global max find.
    // Let's assume for this example a separate kernel found the max and it's passed in a buffer.
    // We'll focus on the second part for clarity.

    // --- Stage 2: Numerically Stable Softmax ---
    // Assume `max_val` was pre-computed and stored in another buffer (or is a uniform)
    let max_val = 0.0; // Placeholder for pre-computed max

    // Subtract max for stability and compute exponent
    let exp_val = exp(input_data[idx] - max_val);
    output_data[idx] = exp_val;

    // --- Stage 3: Summation and Division ---
    // This requires another synchronization point and another reduction.
    // We need to sum all the exp_val's.
    workgroupBarrier();

    // A second kernel would typically perform the sum reduction and then a third would do the division.
    // Running softmax efficiently requires multiple dispatches.
    // The output here is just the numerator. Another kernel would divide by the sum.
}

async function runStableSoftmax(device: GPUDevice, inputData: Float32Array): Promise<Float32Array> {
    const elements = inputData.length;

    // 1. Create GPU buffers and copy data
    const inputBuffer = device.createBuffer({
        size: inputData.byteLength,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
        mappedAtCreation: true,
    });
    new Float32Array(inputBuffer.getMappedRange()).set(inputData);
    inputBuffer.unmap();

    const outputBuffer = device.createBuffer({
        size: inputData.byteLength,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
    });

    const metadataBuffer = device.createBuffer({
        size: 4, // one u32
        usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });

    // This is a simplified example. A full softmax requires multiple passes (max, sum, divide).
    const shaderModule = device.createShaderModule({ code: stable_softmax_wgsl_code });

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

    const bindGroup = device.createBindGroup({
        layout: pipeline.getBindGroupLayout(0),
        entries: [
            { binding: 0, resource: { buffer: inputBuffer } },
            { binding: 1, resource: { buffer: outputBuffer } },
            { binding: 2, resource: { buffer: metadataBuffer } },
        ],
    });

    // 2. Encode and dispatch the command
    const commandEncoder = device.createCommandEncoder();
    device.queue.writeBuffer(metadataBuffer, 0, new Uint32Array([elements]));

    const passEncoder = commandEncoder.beginComputePass();
    passEncoder.setPipeline(pipeline);
    passEncoder.setBindGroup(0, bindGroup);
    const workgroupSize = 256;
    const workgroupCount = Math.ceil(elements / workgroupSize);
    passEncoder.dispatchWorkgroups(workgroupCount);
    passEncoder.end();

    // 3. Create a staging buffer to read data back
    const stagingBuffer = device.createBuffer({
        size: inputData.byteLength,
        usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
    });
    commandEncoder.copyBufferToBuffer(outputBuffer, 0, stagingBuffer, 0, inputData.byteLength);

    // 4. Submit to GPU and wait for results
    device.queue.submit([commandEncoder.finish()]);

    await stagingBuffer.mapAsync(GPUMapMode.READ);
    const result = new Float32Array(stagingBuffer.getMappedRange().slice(0));
    stagingBuffer.unmap();

    // IMPORTANT: Cleanup GPU resources to prevent memory leaks
    inputBuffer.destroy();
    outputBuffer.destroy();
    metadataBuffer.destroy();
    stagingBuffer.destroy();

    return result;
}

import { InferenceSession } from 'onnxruntime-web';
import { simd } from 'wasm-feature-detect';

interface SessionOptions {
    executionProviders: string[];
    graphOptimizationLevel: 'all';
}

async function createInferenceSession(modelPath: string): Promise<InferenceSession> {
    const availableProviders: string[] = ['wasm'];

    // Check for WebGPU support
    if ('gpu' in navigator) {
        try {
            const adapter = await navigator.gpu.requestAdapter();
            if (adapter) {
                availableProviders.unshift('webgpu'); // Add to the front (highest priority)
            }
        } catch (e) {
            console.warn('WebGPU is available but failed to initialize.', e);
        }
    }

    // Check for WASM SIMD support
    if (await simd()) {
        // If WebGPU is not available, SIMD is the next best. Otherwise, it's a fallback.
        const wasmIndex = availableProviders.indexOf('wasm');
        availableProviders.splice(wasmIndex, 0, 'wasm-simd');
    }

    console.log('Attempting to create session with providers:', availableProviders);

    for (const provider of availableProviders) {
        try {
            const options: SessionOptions = {
                executionProviders: [provider],
                graphOptimizationLevel: 'all',
            };
            const session = await InferenceSession.create(modelPath, options);
            console.log(`Successfully created session with ${provider} backend.`);
            return session;
        } catch (error) {
            console.warn(`Failed to create session with ${provider}. Trying next provider.`, error);
        }
    }

    throw new Error('Failed to create inference session with any available provider.');
}

// Usage:
// const mySession = await createInferenceSession('./my-model.onnx');
// const results = await mySession.run(inputs);