Update (base update)

JCNTH · JCNTH · commit 17219f84620a · 2026-07-14T16:29:34.000-07:00
[ghstack-poisoned]
diff --git a/backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp b/backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp
@@ -9,7 +9,7 @@
 #include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
 #include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
 #include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
-#include <executorch/backends/webgpu/runtime/ops/log_softmax/log_softmax_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/softmax/log_softmax_wgsl.h>
 
 #include <webgpu/webgpu.h>
 
diff --git a/backends/webgpu/runtime/ops/log_softmax/log_softmax.wgsl b/backends/webgpu/runtime/ops/log_softmax/log_softmax.wgsl
diff --git a/backends/webgpu/runtime/ops/softmax/log_softmax_wgsl.h b/backends/webgpu/runtime/ops/softmax/log_softmax_wgsl.h
@@ -12,8 +12,8 @@
 
 namespace executorch::backends::webgpu {
 
-// @generated from log_softmax.wgsl - DO NOT EDIT.
-// wgsl-sha256: 0f6d1e9eea8d0ec8246c2c4f0b0be5c03deb4a73790896b8a0b093fc5be16114
+// @generated from softmax.wgsl - DO NOT EDIT.
+// wgsl-sha256: 3e5550d3bf913a4e9e871b1474934278f81b4d635504aceb8ab4e096ef18889a
 inline constexpr const char* kLogSoftmaxWGSL = R"(
 struct Params {
   outer_: u32,
@@ -38,7 +38,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   let oo = t / params.inner_;
   let ii = t % params.inner_;
   let base = oo * params.r_ * params.inner_ + ii;
-  // Online flash-style softmax; -3.4e38 seeds max (Tint overflows -FLT_MAX).
+  // Online (flash-style) softmax; -3.4e38 seeds max (Tint overflows -FLT_MAX).
   var mx: f32 = -3.4e38;
   var ssum: f32 = 0.0;
   for (var r: u32 = 0u; r < params.r_; r = r + 1u) {
diff --git a/backends/webgpu/runtime/ops/softmax/softmax.wgsl b/backends/webgpu/runtime/ops/softmax/softmax.wgsl
@@ -32,8 +32,13 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     }
     ssum = ssum + exp(v - mx);
   }
+  $if LOG:
+    let log_sum = log(ssum);
   for (var r: u32 = 0u; r < params.r_; r = r + 1u) {
     let idx = base + r * params.inner_;
-    out[idx] = exp(inp[idx] - mx) / ssum;
+    $if LOG:
+      out[idx] = (inp[idx] - mx) - log_sum;
+    $else:
+      out[idx] = exp(inp[idx] - mx) / ssum;
   }
 }
diff --git a/backends/webgpu/runtime/ops/softmax/softmax.yaml b/backends/webgpu/runtime/ops/softmax/softmax.yaml
@@ -0,0 +1,8 @@
+softmax:
+  parameter_names_with_default_values:
+    LOG: 0
+  shader_variants:
+    - NAME: softmax
+      LOG: 0
+    - NAME: log_softmax
+      LOG: 1

Original file line number	Diff line number	Diff line change
`@@ -32,8 +32,13 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {`
`32`	`32`	`}`
`33`	`33`	`ssum = ssum + exp(v - mx);`
`34`	`34`	`}`
	`35`	`+ $if LOG:`
	`36`	`+ let log_sum = log(ssum);`
`35`	`37`	`for (var r: u32 = 0u; r < params.r_; r = r + 1u) {`
`36`	`38`	`let idx = base + r * params.inner_;`
`37`		`- out[idx] = exp(inp[idx] - mx) / ssum;`
	`39`	`+ $if LOG:`
	`40`	`+ out[idx] = (inp[idx] - mx) - log_sum;`
	`41`	`+ $else:`
	`42`	`+ out[idx] = exp(inp[idx] - mx) / ssum;`
`38`	`43`	`}`
`39`	`44`	`}`