This is an automated email from the ASF dual-hosted git repository.
guan404ming pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 189cef1249 [Fix][WebGPU] Preserve read-only buffer access modes
(#20113)
189cef1249 is described below
commit 189cef12497f780acd065a1b10d8423830b02614
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Tue Aug 11 07:45:12 2026 -0700
[Fix][WebGPU] Preserve read-only buffer access modes (#20113)
Preserve read-only WebGPU storage-buffer access modes after typed TIR
buffer lowering. Resolve `DeclBuffer` aliases to their underlying
pointer parameters and only emit `read_write` when the buffer is
actually written. This prevents read-only metadata buffers from being
misclassified and triggering WGSL uniform-control-flow validation errors
around `workgroupBarrier`.
---
src/backend/webgpu/codegen/codegen_webgpu.cc | 29 ++++++++++++++-
tests/python/codegen/test_target_codegen_webgpu.py | 42 ++++++++++++++++++++++
2 files changed, 70 insertions(+), 1 deletion(-)
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc
b/src/backend/webgpu/codegen/codegen_webgpu.cc
index 933a81f826..fe67febc45 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -31,7 +31,9 @@
#include <tvm/tirx/transform.h>
#include <algorithm>
+#include <optional>
#include <string>
+#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
@@ -67,6 +69,22 @@ class WebGPUWorkgroupInfoCollector : public StmtExprVisitor {
private:
using StmtExprVisitor::VisitExpr_;
+ static ffi::Optional<Var> GetBufferDataVar(const Expr& data) {
+ if (auto var = data.as<Var>()) {
+ return var;
+ }
+ if (const auto* call = data.as<CallNode>();
+ call && call->op.same_as(tirx::builtin::buffer_data()) &&
call->args.size() == 1) {
+ return call->args[0].as<Var>();
+ }
+ return std::nullopt;
+ }
+
+ Var ResolveBuffer(Var buffer_var) const {
+ auto it = buffer_aliases_.find(buffer_var.get());
+ return it == buffer_aliases_.end() ? buffer_var : it->second;
+ }
+
void VisitExpr_(const VarNode* op) final {
StmtExprVisitor::VisitExpr_(op);
Var buffer_var = ffi::GetRef<Var>(op);
@@ -77,7 +95,15 @@ class WebGPUWorkgroupInfoCollector : public StmtExprVisitor {
void VisitStmt_(const BufferStoreNode* op) final {
StmtExprVisitor::VisitStmt_(op);
- info_.write_access_set.insert(op->buffer.var());
+ info_.write_access_set.insert(ResolveBuffer(op->buffer.var()));
+ }
+
+ void VisitStmt_(const DeclBufferNode* op) final {
+ if (auto source = GetBufferDataVar(op->data)) {
+ buffer_aliases_.insert_or_assign(op->buffer.get(),
ResolveBuffer(source.value()));
+ return;
+ }
+ StmtExprVisitor::VisitStmt_(op);
}
void VisitStmt_(const AttrStmtNode* op) final {
@@ -104,6 +130,7 @@ class WebGPUWorkgroupInfoCollector : public StmtExprVisitor
{
StmtExprVisitor::VisitStmt_(op);
}
WebGPUWorkGroupInfo info_;
+ std::unordered_map<const VarNode*, Var> buffer_aliases_;
};
std::string CodeGenWebGPU::Finish() {
diff --git a/tests/python/codegen/test_target_codegen_webgpu.py
b/tests/python/codegen/test_target_codegen_webgpu.py
new file mode 100644
index 0000000000..96bfc6b90d
--- /dev/null
+++ b/tests/python/codegen/test_target_codegen_webgpu.py
@@ -0,0 +1,42 @@
+# 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.
+
+import tvm
+import tvm.testing
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+
+def test_codegen_buffer_access_modes():
+ """Read-only typed buffer parameters should remain read-only in WGSL."""
+
+ @I.ir_module(s_tir=True)
+ class Module:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer((8,), "float32"), B: T.Buffer((8,), "float32")):
+ for tx in T.thread_binding(8, thread="threadIdx.x"):
+ B[tx] = A[tx]
+
+ executable = tvm.compile(Module, target="webgpu")
+ source = executable.mod.imports[0].inspect_source("wgsl")
+
+ assert "var<storage, read> A_ptr" in source
+ assert "var<storage, read_write> B_ptr" in source
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
