Skip to content

Commit 5a8dae4

Browse files
AharrypotterAharrypotter
authored
[Relax][Frontend][TFLite] Support static hashtable find (#19879)
## Summary This PR extends the TFLite resource hashtable family from #19519 item G from "import + size" to a constant-foldable `HASHTABLE_FIND` subset. It builds on the static `HASHTABLE` / `HASHTABLE_IMPORT` import support added in #19639 and the `HASHTABLE_LOOKUP` converter from #19654. Relax has no string tensor type and no runtime hashtable operator, so this PR targets the only subset that can be lowered today: tables whose keys and values are constants imported in a `CALL_ONCE` init subgraph, queried by a constant string tensor. In that case the lookup is resolved at import time and emitted as a `relax.const`. Runtime string queries remain explicitly guarded instead of being lowered with incorrect semantics. ## Design ### Static String Tensor Decoding The frontend now decodes constant TFLite string tensors through a shared helper `_get_string_tensor_value`. It parses the TFLite string-tensor binary layout: ```text [count(i32)][offset_0 .. offset_count(i32)][utf8 string data] ``` The helper validates the buffer length against the declared count, checks that offsets are in-bounds and monotonically non-decreasing, and verifies the decoded element count matches the tensor shape. This is reusable infrastructure for any future TFLite string support, independent of `HASHTABLE_FIND`. ### Hashtable Import State `HASHTABLE_IMPORT` now stores the actual constant keys and values (numeric or decoded string buffers) in shared conversion state, instead of only recording table metadata (size, key/value dtype). Duplicate keys are rejected, because the constant-fold lookup assumes a unique key mapping. The captured table is keyed by the same `table_id` / handle resolution used by `HASHTABLE` and `HASHTABLE_SIZE`, so a `CALL_ONCE` init subgraph and the main graph agree on the same logical table. ### Constant-Foldable Find `HASHTABLE_FIND` resolves the table handle through the importer-local handle map and the statically imported keys/values. For the supported subset it builds a Python key -> value map, applies the per-element default value, overwrites hits from the table, preserves the query tensor shape, and emits the result as a `relax.const`. The op produces no runtime Relax computation, which matches the fact that both the table and the query are compile-time constants. The supported subset is intentionally narrow and guard-first: - the table must be a constant `string -> int64` table imported via a supported `CALL_ONCE` `HASHTABLE_IMPORT` - the query tensor must be a constant string buffer - the default tensor must be a scalar or match the query shape ### String Graph Input Guard `TensorType.STRING` graph inputs are now rejected in `from_tflite` with a clear `OpNotImplemented` instead of a low-level FFI `unknown dtype string` error, since Relax cannot represent a string tensor. This is the path hit by runtime string queries for `HASHTABLE_FIND` and string-valued `HASHTABLE_LOOKUP`, so the guard gives a frontend-level diagnostic for both. ## Operator Support | Operator | TFLite options | Relax lowering | Supported subset | |---|---|---|---| | `HASHTABLE_IMPORT` | `HashtableImportOptions` | store constant keys/values in importer state | `CALL_ONCE` init, constant keys/values, no duplicate keys | | `HASHTABLE_FIND` | `HashtableFindOptions` | constant-fold to `relax.const` | constant `string -> int64` table + constant string query | ## Not Included - Runtime (non-constant) string queries and runtime hashtable lookup. - `int64 -> string` find, which would require string-typed Relax outputs. - Any general `TensorType.STRING` tensor representation in Relax. - A runtime Relax hashtable operator with string hashing / comparison. - Mutable runtime resource-state threading through Relax functions. These require core Relax support and are out of scope for the frontend. ## Tests The tests manually build minimal TFLite flatbuffers and compare the imported Relax IR with `tvm.ir.assert_structural_equal`. Unsupported patterns use `pytest.raises`. | Test | Coverage | |---|---| | `test_hashtable_call_once_import_find_string_to_int64` | constant `string -> int64` find folds to a `relax.const` | | `test_hashtable_call_once_import_find_string_to_int64_2d_query` | query shape preserved for a 2-D query | | `test_hashtable_call_once_import_find_int64_to_string_unsupported` | `int64 -> string` table rejected | | `test_hashtable_call_once_import_find_runtime_query_unsupported` | runtime string query rejected | | `test_hashtable_call_once_import_duplicate_keys_unsupported` | duplicate static keys rejected | | `test_hashtable_lookup_string_value_unsupported` | string graph input now gives a clean `OpNotImplemented` | Local validation: ```bash python -m ruff format --check \ python/tvm/relax/frontend/tflite/tflite_frontend.py \ tests/python/relax/test_frontend_tflite.py python -m ruff check \ python/tvm/relax/frontend/tflite/tflite_frontend.py \ tests/python/relax/test_frontend_tflite.py python -m pytest \ tests/python/relax/test_frontend_tflite.py \ -k "hashtable or resource or variable" -q python -m pytest \ tests/python/relax/test_frontend_tflite.py -q ``` Result: ```text ruff format --check: 2 files already formatted ruff check: All checks passed 14 passed, 535 deselected 549 passed ``` ## References - Issue #19519 item G: TFLite resource / variable / hashtable operators - PR #19639: TFLite resource variable and static hashtable import support - PR #19654: TFLite `HASHTABLE_LOOKUP` converter
1 parent 730459c commit 5a8dae4

2 files changed

Lines changed: 342 additions & 22 deletions

File tree

‎python/tvm/relax/frontend/tflite/tflite_frontend.py‎

Lines changed: 132 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -721,6 +721,45 @@ def _has_tensor_buffer_data(tensor_wrapper):
721721
and tensor_wrapper.buffer.DataLength() > 0
722722
)
723723

