Skip to content

Commit d73a267

Browse files
authored
[BUILD] TVM_FFI_COLD_CODE / TVM_FFI_PREDICT_FALSE macros and cold-marking of error paths (#589)
## Summary Adds three header-only macros (`TVM_FFI_COLD_CODE`, `TVM_FFI_PREDICT_FALSE`, `TVM_FFI_PREDICT_TRUE`) in `tvm/ffi/base_details.h` and applies them to a small audited set of error-only helpers. No CMake changes. ## What `TVM_FFI_COLD_CODE`: `[[gnu::cold]]` on GCC/Clang, no-op on MSVC. `TVM_FFI_PREDICT_FALSE` / `TVM_FFI_PREDICT_TRUE`: `__builtin_expect` on GCC/Clang, no-op on MSVC. Cold-marked (error / setup / teardown only): - `details::ErrorBuilder` ctors and the `[[noreturn]]` destructor - `TVMFFISegFaultHandler` - `TVMFFIInstallSignalHandler` - `ForwardPyErrorToFFI` C ABI exports stay hot per cross-DSO surface hygiene — `TVMFFIError*` family, `TVMFFIBacktrace`, and `SafeCallContext` setters all remain ordinary entry points. Deleters are not cold-marked: they run on every callback destruction during normal program operation, not on an error path. `TVM_FFI_PREDICT_FALSE` is applied to `TVM_FFI_CHECK_SAFE_CALL`, `TVM_FFI_CHECK`, and ~17 error-check branches in the Python→FFI dispatchers in `tvm_ffi_python_helpers.h`. ## Mechanism GCC and Clang emit cold-marked functions into per-TU `.text.unlikely`. The default GNU linker script's `*(.text.unlikely .text.*_unlikely .text.unlikely.*)` rule gathers them into a contiguous slot inside `.text`. No `-ffunction-sections` flag required. Detailed measurements (binary-size matrix, isolation study, perf table, cold-cluster bounds) posted as a follow-up comment.
1 parent d4cfd86 commit d73a267

7 files changed

Lines changed: 94 additions & 25 deletions

File tree

include/tvm/ffi/base_details.h

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,44 @@
7676
#define TVM_FFI_UNREACHABLE() __builtin_unreachable()
7777
#endif
7878

79+
/*!
80+
* \brief Mark a function as cold so the toolchain places it in a
81+
* separate cold region of `.text`. Apply to functions that only
82+
* run on error / setup / teardown paths.
83+
*
84+
* On GCC and Clang, expands to `[[gnu::cold]]`, which emits the
85+
* function into a per-TU `.text.unlikely` section. The default GNU
86+
* linker script gathers `.text.unlikely.*` into a contiguous slot
87+
* inside `.text`, so cold-marked functions cluster away from hot
88+
* code without any additional CMake flags. On MSVC the attribute
89+
* does not exist and the macro is a no-op.
90+
*/
91+
#if defined(__GNUC__) || defined(__clang__)
92+
#define TVM_FFI_COLD_CODE [[gnu::cold]]
93+
#else
94+
#define TVM_FFI_COLD_CODE
95+
#endif
96+
97+
/*!
98+
* \brief Branch-prediction / layout hint that the condition is unlikely
99+
* to be true. Use on error-checking branches to keep the hot
100+
* fall-through contiguous and push the error-handling block to
101+
* the function tail.
102+
*
103+
* if (TVM_FFI_PREDICT_FALSE(rc != 0)) { ...error... }
104+
*
105+
* On GCC/Clang, expands to `__builtin_expect((cond), 0)`. On MSVC,
106+
* expands to `(cond)` (no equivalent builtin; modern MSVC does its own
107+
* profile-driven block reordering).
108+
*/
109+
#if defined(__GNUC__) || defined(__clang__)
110+
#define TVM_FFI_PREDICT_FALSE(cond) (__builtin_expect(static_cast<bool>(cond), 0))
111+
#define TVM_FFI_PREDICT_TRUE(cond) (__builtin_expect(static_cast<bool>(cond), 1))
112+
#else
113+
#define TVM_FFI_PREDICT_FALSE(cond) (cond)
114+
#define TVM_FFI_PREDICT_TRUE(cond) (cond)
115+
#endif
116+
79117
#define TVM_FFI_STR_CONCAT_(__x, __y) __x##__y
80118
#define TVM_FFI_STR_CONCAT(__x, __y) TVM_FFI_STR_CONCAT_(__x, __y)
81119

include/tvm/ffi/error.h

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -338,11 +338,13 @@ TVM_FFI_INLINE void SetSafeCallRaised(const Error& error) {
338338

339339
class ErrorBuilder {
340340
public:
341+
TVM_FFI_COLD_CODE
341342
explicit ErrorBuilder(std::string kind, std::string backtrace, bool log_before_throw)
342343
: kind_(std::move(kind)),
343344
backtrace_(std::move(backtrace)),
344345
log_before_throw_(log_before_throw) {}
345346

347+
TVM_FFI_COLD_CODE
346348
explicit ErrorBuilder(std::string kind, const TVMFFIByteArray* backtrace, bool log_before_throw)
347349
: ErrorBuilder(std::move(kind), std::string(backtrace->data, backtrace->size),
348350
log_before_throw) {}
@@ -353,7 +355,7 @@ class ErrorBuilder {
353355
#pragma warning(disable : 4722)
354356
#endif
355357
// avoid inline to reduce binary size, error throw path do not need to be fast
356-
[[noreturn]] ~ErrorBuilder() noexcept(false) {
358+
[[noreturn]] TVM_FFI_COLD_CODE ~ErrorBuilder() noexcept(false) {
357359
::tvm::ffi::Error error(std::move(kind_), stream_.str(), std::move(backtrace_));
358360
if (log_before_throw_) {
359361
std::cerr << error.FullMessage();
@@ -456,7 +458,8 @@ TVM_FFI_CHECK_FUNC(_NE, !=)
456458
TVM_FFI_THROW(ErrorKind) << "Check failed: " << #x " " #op " " #y << *__tvm_ffi_log_err << ": "
457459

458460
#define TVM_FFI_CHECK(cond, ErrorKind) \
459-
if (!(cond)) TVM_FFI_THROW(ErrorKind) << "Check failed: (" #cond << ") is false: "
461+
if (TVM_FFI_PREDICT_FALSE(!(cond))) \
462+
TVM_FFI_THROW(ErrorKind) << "Check failed: (" #cond << ") is false: "
460463

461464
#define TVM_FFI_CHECK_LT(x, y, ErrorKind) TVM_FFI_CHECK_BINARY_OP(_LT, <, x, y, ErrorKind)
462465
#define TVM_FFI_CHECK_GT(x, y, ErrorKind) TVM_FFI_CHECK_BINARY_OP(_GT, >, x, y, ErrorKind)

include/tvm/ffi/function.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ namespace ffi {
101101
#define TVM_FFI_CHECK_SAFE_CALL(func) \
102102
{ \
103103
int ret_code = (func); \
104-
if (ret_code != 0) { \
104+
if (TVM_FFI_PREDICT_FALSE(ret_code != 0)) { \
105105
throw ::tvm::ffi::details::MoveFromSafeCallRaised(); \
106106
} \
107107
}

python/tvm_ffi/cython/tvm_ffi_python_helpers.h

Lines changed: 45 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,29 @@
3636
#endif
3737
#endif
3838

39+
// Local mirror of TVM_FFI_COLD_CODE / TVM_FFI_PREDICT_* from
40+
// <tvm/ffi/base_details.h>. The Cython helper deliberately avoids that header
41+
// (keeps the include surface c-headers-only), so we duplicate the macro
42+
// definitions here. Keep these in sync with base_details.h: same expansion on
43+
// GCC/Clang, no-op on MSVC.
44+
#ifndef TVM_FFI_COLD_CODE
45+
#if defined(__GNUC__) || defined(__clang__)
46+
#define TVM_FFI_COLD_CODE [[gnu::cold]]
47+
#else
48+
#define TVM_FFI_COLD_CODE
49+
#endif
50+
#endif
51+
52+
#ifndef TVM_FFI_PREDICT_FALSE
53+
#if defined(__GNUC__) || defined(__clang__)
54+
#define TVM_FFI_PREDICT_FALSE(cond) (__builtin_expect(static_cast<bool>(cond), 0))
55+
#define TVM_FFI_PREDICT_TRUE(cond) (__builtin_expect(static_cast<bool>(cond), 1))
56+
#else
57+
#define TVM_FFI_PREDICT_FALSE(cond) (cond)
58+
#define TVM_FFI_PREDICT_TRUE(cond) (cond)
59+
#endif
60+
#endif
61+
3962
#include <cstring>
4063
#include <exception>
4164
#include <iostream>
@@ -252,7 +275,7 @@ int TVMFFIPyArgSetterInt_(TVMFFIPyArgSetter*, TVMFFIPyCallContext*, PyObject* ar
252275
out->type_index = kTVMFFIInt;
253276
out->v_int64 = PyLong_AsLongLongAndOverflow(arg, &overflow);
254277

255-
if (overflow != 0) {
278+
if (TVM_FFI_PREDICT_FALSE(overflow != 0)) {
256279
PyErr_SetString(PyExc_OverflowError, "Python int too large to convert to int64_t");
257280
return -1;
258281
}
@@ -454,15 +477,15 @@ class TVMFFIPyCallManager {
454477
int* c_api_ret_code, bool release_gil,
455478
const DLPackExchangeAPI** optional_out_ctx_dlpack_api) {
456479
int64_t num_args = PyTuple_Size(py_arg_tuple);
457-
if (num_args == -1) return -1;
480+
if (TVM_FFI_PREDICT_FALSE(num_args == -1)) return -1;
458481
try {
459482
// allocate a call stack
460483
TVMFFIPyCallContext ctx(&call_stack_, num_args);
461484
// Iterate over the arguments and set them
462485
for (int64_t i = 0; i < num_args; ++i) {
463486
PyObject* py_arg = PyTuple_GetItem(py_arg_tuple, i);
464487
TVMFFIAny* c_arg = ctx.packed_args + i;
465-
if (SetArgument(&ctx, py_arg, c_arg) != 0) return -1;
488+
if (TVM_FFI_PREDICT_FALSE(SetArgument(&ctx, py_arg, c_arg) != 0)) return -1;
466489
}
467490
TVMFFIStreamHandle prev_stream = nullptr;
468491
DLPackManagedTensorAllocator prev_tensor_allocator = nullptr;
@@ -471,13 +494,13 @@ class TVMFFIPyCallManager {
471494
c_api_ret_code[0] =
472495
TVMFFIEnvSetStream(ctx.device_type, ctx.device_id, ctx.stream, &prev_stream);
473496
// setting failed, directly return
474-
if (c_api_ret_code[0] != 0) return 0;
497+
if (TVM_FFI_PREDICT_FALSE(c_api_ret_code[0] != 0)) return 0;
475498
}
476499
if (ctx.dlpack_c_exchange_api != nullptr &&
477500
ctx.dlpack_c_exchange_api->managed_tensor_allocator != nullptr) {
478501
c_api_ret_code[0] = TVMFFIEnvSetDLPackManagedTensorAllocator(
479502
ctx.dlpack_c_exchange_api->managed_tensor_allocator, 0, &prev_tensor_allocator);
480-
if (c_api_ret_code[0] != 0) return 0;
503+
if (TVM_FFI_PREDICT_FALSE(c_api_ret_code[0] != 0)) return 0;
481504
}
482505
// call the function
483506
if (release_gil) {
@@ -491,7 +514,8 @@ class TVMFFIPyCallManager {
491514
// restore the original stream
492515
if (ctx.device_type != -1 && prev_stream != ctx.stream) {
493516
// always try recover first, even if error happens
494-
if (TVMFFIEnvSetStream(ctx.device_type, ctx.device_id, prev_stream, nullptr) != 0) {
517+
if (TVM_FFI_PREDICT_FALSE(
518+
TVMFFIEnvSetStream(ctx.device_type, ctx.device_id, prev_stream, nullptr) != 0)) {
495519
// recover failed, set python error
496520
PyErr_SetString(PyExc_RuntimeError, "Failed to recover stream");
497521
return -1;
@@ -502,12 +526,13 @@ class TVMFFIPyCallManager {
502526
prev_tensor_allocator != ctx.dlpack_c_exchange_api->managed_tensor_allocator) {
503527
// note: we cannot set the error value to c_api_ret_code[0] here because it
504528
// will be overwritten by the error value from the function call
505-
if (TVMFFIEnvSetDLPackManagedTensorAllocator(prev_tensor_allocator, 0, nullptr) != 0) {
529+
if (TVM_FFI_PREDICT_FALSE(
530+
TVMFFIEnvSetDLPackManagedTensorAllocator(prev_tensor_allocator, 0, nullptr) != 0)) {
506531
PyErr_SetString(PyExc_RuntimeError, "Failed to recover DLPack managed tensor allocator");
507532
return -1;
508533
}
509534
// return error after
510-
if (c_api_ret_code[0] != 0) return 0;
535+
if (TVM_FFI_PREDICT_FALSE(c_api_ret_code[0] != 0)) return 0;
511536
}
512537
if (optional_out_ctx_dlpack_api != nullptr && ctx.dlpack_c_exchange_api != nullptr) {
513538
*optional_out_ctx_dlpack_api = ctx.dlpack_c_exchange_api;
@@ -540,15 +565,15 @@ class TVMFFIPyCallManager {
540565
TVM_FFI_INLINE int ConstructorCall(void* func_handle, PyObject* py_arg_tuple, TVMFFIAny* result,
541566
int* c_api_ret_code, TVMFFIPyCallContext* parent_ctx) {
542567
int64_t num_args = PyTuple_Size(py_arg_tuple);
543-
if (num_args == -1) return -1;
568+
if (TVM_FFI_PREDICT_FALSE(num_args == -1)) return -1;
544569
try {
545570
// allocate a call stack
546571
TVMFFIPyCallContext ctx(&call_stack_, num_args);
547572
// Iterate over the arguments and set them
548573
for (int64_t i = 0; i < num_args; ++i) {
549574
PyObject* py_arg = PyTuple_GetItem(py_arg_tuple, i);
550575
TVMFFIAny* c_arg = ctx.packed_args + i;
551-
if (SetArgument(&ctx, py_arg, c_arg) != 0) return -1;
576+
if (TVM_FFI_PREDICT_FALSE(SetArgument(&ctx, py_arg, c_arg) != 0)) return -1;
552577
}
553578
c_api_ret_code[0] = TVMFFIFunctionCall(func_handle, ctx.packed_args, num_args, result);
554579
// propagate the call context to the parent context
@@ -577,7 +602,7 @@ class TVMFFIPyCallManager {
577602
try {
578603
TVMFFIPyCallContext ctx(&call_stack_, 1);
579604
TVMFFIAny* c_arg = ctx.packed_args;
580-
if (SetArgument(&ctx, py_arg, c_arg) != 0) return -1;
605+
if (TVM_FFI_PREDICT_FALSE(SetArgument(&ctx, py_arg, c_arg) != 0)) return -1;
581606
if (!(field_flags & kTVMFFIFieldFlagBitSetterIsFunctionObj)) {
582607
auto setter = reinterpret_cast<TVMFFIFieldSetter>(field_setter);
583608
c_api_ret_code[0] = (*setter)(field_ptr, c_arg);
@@ -603,7 +628,7 @@ class TVMFFIPyCallManager {
603628
try {
604629
TVMFFIPyCallContext ctx(&call_stack_, 1);
605630
TVMFFIAny* c_arg = ctx.packed_args;
606-
if (SetArgument(&ctx, py_arg, c_arg) != 0) return -1;
631+
if (TVM_FFI_PREDICT_FALSE(SetArgument(&ctx, py_arg, c_arg) != 0)) return -1;
607632
c_api_ret_code[0] = TVMFFIAnyViewToOwnedAny(c_arg, out);
608633
return 0;
609634
} catch (const std::exception& ex) {
@@ -629,20 +654,20 @@ class TVMFFIPyCallManager {
629654
// find the pre-cached setter
630655
// This class is thread-local, so we don't need to worry about race condition
631656
auto it = arg_dispatch_map_.find(py_type);
632-
if (it != arg_dispatch_map_.end()) {
657+
if (TVM_FFI_PREDICT_TRUE(it != arg_dispatch_map_.end())) {
633658
TVMFFIPyArgSetter setter = it->second;
634659
// if error happens, propagate it back
635-
if (setter(ctx, py_arg, out) != 0) return -1;
660+
if (TVM_FFI_PREDICT_FALSE(setter(ctx, py_arg, out) != 0)) return -1;
636661
} else {
637662
// no dispatch found, query and create a new one.
638663
TVMFFIPyArgSetter setter;
639664
// propagate python error back
640-
if (TVMFFICyArgSetterFactory(py_arg, &setter) != 0) {
665+
if (TVM_FFI_PREDICT_FALSE(TVMFFICyArgSetterFactory(py_arg, &setter) != 0)) {
641666
return -1;
642667
}
643668
// update dispatch table
644669
arg_dispatch_map_.emplace(py_type, setter);
645-
if (setter(ctx, py_arg, out) != 0) return -1;
670+
if (TVM_FFI_PREDICT_FALSE(setter(ctx, py_arg, out) != 0)) return -1;
646671
}
647672
return 0;
648673
}
@@ -706,8 +731,8 @@ class TVMFFIPyCallManager {
706731
TVMFFIPyCallbackContext cb_ctx(&call_stack_, num_args);
707732
// Step 1: Convert each packed arg (borrowed AnyView) to a PyObject*
708733
for (int32_t i = 0; i < num_args; ++i) {
709-
if (SetPyCallbackArg(closure->dlpack_exchange_api, &packed_args[i], &cb_ctx.py_args[i]) !=
710-
0) {
734+
if (TVM_FFI_PREDICT_FALSE(SetPyCallbackArg(closure->dlpack_exchange_api, &packed_args[i],
735+
&cb_ctx.py_args[i]) != 0)) {
711736
ForwardPyErrorToFFI();
712737
return -1;
713738
}
@@ -749,7 +774,7 @@ class TVMFFIPyCallManager {
749774
// The guard's destructor runs AFTER the return value is computed.
750775
TVMFFIPyCallContext ret_ctx(&call_stack_, 1);
751776
TVMFFIAny* view = ret_ctx.packed_args;
752-
if (SetArgument(&ret_ctx, py_result.p, view) != 0) {
777+
if (TVM_FFI_PREDICT_FALSE(SetArgument(&ret_ctx, py_result.p, view) != 0)) {
753778
ForwardPyErrorToFFI();
754779
return -1;
755780
}
@@ -776,7 +801,7 @@ class TVMFFIPyCallManager {
776801
* returned by PyErr_Occurred()) so that set_last_ffi_error can access the
777802
* message and traceback.
778803
*/
779-
static void ForwardPyErrorToFFI() noexcept {
804+
TVM_FFI_COLD_CODE static void ForwardPyErrorToFFI() noexcept {
780805
#if PY_VERSION_HEX >= 0x030C0000
781806
// Python 3.12+: PyErr_Fetch / PyErr_NormalizeException are deprecated.
782807
// PyErr_GetRaisedException returns an already-normalized exception

src/ffi/backtrace.cc

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,7 @@ const TVMFFIByteArray* TVMFFIBacktrace(const char* filename, int lineno, const c
147147
}
148148

149149
#if TVM_FFI_BACKTRACE_ON_SEGFAULT
150+
TVM_FFI_COLD_CODE
150151
void TVMFFISegFaultHandler(int sig) {
151152
// Technically we shouldn't do any allocation in a signal handler, but
152153
// Backtrace may allocate. What's the worst it could do? We're already
@@ -163,6 +164,7 @@ void TVMFFISegFaultHandler(int sig) {
163164
raise(sig);
164165
}
165166

167+
TVM_FFI_COLD_CODE
166168
__attribute__((constructor)) void TVMFFIInstallSignalHandler() {
167169
// this may override already installed signal handlers
168170
std::signal(SIGSEGV, TVMFFISegFaultHandler);

src/ffi/error.cc

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
* \file src/ffi/error.cc
2121
* \brief Error handling implementation
2222
*/
23+
#include <tvm/ffi/base_details.h>
2324
#include <tvm/ffi/c_api.h>
2425
#include <tvm/ffi/error.h>
2526

src/ffi/function.cc

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,7 @@ class GlobalFunctionTable {
8383
};
8484

8585
void Update(const String& name, Function func, bool can_override) {
86-
if (table_.count(name)) {
86+
if (TVM_FFI_PREDICT_FALSE(table_.count(name) != 0)) {
8787
if (!can_override) {
8888
TVM_FFI_THROW(RuntimeError) << "Global Function `" << name << "` is already registered";
8989
}
@@ -93,7 +93,7 @@ class GlobalFunctionTable {
9393

9494
void Update(const TVMFFIMethodInfo* method_info, bool can_override) {
9595
String name(method_info->name.data, method_info->name.size);
96-
if (table_.count(name)) {
96+
if (TVM_FFI_PREDICT_FALSE(table_.count(name) != 0)) {
9797
if (!can_override) {
9898
TVM_FFI_LOG_AND_THROW(RuntimeError)
9999
<< "Global Function `" << name << "` is already registered, possible causes:\n"

0 commit comments

Comments
 (0)