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 c5b0ceff92 [REFACTOR][Runtime] Extract launch parameter decoding
(#20680)
c5b0ceff92 is described below
commit c5b0ceff92edc801772ad2158545c4f60e4d8534
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 12 16:08:01 2026 -0400
[REFACTOR][Runtime] Extract launch parameter decoding (#20680)
Move ThreadWorkLoad and LaunchParamConfig into an internal
launch_param.h so runtime launch decoding no longer depends on
thread/storage scope parsing. Match the twelve supported dimension tags
locally while preserving launch argument layout and options.
Update ROCm, OpenCL, Metal, and Vulkan consumers to include the decoder
directly, remove stale CUDA and OpenCL scope includes, and make the
compiler launch-metadata dependency explicit.
---
src/backend/cuda/runtime/cuda_module.cc | 1 -
src/backend/metal/runtime/metal_module.mm | 2 +-
src/backend/opencl/runtime/opencl_common.h | 1 -
src/backend/opencl/runtime/opencl_module.cc | 1 +
src/backend/rocm/runtime/rocm_module.cc | 2 +-
src/backend/vulkan/runtime/vulkan_wrapped_func.h | 2 +-
src/runtime/launch_param.h | 169 +++++++++++++++++++++++
src/runtime/thread_storage_scope.h | 120 +---------------
src/tirx/transform/split_host_device.cc | 1 +
9 files changed, 178 insertions(+), 121 deletions(-)
diff --git a/src/backend/cuda/runtime/cuda_module.cc
b/src/backend/cuda/runtime/cuda_module.cc
index 3be2043072..e3fa272118 100644
--- a/src/backend/cuda/runtime/cuda_module.cc
+++ b/src/backend/cuda/runtime/cuda_module.cc
@@ -43,7 +43,6 @@
#include "../../../runtime/metadata.h"
#include "../../../runtime/pack_args.h"
-#include "../../../runtime/thread_storage_scope.h"
#include "../../../support/bytes_io.h"
#include "../launch/launch_config.h"
#include "../launch/launch_plan.h"
diff --git a/src/backend/metal/runtime/metal_module.mm
b/src/backend/metal/runtime/metal_module.mm
index afdf11ab67..c547b16a5f 100644
--- a/src/backend/metal/runtime/metal_module.mm
+++ b/src/backend/metal/runtime/metal_module.mm
@@ -38,9 +38,9 @@
#include <unordered_map>
#include <utility>
#include "../../../runtime/file_utils.h"
+#include "../../../runtime/launch_param.h"
#include "../../../runtime/metadata.h"
#include "../../../runtime/pack_args.h"
-#include "../../../runtime/thread_storage_scope.h"
#include "../../../support/bytes_io.h"
#include "metal_common.h"
diff --git a/src/backend/opencl/runtime/opencl_common.h
b/src/backend/opencl/runtime/opencl_common.h
index e197ded776..b2b5ba2c54 100644
--- a/src/backend/opencl/runtime/opencl_common.h
+++ b/src/backend/opencl/runtime/opencl_common.h
@@ -72,7 +72,6 @@
#include "../../../runtime/file_utils.h"
#include "../../../runtime/metadata.h"
#include "../../../runtime/pack_args.h"
-#include "../../../runtime/thread_storage_scope.h"
#include "../../../runtime/workspace_pool.h"
#include "texture.h"
diff --git a/src/backend/opencl/runtime/opencl_module.cc
b/src/backend/opencl/runtime/opencl_module.cc
index aebaf12040..13380f9412 100644
--- a/src/backend/opencl/runtime/opencl_module.cc
+++ b/src/backend/opencl/runtime/opencl_module.cc
@@ -35,6 +35,7 @@
#include <utility>
#include <vector>
+#include "../../../runtime/launch_param.h"
#include "../../../support/bytes_io.h"
#include "opencl_common.h"
#include "source_utils.h"
diff --git a/src/backend/rocm/runtime/rocm_module.cc
b/src/backend/rocm/runtime/rocm_module.cc
index 91a8f1c6c4..5654e184e2 100644
--- a/src/backend/rocm/runtime/rocm_module.cc
+++ b/src/backend/rocm/runtime/rocm_module.cc
@@ -36,9 +36,9 @@
#include <string>
#include <vector>
+#include "../../../runtime/launch_param.h"
#include "../../../runtime/metadata.h"
#include "../../../runtime/pack_args.h"
-#include "../../../runtime/thread_storage_scope.h"
#include "../../../support/bytes_io.h"
#include "rocm_common.h"
diff --git a/src/backend/vulkan/runtime/vulkan_wrapped_func.h
b/src/backend/vulkan/runtime/vulkan_wrapped_func.h
index 495df3a380..b5aa241458 100644
--- a/src/backend/vulkan/runtime/vulkan_wrapped_func.h
+++ b/src/backend/vulkan/runtime/vulkan_wrapped_func.h
@@ -27,9 +27,9 @@
#include <unordered_map>
#include <vector>
+#include "../../../runtime/launch_param.h"
#include "../../../runtime/metadata.h"
#include "../../../runtime/pack_args.h"
-#include "../../../runtime/thread_storage_scope.h"
#include "spirv_shader.h"
#include "vulkan/vulkan_core.h"
#include "vulkan_common.h"
diff --git a/src/runtime/launch_param.h b/src/runtime/launch_param.h
new file mode 100644
index 0000000000..1d3b51b0fc
--- /dev/null
+++ b/src/runtime/launch_param.h
@@ -0,0 +1,169 @@
+/*
+ * 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 launch_param.h
+ * \brief Extract launch parameters configuration from ffi::PackedArgs.
+ */
+#ifndef TVM_RUNTIME_LAUNCH_PARAM_H_
+#define TVM_RUNTIME_LAUNCH_PARAM_H_
+
+#include <tvm/ffi/function.h>
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "metadata.h"
+
+namespace tvm {
+namespace runtime {
+
+/*! \brief workload specification */
+struct ThreadWorkLoad {
+ // work_size layout: [0-2] grid, [3-5] block, [6-8] cluster, [9-11]
preferred_cluster
+ size_t work_size[12];
+ // Dynamic shared memory allocation size in bytes.
+ size_t dyn_shmem_size{0};
+ /*!
+ * \param i The block dimension.
+ * \return i-th block dim
+ */
+ inline size_t block_dim(size_t i) const { return work_size[i + 3]; }
+ /*!
+ * \param i The grid dimension.
+ * \return i-th grid dim
+ */
+ inline size_t grid_dim(size_t i) const { return work_size[i]; }
+ /*!
+ * \param i The cluster dimension.
+ * \return i-th cluster dim
+ */
+ inline size_t cluster_dim(size_t i) const { return work_size[i + 6]; }
+ /*!
+ * \param i The preferred cluster dimension.
+ * \return i-th preferred cluster dim
+ */
+ inline size_t preferred_cluster_dim(size_t i) const { return work_size[i +
9]; }
+};
+
+/*! \brief Launch parameters configuration */
+class LaunchParamConfig {
+ public:
+ void Init(size_t base, const ffi::Array<ffi::String>& launch_param_tags) {
+ base_ = base;
+ use_cluster_launch_ = false;
+ std::vector<bool> filled(12, false);
+ for (size_t i = 0; i < launch_param_tags.size(); ++i) {
+ std::string tag(launch_param_tags[i]);
+ if (tag == launch_param::kUseDynamicSharedMemoryTag) {
+ TVM_FFI_ICHECK_EQ(i, launch_param_tags.size() - 1)
+ << "kUseDynamicSharedMemoryTag should be the last tag in
launch_param_tags.";
+ use_dyn_shared_memory_ = true;
+ } else if (tag == launch_param::kUseProgramaticDependentLaunch) {
+ use_programmatic_dependent_launch_ = true;
+ } else if (tag == launch_param::kUseCooperativeLaunch) {
+ use_cooperative_launch_ = true;
+ } else if (tag == launch_param::kUseRequiredBlockDimension) {
+ use_required_block_dimension_ = true;
+ } else {
+ static constexpr const char* dimension_tags[] = {"blockIdx.x",
+ "blockIdx.y",
+ "blockIdx.z",
+ "threadIdx.x",
+ "threadIdx.y",
+ "threadIdx.z",
+ "clusterCtaIdx.x",
+ "clusterCtaIdx.y",
+ "clusterCtaIdx.z",
+
"preferredClusterCtaIdx.x",
+
"preferredClusterCtaIdx.y",
+
"preferredClusterCtaIdx.z"};
+ uint32_t index = 0;
+ while (index < 12 && tag != dimension_tags[index]) {
+ ++index;
+ }
+ TVM_FFI_ICHECK_LT(index, 12) << "Unsupported launch parameter tag: "
<< tag;
+ if (index >= 6 && index < 9) {
+ use_cluster_launch_ = true;
+ }
+ arg_index_map_.push_back(index);
+ filled[index] = true;
+ }
+ }
+ work_dim_ = 1;
+ for (int i = 0; i < 3; ++i) {
+ if (filled[i] || filled[i + 3]) {
+ work_dim_ = i + 1;
+ }
+ }
+ }
+ // extract workload from arguments.
+ ThreadWorkLoad Extract(ffi::PackedArgs args) const {
+ ThreadWorkLoad w;
+ std::fill(w.work_size, w.work_size + 12, 1);
+ const TVMFFIAny* raw_args = reinterpret_cast<const
TVMFFIAny*>(args.data());
+
+ for (size_t i = 0; i < arg_index_map_.size(); ++i) {
+ // Dynamic shapes can result in 0 dim size. Guard to ensure that the dim
size is at least 1.
+ size_t size = static_cast<size_t>(raw_args[base_ + i].v_int64);
+ if (size > 0) {
+ w.work_size[arg_index_map_[i]] = size;
+ }
+ }
+ if (use_dyn_shared_memory_) {
+ w.dyn_shmem_size = static_cast<size_t>(raw_args[base_ +
arg_index_map_.size()].v_int64);
+ }
+ return w;
+ }
+ // return the work dim
+ size_t work_dim() const { return work_dim_; }
+
+ bool use_programtic_dependent_launch() const { return
use_programmatic_dependent_launch_; }
+
+ bool use_cooperative_launch() const { return use_cooperative_launch_; }
+
+ bool use_required_block_dimension() const { return
use_required_block_dimension_; }
+
+ bool use_cluster_launch() const { return use_cluster_launch_; }
+
+ private:
+ /*! \brief base axis */
+ size_t base_;
+ /*! \brief The worker dimension */
+ size_t work_dim_;
+ /*! \brief The index mapping. */
+ std::vector<uint32_t> arg_index_map_;
+ /*! \brief Whether or not use dynamic shared memory. */
+ bool use_dyn_shared_memory_{false};
+ /*! \brief Whether or not use programmatic dependent launch. */
+ bool use_programmatic_dependent_launch_{false};
+ /*! \brief Whether or not use cooperative launch. */
+ bool use_cooperative_launch_{false};
+ /*! \brief Whether CUDA should use the kernel's statically required block
dimension. */
+ bool use_required_block_dimension_{false};
+ /*! \brief Whether the kernel declares a cluster-to-CTA scope. */
+ bool use_cluster_launch_{false};
+};
+
+} // namespace runtime
+} // namespace tvm
+#endif // TVM_RUNTIME_LAUNCH_PARAM_H_
diff --git a/src/runtime/thread_storage_scope.h
b/src/runtime/thread_storage_scope.h
index 61675d10a4..2db8f68d6f 100644
--- a/src/runtime/thread_storage_scope.h
+++ b/src/runtime/thread_storage_scope.h
@@ -19,17 +19,16 @@
/*!
* \file thread_storage_scope.h
- * \brief Extract launch parameters configuration from ffi::PackedArgs.
+ * \brief Thread and storage scope helpers.
*/
#ifndef TVM_RUNTIME_THREAD_STORAGE_SCOPE_H_
#define TVM_RUNTIME_THREAD_STORAGE_SCOPE_H_
-#include <tvm/ffi/function.h>
+#include <tvm/ffi/error.h>
+#include <cstddef>
+#include <functional>
#include <string>
-#include <vector>
-
-#include "metadata.h"
namespace tvm {
namespace runtime {
@@ -251,117 +250,6 @@ struct ThreadScope {
bool IsClusterCtaIdx() const { return rank == 2; }
};
-/*! \brief workload specification */
-struct ThreadWorkLoad {
- // work_size layout: [0-2] grid, [3-5] block, [6-8] cluster, [9-11]
preferred_cluster
- size_t work_size[12];
- // Dynamic shared memory allocation size in bytes.
- size_t dyn_shmem_size{0};
- /*!
- * \param i The block dimension.
- * \return i-th block dim
- */
- inline size_t block_dim(size_t i) const { return work_size[i + 3]; }
- /*!
- * \param i The grid dimension.
- * \return i-th grid dim
- */
- inline size_t grid_dim(size_t i) const { return work_size[i]; }
- /*!
- * \param i The cluster dimension.
- * \return i-th cluster dim
- */
- inline size_t cluster_dim(size_t i) const { return work_size[i + 6]; }
- /*!
- * \param i The preferred cluster dimension.
- * \return i-th preferred cluster dim
- */
- inline size_t preferred_cluster_dim(size_t i) const { return work_size[i +
9]; }
-};
-
-/*! \brief Launch parameters configuration */
-class LaunchParamConfig {
- public:
- void Init(size_t base, const ffi::Array<ffi::String>& launch_param_tags) {
- base_ = base;
- use_cluster_launch_ = false;
- std::vector<bool> filled(12, false);
- for (size_t i = 0; i < launch_param_tags.size(); ++i) {
- std::string tag(launch_param_tags[i]);
- if (tag == launch_param::kUseDynamicSharedMemoryTag) {
- TVM_FFI_ICHECK_EQ(i, launch_param_tags.size() - 1)
- << "kUseDynamicSharedMemoryTag should be the last tag in
launch_param_tags.";
- use_dyn_shared_memory_ = true;
- } else if (tag == launch_param::kUseProgramaticDependentLaunch) {
- use_programmatic_dependent_launch_ = true;
- } else if (tag == launch_param::kUseCooperativeLaunch) {
- use_cooperative_launch_ = true;
- } else if (tag == launch_param::kUseRequiredBlockDimension) {
- use_required_block_dimension_ = true;
- } else {
- ThreadScope ts = ThreadScope::Create(tag);
- if (ts.IsClusterCtaIdx()) {
- use_cluster_launch_ = true;
- }
- arg_index_map_.push_back(ts.rank * 3 + ts.dim_index);
- filled[ts.rank * 3 + ts.dim_index] = true;
- }
- }
- work_dim_ = 1;
- for (int i = 0; i < 3; ++i) {
- if (filled[i] || filled[i + 3]) {
- work_dim_ = i + 1;
- }
- }
- }
- // extract workload from arguments.
- ThreadWorkLoad Extract(ffi::PackedArgs args) const {
- ThreadWorkLoad w;
- std::fill(w.work_size, w.work_size + 12, 1);
- const TVMFFIAny* raw_args = reinterpret_cast<const
TVMFFIAny*>(args.data());
-
- for (size_t i = 0; i < arg_index_map_.size(); ++i) {
- // Dynamic shapes can result in 0 dim size. Guard to ensure that the dim
size is at least 1.
- size_t size = static_cast<size_t>(raw_args[base_ + i].v_int64);
- if (size > 0) {
- w.work_size[arg_index_map_[i]] = size;
- }
- }
- if (use_dyn_shared_memory_) {
- w.dyn_shmem_size = static_cast<size_t>(raw_args[base_ +
arg_index_map_.size()].v_int64);
- }
- return w;
- }
- // return the work dim
- size_t work_dim() const { return work_dim_; }
-
- bool use_programtic_dependent_launch() const { return
use_programmatic_dependent_launch_; }
-
- bool use_cooperative_launch() const { return use_cooperative_launch_; }
-
- bool use_required_block_dimension() const { return
use_required_block_dimension_; }
-
- bool use_cluster_launch() const { return use_cluster_launch_; }
-
- private:
- /*! \brief base axis */
- size_t base_;
- /*! \brief The worker dimension */
- size_t work_dim_;
- /*! \brief The index mapping. */
- std::vector<uint32_t> arg_index_map_;
- /*! \brief Whether or not use dynamic shared memory. */
- bool use_dyn_shared_memory_{false};
- /*! \brief Whether or not use programmatic dependent launch. */
- bool use_programmatic_dependent_launch_{false};
- /*! \brief Whether or not use cooperative launch. */
- bool use_cooperative_launch_{false};
- /*! \brief Whether CUDA should use the kernel's statically required block
dimension. */
- bool use_required_block_dimension_{false};
- /*! \brief Whether the kernel declares a cluster-to-CTA scope. */
- bool use_cluster_launch_{false};
-};
-
} // namespace runtime
} // namespace tvm
diff --git a/src/tirx/transform/split_host_device.cc
b/src/tirx/transform/split_host_device.cc
index 79e1780a10..097b292106 100644
--- a/src/tirx/transform/split_host_device.cc
+++ b/src/tirx/transform/split_host_device.cc
@@ -46,6 +46,7 @@
#include <optional>
+#include "../../runtime/metadata.h"
#include "../../runtime/thread_storage_scope.h"
#include "ir_utils.h"
