From fb854bd2459a3960561aa0a6549b2ca35eefcdb1 Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Mon, 10 Apr 2023 10:19:18 -0500
Subject: [PATCH] [WebGPU][Bugfix] Replace kDeviceThreadAxis in CodeGenWebGPU

A follow-up from https://github.com/apache/tvm/pull/14495, applying
the same change in CodeGenWebGPU.
---
 src/target/source/codegen_webgpu.cc | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/target/source/codegen_webgpu.cc b/src/target/source/codegen_webgpu.cc
index 32b32063940b..8c34c6144fe3 100644
--- a/src/target/source/codegen_webgpu.cc
+++ b/src/target/source/codegen_webgpu.cc
@@ -185,10 +185,9 @@ runtime::FunctionInfo CodeGenWebGPU::AddFunction(const PrimFunc& f, bool skip_re
   }
 
   // setup thread tags and param access in launch param tags;
-  if (auto opt = f->GetAttr<Array<tir::IterVar>>(tir::attr::kDeviceThreadAxis)) {
-    auto thread_axis = opt.value();
-    for (size_t i = 0; i < thread_axis.size(); ++i) {
-      func_info.launch_param_tags.push_back(thread_axis[i]->thread_tag);
+  if (auto opt = f->GetAttr<Array<String>>(tir::attr::kKernelLaunchParams)) {
+    for (const auto& thread_tag : opt.value()) {
+      func_info.launch_param_tags.push_back(thread_tag);
     }
   }
   os_param_access << "]";