724+
def _get_string_tensor_value(self, tensor_wrapper, op_name):
725+
"""Decode a constant TFLite string tensor buffer."""
726+
if not self._is_tflite_string_type(tensor_wrapper.tensor.Type()):
727+
raise tvm.error.OpNotImplemented(f"{op_name} requires a TensorType.STRING tensor")
728+
if not self._has_tensor_buffer_data(tensor_wrapper):
729+
raise tvm.error.OpNotImplemented(f"{op_name} requires a constant string tensor")
730+
731+
data = bytes(tensor_wrapper.buffer.DataAsNumpy())
732+
if len(data) < 4:
733+
raise tvm.error.OpNotImplemented(f"{op_name} has an invalid string tensor buffer")
734+
735+
count = int(np.frombuffer(data, dtype="<i4", count=1)[0])
736+
if count < 0:
737+
raise tvm.error.OpNotImplemented(f"{op_name} has an invalid string tensor count")
738+
739+
header_size = 4 * (count + 2)
740+
if len(data) < header_size:
741+
raise tvm.error.OpNotImplemented(f"{op_name} has an invalid string tensor offsets")
742+
743+
offsets = np.frombuffer(data, dtype="<i4", count=count + 1, offset=4).astype(np.int64)
744+
if np.any(offsets < header_size) or np.any(offsets > len(data)):
745+
raise tvm.error.OpNotImplemented(f"{op_name} has out-of-bounds string tensor offsets")
746+
if np.any(offsets[:-1] > offsets[1:]):
747+
raise tvm.error.OpNotImplemented(f"{op_name} has non-monotonic string tensor offsets")
748+
749+
try:
750+
values = [
751+
data[int(offsets[i]) : int(offsets[i + 1])].decode("utf-8") for i in range(count)
752+
]
753+
except UnicodeDecodeError as e:
754+
raise tvm.error.OpNotImplemented(f"{op_name} has invalid UTF-8 string data: {e}") from e
755+
shape = self._get_tensor_shape_tuple(tensor_wrapper)
756+
expected_count = math.prod(shape) if shape else 1
757+
if expected_count != count:
758+
raise tvm.error.OpNotImplemented(
759+
f"{op_name} string tensor buffer count does not match its shape"
760+
)
761+
return np.array(values, dtype=object).reshape(shape)
762+
724763
def convert_hashtable(self, op):
725764
"""Convert a TFLite HASHTABLE into an importer-local table handle."""
726765
input_tensors = self.get_input_tensors(op)
@@ -769,21 +808,102 @@ def convert_hashtable_import(self, op):
769808
):
770809
raise tvm.error.OpNotImplemented("HASHTABLE_IMPORT requires constant keys and values")
771810

811+
if self._is_tflite_string_type(table_info["key_dtype"]):
812+
keys = self._get_string_tensor_value(key_tensor, "HASHTABLE_IMPORT")
813+
else:
814+
keys = self.get_tensor_value(key_tensor)
815+
if self._is_tflite_string_type(table_info["value_dtype"]):
816+
values = self._get_string_tensor_value(value_tensor, "HASHTABLE_IMPORT")
817+
else:
818+
values = self.get_tensor_value(value_tensor)
819+
820+
if np.unique(keys).size != keys.size:
821+
raise tvm.error.OpNotImplemented(
822+
"HASHTABLE_IMPORT with duplicate keys is not supported"
823+
)
824+
772825
hashtable_values = self.conversion_state["hashtable_values"]
773826
table_key = table_info["table_key"]
774827
if table_key not in hashtable_values:
775828
hashtable_values[table_key] = {
776829
"size": math.prod(key_shape) if key_shape else 1,
777830
"key_dtype": table_info["key_dtype"],
778831
"value_dtype": table_info["value_dtype"],
832+
"keys": keys,
833+
"values": values,
779834
}
780835
return None
781836

