This is an automated email from the ASF dual-hosted git repository.
tqchen 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 94d8ed3434 [REFACTOR][Support] Expose compiler GPU thread-tag parsing
(#20689)
94d8ed3434 is described below
commit 94d8ed34347908811625e6ec4502e51aee24c270
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 12 19:45:42 2026 -0400
[REFACTOR][Support] Expose compiler GPU thread-tag parsing (#20689)
GPU thread tags describe compiler execution axes. Expose their parser in
a small public support header that consumes ffi::String directly, and
migrate scheduling, lowering, and backend code generation to it,
retiring runtime ThreadScope.
Preserve physical-tag parsing and existing consumer checks, keep
virtual-thread handling local to compiler consumers, and leave
launch-parameter decoding independent.
---
include/tvm/support/gpu_thread_tag.h | 67 ++++++++++++++++++++++
src/backend/cuda/codegen/codegen_cuda.cc | 10 ++--
src/backend/cuda/codegen/llvm/codegen_nvptx.cc | 3 +-
src/backend/metal/codegen/codegen_metal.cc | 6 +-
src/backend/opencl/codegen/codegen_opencl.cc | 4 +-
src/backend/rocm/codegen/llvm/codegen_amdgpu.cc | 3 +-
src/backend/vulkan/codegen/codegen_spirv.cc | 6 +-
src/backend/webgpu/codegen/codegen_webgpu.cc | 10 ++--
src/runtime/thread_storage_scope.h | 45 ---------------
.../analysis/sblock_tensor_access_lca_detector.cc | 10 ++--
.../postproc/rewrite_reduction_block.cc | 3 +-
.../postproc/rewrite_unbound_block.cc | 3 +-
.../meta_schedule/postproc/verify_gpu_code.cc | 3 +-
.../meta_schedule/schedule/cuda/thread_bind.cc | 5 +-
.../schedule_rule/cross_thread_reduction.cc | 7 ++-
src/s_tir/schedule/analysis/analysis.cc | 4 +-
src/s_tir/schedule/primitive/for_kind.cc | 9 +--
src/s_tir/schedule/utils.h | 21 +++++--
src/s_tir/transform/compact_tensor_region.cc | 3 +-
src/s_tir/transform/lift_thread_binding.cc | 11 +---
src/s_tir/transform/loop_partition.cc | 24 ++++----
.../transform/lower_cross_thread_reduction.cc | 40 ++++++-------
.../transform/memhammer_intermediate_stage.cc | 3 +-
src/tirx/transform/lower_thread_allreduce.h | 10 ++--
24 files changed, 170 insertions(+), 140 deletions(-)
diff --git a/include/tvm/support/gpu_thread_tag.h
b/include/tvm/support/gpu_thread_tag.h
new file mode 100644
index 0000000000..800c6cbb27
--- /dev/null
+++ b/include/tvm/support/gpu_thread_tag.h
@@ -0,0 +1,67 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/support/gpu_thread_tag.h
+ * \brief Parse GPU thread tags used by compiler passes and code generation.
+ */
+#ifndef TVM_SUPPORT_GPU_THREAD_TAG_H_
+#define TVM_SUPPORT_GPU_THREAD_TAG_H_
+
+#include <tvm/ffi/error.h>
+#include <tvm/ffi/string.h>
+
+namespace tvm {
+namespace support {
+
+/*! \brief The rank and dimension of a physical GPU thread tag. */
+struct GPUThreadTag {
+ /*! \brief Block, thread, cluster CTA, or preferred cluster CTA rank (0
through 3). */
+ int rank{0};
+ /*! \brief The dimension index, relative to x. */
+ int dim_index{0};
+
+ /*!
+ * \brief Parse a physical GPU thread tag.
+ *
+ * Recognizes blockIdx., threadIdx., clusterCtaIdx., and
preferredClusterCtaIdx.
+ * prefixes. The first character after the prefix determines the dimension;
+ * trailing characters are ignored. Consumers validate supported ranks and
dimensions.
+ * \param tag The thread tag to parse.
+ * \return The parsed rank and dimension.
+ */
+ static GPUThreadTag Parse(const ffi::String& tag) {
+ if (tag.starts_with("blockIdx.")) {
+ return {0, static_cast<int>(tag.c_str()[9] - 'x')};
+ } else if (tag.starts_with("threadIdx.")) {
+ return {1, static_cast<int>(tag.c_str()[10] - 'x')};
+ } else if (tag.starts_with("clusterCtaIdx.")) {
+ return {2, static_cast<int>(tag.c_str()[14] - 'x')};
+ } else if (tag.starts_with("preferredClusterCtaIdx.")) {
+ return {3, static_cast<int>(tag.c_str()[23] - 'x')};
+ } else {
+ TVM_FFI_THROW(InternalError) << "Unknown threadscope " << tag;
+ return {};
+ }
+ }
+};
+
+} // namespace support
+} // namespace tvm
+#endif // TVM_SUPPORT_GPU_THREAD_TAG_H_
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index acf8a51396..8a0961d094 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -28,6 +28,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/function.h>
#include <tvm/ir/prim/op.h>
+#include <tvm/support/gpu_thread_tag.h>
#include <tvm/sym/analyzer.h>
#include <tvm/tirx/index_map.h>
#include <tvm/tirx/op/gpu.h>
@@ -42,7 +43,6 @@
#include <utility>
#include <vector>
-#include "../../../runtime/thread_storage_scope.h"
#include "../../../target/build_common.h"
#include "cuda_fallback_module.h"
#include "literal/cuda_half_t.h"
@@ -207,9 +207,7 @@ class ThreadIdxExtractor : public tirx::StmtExprVisitor {
private:
ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op) final {
if (op->op.same_as(tirx::launch_thread_op())) {
- TVM_FFI_CHECK(
-
std::string(op->args[0].as_or_throw<StringImm>()->value).rfind("vthread", 0) !=
0,
- ValueError)
+
TVM_FFI_CHECK(!op->args[0].as_or_throw<StringImm>()->value.starts_with("vthread"),
ValueError)
<< "Virtual thread launches must be lowered before code generation";
ffi::String tag = op->args[0].as_or_throw<StringImm>()->value;
PrimExpr extent = op->args[1].as_or_throw<PrimExpr>();
@@ -462,10 +460,10 @@ void CodeGenCUDA::Dispatch_(const WhileNode* op) {
void CodeGenCUDA::BindThreadIndex(const PrimVar& var, const ffi::String&
thread_tag) {
TVM_FFI_ICHECK(!var_idmap_.count(var.get()));
- const auto& scope = runtime::ThreadScope::Create(thread_tag);
+ const auto& scope = support::GPUThreadTag::Parse(thread_tag);
TVM_FFI_ICHECK_GE(scope.dim_index, 0);
TVM_FFI_ICHECK_LT(scope.dim_index, 3);
- if (scope.IsClusterCtaIdx()) {
+ if (scope.rank == 2) {
const char dim = static_cast<char>('x' + scope.dim_index);
const std::string sreg = (scope.dim_index == 0 &&
cluster_cta_x_is_linear_rank_)
? "cluster_ctarank"
diff --git a/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
b/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
index f218ebeadf..502aa400f6 100644
--- a/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
+++ b/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
@@ -42,6 +42,7 @@
#include <llvm/Support/raw_ostream.h>
#include <llvm/Target/TargetMachine.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/support/gpu_thread_tag.h>
#include <tvm/tirx/op/gpu.h>
#include <tvm/tirx/op/memory.h>
#include <tvm/tirx/stmt.h>
@@ -149,7 +150,7 @@ class CodeGenNVPTX : public CodeGenLLVM {
// Return the thread index via intrinsics.
llvm::Value* GetThreadIndex(const PrimVar& var, const ffi::String&
thread_tag) final {
- runtime::ThreadScope ts = runtime::ThreadScope::Create(thread_tag);
+ support::GPUThreadTag ts = support::GPUThreadTag::Parse(thread_tag);
llvm::Intrinsic::ID intrin_id = llvm::Intrinsic::nvvm_read_ptx_sreg_tid_x;
if (ts.rank == 1) {
switch (ts.dim_index) {
diff --git a/src/backend/metal/codegen/codegen_metal.cc
b/src/backend/metal/codegen/codegen_metal.cc
index 4e33677960..a8115535d3 100644
--- a/src/backend/metal/codegen/codegen_metal.cc
+++ b/src/backend/metal/codegen/codegen_metal.cc
@@ -28,6 +28,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/function.h>
#include <tvm/runtime/logging.h>
+#include <tvm/support/gpu_thread_tag.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/op/memory.h>
@@ -42,7 +43,6 @@
#include <unordered_map>
#include <utility>
-#include "../../../runtime/thread_storage_scope.h"
#include "../../../target/build_common.h"
#include "metal_fallback_module.h"
@@ -189,7 +189,7 @@ void CodeGenMetal::AddFunction(const GlobalVar& gvar, const
Function& func) {
func->GetAttr<ffi::Array<ffi::String>>(tvm::tirx::attr::kKernelLaunchParams).value();
for (const auto& tag : launch_params) {
if (tag != runtime::launch_param::kUseDynamicSharedMemoryTag) {
- runtime::ThreadScope scope = runtime::ThreadScope::Create(tag);
+ support::GPUThreadTag scope = support::GPUThreadTag::Parse(tag);
work_dim = std::max(work_dim, scope.dim_index + 1);
}
}
@@ -216,7 +216,7 @@ void CodeGenMetal::AddFunction(const GlobalVar& gvar, const
Function& func) {
void CodeGenMetal::BindThreadIndex(const PrimVar& var, const ffi::String&
thread_tag) {
TVM_FFI_ICHECK(!var_idmap_.count(var.get()));
- runtime::ThreadScope ts = runtime::ThreadScope::Create(thread_tag);
+ support::GPUThreadTag ts = support::GPUThreadTag::Parse(thread_tag);
TVM_FFI_ICHECK(ts.rank == 0 || ts.rank == 1) << "Unsupported Metal thread
tag " << thread_tag;
TVM_FFI_ICHECK_GE(ts.dim_index, 0);
TVM_FFI_ICHECK_LT(ts.dim_index, 3);
diff --git a/src/backend/opencl/codegen/codegen_opencl.cc
b/src/backend/opencl/codegen/codegen_opencl.cc
index 7d040e106f..842008b1e3 100644
--- a/src/backend/opencl/codegen/codegen_opencl.cc
+++ b/src/backend/opencl/codegen/codegen_opencl.cc
@@ -26,13 +26,13 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/function.h>
+#include <tvm/support/gpu_thread_tag.h>
#include <tvm/tirx/op/memory.h>
#include <cmath>
#include <string>
#include <vector>
-#include "../../../runtime/thread_storage_scope.h"
#include "../../../target/build_common.h"
#include "../runtime/texture.h"
#include "opencl_fallback_module.h"
@@ -238,7 +238,7 @@ std::string CodeGenOpenCL::Finish() {
void CodeGenOpenCL::BindThreadIndex(const PrimVar& var, const ffi::String&
thread_tag) {
TVM_FFI_ICHECK(!var_idmap_.count(var.get()));
- runtime::ThreadScope ts = runtime::ThreadScope::Create(thread_tag);
+ support::GPUThreadTag ts = support::GPUThreadTag::Parse(thread_tag);
TVM_FFI_ICHECK_GE(ts.dim_index, 0);
TVM_FFI_ICHECK_LT(ts.dim_index, 3);
std::ostringstream os;
diff --git a/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
b/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
index 2a7fcc17b0..541bdaa3a3 100644
--- a/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
+++ b/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
@@ -39,6 +39,7 @@
#include <llvm/Support/raw_ostream.h>
#include <llvm/Target/TargetMachine.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/support/gpu_thread_tag.h>
#include <tvm/tirx/op/memory.h>
#include <tvm/tirx/stmt.h>
#if TVM_LLVM_VERSION < 170
@@ -162,7 +163,7 @@ class CodeGenAMDGPU : public CodeGenLLVM {
// Return the thread index via intrinsics.
llvm::Value* GetThreadIndex(const PrimVar& var, const ffi::String&
thread_tag) final {
- runtime::ThreadScope ts = runtime::ThreadScope::Create(thread_tag);
+ support::GPUThreadTag ts = support::GPUThreadTag::Parse(thread_tag);
llvm::Intrinsic::ID intrin_id = llvm::Intrinsic::amdgcn_workitem_id_x;
if (ts.rank == 1) {
switch (ts.dim_index) {
diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc
b/src/backend/vulkan/codegen/codegen_spirv.cc
index 70674d4292..6bdf2b8465 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.cc
+++ b/src/backend/vulkan/codegen/codegen_spirv.cc
@@ -25,6 +25,7 @@
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/prim/op.h>
+#include <tvm/support/gpu_thread_tag.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/op/abi.h>
#include <tvm/tirx/op/gpu.h>
@@ -206,7 +207,7 @@ void CodeGenSPIRV::InitFuncState() {
spirv::Value CodeGenSPIRV::GetThreadIndex(const PrimVar& var, const
ffi::String& thread_tag,
const PrimExpr& extent) {
- runtime::ThreadScope ts = runtime::ThreadScope::Create(thread_tag);
+ support::GPUThreadTag ts = support::GPUThreadTag::Parse(thread_tag);
TVM_FFI_ICHECK_GE(ts.dim_index, 0);
TVM_FFI_ICHECK_LT(ts.dim_index, 3);
spirv::Value v;
@@ -1021,8 +1022,7 @@ void CodeGenSPIRV::DispatchDeclTensor(const BindNode* op,
const CallNode* buffer
void CodeGenSPIRV::Dispatch_(const RegionStmtNode* op) {
if (op->op.same_as(tirx::launch_thread_op())) {
-
TVM_FFI_CHECK(std::string(op->args[0].as_or_throw<StringImm>()->value).rfind("vthread",
0) != 0,
- ValueError)
+
TVM_FFI_CHECK(!op->args[0].as_or_throw<StringImm>()->value.starts_with("vthread"),
ValueError)
<< "Virtual thread launches must be lowered before code generation";
PrimVar var = op->body_params[0].as_or_throw<PrimVar>();
PrimExpr extent = op->args[1].as_or_throw<PrimExpr>();
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc
b/src/backend/webgpu/codegen/codegen_webgpu.cc
index 37618e86cc..cabdb1e7aa 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -27,6 +27,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/function.h>
#include <tvm/ir/prim/op.h>
+#include <tvm/support/gpu_thread_tag.h>
#include <tvm/support/io.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/function.h>
@@ -139,16 +140,14 @@ class WebGPUWorkgroupInfoCollector : public
StmtExprVisitor {
ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op) final {
if (op->op.same_as(tirx::launch_thread_op())) {
- TVM_FFI_CHECK(
-
std::string(op->args[0].as_or_throw<StringImm>()->value).rfind("vthread", 0) !=
0,
- ValueError)
+
TVM_FFI_CHECK(!op->args[0].as_or_throw<StringImm>()->value.starts_with("vthread"),
ValueError)
<< "Virtual thread launches must be lowered before code generation";
ffi::String tag = op->args[0].as_or_throw<StringImm>()->value;
PrimExpr extent = op->args[1].as_or_throw<PrimExpr>();
auto [it, inserted] = thread_extents_.emplace(tag, extent);
TVM_FFI_CHECK(inserted || analyzer_->CanProveEqual(it->second, extent),
ValueError)
<< "Conflicting launch extents for " << tag;
- runtime::ThreadScope ts = runtime::ThreadScope::Create(tag);
+ support::GPUThreadTag ts = support::GPUThreadTag::Parse(tag);
TVM_FFI_ICHECK_GE(ts.dim_index, 0);
TVM_FFI_ICHECK_LT(ts.dim_index, 3);
if (ts.rank == 1) {
@@ -365,8 +364,7 @@ runtime::FunctionInfo CodeGenWebGPU::AddFunction(const
Function& f, bool skip_re
void CodeGenWebGPU::Dispatch_(const RegionStmtNode* op) {
if (op->op.same_as(tirx::launch_thread_op())) {
-
TVM_FFI_CHECK(std::string(op->args[0].as_or_throw<StringImm>()->value).rfind("vthread",
0) != 0,
- ValueError)
+
TVM_FFI_CHECK(!op->args[0].as_or_throw<StringImm>()->value.starts_with("vthread"),
ValueError)
<< "Virtual thread launches must be lowered before code generation";
PrimVar var = op->body_params[0].as_or_throw<PrimVar>();
PrimExpr extent = op->args[1].as_or_throw<PrimExpr>();
diff --git a/src/runtime/thread_storage_scope.h
b/src/runtime/thread_storage_scope.h
index 2db8f68d6f..04bf808da3 100644
--- a/src/runtime/thread_storage_scope.h
+++ b/src/runtime/thread_storage_scope.h
@@ -205,51 +205,6 @@ struct StorageScope {
}
};
-/*! \brief class to represent thread scope */
-struct ThreadScope {
- /*! \brief The rank of thread scope */
- int rank{0};
- /*! \brief the dimension index under the rank */
- int dim_index{0};
- /*!
- * \brief Create storage scope from string
- * \param s The string to be parsed.
- * \return The storage scope.
- */
- static ThreadScope Create(const std::string& s) {
- ThreadScope r;
- if (s.compare(0, 7, "vthread") == 0) {
- // virtual thread at the same level as local
- r.rank = 1;
- r.dim_index = -1;
- } else if (s.compare(0, 9, "blockIdx.") == 0) {
- r.rank = 0;
- r.dim_index = static_cast<int>(s[9] - 'x');
- } else if (s.compare(0, 10, "threadIdx.") == 0) {
- r.rank = 1;
- r.dim_index = static_cast<int>(s[10] - 'x');
- } else if (s.compare(0, 14, "clusterCtaIdx.") == 0) {
- r.rank = 2;
- r.dim_index = static_cast<int>(s[14] - 'x');
- } else if (s.compare(0, 23, "preferredClusterCtaIdx.") == 0) {
- r.rank = 3;
- r.dim_index = static_cast<int>(s[23] - 'x');
- } else {
- TVM_FFI_THROW(InternalError) << "Unknown threadscope " << s;
- }
- return r;
- }
-
- /*! \brief Whether the thread scope is a virtual thread */
- bool IsVirtualThread() const { return rank == 1 && dim_index == -1; }
- /*! \brief Whether the thread scope is a block */
- bool IsBlockIdx() const { return rank == 0; }
- /*! \brief Whether the thread scope is a thread */
- bool IsThreadIdx() const { return rank == 1 && dim_index != -1; }
- /*! \brief Whether the thread scope is a cluster */
- bool IsClusterCtaIdx() const { return rank == 2; }
-};
-
} // namespace runtime
} // namespace tvm
diff --git a/src/s_tir/analysis/sblock_tensor_access_lca_detector.cc
b/src/s_tir/analysis/sblock_tensor_access_lca_detector.cc
index bcd5c560c5..050e58e20d 100644
--- a/src/s_tir/analysis/sblock_tensor_access_lca_detector.cc
+++ b/src/s_tir/analysis/sblock_tensor_access_lca_detector.cc
@@ -29,6 +29,7 @@
#include <tvm/s_tir/analysis.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
+#include <tvm/support/gpu_thread_tag.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op/region.h>
@@ -102,9 +103,8 @@ class LCADetector : public s_tir::StmtExprVisitor {
auto* current_scope = arena_.make<ScopeInfo>(parent_scope, op, n);
if (tvm::s_tir::GetThreadBinding(op).has_value()) {
- const runtime::ThreadScope& scope =
-
runtime::ThreadScope::Create(tvm::s_tir::GetThreadBinding(op).value());
- if (scope.rank == 0) {
+ const ffi::String tag = tvm::s_tir::GetThreadBinding(op).value();
+ if (!tag.starts_with("vthread") &&
support::GPUThreadTag::Parse(tag).rank == 0) {
blockidx_scopes_.push_back(current_scope);
}
}
@@ -265,8 +265,8 @@ class LCADetector : public s_tir::StmtExprVisitor {
const ScopeInfo* parent_scope = ancestor_scopes_.back();
auto* current_scope = arena_.make<ScopeInfo>(parent_scope, op,
ancestor_scopes_.size());
if (op->op.same_as(tirx::launch_thread_op()) &&
-
std::string(op->args[0].as_or_throw<StringImm>()->value).rfind("vthread", 0) !=
0 &&
-
runtime::ThreadScope::Create(op->args[0].as_or_throw<StringImm>()->value).rank
== 0) {
+ !op->args[0].as_or_throw<StringImm>()->value.starts_with("vthread") &&
+
support::GPUThreadTag::Parse(op->args[0].as_or_throw<StringImm>()->value).rank
== 0) {
blockidx_scopes_.push_back(parent_scope);
}
ancestor_scopes_.push_back(current_scope);
diff --git a/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc
b/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc
index 48a92fbbff..7d7762cd5f 100644
--- a/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc
+++ b/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc
@@ -19,6 +19,7 @@
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/support/gpu_thread_tag.h>
#include "../utils.h"
@@ -54,7 +55,7 @@ struct ReductionBlockFinder : public StmtExprVisitor {
private:
ffi::Optional<VisitInterrupt> Visit_(const ForNode* loop) final {
- runtime::ThreadScope thread_scope = GetThreadScope(loop);
+ support::GPUThreadTag thread_scope = GetThreadBindingTag(loop);
if (IsThreadIdx(thread_scope) || IsBlockIdx(thread_scope)) {
thread_bound_loop_vars_.insert(loop->loop_var.get());
}
diff --git a/src/s_tir/meta_schedule/postproc/rewrite_unbound_block.cc
b/src/s_tir/meta_schedule/postproc/rewrite_unbound_block.cc
index 1955d10c41..92e7700fbb 100644
--- a/src/s_tir/meta_schedule/postproc/rewrite_unbound_block.cc
+++ b/src/s_tir/meta_schedule/postproc/rewrite_unbound_block.cc
@@ -19,6 +19,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/meta_schedule/schedule/cuda/thread_bind.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/support/gpu_thread_tag.h>
#include "../utils.h"
@@ -53,7 +54,7 @@ class UnboundBlockFinder : public StmtExprVisitor {
private:
ffi::Optional<VisitInterrupt> Visit_(const ForNode* loop) final {
- runtime::ThreadScope thread_scope = GetThreadScope(loop);
+ support::GPUThreadTag thread_scope = GetThreadBindingTag(loop);
if (IsBlockIdx(thread_scope)) {
++n_block_idx_;
} else if (IsThreadIdx(thread_scope)) {
diff --git a/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc
b/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc
index 98c3a83fb0..6413e8e258 100644
--- a/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc
+++ b/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc
@@ -21,6 +21,7 @@
#include <tvm/ir/function.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/transform.h>
+#include <tvm/support/gpu_thread_tag.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/transform.h>
@@ -54,7 +55,7 @@ class ThreadExtentChecker : public StmtExprVisitor {
private:
ffi::Optional<VisitInterrupt> Visit_(const ForNode* loop) {
- runtime::ThreadScope thread_scope = GetThreadScope(loop);
+ support::GPUThreadTag thread_scope = GetThreadBindingTag(loop);
if (IsThreadIdx(thread_scope)) {
const auto* p_ext_imm = loop->extent.as<IntImmNode>();
if (auto p_ext = p_ext_imm ? p_ext_imm->value.as<int64_t>() :
std::nullopt;
diff --git a/src/s_tir/meta_schedule/schedule/cuda/thread_bind.cc
b/src/s_tir/meta_schedule/schedule/cuda/thread_bind.cc
index e00709a4b3..5bec958121 100644
--- a/src/s_tir/meta_schedule/schedule/cuda/thread_bind.cc
+++ b/src/s_tir/meta_schedule/schedule/cuda/thread_bind.cc
@@ -19,6 +19,7 @@
#include <tvm/ir/prim/op.h>
#include <tvm/s_tir/meta_schedule/schedule/cuda/thread_bind.h>
#include <tvm/s_tir/schedule/schedule.h>
+#include <tvm/support/gpu_thread_tag.h>
#include <algorithm>
#include <limits>
@@ -35,7 +36,7 @@ using namespace tvm::tirx;
using s_tir::ExprRV;
using s_tir::GetLoopIterType;
using s_tir::GetLoops;
-using s_tir::GetThreadScope;
+using s_tir::GetThreadBindingTag;
using s_tir::HasBeenMultiLevelTiled;
using s_tir::IsBlockIdx;
using s_tir::IsSingleStmt;
@@ -116,7 +117,7 @@ void BindBlockThreadIdx(Schedule sch, SBlockRV block_rv, //
for (int i = 0; i < n; ++i) {
const StmtSRef& loop_sref = loops[i];
const ForNode* loop = TVM_SREF_TO_FOR(loop_sref);
- runtime::ThreadScope thread_scope = GetThreadScope(loop);
+ support::GPUThreadTag thread_scope = GetThreadBindingTag(loop);
if (IsBlockIdx(thread_scope)) {
if (i_block_idx == -1) {
i_block_idx = i;
diff --git a/src/s_tir/meta_schedule/schedule_rule/cross_thread_reduction.cc
b/src/s_tir/meta_schedule/schedule_rule/cross_thread_reduction.cc
index 396781c93a..b5da2d6dbb 100644
--- a/src/s_tir/meta_schedule/schedule_rule/cross_thread_reduction.cc
+++ b/src/s_tir/meta_schedule/schedule_rule/cross_thread_reduction.cc
@@ -18,6 +18,7 @@
*/
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/support/gpu_thread_tag.h>
#include "../utils.h"
@@ -91,7 +92,7 @@ class CrossThreadReductionObj : public ScheduleRuleObj {
// bound the innermost outer loop of `target_sblock` to threadIdx.
Possibly we need to split
// the loop before binding.
// - Otherwise, we search for the extent of "threadIdx.x" and use it as
the split factor.
- if (!InThreadScope(tmp_sch, target_sblock)) {
+ if (!InThreadBinding(tmp_sch, target_sblock)) {
const ffi::Array<s_tir::LoopRV>& split_res =
tmp_sch->Split(tgt_block_innermost_loop, {std::nullopt,
thread_extent});
tmp_sch->Bind(split_res[1], "threadIdx.x");
@@ -134,11 +135,11 @@ class CrossThreadReductionObj : public ScheduleRuleObj {
* \param block The block to be checked
* \return A boolean indicating whether the block is in thread scope.
*/
- bool InThreadScope(const s_tir::Schedule& sch, const s_tir::SBlockRV& block)
{
+ bool InThreadBinding(const s_tir::Schedule& sch, const s_tir::SBlockRV&
block) {
const ffi::Array<s_tir::LoopRV>& axes = sch->GetLoops(block);
for (const s_tir::LoopRV& loop_rv : axes) {
const For& loop = sch->Get(loop_rv);
- runtime::ThreadScope thread_scope = s_tir::GetThreadScope(loop.get());
+ support::GPUThreadTag thread_scope =
s_tir::GetThreadBindingTag(loop.get());
if (s_tir::IsThreadIdx(thread_scope)) {
return true;
}
diff --git a/src/s_tir/schedule/analysis/analysis.cc
b/src/s_tir/schedule/analysis/analysis.cc
index a847fa98ab..f6b5e49bcc 100644
--- a/src/s_tir/schedule/analysis/analysis.cc
+++ b/src/s_tir/schedule/analysis/analysis.cc
@@ -22,6 +22,7 @@
#include <tvm/ir/analysis.h>
#include <tvm/ir/op.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/support/gpu_thread_tag.h>
#include "../ir_comparator.h"
#include "../utils.h"
@@ -708,8 +709,7 @@ ffi::Map<Var, Range> LoopDomainOfSRefTreePath(const
StmtSRef& low_inclusive,
if (const ForNode* loop = p->StmtAs<ForNode>()) {
if (tvm::s_tir::GetThreadBinding(loop).has_value()) {
const ffi::String thread_tag =
tvm::s_tir::GetThreadBinding(loop).value();
- if (CanRelaxStorageUnderThread(extra_relax_scope,
-
runtime::ThreadScope::Create(thread_tag))) {
+ if (CanRelaxStorageUnderThread(extra_relax_scope,
ParseThreadBindingTag(thread_tag))) {
result.Set(loop->loop_var, Range::FromMinExtent(loop->min,
loop->extent));
}
}
diff --git a/src/s_tir/schedule/primitive/for_kind.cc
b/src/s_tir/schedule/primitive/for_kind.cc
index a0eb6fa2d1..3b1c9b7395 100644
--- a/src/s_tir/schedule/primitive/for_kind.cc
+++ b/src/s_tir/schedule/primitive/for_kind.cc
@@ -19,6 +19,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/support/gpu_thread_tag.h>
#include <tvm/tirx/stmt.h>
#include "../utils.h"
@@ -87,7 +88,7 @@ class WrongBlockIterTypeError : public
ScheduleErrorContextObj {
*/
void CheckLoopParallelizableInBlock(const ScheduleState& self, ForKind
for_kind,
const Var& loop_var, const SBlockRealize&
block_realize,
- runtime::ThreadScope thread_scope) {
+ support::GPUThreadTag thread_scope) {
const SBlock& block = block_realize->block;
// Cond 1. The block is required to have affine bindings.
@@ -132,7 +133,7 @@ void CheckLoopParallelizableInBlock(const ScheduleState&
self, ForKind for_kind,
* the operation is not "bind"
*/
void CheckParallelizability(const ScheduleState& self, const For& loop,
ForKind for_kind,
- runtime::ThreadScope thread_scope) {
+ support::GPUThreadTag thread_scope) {
auto walk_fn = [&](const SBlockRealize& realize) ->
ffi::Expected<ffi::WalkResult> {
// If this block doesn't have corresponding StmtSRef in the schedule
state, it must be a
// block inside `tirx.init()`. We don't check the condition for such
blocks.
@@ -175,8 +176,8 @@ void ParallelizeComputation(const ScheduleState& self,
const StmtSRef& loop_sref
// Step 2. Check whether the loop can be parallelized/vectorized/bound with
regard to each
// underlying block.
CheckParallelizability(self, ffi::GetRef<For>(loop), for_kind,
- thread_axis.has_value() ?
runtime::ThreadScope::Create(thread_axis.value())
- : runtime::ThreadScope{-1,
-1});
+ thread_axis.has_value() ?
ParseThreadBindingTag(thread_axis.value())
+ : support::GPUThreadTag{-1,
-1});
// Step 3. Loop update and IR replacement
ffi::ObjectPtr<ForNode> new_loop = ffi::make_object<ForNode>(*loop);
diff --git a/src/s_tir/schedule/utils.h b/src/s_tir/schedule/utils.h
index 08adea2731..b6fe4b6342 100644
--- a/src/s_tir/schedule/utils.h
+++ b/src/s_tir/schedule/utils.h
@@ -32,6 +32,7 @@
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/utils.h>
+#include <tvm/support/gpu_thread_tag.h>
#include <tvm/sym/analyzer.h>
#include <tvm/sym/int_set.h>
#include <tvm/sym/iter_affine_map.h>
@@ -98,7 +99,7 @@ inline ffi::Array<StmtSRef> SBlockRVs2StmtSRefs(const
Schedule& sch,
* \return A boolean indicating the result
*/
inline bool CanRelaxStorageUnderThread(const runtime::StorageScope&
storage_scope,
- const runtime::ThreadScope&
thread_scope) {
+ const support::GPUThreadTag&
thread_scope) {
if (storage_scope.rank == runtime::StorageRank::kWarp) {
// for warp memory, we only relax threadIdx.x
return thread_scope.rank == 1 && thread_scope.dim_index == 0;
@@ -160,16 +161,24 @@ inline IterVar IterVarFromLoop(const For& loop,
ffi::String name, IterVarType it
PrimVar(std::move(name), loop->loop_var.ty()), iter_var_type);
}
+/*! \brief Interpret physical and virtual thread bindings for S-TIR
scheduling. */
+inline support::GPUThreadTag ParseThreadBindingTag(const ffi::String& tag) {
+ if (tag.starts_with("vthread")) {
+ return {1, -1};
+ }
+ return support::GPUThreadTag::Parse(tag);
+}
+
/*!
* \brief Get the thread scope bound to the specific loop
* \param loop The loop to be inspected
* \return The thread scope bound to the loop
*/
-inline runtime::ThreadScope GetThreadScope(const ForNode* loop) {
+inline support::GPUThreadTag GetThreadBindingTag(const ForNode* loop) {
if (tvm::s_tir::GetThreadBinding(loop).has_value()) {
- return
runtime::ThreadScope::Create(tvm::s_tir::GetThreadBinding(loop).value());
+ return ParseThreadBindingTag(tvm::s_tir::GetThreadBinding(loop).value());
}
- return runtime::ThreadScope{-1, -1};
+ return support::GPUThreadTag{-1, -1};
}
/*!
@@ -177,7 +186,7 @@ inline runtime::ThreadScope GetThreadScope(const ForNode*
loop) {
* \param thread_scope The thread scope to be checked
* \return True if the thread scope is blockIdx
*/
-inline bool IsBlockIdx(const runtime::ThreadScope& thread_scope) {
+inline bool IsBlockIdx(const support::GPUThreadTag& thread_scope) {
return thread_scope.rank == 0; // The rank of blockIdx is 0
}
@@ -186,7 +195,7 @@ inline bool IsBlockIdx(const runtime::ThreadScope&
thread_scope) {
* \param thread_scope The thread scope to be checked
* \return True if the thread scope is threadIdx
*/
-inline bool IsThreadIdx(const runtime::ThreadScope& thread_scope) {
+inline bool IsThreadIdx(const support::GPUThreadTag& thread_scope) {
return thread_scope.rank == 1 && thread_scope.dim_index >= 0;
}
diff --git a/src/s_tir/transform/compact_tensor_region.cc
b/src/s_tir/transform/compact_tensor_region.cc
index cb721587f8..f6567adafc 100644
--- a/src/s_tir/transform/compact_tensor_region.cc
+++ b/src/s_tir/transform/compact_tensor_region.cc
@@ -30,6 +30,7 @@
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
+#include <tvm/support/gpu_thread_tag.h>
#include <tvm/sym/int_set.h>
#include <tvm/tirx/op/memory.h>
#include <tvm/tirx/op/region.h>
@@ -459,7 +460,7 @@ class TensorAccessRegionCollector : public StmtExprVisitor {
}
// When there is warp memory
// threadIdx.x must be set to be warp index.
- return CanRelaxStorageUnderThread(scope,
runtime::ThreadScope::Create((iter->thread_tag)));
+ return CanRelaxStorageUnderThread(scope,
ParseThreadBindingTag((iter->thread_tag)));
}
/*!
diff --git a/src/s_tir/transform/lift_thread_binding.cc
b/src/s_tir/transform/lift_thread_binding.cc
index 9661860d8f..6874025352 100644
--- a/src/s_tir/transform/lift_thread_binding.cc
+++ b/src/s_tir/transform/lift_thread_binding.cc
@@ -28,7 +28,6 @@
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/stmt.h>
-#include "../../runtime/thread_storage_scope.h"
#include "../../tirx/transform/ir_utils.h"
namespace tvm {
@@ -105,15 +104,7 @@ FindLoopLCA(const Stmt& root) {
sorted_thread_tags.push_back(kv.first);
}
std::sort(sorted_thread_tags.begin(), sorted_thread_tags.end(),
- [](const std::string& lhs, const std::string& rhs) {
- return lhs.size() > rhs.size();
- runtime::ThreadScope lhs_scope =
runtime::ThreadScope::Create(lhs);
- runtime::ThreadScope rhs_scope =
runtime::ThreadScope::Create(rhs);
- if (lhs_scope.rank != rhs_scope.rank) {
- return lhs_scope.rank < rhs_scope.rank;
- }
- return lhs_scope.dim_index < rhs_scope.dim_index;
- });
+ [](const std::string& lhs, const std::string& rhs) { return
lhs.size() > rhs.size(); });
for (const auto& thread_tag : sorted_thread_tags) {
Stmt lca = finder->lca[thread_tag].back();
const IterVar& iter = finder->iters.at(thread_tag);
diff --git a/src/s_tir/transform/loop_partition.cc
b/src/s_tir/transform/loop_partition.cc
index dde8221b26..23273ab176 100644
--- a/src/s_tir/transform/loop_partition.cc
+++ b/src/s_tir/transform/loop_partition.cc
@@ -32,6 +32,7 @@
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
+#include <tvm/support/gpu_thread_tag.h>
#include <tvm/sym/analyzer.h>
#include <tvm/sym/bound.h>
#include <tvm/tirx/analysis.h>
@@ -43,7 +44,6 @@
#include <unordered_map>
#include <unordered_set>
-#include "../../runtime/thread_storage_scope.h"
#include "../../sym/interval_set.h"
#include "ir_utils.h"
@@ -115,7 +115,7 @@ using ExpressionSet = std::unordered_set<PrimExpr,
ffi::ObjectPtrHash, ffi::Obje
static bool IsVirtualThread(const ForNode* op) {
if (!tvm::s_tir::GetThreadBinding(op).has_value()) return false;
const auto tag = tvm::s_tir::GetThreadBinding(op).value();
- return std::string(tag).rfind("vthread", 0) == 0;
+ return tag.starts_with("vthread");
}
static bool HasPartitionHint(const ForNode* op) {
@@ -171,7 +171,7 @@ class CandidateSelector final : public StmtExprVisitor {
return StmtExprVisitor::Visit_(op);
}
if (tvm::s_tir::GetThreadBinding(op).has_value() &&
-
runtime::ThreadScope::Create(tvm::s_tir::GetThreadBinding(op).value()).rank !=
0) {
+
support::GPUThreadTag::Parse(tvm::s_tir::GetThreadBinding(op).value()).rank !=
0) {
return StmtExprVisitor::Visit_(op);
}
// partition const loop when sets partition_const_loop_
@@ -190,10 +190,10 @@ class CandidateSelector final : public StmtExprVisitor {
ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op) final {
if (op->op.same_as(tirx::launch_thread_op()) &&
-
std::string(op->args[0].as_or_throw<StringImm>()->value).rfind("vthread", 0) !=
0) {
+ !op->args[0].as_or_throw<StringImm>()->value.starts_with("vthread")) {
Var var = op->body_params[0].as_or_throw<PrimVar>();
- runtime::ThreadScope scope =
-
runtime::ThreadScope::Create(op->args[0].as_or_throw<StringImm>()->value);
+ support::GPUThreadTag scope =
+
support::GPUThreadTag::Parse(op->args[0].as_or_throw<StringImm>()->value);
PrimExpr extent = op->args[1].as_or_throw<PrimExpr>();
if ((scope.rank == 0) && (!IsConstInt(extent) || partition_const_loop_))
{
record_.insert({var.get(), false});
@@ -307,7 +307,7 @@ class PartitionFinder : public StmtExprVisitor {
ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op) final {
// handle thread_axis
if (op->op.same_as(tirx::launch_thread_op()) &&
-
std::string(op->args[0].as_or_throw<StringImm>()->value).rfind("vthread", 0) !=
0) {
+ !op->args[0].as_or_throw<StringImm>()->value.starts_with("vthread")) {
const VarNode* var = op->body_params[0].as_or_throw<PrimVar>().get();
PrimExpr extent = op->args[1].as_or_throw<PrimExpr>();
IntSet dom = IntSet::FromRange(Range(IntImm(extent.ty(), 0), extent));
@@ -476,7 +476,7 @@ class ThreadPartitionInserter : public StmtExprMutator {
UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) final {
if (op->op.same_as(tirx::launch_thread_op()) &&
-
std::string(op->args[0].as_or_throw<StringImm>()->value).rfind("vthread", 0) !=
0) {
+ !op->args[0].as_or_throw<StringImm>()->value.starts_with("vthread")) {
innermost_thread_scope_ = true;
Stmt stmt =
StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
@@ -537,7 +537,7 @@ class LoopPartitioner : public StmtExprMutator {
// Relax threadIdx ranges to avoid introducing divergent partition
branches.
bool relax = thread_scope &&
-
runtime::ThreadScope::Create(tvm::s_tir::GetThreadBinding(op).value()).rank ==
1;
+
support::GPUThreadTag::Parse(tvm::s_tir::GetThreadBinding(op).value()).rank ==
1;
auto& ranges = relax ? relax_map_ : hint_map_;
ranges.insert({op->loop_var.get(), IntSet::Interval(op->min, op->min +
op->extent - 1)});
Stmt res = StmtExprMutator::Mutate_(op, InplaceMode::kDisallow)
@@ -548,7 +548,7 @@ class LoopPartitioner : public StmtExprMutator {
UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) final {
if (!op->op.same_as(tirx::launch_thread_op()) ||
-
std::string(op->args[0].as_or_throw<StringImm>()->value).rfind("vthread", 0) ==
0) {
+ op->args[0].as_or_throw<StringImm>()->value.starts_with("vthread")) {
return StmtExprMutator::Mutate_(op, inplace_mode);
}
@@ -561,8 +561,8 @@ class LoopPartitioner : public StmtExprMutator {
}
// normal path when loop parittion fails.
- runtime::ThreadScope scope =
-
runtime::ThreadScope::Create(op->args[0].as_or_throw<StringImm>()->value);
+ support::GPUThreadTag scope =
+
support::GPUThreadTag::Parse(op->args[0].as_or_throw<StringImm>()->value);
Stmt res{ffi::UnsafeInit{}};
if (scope.rank == 1) {
// threadIdx should be put into relax map, in case of divergence.
diff --git a/src/s_tir/transform/lower_cross_thread_reduction.cc
b/src/s_tir/transform/lower_cross_thread_reduction.cc
index d8b8a5301e..e32d5fd3d3 100644
--- a/src/s_tir/transform/lower_cross_thread_reduction.cc
+++ b/src/s_tir/transform/lower_cross_thread_reduction.cc
@@ -30,35 +30,36 @@
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
+#include <tvm/support/gpu_thread_tag.h>
#include <tvm/sym/analyzer.h>
#include <tvm/te/operation.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op/gpu.h>
#include <tvm/tirx/stmt.h>
-#include "../../runtime/thread_storage_scope.h"
#include "../../support/utils.h"
#include "../../tirx/transform/ir_utils.h"
#include "../schedule/analysis.h"
+#include "../schedule/utils.h"
namespace tvm {
namespace s_tir {
using namespace tvm::prim;
using namespace tvm::tirx;
-using runtime::ThreadScope;
+using support::GPUThreadTag;
using support::StartsWith;
-// Implement a hash and equality function for ThreadScope so that
-// ThreadScope can serve as map key class
-struct ThreadScopeHash {
- size_t operator()(const ThreadScope& scope) const {
+// Implement a hash and equality function for GPUThreadTag so that
+// GPUThreadTag can serve as map key class
+struct GPUThreadTagHash {
+ size_t operator()(const GPUThreadTag& scope) const {
return static_cast<size_t>(scope.rank * 30 + scope.dim_index);
}
};
-struct ThreadScopeEqual {
- bool operator()(const ThreadScope& a, const ThreadScope& b) const {
+struct GPUThreadTagEqual {
+ bool operator()(const GPUThreadTag& a, const GPUThreadTag& b) const {
return a.rank == b.rank && a.dim_index == b.dim_index;
}
};
@@ -72,8 +73,7 @@ bool IsBoundToThreadIdx(const ForNode* loop) {
if (!tvm::s_tir::GetThreadBinding(loop).has_value()) {
return false;
}
- runtime::ThreadScope scope =
- runtime::ThreadScope::Create(tvm::s_tir::GetThreadBinding(loop).value());
+ support::GPUThreadTag scope =
ParseThreadBindingTag(tvm::s_tir::GetThreadBinding(loop).value());
return scope.rank == 1 && scope.dim_index >= 0;
}
@@ -646,7 +646,7 @@ class CrossThreadReductionTransformer : public
StmtExprMutator {
// 2. it writes to non-local memory,
// 3. at least one of the reduction thread vars of the cross-thread reduction
// is free to this block (i.e., not bound to the block).
- std::vector<std::pair<ThreadScope, Range>> NeedCrossThreadBroadcast(
+ std::vector<std::pair<GPUThreadTag, Range>> NeedCrossThreadBroadcast(
const SBlockRealizeNode* realize) {
SBlock block = realize->block;
@@ -659,7 +659,7 @@ class CrossThreadReductionTransformer : public
StmtExprMutator {
// Find out the reduction threads for the read-tensors which are produced
by
// cross-thread reduction.
- std::unordered_map<ThreadScope, Range, ThreadScopeHash, ThreadScopeEqual>
thread2range;
+ std::unordered_map<GPUThreadTag, Range, GPUThreadTagHash,
GPUThreadTagEqual> thread2range;
for (TensorRegion read_region : block->reads) {
auto tensor_it =
crt_tensor2threads_.find(read_region->source.as_or_throw<tvm::tirx::TensorVar>().get());
@@ -674,11 +674,11 @@ class CrossThreadReductionTransformer : public
StmtExprMutator {
// Erase those threads which are not free to this block.
for (const ForNode* loop : loop_stack_) {
if (tvm::s_tir::GetThreadBinding(loop).has_value()) {
- ThreadScope scope =
ThreadScope::Create(tvm::s_tir::GetThreadBinding(loop).value());
+ GPUThreadTag scope =
ParseThreadBindingTag(tvm::s_tir::GetThreadBinding(loop).value());
thread2range.erase(scope);
}
}
- std::vector<std::pair<ThreadScope, Range>> unbound_thread2range_list;
+ std::vector<std::pair<GPUThreadTag, Range>> unbound_thread2range_list;
for (auto [scope, range] : thread2range) {
unbound_thread2range_list.emplace_back(scope, range);
}
@@ -812,7 +812,7 @@ class CrossThreadReductionTransformer : public
StmtExprMutator {
bool is_block_idx = false;
bool is_thread_idx = false;
if (tvm::s_tir::GetThreadBinding(loop).has_value()) {
- ThreadScope scope =
ThreadScope::Create(tvm::s_tir::GetThreadBinding(loop).value());
+ GPUThreadTag scope =
ParseThreadBindingTag(tvm::s_tir::GetThreadBinding(loop).value());
if (scope.rank == 1 && scope.dim_index >= 0) {
is_thread_idx = true;
++thread_idx_depth;
@@ -906,12 +906,12 @@ class CrossThreadReductionTransformer : public
StmtExprMutator {
// Step 5. Record the reduction thread dims for the write-back tensors.
// The information is used for consumer block broadcasting detection.
- std::vector<std::pair<ThreadScope, Range>> reduction_threads;
+ std::vector<std::pair<GPUThreadTag, Range>> reduction_threads;
reduction_threads.reserve(reduction_loops.size());
for (const ForNode* loop : reduction_loops) {
if (tvm::s_tir::GetThreadBinding(loop).has_value()) {
reduction_threads.emplace_back(
- ThreadScope::Create(tvm::s_tir::GetThreadBinding(loop).value()),
+ ParseThreadBindingTag(tvm::s_tir::GetThreadBinding(loop).value()),
Range::FromMinExtent(loop->min, loop->extent));
}
}
@@ -922,7 +922,7 @@ class CrossThreadReductionTransformer : public
StmtExprMutator {
Stmt MakeCrossThreadBroadcast(
const SBlockRealizeNode* realize,
- const std::vector<std::pair<ThreadScope, Range>>& unbound_thread2range) {
+ const std::vector<std::pair<GPUThreadTag, Range>>& unbound_thread2range)
{
// Step 1. Generate loop var for each unbound thread.
// Update the block predicate with clauses of `thread_var == min`.
PrimExpr predicate = realize->predicate;
@@ -973,7 +973,7 @@ class CrossThreadReductionTransformer : public
StmtExprMutator {
// Part 2. Check if the block needs all-thread broadcasting rewrite.
// We only check this when cross-thread reduction was detected.
- std::vector<std::pair<ThreadScope, Range>> unbound_thread2range =
+ std::vector<std::pair<GPUThreadTag, Range>> unbound_thread2range =
NeedCrossThreadBroadcast(realize);
if (!unbound_thread2range.empty()) {
return MakeCrossThreadBroadcast(realize, unbound_thread2range);
@@ -993,7 +993,7 @@ class CrossThreadReductionTransformer : public
StmtExprMutator {
int block_idx_depth = 0;
int thread_idx_depth = 0;
- std::unordered_map<const VarNode*, std::vector<std::pair<ThreadScope,
Range>>>
+ std::unordered_map<const VarNode*, std::vector<std::pair<GPUThreadTag,
Range>>>
crt_tensor2threads_;
};
diff --git a/src/s_tir/transform/memhammer_intermediate_stage.cc
b/src/s_tir/transform/memhammer_intermediate_stage.cc
index 9e95051a46..4218623606 100644
--- a/src/s_tir/transform/memhammer_intermediate_stage.cc
+++ b/src/s_tir/transform/memhammer_intermediate_stage.cc
@@ -19,6 +19,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/support/gpu_thread_tag.h>
#include "memhammer_rewrite_rule.h"
@@ -277,7 +278,7 @@ std::pair<Stmt, SeqStmt> InsertCacheStage(Stmt stmt, bool
is_write_cache, ffi::S
for (const For& loop : outer_loops) {
if (tvm::s_tir::GetThreadBinding(loop).has_value()) {
const ffi::String thread_tag =
tvm::s_tir::GetThreadBinding(loop).value();
- auto thread_scope = runtime::ThreadScope::Create(thread_tag);
+ auto thread_scope = ParseThreadBindingTag(thread_tag);
if
(CanRelaxStorageUnderThread(runtime::StorageScope::Create(storage_scope),
thread_scope)) {
if (is_write_cache && thread_scope.dim_index == 0) {
// writing C_reindex_m16n8k8_matrixC_shared_dyn is warp execution
diff --git a/src/tirx/transform/lower_thread_allreduce.h
b/src/tirx/transform/lower_thread_allreduce.h
index 122d42c2e1..069a0755af 100644
--- a/src/tirx/transform/lower_thread_allreduce.h
+++ b/src/tirx/transform/lower_thread_allreduce.h
@@ -26,6 +26,7 @@
#include <tvm/ir/function.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/prim/op.h>
+#include <tvm/support/gpu_thread_tag.h>
#include <tvm/sym/analyzer.h>
#include <tvm/target/target.h>
#include <tvm/tirx/op/gpu.h>
@@ -38,7 +39,6 @@
#include <map>
#include <unordered_set>
-#include "../../runtime/thread_storage_scope.h"
#include "ir_utils.h"
#include "update_pointer_storage_scope.h"
@@ -78,7 +78,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) final {
if (!op->op.same_as(tirx::launch_thread_op()) ||
-
std::string(op->args[0].as_or_throw<StringImm>()->value).rfind("vthread", 0) ==
0) {
+ op->args[0].as_or_throw<StringImm>()->value.starts_with("vthread")) {
return DialectMutator::Mutate_(op, inplace_mode);
}
thread_extents_.push_back(op);
@@ -207,7 +207,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
private:
// Thread entry
struct ThreadEntry {
- runtime::ThreadScope scope;
+ support::GPUThreadTag scope;
ffi::Optional<PrimVar> var;
int extent;
// comparator
@@ -288,7 +288,9 @@ class ThreadAllreduceBuilder final : public DialectMutator {
ThreadEntry e;
PrimVar var = launch->body_params[0].as_or_throw<PrimVar>();
ffi::String tag = launch->args[0].as_or_throw<StringImm>()->value;
- e.scope = runtime::ThreadScope::Create(tag);
+ TVM_FFI_ICHECK(!tag.starts_with("vthread"))
+ << "vthread do not work with cross thread reduction";
+ e.scope = support::GPUThreadTag::Parse(tag);
e.var = var;
TVM_FFI_ICHECK_LE(e.scope.rank, 1);
TVM_FFI_ICHECK_GE(e.scope.dim_index, 0) << "vthread do not work with
cross thread reduction";
