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 d310e7befa [Runtime] Expose device selection through the backend C API
(#20661)
d310e7befa is described below
commit d310e7befaed1c46e47845296f3ed02e529882b3
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Oct 11 12:21:11 2026 -0400
[Runtime] Expose device selection through the backend C API (#20661)
---
include/tvm/runtime/c_backend_api.h | 35 ++++++++++++++++++++++
include/tvm/runtime/device_api.h | 18 +++++++----
src/backend/cuda/codegen/codegen_cuda_host.cc | 27 +++++++++--------
src/runtime/device_api.cc | 32 +++++++++++++++-----
src/runtime/module.cc | 2 ++
src/target/llvm/codegen_cpu.cc | 2 ++
src/tirx/transform/make_packed_api.cc | 14 +++++----
.../s_tir/script/test_s_tir_script_printer.py | 8 ++---
8 files changed, 102 insertions(+), 36 deletions(-)
diff --git a/include/tvm/runtime/c_backend_api.h
b/include/tvm/runtime/c_backend_api.h
index 07d1d305de..d167e98a28 100644
--- a/include/tvm/runtime/c_backend_api.h
+++ b/include/tvm/runtime/c_backend_api.h
@@ -34,6 +34,41 @@
extern "C" {
#endif
+/*!
+ * \brief Select the runtime device for the calling host thread.
+ *
+ * Dispatches to DeviceAPI::SetDevice with the given local DLPack device type
+ * and runtime device index (DLDevice::device_id). This does not save or
restore
+ * the previous selection. Backends without device-selection state may treat
+ * selection as a no-op. Valid indices and index validation are
backend-specific.
+ *
+ * \param device_type Local DLPack device type, not an RPC session-encoded
type.
+ * \param device_index Nonnegative runtime device index, not a device type or
+ * target configuration index.
+ * \return 0 on success, -1 on invalid arguments, unavailable backend, or
backend
+ * failure. Failure records the raised TVM FFI error on the calling thread; no
+ * C++ exception crosses this C entry point.
+ */
+TVM_RUNTIME_DLL int TVMBackendSetDevice(int device_type, int device_index);
+
+/*!
+ * \brief Query the selected runtime device index on the calling host thread.
+ *
+ * Dispatches to DeviceAPI::GetDeviceIndex without changing device selection.
+ * Getter support is optional and independent of setter support: the default
+ * implementation reports NotImplementedError, including for backends without
+ * device-selection state. RPC devices are not supported.
+ *
+ * \param device_type Local DLPack device type, not an RPC session-encoded
type.
+ * \param out_device_index Non-null pointer receiving the runtime device index
+ * (DLDevice::device_id), not a device type or target configuration index.
Written
+ * only on success; its previous value is unchanged on failure.
+ * \return 0 on success, -1 on invalid arguments, unavailable backend,
unsupported
+ * query, or backend failure. Failure records the raised TVM FFI error on the
+ * calling thread; no C++ exception crosses this C entry point.
+ */
+TVM_RUNTIME_DLL int TVMBackendGetDeviceIndex(int device_type, int*
out_device_index);
+
/*!
* \brief Backend function to allocate temporal workspace.
*
diff --git a/include/tvm/runtime/device_api.h b/include/tvm/runtime/device_api.h
index ed2d2806af..70e40f9d48 100644
--- a/include/tvm/runtime/device_api.h
+++ b/include/tvm/runtime/device_api.h
@@ -136,6 +136,18 @@ class TVM_RUNTIME_DLL DeviceAPI {
* \param dev The device to be set.
*/
virtual void SetDevice(Device dev) = 0;
+ /*!
+ * \brief Get the currently selected runtime device index on the calling
host thread.
+ *
+ * The device type is determined by this DeviceAPI instance. The result is a
+ * runtime device index (DLDevice::device_id), not a device type or a target
+ * configuration index. Implementations must query the current selection
without
+ * changing it. The default implementation throws NotImplementedError;
support
+ * for SetDevice does not imply support for this query. RPC devices are
excluded.
+ *
+ * \return The selected device index, on success.
+ */
+ virtual int GetDeviceIndex();
/*!
* \brief Get attribute of specified device.
* \param dev The device device
@@ -415,12 +427,6 @@ inline Device AddRPCSessionMask(Device dev, int
session_table_index) {
*/
TVM_RUNTIME_DLL bool RuntimeEnabled(const ffi::String& target);
-/*! \brief namespace for constant symbols */
-namespace symbol {
-/*! \brief global function to set device */
-constexpr const char* tvm_set_device = "__tvm_set_device";
-} // namespace symbol
-
} // namespace runtime
} // namespace tvm
diff --git a/src/backend/cuda/codegen/codegen_cuda_host.cc
b/src/backend/cuda/codegen/codegen_cuda_host.cc
index 287e170995..10126ebd46 100644
--- a/src/backend/cuda/codegen/codegen_cuda_host.cc
+++ b/src/backend/cuda/codegen/codegen_cuda_host.cc
@@ -199,19 +199,6 @@ class CodeGenCUDAHost : public CodeGenCHost {
TVM_FFI_THROW(ValueError)
<< "cuda_host requires tensormap_encode_tiled instead of a packed
tensor-map encoder";
}
- if (name == "__tvm_set_device") {
- std::string args =
- "((TVMFFIAny*)" + PrintExpr(op->args[1]) + " + " +
PrintExpr(op->args[2]) + ")";
- std::string guard = name_supply_->FreshName("cuda_device_guard");
- PrintIndent();
- stream << "if (" << args << "[0].v_int64 != kDLCUDA) {\n"
- << " TVMFFIErrorSetRaisedFromCStr(\"ValueError\", "
- << "\"cuda_host requires a CUDA device\");\n return -1;\n}\n";
- PrintIndent();
- stream << "tvm::ffi::CUDADeviceGuard " << guard << "(" << args <<
"[1].v_int64);\n";
- os << "0";
- return;
- }
std::string function_name = name;
TVM_FFI_CHECK(function_name.rfind("runtime.", 0) != 0 &&
function_name.rfind("device_api.", 0) != 0 &&
@@ -228,6 +215,20 @@ class CodeGenCUDAHost : public CodeGenCHost {
}
if (op->op.same_as(tirx::call_extern_op()) ||
op->op.same_as(tirx::call_pure_extern_op())) {
const auto& symbol = op->args[0].as_or_throw<StringImm>()->value;
+ if (op->op.same_as(tirx::call_extern_op()) && symbol ==
"TVMBackendSetDevice") {
+ TVM_FFI_CHECK_EQ(op->args.size(), 3, ValueError);
+ std::string device_type = PrintExpr(op->args[1]);
+ std::string device_index = PrintExpr(op->args[2]);
+ std::string guard = name_supply_->FreshName("cuda_device_guard");
+ PrintIndent();
+ stream << "if (" << device_type << " != kDLCUDA) {\n"
+ << " TVMFFIErrorSetRaisedFromCStr(\"ValueError\", "
+ << "\"cuda_host requires a CUDA device\");\n return -1;\n}\n";
+ PrintIndent();
+ stream << "tvm::ffi::CUDADeviceGuard " << guard << "(" << device_index
<< ");\n";
+ os << "0";
+ return;
+ }
TVM_FFI_CHECK(std::string(symbol).rfind("TVMBackend", 0) != 0,
ValueError)
<< "cuda_host does not provide TVM runtime operation: " << symbol;
}
diff --git a/src/runtime/device_api.cc b/src/runtime/device_api.cc
index e573b8fb4e..c46d2ca9af 100644
--- a/src/runtime/device_api.cc
+++ b/src/runtime/device_api.cc
@@ -97,6 +97,10 @@ DeviceAPI* DeviceAPI::Get(Device dev, bool allow_missing) {
return DeviceAPIManager::Get(static_cast<int>(dev.device_type),
allow_missing);
}
+int DeviceAPI::GetDeviceIndex() {
+ TVM_FFI_THROW(NotImplementedError) << "Device does not support
GetDeviceIndex.";
+}
+
void* DeviceAPI::AllocWorkspace(Device dev, size_t size, DLDataType type_hint)
{
return AllocDataSpace(dev, size, kTempAllocaAlignment, type_hint);
}
@@ -204,13 +208,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
- .def_packed(tvm::runtime::symbol::tvm_set_device,
- [](tvm::ffi::PackedArgs args, tvm::ffi::Any* ret) {
- DLDevice dev;
- dev.device_type =
static_cast<DLDeviceType>(args[0].cast<int>());
- dev.device_id = args[1].cast<int>();
- DeviceAPIManager::Get(dev)->SetDevice(dev);
- })
.def_packed("runtime.GetDeviceAttr",
[](tvm::ffi::PackedArgs args, tvm::ffi::Any* ret) {
DLDevice dev;
@@ -241,6 +238,27 @@ TVM_FFI_STATIC_INIT_BLOCK() {
using namespace tvm::runtime;
+int TVMBackendSetDevice(int device_type, int device_index) {
+ TVM_FFI_SAFE_CALL_BEGIN();
+ TVM_FFI_CHECK(device_type > 0 && device_type < TVMDeviceExtType_End,
ValueError)
+ << "Expected a local device type, got " << device_type;
+ TVM_FFI_CHECK_GE(device_index, 0, ValueError) << "Expected a nonnegative
device index";
+ DLDevice dev{static_cast<DLDeviceType>(device_type), device_index};
+ DeviceAPI::Get(dev)->SetDevice(dev);
+ TVM_FFI_SAFE_CALL_END();
+}
+
+int TVMBackendGetDeviceIndex(int device_type, int* out_device_index) {
+ TVM_FFI_SAFE_CALL_BEGIN();
+ TVM_FFI_CHECK(out_device_index != nullptr, ValueError) << "out_device_index
must not be null";
+ TVM_FFI_CHECK(device_type > 0 && device_type < TVMDeviceExtType_End,
ValueError)
+ << "Expected a local device type, got " << device_type;
+ DLDevice dev{static_cast<DLDeviceType>(device_type), 0};
+ int device_index = DeviceAPI::Get(dev)->GetDeviceIndex();
+ *out_device_index = device_index;
+ TVM_FFI_SAFE_CALL_END();
+}
+
void* TVMBackendAllocWorkspace(int device_type, int device_id, uint64_t size,
int dtype_code_hint,
int dtype_bits_hint) {
DLDevice dev;
diff --git a/src/runtime/module.cc b/src/runtime/module.cc
index 7892f3eb4e..427766a158 100644
--- a/src/runtime/module.cc
+++ b/src/runtime/module.cc
@@ -80,6 +80,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TVM_INIT_CONTEXT_FUNC(TVMFFIErrorSetRaisedFromCStr);
TVM_INIT_CONTEXT_FUNC(TVMFFIEnvModLookupFromImports);
TVM_INIT_CONTEXT_FUNC(TVMFFIHandleInitOnce);
+ TVM_INIT_CONTEXT_FUNC(TVMBackendSetDevice);
+ TVM_INIT_CONTEXT_FUNC(TVMBackendGetDeviceIndex);
TVM_INIT_CONTEXT_FUNC(TVMBackendAllocWorkspace);
TVM_INIT_CONTEXT_FUNC(TVMBackendFreeWorkspace);
TVM_INIT_CONTEXT_FUNC(TVMBackendParallelLaunch);
diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index ddeb4f5f38..3ef5543702 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -493,6 +493,8 @@ void CodeGenCPU::InitGlobalContext(bool dynamic_lookup) {
gv_tvm_parallel_barrier_ =
InitContextPtr(llvmGetPointerTo(ftype_tvm_parallel_barrier_, 0),
"__TVMBackendParallelBarrier");
// Mark as context functions
+ gv_func_map_["TVMBackendSetDevice"] = nullptr;
+ gv_func_map_["TVMBackendGetDeviceIndex"] = nullptr;
gv_func_map_["TVMBackendAllocWorkspace"] = nullptr;
gv_func_map_["TVMBackendFreeWorkspace"] = nullptr;
}
diff --git a/src/tirx/transform/make_packed_api.cc
b/src/tirx/transform/make_packed_api.cc
index 1de656c2ff..64242216ec 100644
--- a/src/tirx/transform/make_packed_api.cc
+++ b/src/tirx/transform/make_packed_api.cc
@@ -269,11 +269,15 @@ Function MakePackedAPI(Function func) {
// Set device context
if (need_set_device) {
if (runtime::DeviceAPI::NeedSetDevice(target_device_type)) {
- Stmt set_device = Evaluate(Call(PrimType::Int(32),
tirx::call_packed_op(),
-
{StringImm(runtime::symbol::tvm_set_device), device_type,
- device_id.as_or_throw<PrimExpr>()})
- .as_or_throw<PrimExpr>());
- body = SeqStmt({set_device, body});
+ PrimVar status("set_device_status", PrimType::Int(32));
+ Stmt set_device =
+ Bind(status, Call(PrimType::Int(32), tirx::call_extern_op(),
+ {StringImm("TVMBackendSetDevice"), device_type,
device_id}));
+ Stmt check_status = If(
+ status != IntImm::Int32(0),
+ SeqStmt({Evaluate(
+ Call(PrimType::Int(32), tirx::throw_last_error_op(),
{}).as_or_throw<PrimExpr>())}));
+ body = SeqStmt({set_device, check_status, body});
}
}
diff --git a/tests/python/s_tir/script/test_s_tir_script_printer.py
b/tests/python/s_tir/script/test_s_tir_script_printer.py
index 1a56098e70..d96a08250c 100644
--- a/tests/python/s_tir/script/test_s_tir_script_printer.py
+++ b/tests/python/s_tir/script/test_s_tir_script_printer.py
@@ -1161,11 +1161,9 @@ def opt_conv_tensorcore_mod_host():
assert dev_id == T.abi_field_get(arg2, 0, 9, ty="int32"), (
"Argument arg2.device_id has an unsatisfied constraint"
)
- T.evaluate(T.abi_field_set(stack_value, 0, 12, T.cast(2, "int64")))
- stack_tcode[0] = 0
- T.evaluate(T.abi_field_set(stack_value, 1, 12, T.cast(dev_id,
"int64")))
- stack_tcode[1] = 0
- T.evaluate(T.call_packed_lowered("__tvm_set_device", stack_value, 0,
2))
+ set_device_status = T.call_extern("TVMBackendSetDevice", 2, dev_id,
ty="int32")
+ if set_device_status != 0:
+ T.evaluate(T.throw_last_error())
with T.compute_scope("default_function_compute_"):
T.evaluate(T.abi_field_set(stack_value, 0, 12, A))
stack_tcode[0] = 3
