fasthtml live wgsl editing experiment works

Author: austinvhuang

Diff for: examples/web/run.cpp

@@ -4,7 +4,6 @@
 #include <memory>
 #include "gpu.h"
 
-// #include <webgpu/webgpu.h>
 #include "emscripten/emscripten.h"
 
 using namespace gpu; // createContext, createTensor, createKernel,
@@ -33,10 +32,6 @@ int main(int argc, char **argv) {
   printf("\nHello gpu.cpp!\n");
   printf("--------------\n\n");
 
-  // const WGPUInstanceDescriptor descriptor = { };
-  // std::unique_ptr<WGPUInstanceDescriptor> descriptor = std::make_unique<WGPUInstanceDescriptor>();
-
-  // WGPUInstance instance = wgpuCreateInstance({});
   Context ctx = createContext({});
   static constexpr size_t N = 5000;
   std::array<float, N> inputArr, outputArr;

Diff for: experimental/fasthtml/Makefile

@@ -1,23 +1,29 @@
 GPUCPP=../..
-FLAGS=-std=c++17 -s USE_WEBGPU=1 -s ASYNCIFY=1 -I$(GPUCPP)
+COMMON_FLAGS=-std=c++17 -s USE_WEBGPU=1 -s ASYNCIFY=1 -I$(GPUCPP)
+# Note - no spaces after comma
+# enable exceptions to recover from WGSL failure
+JS_FLAGS=-s EXPORTED_RUNTIME_METHODS=['UTF8ToString','setValue','addFunction'] -s EXPORTED_FUNCTIONS=['_malloc','_free','_executeKernel'] -s NO_DISABLE_EXCEPTION_CATCHING 
+WASM_FLAGS=-s STANDALONE_WASM -s ERROR_ON_UNDEFINED_SYMBOLS=0 -s EXPORTED_FUNCTIONS=['_executeKernel'] -s EXPORTED_RUNTIME_METHODS=['ccall','cwrap'] -DSTANDALONE_WASM
+MODULARIZE_FLAGS=-s EXPORT_NAME='createModule' -s MODULARIZE=1 --bind
+NO_MODULARIZE_FLAGS=-s EXPORTED_FUNCTIONS=['_executeKernel'] -s EXPORTED_RUNTIME_METHODS=['ccall','cwrap'] --bind
 
-.PHONY: default cmake check-emsdk browser clean
+.PHONY: default cmake check-emsdk browser clean server
 
 default: server
 
-build/run.html: check-emsdk run.cpp
-	em++ run.cpp -o build/run.html \
-		$(FLAGS) --shell-file ./custom_shell.html
+# build/run.html: check-emsdk run.cpp
+#	em++ run.cpp -o build/run.html \
+#		$(COMMON_FLAGS) $(JS_FLAGS) --shell-file ./custom_shell.html
 
 build/run.js: check-emsdk run.cpp
-	em++ run.cpp -o build/run.js --shell-file ./custom_shell.html \
-		$(FLAGS)
+	em++ run.cpp -o build/run.js \
+		$(COMMON_FLAGS) $(JS_FLAGS) $(MODULARIZE_FLAGS)
 
 build/run.wasm: check-emsdk run.cpp
 	em++ run.cpp -o build/run.wasm \
-		$(FLAGS)
+		$(COMMON_FLAGS) $(WASM_FLAGS)
 
-server: build/run.wasm
+server: build/run.js
 	python3 run.py
 
 clean:

Diff for: experimental/fasthtml/components/code_editor.py

@@ -8,6 +8,7 @@
 let completionTippy;
 let currentCompletion = '';
 
+
 function initEditor() {
     editor = ace.edit("editor");
     editor.setTheme("ace/theme/monokai");
@@ -23,20 +24,16 @@
     editor.setKeyboardHandler("ace/keyboard/vim");
     
     editor.setValue(
-"// Puzzle 1 : Map\\n\
-// Implement a kernel that adds 10 to each position of vector\\n\
-// a and stores it in vector out. You have 1 thread per position.\\n\
-//\\n\
-// Warning: You are in vim mode.\\n\
+"// Start editing here to see the results.\\n\// Warning: You are in vim mode.\\n\
 \\n\
-@group(0) @binding(0) var<storage, read_write> a: array<f32>;\\n\
+@group(0) @binding(0) var<storage, read_write> input: array<f32>;\\n\
 @group(0) @binding(1) var<storage, read_write> output : array<f32>;\\n\
 @compute @workgroup_size(256)\\n\
 fn main(\\n\
   @builtin(local_invocation_id) LocalInvocationID: vec3<u32>) {\\n\
     let local_idx = LocalInvocationID.x;\\n\
-    if (local_idx < arrayLength(&a)) {\\n\
-      output[local_idx] = a[local_idx] + 10;\\n\
+    if (local_idx < arrayLength(&input)) {\\n\
+      output[local_idx] = input[local_idx] + 1;\\n\
     }\\n\
   }\\n\
 ");
@@ -54,6 +51,24 @@
         if (delta.action === 'insert' && (delta.lines[0] === '.' || delta.lines[0] === ' ')) {
             showCompletionSuggestion();
         }
+
+        // Recover from errors TODO(avh): only do this if there's an error
+        createModule().then((Module) => {
+            // Keep your existing Module setup
+            Module.print = window.customPrint;
+            Module.printErr = window.customPrint;
+            window.Module = Module;
+            console.log("Module ready");
+        });
+
+        if (window.Module && window.Module.executeKernel) {
+            console.log("Executing kernel");
+            window.terminal.clear();
+            window.Module.executeKernel(editor.getValue());
+        } else {
+            console.log("Module not ready");
+        }
+
     });
 
     completionTippy = tippy(document.getElementById('editor'), {

Diff for: experimental/fasthtml/custom_shell.html

@@ -1,62 +1,72 @@
+#include "gpu.h"
 #include <array>
 #include <cstdio>
+#include <emscripten/emscripten.h>
 #include <future>
 #include <memory>
-#include "gpu.h"
+#include <string>
 
-// #include <webgpu/webgpu.h>
-#include "emscripten/emscripten.h"
-
-using namespace gpu; // createContext, createTensor, createKernel,
-                     // createShader, dispatchKernel, wait, toCPU
-                     // Tensor, Kernel, Context, Shape, kf32
-
-static const char *kGelu = R"(
-const GELU_SCALING_FACTOR: f32 = 0.7978845608028654; // sqrt(2.0 / PI)
-@group(0) @binding(0) var<storage, read_write> inp: array<{{precision}}>;
-@group(0) @binding(1) var<storage, read_write> out: array<{{precision}}>;
-@group(0) @binding(1) var<storage, read_write> dummy: array<{{precision}}>;
-@compute @workgroup_size({{workgroupSize}})
-fn main(
-    @builtin(global_invocation_id) GlobalInvocationID: vec3<u32>) {
-    let i: u32 = GlobalInvocationID.x;
-    if (i < arrayLength(&inp)) {
-        let x: f32 = inp[i];
-        out[i] = select(0.5 * x * (1.0 + tanh(GELU_SCALING_FACTOR 
-                 * (x + .044715 * x * x * x))), x, x > 10.0);
-    }
-}
-)";
+using namespace gpu;
 
-int main(int argc, char **argv) {
-  printf("\033[2J\033[1;1H");
-  printf("\nHello gpu.cpp!\n");
-  printf("--------------\n\n");
+EM_JS(void, js_print, (const char *str), {
+  if (typeof window != 'undefined' && window.customPrint) {
+    window.customPrint(UTF8ToString(str));
+  } else {
+    console.log("window.customPrint is not defined.");
+    console.log(UTF8ToString(str));
+  }
+});
 
-  // const WGPUInstanceDescriptor descriptor = { };
-  // std::unique_ptr<WGPUInstanceDescriptor> descriptor = std::make_unique<WGPUInstanceDescriptor>();
+extern "C" {
 
-  // WGPUInstance instance = wgpuCreateInstance({});
+EMSCRIPTEN_KEEPALIVE
+void executeKernel(const char *kernelCode) {
   Context ctx = createContext({});
   static constexpr size_t N = 5000;
   std::array<float, N> inputArr, outputArr;
+
   for (int i = 0; i < N; ++i) {
-    inputArr[i] = static_cast<float>(i) / 10.0; // dummy input data
+    inputArr[i] = static_cast<float>(i);
   }
+
   Tensor input = createTensor(ctx, Shape{N}, kf32, inputArr.data());
   Tensor output = createTensor(ctx, Shape{N}, kf32);
+
   std::promise<void> promise;
   std::future<void> future = promise.get_future();
-  Kernel op = createKernel(ctx, {kGelu, 256, kf32},
-                           Bindings{input, output},
-                           {cdiv(N, 256), 1, 1});
-  dispatchKernel(ctx, op, promise);
-  wait(ctx, future);
+
+  try {
+    Kernel op = createKernel(ctx, {kernelCode, 256, kf32},
+                             Bindings{input, output}, {cdiv(N, 256), 1, 1});
+
+    dispatchKernel(ctx, op, promise);
+    wait(ctx, future);
+  } catch (const std::exception &e) {
+    js_print("Invalid kernel code.");
+    exit(1);
+  }
+
   toCPU(ctx, output, outputArr.data(), sizeof(outputArr));
+
+  char buffer[1024];
   for (int i = 0; i < 12; ++i) {
-    printf("  gelu(%.2f) = %.2f\n", inputArr[i], outputArr[i]);
+    snprintf(buffer, sizeof(buffer), "  kernel(%.1f) = %.4f", inputArr[i],
+             outputArr[i]);
+    js_print(buffer);
   }
-  printf("  ...\n\n");
-  printf("Computed %zu values of GELU(x)\n\n", N);
-  return 0;
+  snprintf(buffer, sizeof(buffer), "  ...");
+  js_print(buffer);
+  snprintf(buffer, sizeof(buffer), "Computed %zu values", N);
+  js_print(buffer);
+}
+}
+
+#ifndef STANDALONE_WASM
+#include "emscripten/bind.h"
+EMSCRIPTEN_BINDINGS(module) {
+  emscripten::function("executeKernel", emscripten::optional_override(
+                                            [](const std::string &kernelCode) {
+                                              executeKernel(kernelCode.c_str());
+                                            }));
 }
+#endif

Diff for: experimental/fasthtml/run.cpp

@@ -9,6 +9,12 @@
 ace_editor = Script(src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.12/ace.js")
 gpucpp_runtime = Script(src="/build/run.js")    
 gpucpp_wasm = Script(src="/build/run.wasm")    
+tippy_css = Link(rel="stylesheet", href="https://unpkg.com/tippy.js@6/dist/tippy.css")
+tippy_js = Script(src="https://unpkg.com/@popperjs/core@2")
+tippy_js2 = Script(src="https://unpkg.com/tippy.js@6")
+xterm_css = Link(rel="stylesheet", href="https://cdn.jsdelivr.net/npm/xterm/css/xterm.css")
+xterm_js = Script(src="https://cdn.jsdelivr.net/npm/xterm/lib/xterm.js")
+xterm_fit_js = Script(src="https://cdn.jsdelivr.net/npm/xterm-addon-fit/lib/xterm-addon-fit.js")
 
 global_style = Style("""
 #editor {
@@ -17,11 +23,53 @@
 }
 """)
 
+terminal_init = Script("""
+    console.log("Terminal initialized");
+    const terminal = new Terminal();
+    const fitAddon = new FitAddon.FitAddon();
+    terminal.loadAddon(fitAddon);
+    // terminal.open(document.getElementById('output'));
+    window.terminal = terminal;
+    fitAddon.fit();
+    console.log("Terminal initialized");
+""");
+
+print_script = Script("""
+window.customPrint = function(text) {
+  console.log(text);
+  if (window.terminal) {
+    window.terminal.writeln(text);
+  } else {
+    console.warn("Terminal not initialized");
+  }
+};
+createModule().then((Module) => {
+  Module.print = window.customPrint;
+  Module.printErr = window.customPrint;
+  window.Module = Module;
+  console.log("Module ready");
+});
+"""),
+
+bind_terminal = Script("window.terminal.open(document.getElementById('output'));")
+
+# TODO(avh): Global state handling of terminal binding, module creation, etc.
+# could be improved
+
 HDRS = (
         picolink,
     ace_editor,
+    xterm_css,
+    xterm_js,
+    xterm_fit_js,
+    terminal_init,
     gpucpp_runtime,
+    print_script,
+    bind_terminal,
     global_style,
+    tippy_css,
+    tippy_js,
+    tippy_js2,
     *Socials(
         title="gpu.cpp gpu puzzles",
         description="",
@@ -39,9 +87,9 @@
 
 rt = app.route
 
-@app.get("/build/{fname:path}.{ext:static}")
-async def build(fname: str, ext: str):
-    return FileResponse(f"build/{fname}.{ext}")
+@app.get("/build/run.js")
+async def serve_wasm(fname: str, ext: str):
+    return FileResponse(f"build/run.js")
 
 @app.get("/build/run.wasm")
 async def serve_wasm(fname: str, ext: str):
@@ -56,10 +104,12 @@ def page():
             ),
             Div(
                 "Output",
+                id="output",
                 style="width: 34vw; height:100vh; background-color: #444; float: right;",
             ),
         ),
-        style="height: 100vh; overflow: hidden;")
+        bind_terminal,
+        style="height: 100vh; overflow: hidden;"),
 
 
 @rt("/")