From 5ae9c127eb9af4b34b5d5d7ca5004bc7baf57456 Mon Sep 17 00:00:00 2001
From: Ahmed Harmouche <ahmedharmouche92@gmail.com>
Date: Sun, 10 Dec 2023 10:38:20 +0100
Subject: [PATCH] Refactor f16tof32GPU into an npm package, load it from unpkg

---
 f16_to_f32.js | 64 ---------------------------------------------------
 index.html    | 16 ++++---------
 2 files changed, 5 insertions(+), 75 deletions(-)
 delete mode 100644 f16_to_f32.js

diff --git a/f16_to_f32.js b/f16_to_f32.js
deleted file mode 100644
index fb91ea0..0000000
--- a/f16_to_f32.js
+++ /dev/null
@@ -1,64 +0,0 @@
-const f16tof32 = `
-fn u16_to_f16(x: u32) -> f32 {
-    let sign = f32((x >> 15) & 0x1);
-    let exponent = f32((x >> 10) & 0x1F);
-    let fraction = f32(x & 0x3FF);
-
-    let sign_multiplier = select(1.0, -1.0, sign == 1.0);
-    if (exponent == 0.0) {
-        return sign_multiplier * 6.103515625e-5 * (fraction / 1024.0);
-    } else {
-        return sign_multiplier * exp2(exponent - 15.0) * (1.0 + fraction / 1024.0);
-    }
-}
-
-@group(0) @binding(0) var<storage,read_write> data0: array<u32>;
-@group(0) @binding(1) var<storage,read_write> data1: array<f32>;
-@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let gidx = gid.x;
-    let outgidx = gidx*2;
-
-    if (gidx >= arrayLength(&data0)) {
-        return;
-    }
-
-    let oo = data0[gidx];
-    let oo1 = (oo >> 16);
-    let oo2 = oo & 0xFFFFu;
-
-    let f1 = u16_to_f16(oo2);
-    let f2 = u16_to_f16(oo1);
-    
-    data1[outgidx] = f1;
-    data1[outgidx + 1] = f2;
-}`;
-
-window.f16tof32GPU = async(device, inf16) => {
-    const input = device.createBuffer({size: inf16.length, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST });
-    const output = device.createBuffer({size: inf16.length*2, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST });
-
-    const gpuWriteBuffer = device.createBuffer({size: input.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE });
-    const gpuReadBuffer = device.createBuffer({ size: output.size, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ });
-    const commandEncoder = device.createCommandEncoder();
-    await gpuWriteBuffer.mapAsync(GPUMapMode.WRITE);
-
-    const alignedUint32View = new Uint32Array(inf16.buffer, inf16.byteOffset, inf16.length / 4);
-    new Uint32Array(gpuWriteBuffer.getMappedRange()).set(alignedUint32View);
-
-    gpuWriteBuffer.unmap();
-    commandEncoder.copyBufferToBuffer(gpuWriteBuffer, 0, input, 0, gpuWriteBuffer.size);
-    const pipeline = await device.createComputePipelineAsync({layout: "auto", compute: { module: device.createShaderModule({ code: f16tof32 }), entryPoint: "main" }});
-
-    addComputePass(device, commandEncoder, pipeline, [input, output], [Math.ceil(inf16.length/(4*256)), 1, 1]);
-
-    commandEncoder.copyBufferToBuffer(output, 0, gpuReadBuffer, 0, output.size);
-    const gpuCommands = commandEncoder.finish();
-    device.queue.submit([gpuCommands]);
-
-    await gpuReadBuffer.mapAsync(GPUMapMode.READ);
-    const resultBuffer = new Float32Array(gpuReadBuffer.size/4);
-    resultBuffer.set(new Float32Array(gpuReadBuffer.getMappedRange()));
-    gpuReadBuffer.unmap();
-
-    return resultBuffer;
-}
diff --git a/index.html b/index.html
index 4c61524..bde471c 100644
--- a/index.html
+++ b/index.html
@@ -68,7 +68,10 @@
         import ClipTokenizer from 'https://softwired.nyc3.cdn.digitaloceanspaces.com/sd/clip_tokenizer.js';
         window.clipTokenizer = new ClipTokenizer();
     </script>
-    <script src="./f16_to_f32.js"></script>
+    <script type="module">
+        import { f16tof32GPU } from 'https://unpkg.com/f16-to-f32-gpu@0.1.0/src/index.js';
+        window.f16tof32GPU = f16tof32GPU;
+    </script>
     <script src="./net.js"></script>
 </head>
 
@@ -351,16 +354,7 @@ <h1 id="sdTitle">StableDiffusion by <a href="https://github.com/tinygrad/tinygra
                 let chunkStartF16 = 8 + metadataLength + (decodeChunkSize * i);
                 let chunkEndF16 = chunkStartF16 + decodeChunkSize;
                 let chunk = combinedBuffer.subarray(chunkStartF16, chunkEndF16);
-
-                if (chunk.byteLength %4 != 0) {
-                    const paddingBytes = 4 - (chunk.byteLength % 4);
-                    const alignedBuffer = new ArrayBuffer(chunk.byteLength + paddingBytes);
-                    const alignedView = new Uint8Array(alignedBuffer);
-                    alignedView.set(new Uint8Array(chunk));
-                    chunk = alignedView;
-                }
-
-                let result = await f16tof32GPU(device, chunk);
+                let result = await f16tof32GPU(chunk);
                 let resultUint8 = new Uint8Array(result.buffer);
                 let chunkStartF32 = 8 + metadataLength + (decodeChunkSize * i * 2);
                 let chunkEndF32 = chunkStartF32 + resultUint8.byteLength;