[PYTHON] Further streamline number handling (#242)

This PR further streamlines number handling by introducing two custom protocols and move Integral and Real handling to more conservative path.

Author: tqchen

python/tvm_ffi/_convert.py

@@ -129,6 +129,10 @@ def convert(value: Any) -> Any: # noqa: PLR0911,PLR0912
  return value
  elif hasattr(value, "__dlpack_device__"):
  return value
+ elif hasattr(value, "__tvm_ffi_int__"):
+ return value
+ elif hasattr(value, "__tvm_ffi_float__"):
+ return value
  else:
  # in this case, it is an opaque python object
  return core._convert_to_opaque_object(value)

python/tvm_ffi/cython/function.pxi

@@ -17,7 +17,7 @@
 import ctypes
 import threading
 import os
-from numbers import Real, Integral
+from numbers import Integral, Real
 from typing import Any, Callable
 
 
@@ -276,7 +276,33 @@ cdef int TVMFFIPyArgSetterDLPack_(
  return 0
 
 
-cdef int TVMFFIPyArgSetterFFIObjectCompatible_(
+cdef int TVMFFIPyArgSetterIntegral_(
+ TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx,
+ PyObject* py_arg, TVMFFIAny* out
+) except -1:
+ """Setter for Integral"""
+ cdef object arg = <object>py_arg
+ out.type_index = kTVMFFIInt
+ # keep it in cython so it will also check for fallback cases
+ # where the arg is not exactly the int class
+ out.v_int64 = <long long>arg
+ return 0
+
+
+cdef int TVMFFIPyArgSetterReal_(
+ TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx,
+ PyObject* py_arg, TVMFFIAny* out
+) except -1:
+ """Setter for Real"""
+ cdef object arg = <object>py_arg
+ out.type_index = kTVMFFIFloat
+ # keep it in cython so it will also check for fallback cases
+ # where the arg is not exactly the float class
+ out.v_float64 = <double>arg
+ return 0
+
+
+cdef int TVMFFIPyArgSetterFFIObjectProtocol_(
  TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx,
  PyObject* py_arg, TVMFFIAny* out
 ) except -1:
@@ -608,6 +634,7 @@ cdef int TVMFFIPyArgSetterDTypeFromNumpy_(
  out.v_dtype = NUMPY_DTYPE_TO_DL_DATA_TYPE[py_obj]
  return 0
 
+
 cdef int TVMFFIPyArgSetterDLPackDataTypeProtocol_(
  TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx,
  PyObject* py_arg, TVMFFIAny* out
@@ -621,6 +648,29 @@ cdef int TVMFFIPyArgSetterDLPackDataTypeProtocol_(
  out.v_dtype.lanes = <long long>dltype_data_type[2]
  return 0
 
+
+cdef int TVMFFIPyArgSetterIntProtocol_(
+ TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx,
+ PyObject* py_arg, TVMFFIAny* out
+) except -1:
+ """Setter for class with __tvm_ffi_int__() method"""
+ cdef object arg = <object>py_arg
+ out.type_index = kTVMFFIInt
+ out.v_int64 = <long long>(arg.__tvm_ffi_int__())
+ return 0
+
+
+cdef int TVMFFIPyArgSetterFloatProtocol_(
+ TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx,
+ PyObject* py_arg, TVMFFIAny* out
+) except -1:
+ """Setter for class with __tvm_ffi_float__() method"""
+ cdef object arg = <object>py_arg
+ out.type_index = kTVMFFIFloat
+ out.v_float64 = <double>(arg.__tvm_ffi_float__())
+ return 0
+
+
 cdef _DISPATCH_TYPE_KEEP_ALIVE = set()
 cdef _DISPATCH_TYPE_KEEP_ALIVE_LOCK = threading.Lock()
 
@@ -668,7 +718,7 @@ cdef int TVMFFIPyArgSetterFactory_(PyObject* value, TVMFFIPyArgSetter* out) exce
  # can directly map to tvm ffi object
  # usually used for solutions that takes subclass of ffi.Object
  # as a member variable
- out.func = TVMFFIPyArgSetterFFIObjectCompatible_
+ out.func = TVMFFIPyArgSetterFFIObjectProtocol_
  return 0
  if os.environ.get("TVM_FFI_SKIP_C_DLPACK_EXCHANGE_API", "0") != "1":
  # Check for DLPackExchangeAPI struct (new approach)
@@ -698,10 +748,15 @@ cdef int TVMFFIPyArgSetterFactory_(PyObject* value, TVMFFIPyArgSetter* out) exce
  out.func = TVMFFIPyArgSetterBool_
  return 0
  if isinstance(arg, Integral):
- out.func = TVMFFIPyArgSetterInt_
+ # must occur before Real check
+ # cannot simply use TVMFFIPyArgSetterInt
+ # because Integral may not be exactly the int class
+ out.func = TVMFFIPyArgSetterIntegral_
  return 0
  if isinstance(arg, Real):
- out.func = TVMFFIPyArgSetterFloat_
+ # cannot simply use TVMFFIPyArgSetterFloat
+ # because Real may not be exactly the float class
+ out.func = TVMFFIPyArgSetterReal_
  return 0
  # dtype is a subclass of str, so this check must occur before str
  if isinstance(arg, _CLASS_DTYPE):
@@ -760,6 +815,12 @@ cdef int TVMFFIPyArgSetterFactory_(PyObject* value, TVMFFIPyArgSetter* out) exce
  # then it is a DLPack device protocol
  out.func = TVMFFIPyArgSetterDLPackDeviceProtocol_
  return 0
+ if hasattr(arg_class, "__tvm_ffi_int__"):
+ out.func = TVMFFIPyArgSetterIntProtocol_
+ return 0
+ if hasattr(arg_class, "__tvm_ffi_float__"):
+ out.func = TVMFFIPyArgSetterFloatProtocol_
+ return 0
  if isinstance(arg, Exception):
  out.func = TVMFFIPyArgSetterException_
  return 0

tests/python/test_function.py

@@ -350,3 +350,36 @@ def __dlpack_device__(self) -> tuple[int, int]:
  x = DLPackDeviceProtocol(device)
  y = fecho(x)
  assert y == device
+
+
+def test_integral_float_variants_passing() -> None:
+ fecho = tvm_ffi.get_global_func("testing.echo")
+ y = fecho(np.int32(1))
+ assert isinstance(y, int)
+ assert y == 1
+
+ y = fecho(np.float64(2.0))
+ assert isinstance(y, float)
+ assert y == 2.0
+
+ class IntProtocol:
+ def __init__(self, value: int) -> None:
+ self.value = value
+
+ def __tvm_ffi_int__(self) -> int:
+ return self.value
+
+ y = fecho(IntProtocol(10))
+ assert isinstance(y, int)
+ assert y == 10
+
+ class FloatProtocol:
+ def __init__(self, value: float) -> None:
+ self.value = value
+
+ def __tvm_ffi_float__(self) -> float:
+ return self.value
+
+ y = fecho(FloatProtocol(10))
+ assert isinstance(y, float)
+ assert y == 10