chore: import upstream snapshot with attribution

web/emcc/webgpu_runtime.cc

@@ -0,0 +1,260 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*
+ * \file webgpu_runtime.cc
+ * \brief WebGPU runtime based on the TVM JS.
+ */
+
+// configurations for tvm logging
+#define TVM_LOG_DEBUG 0
+#define TVM_LOG_CUSTOMIZE 1
+#define TVM_FFI_ALWAYS_LOG_BEFORE_THROW 1
+
+#include <tvm/ffi/extra/json.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/runtime/device_api.h>
+
+#include <iostream>
+#include <string>
+
+#include "../../src/runtime/file_utils.h"
+#include "../../src/runtime/metadata.h"
+#include "../../src/runtime/workspace_pool.h"
+#include "../../src/support/bytes_io.h"
+
+namespace tvm {
+namespace runtime {
+
+/*! \brief Thread local workspace */
+class WebGPUThreadEntry {
+ public:
+  /*! \brief thread local pool*/
+  WorkspacePool pool;
+  /*! \brief constructor */
+  WebGPUThreadEntry();
+  // get the threadlocal workspace
+  static WebGPUThreadEntry* ThreadLocal();
+};
+
+// All the implementations are redirectly to the JS side.
+class WebGPUDeviceAPI : public DeviceAPI {
+ public:
+  WebGPUDeviceAPI() {
+    auto fp = tvm::ffi::Function::GetGlobal("wasm.WebGPUDeviceAPI");
+    TVM_FFI_ICHECK(fp.has_value()) << "Cannot find wasm.WebGPUContext in the env";
+    auto getter = ffi::TypedFunction<ffi::Function(std::string)>(*fp);
+    alloc_space_ = getter("deviceAllocDataSpace");
+    free_space_ = getter("deviceFreeDataSpace");
+    copy_to_gpu_ = getter("deviceCopyToGPU");
+    copy_from_gpu_ = getter("deviceCopyFromGPU");
+    copy_within_gpu_ = getter("deviceCopyWithinGPU");
+  }
+
+  void SetDevice(Device dev) final {}
+  void GetAttr(Device dev, DeviceAttrKind kind, ffi::Any* rv) final {
+    if (kind == kExist) {
+      *rv = 1;
+    }
+  }
+
+  void* AllocDataSpace(Device dev, size_t nbytes, size_t alignment, DLDataType type_hint) final {
+    double ptr_number = alloc_space_(nbytes);
+    return reinterpret_cast<void*>(static_cast<int64_t>(ptr_number));
+  }
+
+  void FreeDataSpace(Device dev, void* ptr) final { return free_space_(ptr); }
+
+ protected:
+  void CopyDataFromTo(const void* from, size_t from_offset, void* to, size_t to_offset, size_t size,
+                      Device dev_from, Device dev_to, DLDataType type_hint,
+                      TVMStreamHandle stream) final {
+    if (static_cast<int>(dev_from.device_type) == kDLWebGPU &&
+        static_cast<int>(dev_to.device_type) == kDLWebGPU) {
+      TVM_FFI_ICHECK_EQ(dev_from.device_id, dev_to.device_id);
+      copy_within_gpu_(const_cast<void*>(from), from_offset, to, to_offset, size);
+    } else if (static_cast<int>(dev_from.device_type) == kDLWebGPU &&
+               dev_to.device_type == kDLCPU) {
+      void* to_ptr = static_cast<uint8_t*>(to) + to_offset;
+      copy_from_gpu_(const_cast<void*>(from), from_offset, to_ptr, size);
+    } else if (dev_from.device_type == kDLCPU &&
+               static_cast<int>(dev_to.device_type) == kDLWebGPU) {
+      void* from_ptr = static_cast<uint8_t*>(const_cast<void*>(from)) + from_offset;
+      copy_to_gpu_(from_ptr, to, to_offset, size);
+    } else {
+      TVM_FFI_THROW(InternalError) << "expect copy from/to WebGPU or between WebGPU";
+    }
+  }
+
+ public:
+  TVMStreamHandle CreateStream(Device dev) final {
+    TVM_FFI_THROW(InternalError) << "Not implemented";
+  }
+
+  void FreeStream(Device dev, TVMStreamHandle stream) final {
+    TVM_FFI_THROW(InternalError) << "Not implemented";
+  }
+
+  void SyncStreamFromTo(Device dev, TVMStreamHandle event_src, TVMStreamHandle event_dst) {
+    TVM_FFI_THROW(InternalError) << "Not implemented";
+  }
+
+  void StreamSync(Device dev, TVMStreamHandle stream) final {
+    static auto func = tvm::ffi::Function::GetGlobal("__asyncify.WebGPUWaitForTasks");
+    TVM_FFI_ICHECK(func.has_value()) << "Stream sync inside c++ only supported in asyncify mode";
+    (*func)();
+  }
+
+  void* AllocWorkspace(Device dev, size_t size, DLDataType type_hint) final {
+    return WebGPUThreadEntry::ThreadLocal()->pool.AllocWorkspace(dev, size);
+  }
+
+  void FreeWorkspace(Device dev, void* data) final {
+    WebGPUThreadEntry::ThreadLocal()->pool.FreeWorkspace(dev, data);
+  }
+
+  static WebGPUDeviceAPI* Global() {
+    static WebGPUDeviceAPI* inst = new WebGPUDeviceAPI();
+    return inst;
+  }
+
+ private:
+  // NOTE: js return number as double.
+  ffi::TypedFunction<double(int64_t nbytes)> alloc_space_;
+  ffi::TypedFunction<void(void* ptr)> free_space_;
+  ffi::TypedFunction<void(void* from, void* to, int64_t to_offset, int64_t nbytes)> copy_to_gpu_;
+  ffi::TypedFunction<void(void* from, int64_t from_offset, void* to, int64_t nbytes)>
+      copy_from_gpu_;
+  ffi::TypedFunction<void(void* from, int64_t from_offset, void* to, int64_t to_offset,
+                          int64_t nbytes)>
+      copy_within_gpu_;
+};
+
+WebGPUThreadEntry::WebGPUThreadEntry()
+    : pool(static_cast<DLDeviceType>(kDLWebGPU), WebGPUDeviceAPI::Global()) {}
+
+WebGPUThreadEntry* WebGPUThreadEntry::ThreadLocal() {
+  static thread_local WebGPUThreadEntry inst;
+  return &inst;
+}
+
+class WebGPUModuleNode final : public ffi::ModuleObj {
+ public:
+  explicit WebGPUModuleNode(std::unordered_map<std::string, std::string> smap,
+                            ffi::Map<ffi::String, FunctionInfo> fmap)
+      : smap_(smap), fmap_(fmap) {
+    auto fp = tvm::ffi::Function::GetGlobal("wasm.WebGPUCreateShader");
+    TVM_FFI_ICHECK(fp.has_value());
+    create_shader_ = *fp;
+  }
+
+  const char* kind() const final { return "webgpu"; }
+
+  ffi::Optional<ffi::Function> GetFunction(const ffi::String& name) final {
+    // special function
+    if (name == "webgpu.get_fmap") {
+      return ffi::Function([this](ffi::PackedArgs args, ffi::Any* rv) {
+        namespace json = ::tvm::ffi::json;
+        json::Object obj;
+        for (const auto& kv : fmap_) {
+          obj.Set(kv.first, kv.second->SaveToJSON());
+        }
+        *rv = std::string(json::Stringify(obj));
+      });
+    } else if (name == "webgpu.get_shader") {
+      return ffi::Function([this](ffi::PackedArgs args, ffi::Any* rv) {
+        auto name = args[0].cast<std::string>();
+        auto it = smap_.find(name);
+        TVM_FFI_ICHECK(it != smap_.end()) << "Cannot find code " << name;
+        *rv = it->second;
+      });
+    } else if (name == "webgpu.update_prebuild") {
+      return ffi::Function([this](ffi::PackedArgs args, ffi::Any* rv) {
+        auto name = args[0].cast<std::string>();
+        ffi::Function func = args[1].cast<ffi::Function>();
+        prebuild_[name] = func;
+      });
+    }
+    // check prebuild cache
+    auto prebuild_it = prebuild_.find(name);
+    if (prebuild_it != prebuild_.end()) {
+      return prebuild_it->second;
+    }
+
+    auto it = smap_.find(name);
+    if (it != smap_.end()) {
+      auto opt_info = fmap_.Get(name);
+      TVM_FFI_ICHECK(opt_info.has_value());
+      FunctionInfo orig_info = opt_info.value();
+      FunctionInfo info(name, orig_info->arg_types, orig_info->launch_param_tags,
+                        orig_info->arg_extra_tags);
+      namespace json = ::tvm::ffi::json;
+      std::string info_json = std::string(json::Stringify(info->SaveToJSON()));
+      return create_shader_(info_json, it->second);
+    } else {
+      return std::nullopt;
+    }
+  }
+
+  int GetPropertyMask() const final { return ffi::Module::kBinarySerializable; };
+
+  ffi::Bytes SaveToBytes() const final { TVM_FFI_THROW(InternalError) << "Not implemented"; }
+
+  ffi::String InspectSource(const ffi::String& format) const final {
+    // can only return source code.
+    return source_;
+  }
+
+ private:
+  // code table
+  std::unordered_map<std::string, std::string> smap_;
+  // function information table.
+  ffi::Map<ffi::String, FunctionInfo> fmap_;
+  // The source
+  std::string source_;
+  // prebuild_ functions
+  std::unordered_map<std::string, ffi::Function> prebuild_;
+  // Callback to get the GPU function.
+  ffi::TypedFunction<ffi::Function(std::string finfo, std::string shader)> create_shader_;
+};
+
+ffi::Module WebGPUModuleLoadFromBytes(const ffi::Bytes& bytes) {
+  support::BytesInStream stream(bytes);
+  std::unordered_map<std::string, std::string> smap;
+
+  ffi::Map<ffi::String, FunctionInfo> fmap;
+  TVM_FFI_ICHECK(stream.Read(&fmap));
+  stream.Read(&smap);
+  return ffi::Module(ffi::make_object<WebGPUModuleNode>(smap, fmap));
+}
+
+// for now webgpu is hosted via a vulkan module.
+TVM_FFI_STATIC_INIT_BLOCK() {
+  namespace refl = tvm::ffi::reflection;
+  refl::GlobalDef()
+      .def("ffi.Module.load_from_bytes.webgpu", WebGPUModuleLoadFromBytes)
+      .def_packed("device_api.webgpu", [](ffi::PackedArgs args, ffi::Any* rv) {
+        DeviceAPI* ptr = WebGPUDeviceAPI::Global();
+        *rv = static_cast<void*>(ptr);
+      });
+}
+
+}  // namespace runtime
+}  // namespace tvm