Skip to content

Commit

Permalink
Update cython bindings to use kTVMArgBool
Browse files Browse the repository at this point in the history
  • Loading branch information
Lunderberg committed Apr 22, 2024
1 parent 09bbb80 commit b0bc242
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 5 deletions.
5 changes: 4 additions & 1 deletion python/tvm/_ffi/_cython/base.pxi
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
# under the License.

from ..base import raise_last_ffi_error
from libcpp cimport bool as bool_t
from libcpp.vector cimport vector
from cpython.version cimport PY_MAJOR_VERSION
from cpython cimport pycapsule
Expand All @@ -38,7 +39,8 @@ cdef enum TVMArgTypeCode:
kTVMBytes = 12
kTVMNDArrayHandle = 13
kTVMObjectRefArg = 14
kTVMExtBegin = 15
kTVMArgBool = 15
kTVMExtBegin = 16

cdef extern from "tvm/runtime/c_runtime_api.h":
ctypedef struct DLDataType:
Expand Down Expand Up @@ -66,6 +68,7 @@ cdef extern from "tvm/runtime/c_runtime_api.h":

ctypedef struct TVMValue:
int64_t v_int64
bool_t v_bool
double v_float64
void* v_handle
const char* v_str
Expand Down
8 changes: 4 additions & 4 deletions python/tvm/_ffi/_cython/packed_func.pxi
Original file line number Diff line number Diff line change
Expand Up @@ -121,10 +121,8 @@ cdef inline int make_arg(object arg,
elif isinstance(arg, bool):
# A python `bool` is a subclass of `int`, so this check
# must occur before `Integral`.
arg = _FUNC_CONVERT_TO_OBJECT(arg)
value[0].v_handle = (<ObjectBase>arg).chandle
tcode[0] = kTVMObjectHandle
temp_args.append(arg)
value[0].v_bool = arg
tcode[0] = kTVMArgBool
elif isinstance(arg, Integral):
value[0].v_int64 = arg
tcode[0] = kInt
Expand Down Expand Up @@ -216,6 +214,8 @@ cdef inline object make_ret(TVMValue value, int tcode):
return make_ret_object(value.v_handle)
elif tcode == kTVMNullptr:
return None
elif tcode == kTVMArgBool:
return value.v_bool
elif tcode == kInt:
return value.v_int64
elif tcode == kFloat:
Expand Down

0 comments on commit b0bc242

Please sign in to comment.