refine state handling in fasthtml experiment, begin work on shadertui+fasthtml (WIP)

Author: austinvhuang

Diff for: experimental/fasthtml/Makefile

@@ -11,10 +11,6 @@ NO_MODULARIZE_FLAGS=-s EXPORTED_FUNCTIONS=['_executeKernel'] -s EXPORTED_RUNTIME
 
 default: server
 
-# 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 \
 		$(COMMON_FLAGS) $(JS_FLAGS) $(MODULARIZE_FLAGS)
@@ -23,6 +19,10 @@ build/run.wasm: check-emsdk run.cpp
 	em++ run.cpp -o build/run.wasm \
 		$(COMMON_FLAGS) $(WASM_FLAGS)
 
+watch:
+	@echo "Watching for changes..."
+	ls run.cpp | entr -c make build/run.js
+
 server: build/run.js
 	python3 run.py
 

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

@@ -2,46 +2,33 @@
 
 from fasthtml.common import *
 from .toolbar import Toolbar
+import json
 
-editor_script = Script("""
+def editor_script(initial_content: str) -> Script:
+    return Script("""
 let editor;
 let completionTippy;
 let currentCompletion = '';
 
-
 function initEditor() {
     editor = ace.edit("editor");
     editor.setTheme("ace/theme/monokai");
     editor.session.setMode("ace/mode/javascript");
     editor.setOptions({
         fontSize: "14px",
         showPrintMargin: false,
-        // disable showing errors in gutter Ace WGSL parser is out of date
+        // disable showing errors in gutter, Ace's WGSL parser is out of date
         showGutter: false,
         highlightActiveLine: true,
         wrap: true,
     });
     editor.setKeyboardHandler("ace/keyboard/vim");
     
-    editor.setValue(
-"// Start editing here to see the results.\\n\// Warning: You are in vim mode.\\n\
-\\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(&input)) {\\n\
-      output[local_idx] = input[local_idx] + 1;\\n\
-    }\\n\
-  }\\n\
-");
-    
+    editor.setValue(""" + json.dumps(initial_content) + ");" + 
+"""
     window.addEventListener('resize', function() {
         editor.resize();
     });
-
     document.getElementById('language').addEventListener('change', function(e) {
         let mode = "ace/mode/" + e.target.value;
         editor.session.setMode(mode);
@@ -149,7 +136,7 @@
 document.addEventListener('DOMContentLoaded', initEditor);
 """)
 
-def CodeEditor():
+def CodeEditor(initial_content: str):
     return (
         Div(
             Toolbar(),
@@ -167,5 +154,5 @@ def CodeEditor():
             # cls="flex flex-col h-screen w-full", style="height: 100vh; overflow: hidden;"
             style="height: 100vh; overflow: hidden;"
         ),
-        editor_script
+        editor_script(initial_content)
     )

Diff for: experimental/fasthtml/run.cpp

@@ -49,17 +49,23 @@ void executeKernel(const char *kernelCode) {
   toCPU(ctx, output, outputArr.data(), sizeof(outputArr));
 
   char buffer[1024];
-  for (int i = 0; i < 12; ++i) {
-    snprintf(buffer, sizeof(buffer), "  kernel(%.1f) = %.4f", inputArr[i],
-             outputArr[i]);
+  for (int i = 0; i < 10; ++i) {
+    snprintf(buffer, sizeof(buffer), "  [%d] kernel(%.1f) = %.4f", i,
+             inputArr[i], outputArr[i]);
     js_print(buffer);
   }
   snprintf(buffer, sizeof(buffer), "  ...");
   js_print(buffer);
+  for (int i = N - 10; i < N; ++i) {
+    snprintf(buffer, sizeof(buffer), "  [%d] kernel(%.1f) = %.4f", i,
+             inputArr[i], outputArr[i]);
+    js_print(buffer);
+  }
   snprintf(buffer, sizeof(buffer), "Computed %zu values", N);
   js_print(buffer);
-}
-}
+} // executeKernel
+
+} // extern "C"
 
 #ifndef STANDALONE_WASM
 #include "emscripten/bind.h"

Diff for: experimental/fasthtml/run.py

@@ -7,34 +7,42 @@
 TARGET = os.getenv("TARGET", "debug")
 
 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")    
+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_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")
+xterm_fit_js = Script(
+    src="https://cdn.jsdelivr.net/npm/xterm-addon-fit/lib/xterm-addon-fit.js"
+)
 
-global_style = Style("""
+global_style = Style(
+    """
 #editor {
     height: 50vh;
     width: 50vw;
 }
-""")
+"""
+)
 
-terminal_init = Script("""
-    console.log("Terminal initialized");
+terminal_init = \
+    """
     const terminal = new Terminal();
     const fitAddon = new FitAddon.FitAddon();
     terminal.loadAddon(fitAddon);
     // terminal.open(document.getElementById('output'));
     window.terminal = terminal;
-    fitAddon.fit();
+    // fitAddon.fit();
     console.log("Terminal initialized");
-""");
+"""
 
-print_script = Script("""
+print_script = (
+    Script(
+        """
 window.customPrint = function(text) {
   console.log(text);
   if (window.terminal) {
@@ -48,24 +56,43 @@
   Module.printErr = window.customPrint;
   window.Module = Module;
   console.log("Module ready");
+  // window.Module.executeKernel(editor.getValue());
 });
-"""),
-
-bind_terminal = Script("window.terminal.open(document.getElementById('output'));")
+"""
+    ),
+)
 
-# TODO(avh): Global state handling of terminal binding, module creation, etc.
+bind_terminal = """
+    window.terminal.open(document.getElementById('output'));
+    fitAddon.fit();
+"""
+
+gelu_kernel = """// Start editing here to see the results.
+// Warning: You are in vim mode.
+@group(0) @binding(0) var<storage, read_write> input : array<f32>;
+@group(0) @binding(1) var<storage, read_write> output : array<f32>;
+@compute @workgroup_size(256)
+fn main(
+    @builtin(global_invocation_id) GlobalInvocationID: vec3<u32>) {
+        let i: u32 = GlobalInvocationID.x;
+        if (i < arrayLength(&input)) {
+            output[i] = input[i] + 1;
+        }
+    }
+  """
+
+# TODO(avh) : Global state handling of terminal binding, module creation, etc.
 # could be improved
 
 HDRS = (
-        picolink,
+    picolink,
     ace_editor,
     xterm_css,
     xterm_js,
     xterm_fit_js,
-    terminal_init,
+    Script(terminal_init),
     gpucpp_runtime,
     print_script,
-    bind_terminal,
     global_style,
     tippy_css,
     tippy_js,
@@ -78,7 +105,7 @@
         image="",
         url="https://gpucpp.answer.ai",
     ),
-    )
+)
 
 if TARGET == "release":
     app = FastHTML(hdrs=HDRS)
@@ -87,29 +114,36 @@
 
 rt = app.route
 
+
 @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):
     return FileResponse(f"build/run.wasm")
 
+
 def page():
-    return Title("GPU Puzzles"), Body(
-        Div(
+    return (
+        Title("GPU Puzzles"),
+        Body(
             Div(
-                CodeEditor(),
-                style="width: 66vw; height:100vh; background-color: #333; float: left;",
-            ),
-            Div(
-                "Output",
-                id="output",
-                style="width: 34vw; height:100vh; background-color: #444; float: right;",
+                Div(
+                    CodeEditor(initial_content=gelu_kernel),
+                    style="width: 66vw; height:100vh; background-color: #333; float: left;",
+                ),
+                Div(
+                    "Output",
+                    id="output",
+                    style="width: 34vw; height:100vh; background-color: #444; float: right;",
+                ),
             ),
+            Script(bind_terminal),
+            style="height: 100vh; overflow: hidden;",
         ),
-        bind_terminal,
-        style="height: 100vh; overflow: hidden;"),
+    )
 
 
 @rt("/")

Diff for: experimental/fasthtml/shadertui_web.cpp

@@ -0,0 +1,99 @@
+#include "gpu.h"
+#include <array>
+#include <cstdio>
+#include <emscripten/emscripten.h>
+#include <future>
+#include <chrono>
+#include <string>
+#include <algorithm>
+
+using namespace gpu;
+
+static constexpr size_t kRows = 64;
+static constexpr size_t kCols = 96;
+
+template <size_t rows, size_t cols>
+void rasterize(const std::array<float, rows * cols> &values,
+               std::array<char, rows *(cols + 1)> &raster) {
+  static const char intensity[] = " .`'^-+=*x17X$8#%@";
+  for (size_t i = 0; i < rows; ++i) {
+    for (size_t j = 0; j < cols; ++j) {
+      size_t index =
+          std::min(sizeof(intensity) - 2,
+                   std::max(0ul, static_cast<size_t>(values[i * cols + j] *
+                                                     (sizeof(intensity) - 2))));
+      raster[i * (cols + 1) + j] = intensity[index];
+    }
+    raster[i * (cols + 1) + cols] = '\n';
+  }
+}
+
+float getCurrentTimeInMilliseconds(
+    std::chrono::time_point<std::chrono::high_resolution_clock> &zeroTime) {
+  std::chrono::duration<float> duration =
+      std::chrono::high_resolution_clock::now() - zeroTime;
+  return duration.count() * 1000.0f;  // Convert to milliseconds
+}
+
+EM_JS(void, js_print, (const char *str), {
+  if (typeof window != 'undefined' && window.customPrint) {
+    window.customPrint(UTF8ToString(str));
+  } else {
+    console.log(UTF8ToString(str));
+  }
+});
+
+extern "C" {
+EMSCRIPTEN_KEEPALIVE
+void executeShader(const char *shaderCode) {
+  Context ctx = createContext();
+  
+  std::array<float, kRows * kCols> screenArr;
+  Tensor screen = createTensor(ctx, {kRows, kCols}, kf32, screenArr.data());
+
+  struct Params {
+    float time;
+    uint32_t screenWidth;
+    uint32_t screenHeight;
+  } params = {0.0f, kCols, kRows};
+
+  KernelCode shader{shaderCode, Shape{16, 16, 1}};
+  Kernel renderKernel =
+      createKernel(ctx, shader, Bindings{screen},
+                   cdiv({kCols, kRows, 1}, shader.workgroupSize), params);
+
+  std::array<char, kRows *(kCols + 1)> raster;
+
+  auto start = std::chrono::high_resolution_clock::now();
+  
+  // Render a few frames
+  for (int frame = 0; frame < 10; ++frame) {
+    params.time = getCurrentTimeInMilliseconds(start);
+    toGPU(ctx, params, renderKernel);
+    
+    std::promise<void> promise;
+    std::future<void> future = promise.get_future();
+    dispatchKernel(ctx, renderKernel, promise);
+    wait(ctx, future);
+    
+    toCPU(ctx, screen, screenArr.data());
+    rasterize<kRows, kCols>(screenArr, raster);
+    
+    // Print the rasterized frame
+    js_print(raster.data());
+    
+    // Add a small delay between frames
+    emscripten_sleep(100);
+  }
+}
+}
+
+#ifndef STANDALONE_WASM
+#include <emscripten/bind.h>
+EMSCRIPTEN_BINDINGS(module) {
+  emscripten::function("executeShader", emscripten::optional_override(
+                                            [](const std::string &shaderCode) {
+                                              executeShader(shaderCode.c_str());
+                                            }));
+}
+#endif