gemini-code-assist[bot] commented on code in PR #178:
URL: https://github.com/apache/tvm-ffi/pull/178#discussion_r2445809920
##########
python/tvm_ffi/cython/function.pxi:
##########
@@ -573,6 +573,18 @@ cdef int TVMFFIPyArgSetterDTypeFromNumpy_(
out.v_dtype = NUMPY_DTYPE_TO_DTYPE[py_obj]
return 0
+cdef int TVMFFIPyArgSetterDLPackDataTypeProtocol_(
+ TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx,
+ PyObject* py_arg, TVMFFIAny* out
+) except -1:
+ """Setter for dtype protocol"""
+ cdef object arg = <object>py_arg
+ cdef tuple dltype_data_type = arg.__dlpack_data_type__()
+ out.type_index = kTVMFFIDataType
+ out.v_dtype.code = <long long>dltype_data_type[0]
+ out.v_dtype.bits = <long long>dltype_data_type[1]
+ out.v_dtype.lanes = <long long>dltype_data_type[2]
Review Comment:
The `DLDataType` struct defines `code` and `bits` as `uint8_t`, and `lanes`
as `uint16_t`. Casting the Python integers from `dltype_data_type` to `long
long` before assigning them to these smaller integer types is misleading and
potentially unsafe. It's better to cast to the specific C types for type safety
and clarity.
```
out.v_dtype.code = <uint8_t>dltype_data_type[0]
out.v_dtype.bits = <uint8_t>dltype_data_type[1]
out.v_dtype.lanes = <uint16_t>dltype_data_type[2]
```
##########
python/tvm_ffi/_dtype.py:
##########
@@ -60,15 +60,39 @@ class dtype(str):
"""
- __slots__ = ["__tvm_ffi_dtype__"]
- __tvm_ffi_dtype__: core.DataType
+ __slots__ = ["_tvm_ffi_dtype"]
+ _tvm_ffi_dtype: core.DataType
_NUMPY_DTYPE_TO_STR: ClassVar[dict[Any, str]] = {}
def __new__(cls, content: Any) -> dtype:
content = str(content)
val = str.__new__(cls, content)
- val.__tvm_ffi_dtype__ = core.DataType(content)
+ val._tvm_ffi_dtype = core.DataType(content)
+ return val
+
+ @staticmethod
+ def from_dlpack_data_type(dltype_data_type: tuple[int, int, int]) -> dtype:
+ """Create a dtype from a DLPack data type tuple.
+
+ Parameters
+ ----------
+ dltype_data_type
+ The DLPack data type tuple (type_code, bits, lanes).
+
+ Returns
+ -------
+ The created dtype.
+
+ """
+ cdtype = core._create_dtype_from_tuple(
+ core.DataType,
+ dltype_data_type[0],
+ dltype_data_type[1],
+ dltype_data_type[2],
+ )
+ val = str.__new__(dtype, str(cdtype))
+ val._tvm_ffi_dtype = cdtype
return val
Review Comment:
This logic for creating a `dtype` instance from a `core.DataType` is
duplicated in the `with_lanes` method. To improve maintainability, consider
extracting this into a private class method.
For example:
```python
@classmethod
def _from_core_datatype(cls, cdtype: core.DataType) -> "dtype":
"""Create a dtype from a core.DataType instance."""
val = str.__new__(cls, str(cdtype))
val._tvm_ffi_dtype = cdtype
return val
```
Then `from_dlpack_data_type` and `with_lanes` can be simplified:
```python
# In from_dlpack_data_type
return dtype._from_core_datatype(cdtype)
# In with_lanes
return type(self)._from_core_datatype(cdtype)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
