[tvm] branch unity updated: [Unity][WebGPU] Try F16 support for WebGPU Backend (#14904)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/unity by this push:
     new 7d2579e62e [Unity][WebGPU] Try F16 support for  WebGPU Backend (#14904)
7d2579e62e is described below

commit 7d2579e62e80c157167bb377e5cd58f8c3cb48a8
Author: jzm-intel <zhaoming.ji...@intel.com>
AuthorDate: Tue May 23 20:00:35 2023 +0800

    [Unity][WebGPU] Try F16 support for  WebGPU Backend (#14904)
---
 src/target/source/codegen_webgpu.cc | 20 +++++++++++++++-----
 src/target/source/codegen_webgpu.h  |  8 ++++++++
 web/src/webgpu.ts                   |  9 ++++++++-
 3 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/src/target/source/codegen_webgpu.cc 
b/src/target/source/codegen_webgpu.cc
index e933d58e17..538acb0a35 100644
--- a/src/target/source/codegen_webgpu.cc
+++ b/src/target/source/codegen_webgpu.cc
@@ -99,7 +99,11 @@ class WebGPUWorkgroupInfoCollector : public StmtExprVisitor {
 };
 
 std::string CodeGenWebGPU::Finish() {
-  return decl_stream.str() + this->fwd_decl_stream.str() + stream.str();
+  // Using f16 requires enable directive
+  if (enable_fp16_) {
+    header_stream << "enable f16;\n\n";
+  }
+  return header_stream.str() + decl_stream.str() + this->fwd_decl_stream.str() 
+ stream.str();
 }
 
 void CodeGenWebGPU::InitFuncState(const PrimFunc& f) {
@@ -134,9 +138,9 @@ runtime::FunctionInfo CodeGenWebGPU::AddFunction(const 
PrimFunc& f, bool skip_re
   ICHECK(global_symbol.defined())
       << "CodeGenWebGPU: Expect PrimFunc to have the global_symbol attribute";
 
-  decl_stream << "//----------------------------------------\n"
-              << "// Function: " << global_symbol.value() << "\n"
-              << "//----------------------------------------\n";
+  header_stream << "//----------------------------------------\n"
+                << "// Function: " << global_symbol.value() << "\n"
+                << "//----------------------------------------\n";
   runtime::FunctionInfo func_info;
   func_info.name = global_symbol.value();
 
@@ -299,6 +303,10 @@ void CodeGenWebGPU::PrintType(DataType t, std::ostream& 
os) {  // NOLINT(*)
 
   if (t.is_float()) {
     ICHECK(t.bits() == 16 || t.bits() == 32) << "CodeGenWebGPU: only support 
f16 or f32";
+    if (t.bits() == 16) {
+      // Using f16 requires enable directive
+      enable_fp16_ = true;
+    }
     os << "f" << t.bits();
   } else if (t.is_uint()) {
     ICHECK(t.bits() != 64) << "CodeGenWebGPU: do not support u64";
@@ -434,6 +442,8 @@ void CodeGenWebGPU::VisitExpr_(const FloatImmNode* op, 
std::ostream& os) {  // N
   if (op->dtype.bits() == 32) {
     temp << 'f';
   } else if (op->dtype.bits() == 16) {
+    // Using f16 requires enable directive
+    enable_fp16_ = true;
     temp << 'h';
   } else {
     LOG(FATAL) << "Unsupported floating point bits " << op->dtype.bits();
@@ -668,7 +678,7 @@ class WebGPUSourceModuleNode final : public 
runtime::ModuleNode {
   }
 
  private:
-  // function information table.
+  // function shader code table.
   std::unordered_map<std::string, std::string> smap_;
   // function information table.
   std::unordered_map<std::string, runtime::FunctionInfo> fmap_;
diff --git a/src/target/source/codegen_webgpu.h 
b/src/target/source/codegen_webgpu.h
index ff99f4608a..f12cd3430d 100644
--- a/src/target/source/codegen_webgpu.h
+++ b/src/target/source/codegen_webgpu.h
@@ -83,6 +83,14 @@ class CodeGenWebGPU final : public CodeGenC {
    * \brief Storage type of bool values.
    */
   DataType boolean_storage_type_{DataType::Int(8)};
+
+  // whether enable fp16
+  bool enable_fp16_{false};
+
+  /*! \brief the header stream for function label and enable directive if any, 
goes before any other
+   * declaration */
+  std::ostringstream header_stream;
+
   Target target_;
 };
 }  // namespace codegen
diff --git a/web/src/webgpu.ts b/web/src/webgpu.ts
index 0459e32415..0d7de70655 100644
--- a/web/src/webgpu.ts
+++ b/web/src/webgpu.ts
@@ -74,13 +74,20 @@ export async function detectGPUDevice(): 
Promise<GPUDeviceDetectOutput | undefin
       );
     }
 
+    let requiredFeatures = [] as string[];
+    // Always require f16 if available
+    if (adapter.features.has("shader-f16")) {
+      requiredFeatures.push("shader-f16");
+    }
+
     const adapterInfo = await adapter.requestAdapterInfo();
     const device = await adapter.requestDevice({
       requiredLimits: {
         maxBufferSize: requiedMaxBufferSize,
         maxStorageBufferBindingSize: requiredMaxStorageBufferBindingSize,
         maxComputeWorkgroupStorageSize: requiredMaxComputeWorkgroupStorageSize,
-      }
+      },
+      requiredFeatures
     });
     return {
       adapter: adapter,