782837
def convert_hashtable_find(self, op):
783-
"""Reject HASHTABLE_FIND until Relax can represent TFLite string tensors."""
784-
raise tvm.error.OpNotImplemented(
785-
"HASHTABLE_FIND requires TensorType.STRING support in Relax TFLite frontend"
838+
"""Convert the constant-foldable string-to-int64 HASHTABLE_FIND subset."""
839+
from tflite.TensorType import TensorType
840+
841+
input_tensors = self.get_input_tensors(op)
842+
output_tensors = self.get_output_tensors(op)
843+
if len(input_tensors) != 3 or len(output_tensors) != 1:
844+
raise tvm.error.OpNotImplemented(
845+
"HASHTABLE_FIND expects table, query, and default inputs with one output"
846+
)
847+
848+
table_tensor, query_tensor, default_tensor = input_tensors
849+
output_tensor = output_tensors[0]
850+
table_info = self._get_hashtable_info_for_handle(table_tensor, "HASHTABLE_FIND")
851+
table_key = table_info["table_key"]
852+
hashtable_values = self.conversion_state["hashtable_values"]
853+
if table_key not in hashtable_values:
854+
raise tvm.error.OpNotImplemented(
855+
"HASHTABLE_FIND requires a table initialized by a supported CALL_ONCE subgraph"
856+
)
857+
table_values = hashtable_values[table_key]
858+
859+
if (
860+
query_tensor.tensor.Type() != table_values["key_dtype"]
861+
or default_tensor.tensor.Type() != table_values["value_dtype"]
862+
or output_tensor.tensor.Type() != table_values["value_dtype"]
863+
):
864+
raise tvm.error.OpNotImplemented("HASHTABLE_FIND key/value dtypes mismatch")
865+
866+
if not (
867+
self._is_tflite_string_type(table_values["key_dtype"])
868+
and table_values["value_dtype"] == TensorType.INT64
869+
):
870+
raise tvm.error.OpNotImplemented(
871+
"HASHTABLE_FIND only supports constant string -> int64 tables"
872+
)
873+
if not self._has_tensor_buffer_data(query_tensor):
874+
raise tvm.error.OpNotImplemented(
875+
"HASHTABLE_FIND with runtime string queries is not supported"
876+
)
877+
if not self._has_tensor_buffer_data(default_tensor):
878+
raise tvm.error.OpNotImplemented("HASHTABLE_FIND requires constant default values")
879+
880+
query_shape = self._get_tensor_shape_tuple(query_tensor)
881+
output_shape = self._get_tensor_shape_tuple(output_tensor)
882+
if output_shape != query_shape:
883+
raise tvm.error.OpNotImplemented("HASHTABLE_FIND output shape must match query shape")
884+
885+
query_values = self._get_string_tensor_value(query_tensor, "HASHTABLE_FIND")
886+
default_values = self.get_tensor_value(default_tensor)
887+
default_shape = self._get_tensor_shape_tuple(default_tensor)
888+
if default_shape == () or default_values.size == 1:
889+
result = np.full(output_shape, int(default_values.item()), dtype=np.int64)
890+
elif default_shape == query_shape:
891+
result = default_values.astype(np.int64).copy()
892+
else:
893+
raise tvm.error.OpNotImplemented(
894+
"HASHTABLE_FIND default value must be scalar or match query shape"
895+
)
896+
897+
table_map = dict(
898+
zip(
899+
table_values["keys"].reshape(-1).tolist(),
900+
table_values["values"].reshape(-1).astype(np.int64).tolist(),
901+
)
786902
)
903+
for index, key in np.ndenumerate(query_values):
904+
if key in table_map:
905+
result[index] = table_map[key]
906+
return relax.const(result.astype(np.int64), "int64")
787907

788908
def convert_hashtable_lookup(self, op):
789909
"""Convert TFLite HASHTABLE_LOOKUP for non-string value tensors."""
@@ -8834,6 +8954,15 @@ def func(self, data):
88348954
dtype = "float32"
88358955
if shape is not None:
88368956
shape = tuple(shape) + (2,)
8957+
if dtype == "string":
8958+
# Relax has no string tensor type, so TFLite TensorType.STRING graph
8959+
# inputs cannot be represented. This also covers runtime string queries
8960+
# for ops like HASHTABLE_FIND, whose constant-foldable subset is handled
8961+
# in the op converter.
8962+
raise tvm.error.OpNotImplemented(
8963+
"Relax TFLite frontend does not support TensorType.STRING graph inputs "
8964+
"(e.g. runtime string queries)"
8965+
)
88378966
input_var = relax.Var(
88388967
name_hint=model_input_name,
88398968
ty=relax.TensorType(shape=shape, dtype=dtype),

0 commit comments

Comments
 (0)