1D workgroup specification overload for KernelCode constructor accepting totalWorkgroups, add webprint.h to gpu puzzles project

Author: austinvhuang

experimental/fasthtml/gpu_puzzles/webprint.h

@@ -0,0 +1,56 @@
+#ifndef WEBPRINT_H
+#define WEBPRINT_H
+
+
+#include <emscripten/emscripten.h>
+
+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));
+  }
+});
+
+
+// need to allow printf with variable arguments
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wformat-security"
+template<typename... Args>
+void wprintf(const char *str, Args... args) {
+  char buffer[1024];
+  snprintf(buffer, sizeof(buffer), str, args...);
+  js_print(buffer);
+}
+#pragma clang diagnostic pop
+
+void printVec(const std::vector<float> &vec, const char *name = "") {
+    char buffer[1024];
+    size_t pos = 0;
+    pos += snprintf(buffer + pos, sizeof(buffer) - pos, "[ ");
+    for (size_t i = 0; i < vec.size(); ++i) {
+        pos += snprintf(buffer + pos, sizeof(buffer) - pos, "%.1f", vec[i]);
+        if (i != vec.size() - 1) {
+            pos += snprintf(buffer + pos, sizeof(buffer) - pos, ", ");
+        }
+    }
+    snprintf(buffer + pos, sizeof(buffer) - pos, " ]");
+    wprintf("%s %s", name, buffer);
+}
+
+
+void printVecBuf(const std::vector<float> &vec, const char *name, char *buffer, size_t& pos) {
+    pos += snprintf(buffer + pos, sizeof(buffer) - pos, "%s", name);
+    pos += snprintf(buffer + pos, sizeof(buffer) - pos, "[ ");
+    for (size_t i = 0; i < vec.size(); ++i) {
+        pos += snprintf(buffer + pos, sizeof(buffer) - pos, "%2.0f", vec[i]);
+        if (i != vec.size() - 1) {
+            pos += snprintf(buffer + pos, sizeof(buffer) - pos, ", ");
+        }
+    }
+    pos += snprintf(buffer + pos, sizeof(buffer) - pos, " ]\n\r");
+}
+
+
+#endif // WEBPRINT_H

gpu.h

@@ -307,38 +307,59 @@ struct KernelCode {
   }
 
   /**
-   * @brief Overload of the constructor to create a code object from a
-   * template string and workgroup size. Unlike the main factory function,
-   * this overload takes a single size_t workgroupSize parameter instead of a
-   * 3D shape for the workgroup size and instantiates a 3D shape with the
-   * workgroupSize in the x dimension and 1 in the y and z dimensions.
+   * @brief Overload of the constructor to create a code object from a template
+   * string and workgroup size. This overload takes a single size_t
+   * workgroupSize parameter instead of a 3D shape for the workgroup size and
+   * instantiates a 3D shape with the workgroupSize in the x dimension and 1 in
+   * the y and z dimensions.
+   *
+   * @param[in] pData Shader template string with placeholders @param[in]
+   * workgroupSize 3D Workgroup size 
+   * @param[in] precision Data type precision for the shader
+   *
+   * @code KernelCode code = {kPuzzle1, 256, kf32}; @endcode
+   */
+  inline KernelCode(const std::string &pData, const Shape &workgroupSize =
+      {256, 1, 1}, NumType precision = kf32) : data(pData),
+  workgroupSize(workgroupSize), precision(precision) { if (precision == kf16) {
+    data = "enable f16;\n" + data; } replaceAll(data, "{{workgroupSize}}",
+        toString(workgroupSize)); replaceAll(data, "{{precision}}",
+        toString(precision)); LOG(kDefLog, kInfo, "Shader code:\n%s",
+        data.c_str()); }
+
+
+  /**
+   * @brief Overload of the constructor, adding totalWorkgroups parameter to
+   * perform a string replacement for the total number of workgroups in the
+   * kernel code.
    *
    * @param[in] pData Shader template string with placeholders
-   * @param[in] workgroupSize Workgroup size in the x dimension
+   * @param[in] workgroupSize 3D Workgroup size
    * @param[in] precision Data type precision for the shader
+   * @param[in] totalWorkgroups Total number of workgroups in the kernel
    *
    * @code
-   * KernelCode code = {kPuzzle1, 256, kf32};
+   * KernelCode code = {kPuzzle1, {256, 1, 1}, kf32, {2, 2, 1}};
    * @endcode
    */
-
   inline KernelCode(const std::string &pData,
-                    const Shape &workgroupSize = {256, 1, 1},
-                    NumType precision = kf32)
+                    const Shape &workgroupSize,
+                    NumType precision,
+                    const Shape &totalWorkgroups)
       : data(pData), workgroupSize(workgroupSize), precision(precision) {
     if (precision == kf16) {
       data = "enable f16;\n" + data;
     }
     replaceAll(data, "{{workgroupSize}}", toString(workgroupSize));
     replaceAll(data, "{{precision}}", toString(precision));
+    replaceAll(data, "{{totalWorkgroups}}", toString(totalWorkgroups));
     LOG(kDefLog, kInfo, "Shader code:\n%s", data.c_str());
   }
 
 
   /**
-   * @brief Overload of the constructor, adding totalWorkgroups parameter to
-   * perform a string replacement for the total number of workgroups in the
-   * kernel code.
+   * @brief Overload of the constructor, adding totalWorkgroups parameter as
+   * well as the size_t 1D workgroupSize parameter.
    *
    * @param[in] pData Shader template string with placeholders
    * @param[in] workgroupSize Workgroup size in the x dimension
@@ -350,14 +371,14 @@ struct KernelCode {
    * @endcode
    */
   inline KernelCode(const std::string &pData,
-                    const Shape &workgroupSize,
+                    const size_t &workgroupSize,
                     NumType precision,
                     const Shape &totalWorkgroups)
-      : data(pData), workgroupSize(workgroupSize), precision(precision) {
+      : data(pData), workgroupSize({workgroupSize, 1, 1}), precision(precision) {
     if (precision == kf16) {
       data = "enable f16;\n" + data;
     }
-    replaceAll(data, "{{workgroupSize}}", toString(workgroupSize));
+    replaceAll(data, "{{workgroupSize}}", toString({workgroupSize, 1, 1}));
     replaceAll(data, "{{precision}}", toString(precision));
     replaceAll(data, "{{totalWorkgroups}}", toString(totalWorkgroups));
     LOG(kDefLog, kInfo, "Shader code:\n%s", data.c_str());