From 69952ca29eb51aaa5263ceefeee2de91cdcd65e0 Mon Sep 17 00:00:00 2001
From: Akaash Parthasarathy <akaashrp@gmail.com>
Date: Sun, 9 Aug 2026 06:56:49 -0400
Subject: [PATCH] [FIX][WEBGPU] Preserve read-only buffer access modes

---
 src/backend/webgpu/codegen/codegen_webgpu.cc  | 29 ++++++++++++-
 .../codegen/test_target_codegen_webgpu.py     | 42 +++++++++++++++++++
 2 files changed, 70 insertions(+), 1 deletion(-)
 create mode 100644 tests/python/codegen/test_target_codegen_webgpu.py

diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc b/src/backend/webgpu/codegen/codegen_webgpu.cc
index 933a81f826fd..fe67febc4577 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 000000000000..96bfc6b90d85
--- /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()