gemini-code-assist[bot] commented on code in PR #665:
URL: https://github.com/apache/tvm-ffi/pull/665#discussion_r3574238593
##########
include/tvm/ffi/type_traits.h:
##########
@@ -536,6 +537,88 @@ struct TypeTraits<void*> : public TypeTraitsBase {
}
};
+namespace details {
+
+TVM_FFI_INLINE static std::optional<std::string_view> TryGetStringViewFromAny(
+ const TVMFFIAny* src) {
+ if (src->type_index == TypeIndex::kTVMFFIRawStr) {
+ return std::string_view(src->v_c_str);
+ }
+ if (src->type_index == TypeIndex::kTVMFFISmallStr) {
+ TVMFFIByteArray bytes = TVMFFISmallBytesGetContentByteArray(src);
+ return std::string_view(bytes.data, bytes.size);
+ }
+ if (src->type_index == TypeIndex::kTVMFFIStr) {
+ TVMFFIByteArray* bytes = TVMFFIBytesGetByteArrayPtr(src->v_obj);
+ return std::string_view(bytes->data, bytes->size);
+ }
+ return std::nullopt;
+}
+
+TVM_FFI_INLINE static std::optional<DLDeviceType>
TryParseDLDeviceType(std::string_view name) {
+ if (name == "llvm" || name == "cpu" || name == "c" || name == "test") return
kDLCPU;
+ if (name == "cuda" || name == "nvptx") return kDLCUDA;
+ if (name == "cl" || name == "opencl") return kDLOpenCL;
+ if (name == "vulkan") return kDLVulkan;
+ if (name == "metal") return kDLMetal;
+ if (name == "vpi") return kDLVPI;
+ if (name == "rocm") return kDLROCM;
+ if (name == "ext_dev") return kDLExtDev;
+ if (name == "hexagon") return kDLHexagon;
+ if (name == "webgpu") return kDLWebGPU;
+ if (name == "maia") return kDLMAIA;
+ if (name == "trn") return kDLTrn;
+ return std::nullopt;
+}
+
+TVM_FFI_INLINE static std::optional<int32_t>
TryParseDLDeviceIndex(std::string_view index) {
+ if (index.empty()) return std::nullopt;
+ int64_t sign = 1;
+ size_t pos = 0;
+ if (index[0] == '-') {
+ sign = -1;
+ pos = 1;
+ }
+ if (pos >= index.size()) return std::nullopt;
+ int64_t value = 0;
+ for (; pos < index.size(); ++pos) {
+ char ch = index[pos];
+ if (ch < '0' || ch > '9') return std::nullopt;
+ value = value * 10 + (ch - '0');
+ if (value > std::numeric_limits<int32_t>::max()) return std::nullopt;
+ }
Review Comment:
The loop check `value > std::numeric_limits<int32_t>::max()` prematurely
rejects the valid `int32_t` minimum value `-2147483648` when the sign is
negative (since `2147483648` is greater than `2147483647`). To support the full
range of `int32_t` including its minimum value, the absolute value limit check
inside the loop should allow up to `2147483648LL`.
```c
for (; pos < index.size(); ++pos) {
char ch = index[pos];
if (ch < '0' || ch > '9') return std::nullopt;
value = value * 10 + (ch - '0');
if (value > 2147483648LL) return std::nullopt;
}
```
##########
python/tvm_ffi/dataclasses/enum.py:
##########
@@ -593,511 +259,256 @@ def __init__(self, enum_cls: type, name: str, *,
default: Any = core.MISSING) ->
self._name = name
self._default = default
- def _ordinal_of(self, variant: object) -> int:
+ def _check(self, variant: object) -> None:
if not isinstance(variant, self._enum_cls):
- raise TypeError(
- f"{self._enum_cls.__name__}.def_attr({self._name!r}) expects a
"
- f"{self._enum_cls.__name__} variant, got
{type(variant).__name__}"
- )
- return int(variant._value) # type: ignore[attr-defined]
-
- def _column(self, *, create: bool) -> Any | None:
- """Return the ``List[Any]`` column for this attribute; create if
missing.
+ raise TypeError(f"expected {self._enum_cls.__name__}, got
{type(variant).__name__}")
- Returns ``None`` when ``create`` is false and the column doesn't exist.
- """
- attrs = _attrs_dict(self._enum_cls) if not create else
_ensure_attrs_dict(self._enum_cls)
- if attrs is None:
+ def _column(self, create: bool) -> Any:
+ state = _state(self._enum_cls, create=create)
+ if state is None:
return None
- if self._name in attrs:
- return attrs[self._name]
- if not create:
- return None
- col = List([])
- attrs[self._name] = col
- return col
+ attrs = state.attrs
+ column = attrs.get(self._name)
+ if column is None and create:
+ column = Dict({})
+ attrs[self._name] = column
+ return column
def __setitem__(self, variant: object, value: Any) -> None:
- if value is None:
- raise TypeError(
- f"{self._enum_cls.__name__}.def_attr({self._name!r}): "
- f"None is reserved as the 'unset' sentinel for extensible "
- f"attributes and cannot be written explicitly."
- )
- ordinal = self._ordinal_of(variant)
- col = self._column(create=True)
- assert col is not None # create=True always materialises the column.
- while len(col) <= ordinal:
- col.append(None)
- col[ordinal] = value
+ self._check(variant)
+ self._column(True)[variant] = value
def __getitem__(self, variant: object) -> Any:
- ordinal = self._ordinal_of(variant)
- col = self._column(create=False)
- if col is not None and ordinal < len(col):
- v = col[ordinal]
- if v is not None:
- return v
- if self._default is core.MISSING:
- raise KeyError(
- f"{self._enum_cls.__name__}.{variant._name} has no " # type:
ignore[attr-defined]
- f"extensible attribute {self._name!r} set"
- )
- return self._default
+ self._check(variant)
+ column = self._column(False)
+ if column is not None and variant in column:
+ return column[variant]
+ if self._default is not core.MISSING:
+ return self._default
+ raise KeyError(f"{variant!r} has no extensible attribute
{self._name!r}")
def __contains__(self, variant: object) -> bool:
if not isinstance(variant, self._enum_cls):
return False
- try:
- ordinal = self._ordinal_of(variant)
- except TypeError:
- return False
- col = self._column(create=False)
- return col is not None and ordinal < len(col) and col[ordinal] is not
None
+ column = self._column(False)
+ return column is not None and variant in column
def get(self, variant: object, default: Any = None) -> Any:
- """Return the value for *variant*, or *default* if unset or foreign."""
- if not isinstance(variant, self._enum_cls):
- return default
- try:
- return self[variant]
- except KeyError:
+ """Return a stored value or *default*."""
+ if variant not in self:
return default
+ return self[variant]
@property
def name(self) -> str:
- """The extensible-attribute name passed to :meth:`Enum.def_attr`."""
+ """Return the extensible-attribute name."""
return self._name
-class IntEnum(Enum):
- """Enum variant base with a user-visible ``value: int`` field."""
-
- __slots__ = ()
- value: int
-
- def __init_subclass__(
- cls,
- *,
- type_key: str | None = None,
- frozen: bool = True,
- init: bool = True,
- repr: bool = True,
- **kwargs: Any,
- ) -> None:
- _prepare_payload_enum_subclass(cls, value_type=int,
base_name="IntEnum")
- super().__init_subclass__(
- type_key=type_key,
- frozen=frozen,
- init=init,
- repr=repr,
- **kwargs,
- )
+def _install_payload_behavior(cls: type, payload_type: type, user_repr: bool)
-> None:
+ def eq(self: Enum, other: object) -> Any:
+ if isinstance(other, type(self)):
+ return self.value == other.value
+ if isinstance(other, payload_type):
+ return self.value == other
+ return NotImplemented
+ def ne(self: Enum, other: object) -> Any:
+ result = eq(self, other)
+ return NotImplemented if result is NotImplemented else not result
-class StrEnum(Enum):
- """Enum variant base with a user-visible ``value: str`` field."""
+ def name(self: Enum) -> str:
+ return str(self._str_index)
- __slots__ = ()
- value: str
+ defaults = {
+ "__eq__": eq,
+ "__ne__": ne,
+ "__str__": lambda self: str(self.value),
+ "__hash__": lambda self: hash(self.value),
+ }
+ for key, value in defaults.items():
+ if key not in cls.__dict__:
+ setattr(cls, key, value)
+ if "name" not in cls.__dict__:
+ cls.name = property(name) # ty: ignore[unresolved-attribute]
+ if not user_repr:
+ cls.__repr__ = lambda self: f"{type(self).__name__}.{self.name}" #
type: ignore[attr-defined]
- def __init_subclass__(
- cls,
- *,
- type_key: str | None = None,
- frozen: bool = True,
- init: bool = True,
- repr: bool = True,
- **kwargs: Any,
- ) -> None:
- _prepare_payload_enum_subclass(cls, value_type=str,
base_name="StrEnum")
- super().__init_subclass__(
- type_key=type_key,
- frozen=frozen,
- init=init,
- repr=repr,
- **kwargs,
- )
+@_enum_c_class("ffi.IntEnum", init=False)
+class IntEnum(Enum):
+ """Enum whose ``value`` aliases its signed 64-bit integer index."""
-# ---------------------------------------------------------------------------
-# TypeAttr accessors
-# ---------------------------------------------------------------------------
+ if not TYPE_CHECKING:
+ __slots__ = ()
+ value: int
+ if TYPE_CHECKING:
-def _entries_dict(cls: type) -> Any:
- type_info = getattr(cls, "__tvm_ffi_type_info__", None)
- if type_info is None:
- return None
- return core._lookup_type_attr(type_info.type_index, ENUM_ENTRIES_ATTR)
+ @property
+ def value(self) -> int:
+ """Return the canonical integer index."""
+ ...
+ def __init_subclass__(cls, **kwargs: Any) -> None:
+ _reject_value_field(cls, "IntEnum")
+ super().__init_subclass__(**kwargs)
-def _attrs_dict(cls: type) -> Any:
- type_info = getattr(cls, "__tvm_ffi_type_info__", None)
- if type_info is None:
- return None
- return core._lookup_type_attr(type_info.type_index, ENUM_ATTRS_ATTR)
+@_enum_c_class("ffi.StrEnum", init=False)
+class StrEnum(Enum):
+ """Enum whose ``name`` and ``value`` alias its string index."""
-def _value_entries_dict(cls: type) -> Any:
- """Return the live ``Dict[Any, Enum]`` payload-to-variant index, or
None."""
- type_info = getattr(cls, "__tvm_ffi_type_info__", None)
- if type_info is None:
- return None
- return core._lookup_type_attr(type_info.type_index,
ENUM_VALUE_ENTRIES_ATTR)
+ if not TYPE_CHECKING:
+ __slots__ = ()
+ value: str
+ if TYPE_CHECKING:
-def _ordered_entries(cls: type) -> list[Any]:
- """Return all variants ordered by ordinal."""
- entries = _entries_dict(cls)
- if entries is None:
- return []
- ordered: list[Any] = [None] * len(entries)
- for inst in entries.values():
- ordered[int(inst._value)] = inst
- return ordered
-
-
-def _ensure_entries_dict(cls: type) -> Any:
- """Return the live ``__ffi_enum_entries__`` dict, registering it if
absent."""
- type_info = cls.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute]
- entries = core._lookup_type_attr(type_info.type_index, ENUM_ENTRIES_ATTR)
- if entries is not None:
- return entries
- entries = Dict({})
- core._register_type_attr(type_info.type_index, ENUM_ENTRIES_ATTR, entries)
- # Re-read so mutations go through the ref owned by the registry.
- return core._lookup_type_attr(type_info.type_index, ENUM_ENTRIES_ATTR)
-
-
-def _ensure_attrs_dict(cls: type) -> Any:
- """Return the live ``__ffi_enum_attrs__`` dict, registering it if
absent."""
- type_info = cls.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute]
- attrs = core._lookup_type_attr(type_info.type_index, ENUM_ATTRS_ATTR)
- if attrs is not None:
- return attrs
- attrs = Dict({})
- core._register_type_attr(type_info.type_index, ENUM_ATTRS_ATTR, attrs)
- return core._lookup_type_attr(type_info.type_index, ENUM_ATTRS_ATTR)
-
-
-def _ensure_value_entries_dict(cls: type) -> Any:
- """Return the live ``__ffi_enum_value_entries__`` dict, registering it if
absent."""
- type_info = cls.__tvm_ffi_type_info__ # ty: ignore[unresolved-attribute]
- entries = core._lookup_type_attr(type_info.type_index,
ENUM_VALUE_ENTRIES_ATTR)
- if entries is not None:
- return entries
- entries = Dict({})
- core._register_type_attr(type_info.type_index, ENUM_VALUE_ENTRIES_ATTR,
entries)
- return core._lookup_type_attr(type_info.type_index,
ENUM_VALUE_ENTRIES_ATTR)
-
-
-# ---------------------------------------------------------------------------
-# Class-body scanning + entry materialisation
-# ---------------------------------------------------------------------------
-
-
-def _collect_entry_declarations(
- cls: type,
-) -> tuple[list[str], dict[str, _EnumEntry]]:
- """Scan *cls.__dict__* for variant declarations.
+ @property
+ def value(self) -> str:
+ """Return the canonical string index."""
+ ...
- Returns ``(binders, python_entries)`` in declaration order:
+ def __init_subclass__(cls, **kwargs: Any) -> None:
+ _reject_value_field(cls, "StrEnum")
+ super().__init_subclass__(**kwargs)
- * *binders* — names annotated as ``ClassVar[Cls]`` with no assigned value.
- Each either binds to an existing C++-registered entry with the same
- name or registers a new blank Python entry.
- * *python_entries* — names assigned an ``entry(...)`` sentinel (with or
- without a ``ClassVar`` annotation). Each registers a new Python entry
- using the captured args/kwargs.
- Matched assignments are removed from ``cls.__dict__`` so that
- ``@c_class`` / ``@py_class`` don't misinterpret them as field defaults or
- class constants.
- """
+def _collect_declarations(
+ cls: type, payload_type: type | None
+) -> tuple[list[str], dict[str, _EnumEntry]]:
annotations = _own_annotations(cls)
- dict_keys = set(cls.__dict__.keys())
-
- binders: list[str] = []
- ordinary_fields: set[str] = set()
- payload_value_type = getattr(cls, "__ffi_enum_payload_value_type__", None)
- for name, ann in annotations.items():
- if name.startswith("_"):
- continue
- if _is_class_var(ann):
- if name not in dict_keys:
- binders.append(name)
- else:
- ordinary_fields.add(name)
-
- python_entries: dict[str, _EnumEntry] = {}
+ binders = [
+ name
+ for name, annotation in annotations.items()
+ if not name.startswith("_") and _is_class_var(annotation) and name not
in cls.__dict__
+ ]
+ fields = {name for name, annotation in annotations.items() if not
_is_class_var(annotation)}
+ declarations: dict[str, _EnumEntry] = {}
for name, value in list(cls.__dict__.items()):
- if name.startswith("_"):
+ if name.startswith("_") or name in fields:
continue
if isinstance(value, _EnumEntry):
- python_entries[name] = value
- try:
- delattr(cls, name)
- except AttributeError:
- pass
- elif payload_value_type is not None and name not in ordinary_fields:
- if isinstance(value, (staticmethod, classmethod, property)) or
callable(value):
- continue
- python_entries[name] = _EnumEntry(value=value)
- try:
- delattr(cls, name)
- except AttributeError:
- pass
-
- return binders, python_entries
-
-
-def _resolve_entries(
- cls: type,
- binders: list[str],
- python_entries: dict[str, _EnumEntry],
- *,
- type_key: str,
- cxx_backed: bool,
-) -> None:
- """Materialise *binders* and *python_entries* into class-attribute
singletons.
-
- Processing order matches declaration order: ``binders`` first (because
- their annotations appear before any class-body assignments), then
- ``python_entries`` in their class-body order. Each newly registered
- entry gets a dense ordinal equal to the current entries-dict size, so
- ordinals stay compact and stable across registrations.
-
- A cxx-backed enum (``type_key`` was already registered in the FFI type
- system before this Python subclass was created) supports mixing C++ and
- Python entries: bare ``ClassVar[Cls]`` binders must name an existing
- C++-registered entry, but ``entry(...)``/``auto()`` sentinels may add
- fresh Python-side entries whose ordinals extend past the C++ entries.
- """
- entries = _ensure_entries_dict(cls)
- payload_value_type = getattr(cls, "__ffi_enum_payload_value_type__", None)
-
- for name in binders:
- if name in entries:
- # Already materialised — either C++-registered or previously bound.
- instance = entries[name]
- setattr(cls, name, instance)
- _index_payload(cls, instance, payload_value_type)
+ declarations[name] = value
+ elif payload_type is not None and not (
+ isinstance(value, (staticmethod, classmethod, property)) or
callable(value)
+ ):
+ declarations[name] = _EnumEntry(value=value)
+ else:
continue
- if cxx_backed:
- raise _cxx_backed_unknown_binder_error(cls, name, type_key,
entries)
- ordinal = len(entries)
- instance = _instantiate(cls, args=(), kwargs={}, ordinal=ordinal,
name=name)
- entries[name] = instance
- setattr(cls, name, instance)
- _index_payload(cls, instance, payload_value_type)
-
- for name, e in python_entries.items():
- if name in entries:
- raise RuntimeError(
- f"Duplicate enum entry {name!r} for {cls.__name__}: already "
- f"registered as ordinal {int(entries[name]._value)}."
- )
- if "_value" in e.kwargs or "_name" in e.kwargs:
- raise TypeError(
- f"{cls.__name__}.{name}: `_value` and `_name` are
auto-assigned "
- f"and must not appear in entry(...) arguments."
- )
- ordinal = len(entries)
- instance = _instantiate_entry(
- cls, entry=e, ordinal=ordinal, name=name, cxx_backed=cxx_backed
- )
- entries[name] = instance
- setattr(cls, name, instance)
- _index_payload(cls, instance, payload_value_type)
+ delattr(cls, name)
+ return binders, declarations
-def _index_payload(cls: type, instance: Any, payload_value_type: type | None)
-> None:
- """Record ``(instance.value → instance)`` in the value-entries column.
-
- No-op for non-payload enums. For payload enums, the first writer of a
- given payload wins — matches the "first-match" semantics of the linear
- scan that FFI converters fall back to when this column is incomplete.
- """
- if payload_value_type is None:
- return
- payload = getattr(instance, "value", None)
- if payload is None:
- return
- value_entries = _ensure_value_entries_dict(cls)
- if payload not in value_entries:
- value_entries[payload] = instance
-
-
-def _cxx_backed_unknown_binder_error(
+def _resolve(
cls: type,
- name: str,
+ binders: list[str],
+ declarations: dict[str, _EnumEntry],
type_key: str,
- entries: Any,
-) -> RuntimeError:
- """Build a descriptive error for an unbindable bare ``ClassVar`` binder.
-
- A bare ``ClassVar[Cls]`` annotation on a cxx-backed enum means "bind to
- an existing C++ entry with this name" — if the C++ registry has no such
- entry, the declaration is almost always a typo. For adding a *new*
- Python-side variant on a cxx-backed enum, use ``entry(...)`` or
- ``auto()`` instead.
- """
- known = list(entries.keys()) if entries is not None else []
- known_str = ", ".join(repr(k) for k in known) if known else "<none>"
- return RuntimeError(
- f"Cannot bind enum variant {name!r} on {cls.__name__}: the FFI "
- f"type {type_key!r} is already registered in C++ with entries "
- f"[{known_str}], but has no C++ entry named {name!r}. "
- f"Bare ``ClassVar[{cls.__name__}]`` binders on a C++-backed enum "
- f"must name an entry already registered in C++; they cannot "
- f"introduce new variants from Python. "
- f"If this was a typo, double-check the spelling against the known "
- f"entries above (`{name}: ClassVar[{cls.__name__}]`); if you meant "
- f"to add a new Python-side variant, use `{name} = auto()` or "
- f"`{name} = entry(...)` instead."
- )
-
-
-def _instantiate(
- cls: type,
- *,
- args: tuple[Any, ...],
- kwargs: dict[str, Any],
- ordinal: int,
- name: str,
-) -> Any:
- """Construct a subclass instance with auto-assigned
``_value``/``_name``."""
- merged = dict(kwargs)
- merged["_value"] = ordinal
- merged["_name"] = name
- return cls(*args, **merged)
-
-
-def _instantiate_entry(
+ cxx_backed: bool,
+ payload_type: type | None,
+) -> None:
+ state = _state(cls, create=True)
+ assert state is not None
+ entries, indexes = state.entries, state.indexes
+ for alias in binders:
+ instance = indexes.get(alias)
+ if instance is None:
+ if cxx_backed:
+ raise RuntimeError(f"{type_key!r} has no enum at string index
{alias!r}")
+ if payload_type is int:
+ raise TypeError(f"{cls.__name__}.{alias} requires an integer
value")
+ instance = _create(cls, len(entries), alias, _EnumEntry(),
cxx_backed)
+ _register(entries, indexes, instance)
+ setattr(cls, alias, instance)
+
+ for alias, declaration in declarations.items():
+ int_index, str_index, fields = _indices(cls, alias, len(entries),
declaration, payload_type)
+ lookup = int_index if payload_type is int else str_index
+ instance = indexes.get(lookup)
+ if instance is None:
+ if int_index in indexes or str_index in indexes:
+ raise ValueError(f"duplicate enum index in
{cls.__name__}.{alias}")
+ instance = _create(cls, int_index, str_index, fields, cxx_backed)
+ _register(entries, indexes, instance)
+ setattr(cls, alias, instance)
+
+
+def _indices(
cls: type,
- *,
- entry: _EnumEntry,
+ alias: str,
ordinal: int,
- name: str,
- cxx_backed: bool,
+ declaration: _EnumEntry,
+ payload_type: type | None,
+) -> tuple[int, str, _EnumEntry]:
+ kwargs = dict(declaration.kwargs)
+ if "_int_index" in kwargs or "_str_index" in kwargs:
+ raise TypeError(f"{cls.__name__}.{alias}: enum indices are inferred")
+ if payload_type is None:
+ return ordinal, alias, declaration
+ if "value" not in kwargs:
+ raise TypeError(f"{cls.__name__}.{alias} requires a
{payload_type.__name__} value")
+ value = kwargs.pop("value")
+ if type(value) is not payload_type:
+ raise TypeError(f"{cls.__name__}.{alias} value must be
{payload_type.__name__}")
+ if payload_type is int:
+ _normalize_index(value)
+ indices = value, alias
Review Comment:
There are two issues in this block:
1. Using `type(value) is not payload_type` is overly restrictive and rejects
subclasses of the payload type, such as `bool` for `IntEnum` (since `type(True)
is int` is `False`). Using `isinstance(value, payload_type)` is more Pythonic
and correctly supports subclasses.
2. The return value of `_normalize_index(value)` is ignored on line 456,
meaning `value` remains a `bool` (e.g., `True` or `False`) instead of being
converted to an `int` (`1` or `0`). This can lead to duplicate keys or type
mismatches in the index map since `True` and `1` share the same hash in Python.
We should use `isinstance` and assign the normalized index back to `value`.
```suggestion
if not isinstance(value, payload_type):
raise TypeError(f"{cls.__name__}.{alias} value must be
{payload_type.__name__}")
if payload_type is int:
value = _normalize_index(value)
indices = value, alias
```
--
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]
