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
0 commit comments