feat: compile TFLite SVDF with its state

scripts/crossrepo_contract.py

@@ -11,6 +11,7 @@
 
 import argparse
 import copy
+import json
 import math
 import re
 import subprocess
@@ -41,6 +42,7 @@
 from tigris.emitters.binary.reader import read_binary_plan
 from tigris.emitters.binary.writer import emit_binary
 from tigris.fixtures import build_tcn
+from tigris.frontends.tflite import STATE_KEY
 from tigris.graph.ir import Stage
 
 # The byte-level line-buffer flag decoder already exists in the compiler's
@@ -844,6 +846,55 @@ def build(body):
 
 
 
+def _svdf_case() -> ContractCase:
+    """A float SVDF with a fused Relu, run once from its zero initial state, so
+    each filter's time projection sees only the newest slot of its memory."""
+    rng = np.random.default_rng(85)
+    batch, features, units, rank, memory = 2, 6, 3, 2, 4
+    filters = units * rank
+    feature = rng.normal(0.0, 0.5, (filters, features)).astype(np.float32)
+    time = rng.normal(0.0, 0.5, (filters, memory)).astype(np.float32)
+    bias = rng.normal(0.0, 0.3, (units,)).astype(np.float32)
+    x = rng.uniform(-2.0, 2.0, (batch, features)).astype(np.float32)
+    shape_x, shape_y, shape_s = [batch, features], [batch, units], [batch, filters * memory]
+    weights = [numpy_helper.from_array(feature, "feature"), numpy_helper.from_array(time, "time"),
+               numpy_helper.from_array(bias, "bias"),
+               numpy_helper.from_array(np.zeros(shape_s, np.float32), "state_initial")]
+    compile_graph = helper.make_graph(
+        [helper.make_node("Svdf", ["x", "feature", "time", "bias", "state_in"], ["y", "state_out"],
+                          domain="tigris", rank=rank, activation="relu")],
+        "svdf",
+        [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape_x),
+         helper.make_tensor_value_info("state_in", TensorProto.FLOAT, shape_s)],
+        [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape_y),
+         helper.make_tensor_value_info("state_out", TensorProto.FLOAT, shape_s)],
+        weights,
+        value_info=[helper.make_tensor_value_info("y", TensorProto.FLOAT, shape_y),
+                    helper.make_tensor_value_info("state_out", TensorProto.FLOAT, shape_s)])
+    compile_model = helper.make_model(compile_graph, opset_imports=[
+        helper.make_opsetid("", 17), helper.make_opsetid("tigris", 1)])
+    compile_model.ir_version = 9
+    helper.set_model_props(compile_model, {STATE_KEY: json.dumps(
+        [{"input": "state_in", "output": "state_out", "initial": "state_initial"}])})
+    reference = _model(
+        "svdf_reference",
+        [helper.make_node("MatMul", ["x", "feature_t"], ["projection"]),
+         helper.make_node("Mul", ["projection", "newest"], ["weighted"]),
+         helper.make_node("Reshape", ["weighted", "grouped_shape"], ["grouped"]),
+         helper.make_node("ReduceSum", ["grouped", "rank_axis"], ["summed"], keepdims=0),
+         helper.make_node("Add", ["summed", "bias"], ["biased"]),
+         helper.make_node("Relu", ["biased"], ["y"])],
+        [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape_x)],
+        [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape_y)],
+        [numpy_helper.from_array(np.ascontiguousarray(feature.T), "feature_t"),
+         numpy_helper.from_array(np.ascontiguousarray(time[:, -1]), "newest"),
+         numpy_helper.from_array(np.asarray([batch, units, rank], np.int64), "grouped_shape"),
+         numpy_helper.from_array(np.asarray([2], np.int64), "rank_axis"),
+         numpy_helper.from_array(bias, "bias")])
+    return ContractCase(name="svdf_relu", compile_model=compile_model, reference_model=reference,
+                        inputs={"x": x}, expected_operators=("Svdf",))
+
+
 def _runtime_index_case(kind: str, quantized: bool, variant: int = 0, arg: str | None = None,
                         cast: bool = False) -> ContractCase:
     case = _movement_case(kind, quantized, variant)
@@ -6918,6 +6969,7 @@ def _run_gate(runtime: Path, work_dir: Path) -> None:
         _convtranspose_2d_tiled_case(),
         _qdq_convtranspose_2d_tiled_case(),
         _convtranspose_2d_partial_edge_case(),
+        _svdf_case(),
     ]
     covered_operators = {
         operator for case in cases for operator in case.expected_operators

scripts/gen_tflite_fixtures.py

@@ -469,6 +469,104 @@ def _embedding_lookup(model: bytes) -> bytes:
 
 REWRITES["embedding_lookup"] = _embedding_lookup
 REWRITES["embedding_lookup_runtime"] = _embedding_lookup
+def _one_operator(code, options_type, options, tensors, op_inputs, op_outputs, inputs, outputs):
+    """A model of one builtin operator. `tensors` are (shape, type, data, scale,
+    zero point, variable) records; data None marks an activation."""
+    tree = schema.ModelT()
+    tree.version = 3
+    tree.buffers = [schema.BufferT()]
+    graph = schema.SubGraphT()
+    graph.tensors = []
+    for i, (shape, kind, data, scale, zero_point, variable) in enumerate(tensors):
+        tensor = schema.TensorT()
+        tensor.name = f"t{i}".encode()
+        tensor.shape = np.asarray(shape, np.int32)
+        tensor.type = kind
+        tensor.isVariable = variable
+        tensor.buffer = len(tree.buffers)
+        buffer = schema.BufferT()
+        if data is not None:
+            buffer.data = np.frombuffer(np.ascontiguousarray(data).tobytes(), np.uint8)
+        tree.buffers.append(buffer)
+        if scale is not None:
+            quant = schema.QuantizationParametersT()
+            quant.scale = np.asarray([scale], np.float32)
+            quant.zeroPoint = np.asarray([zero_point], np.int64)
+            tensor.quantization = quant
+        graph.tensors.append(tensor)
+    op = schema.OperatorT()
+    op.opcodeIndex = 0
+    op.inputs = np.asarray(op_inputs, np.int32)
+    op.outputs = np.asarray(op_outputs, np.int32)
+    op.builtinOptionsType = options_type
+    op.builtinOptions = options
+    graph.operators = [op]
+    graph.inputs = np.asarray(inputs, np.int32)
+    graph.outputs = np.asarray(outputs, np.int32)
+    tree.subgraphs = [graph]
+    opcode = schema.OperatorCodeT()
+    opcode.builtinCode = code
+    opcode.deprecatedBuiltinCode = min(code, 127)
+    opcode.version = 1
+    tree.operatorCodes = [opcode]
+    builder = flatbuffers.Builder(4096)
+    builder.Finish(tree.Pack(builder), file_identifier=b"TFL3")
+    return bytes(builder.Output())
+
+
+def _quantize(values, scale, dtype):
+    info = np.iinfo(dtype)
+    return np.clip(np.round(values / scale), info.min, info.max).astype(dtype)
+
+
+def _svdf(batch, features, units, rank, memory, activation, quantized=False):
+    """SVDF on a [batch, features] input, its state a variable tensor; seeded
+    weights. The int8 form keeps its state and time weights in int16. TFLite
+    Micro's Prepare dereferences the bias, so every case has one."""
+    rng = np.random.default_rng(features * 100 + units * 10 + rank)
+    filters = units * rank
+    feature = rng.normal(0.0, 0.4, (filters, features)).astype(np.float32)
+    time = rng.normal(0.0, 0.4, (filters, memory)).astype(np.float32)
+    offsets = rng.normal(0.0, 0.3, (units,)).astype(np.float32)
+    options = schema.SVDFOptionsT()
+    options.rank = rank
+    options.fusedActivationFunction = activation
+    T = schema.TensorType
+    if not quantized:
+        tensors = [((batch, features), T.FLOAT32, None, None, 0, False),
+                   ((filters, features), T.FLOAT32, feature, None, 0, False),
+                   ((filters, memory), T.FLOAT32, time, None, 0, False),
+                   ((units,), T.FLOAT32, offsets, None, 0, False),
+                   ((batch, filters * memory), T.FLOAT32, None, None, 0, True),
+                   ((batch, units), T.FLOAT32, None, None, 0, False)]
+    else:
+        x_scale, state_scale, out_scale = 6.0 / 255, 8.0 / 32767, 0.06
+        f_scale = float(np.abs(feature).max()) / 127
+        t_scale = float(np.abs(time).max()) / 32767
+        b_scale = float(np.float32(state_scale) * np.float32(t_scale))
+        tensors = [((batch, features), T.INT8, None, x_scale, 2, False),
+                   ((filters, features), T.INT8, _quantize(feature, f_scale, np.int8), f_scale, 0, False),
+                   ((filters, memory), T.INT16, _quantize(time, t_scale, np.int16), t_scale, 0, False),
+                   ((units,), T.INT32, _quantize(offsets, b_scale, np.int32), b_scale, 0, False),
+                   ((batch, filters * memory), T.INT16, None, state_scale, 0, True),
+                   ((batch, units), T.INT8, None, out_scale, -3, False)]
+    op_inputs = [0, 1, 2, 3, 4]
+    return _one_operator(schema.BuiltinOperator.SVDF, schema.BuiltinOptions.SVDFOptions,
+                         options, tensors, op_inputs, [5], [0], [5])
+
+
+# Models no converter writes, built operator by operator.
+_RELU = schema.ActivationFunctionType.RELU
+_NONE = schema.ActivationFunctionType.NONE
+HANDMADE = {
+    "float_svdf": lambda: _svdf(1, 8, 4, 1, 5, _NONE),
+    "float_svdf_rank2_relu": lambda: _svdf(2, 6, 3, 2, 4, _RELU),
+    "svdf": lambda: _svdf(1, 8, 4, 2, 5, _NONE, quantized=True),
+    "svdf_batch": lambda: _svdf(2, 6, 3, 1, 4, _NONE, quantized=True),
+}
+# TFLite's float SVDF sums in another order than TFLite Micro, so the two
+# differ in the last bits; TFLite Micro's outputs are recorded for these.
+SUMMATION_ORDER = {"float_svdf", "float_svdf_rank2_relu"}
 # The tier-1 cases again, converted without quantization.
 _FLOAT_TIER1 = (
     "max_pool_valid", "max_pool_same", "avg_pool_valid", "avg_pool_same", "concat_channels",
@@ -540,13 +638,16 @@ def _inputs(details, rng, value_range=None, index_range=None):
 
 def generate(name: str) -> bool:
     """Writes the case; True when TFLite Micro deviates from the reference kernels."""
-    shapes, fn = CASES[name]
     rng = np.random.default_rng(sum(name.encode()))
     ranges = RANGES.get(name)
-    model = _convert(fn, shapes, ranges or [(-3.0, 3.0)] * len(shapes), rng,
-                     float_io=name in FLOAT_BOUNDARIES, quantize=name not in FLOAT_MODELS,
-                     trackable=TRACKABLES.get(name.removeprefix("float_")),
-                     index_inputs=INDEX_INPUTS.get(name.removeprefix("float_")))
+    if name in HANDMADE:
+        model = HANDMADE[name]()
+    else:
+        shapes, fn = CASES[name]
+        model = _convert(fn, shapes, ranges or [(-3.0, 3.0)] * len(shapes), rng,
+                         float_io=name in FLOAT_BOUNDARIES, quantize=name not in FLOAT_MODELS,
+                         trackable=TRACKABLES.get(name.removeprefix("float_")),
+                         index_inputs=INDEX_INPUTS.get(name.removeprefix("float_")))
     if name in REWRITES:
         model = REWRITES[name](model)
     micro_interpreter = micro.Interpreter.from_bytes(model, arena_size=1024 * 1024)
@@ -597,7 +698,8 @@ def generate(name: str) -> bool:
                 continue
             reference_outputs.append([reference.get_tensor(d["index"]).copy()
                                       for d in reference.get_output_details()])
-    has_reference = reference is not None and name not in BROKEN_REFERENCE
+    has_reference = (reference is not None and name not in BROKEN_REFERENCE
+                     and name not in SUMMATION_ORDER)
     deviates = has_reference and any(not np.array_equal(m, r) for ms, rs in
                                      zip(micro_outputs, reference_outputs) for m, r in zip(ms, rs))
     outputs = reference_outputs if deviates else micro_outputs
@@ -637,7 +739,7 @@ def main(names: list[str]) -> None:
         for name in REFERENCE_MODELS:
             generate_model(name)
         return
-    deviations = [name for name in names or sorted(CASES) if generate(name)]
+    deviations = [name for name in names or sorted({*CASES, *HANDMADE}) if generate(name)]
     if deviations:
         print("recorded from the reference kernels: " + ", ".join(deviations))
     print(f"tensorflow {tf.__version__}, tflite-micro {version('tflite-micro')}")

scripts/generate_schema_package.py

@@ -33,6 +33,8 @@ def schema_package() -> dict[str, object]:
             "cumsum_options": defs.OP_ATTR_CUMSUM_OPTIONS,
             "movement": defs.OP_ATTR_MOVEMENT,
             "comparison_requant": defs.OP_ATTR_COMPARISON_REQUANT,
+            "constants": defs.OP_ATTR_CONSTANTS,
+            "svdf": defs.OP_ATTR_SVDF,
             "epsilon": defs.OP_ATTR_EPSILON,
             "pads": defs.OP_ATTR_PADS,
             "pool_rounding": defs.OP_ATTR_POOL_ROUNDING,
@@ -66,6 +68,7 @@ def schema_package() -> dict[str, object]:
                 "stage table authoritative; operator byte is canonical low-byte hint"
             ),
             "state_v10": "only a plan with a state section is written at schema 10",
+            "state_int16": "int16 (ONNX dtype 5) is stored only on state tensors",
         },
         "section_alignment": defs.PLAN_SECTION_ALIGNMENT,
         "tensor_flags": {

src/tigris/analysis/validation.py

@@ -17,6 +17,7 @@
 from tigris.analysis.broadcast import DENSE, PERIODIC, stored_operands
 from tigris.capabilities import KERNEL_CAPABILITIES, effective_operators
 from tigris.emitters.binary.defs import OP_TYPE_MAP
+from tigris.graph.ir import state_tensor_names
 from tigris.dtypes import check_dtype_signatures
 from tigris.graph.ir import (
     AnalyzedGraph,
@@ -96,7 +97,7 @@ def validate_execution_dtype(ag: AnalyzedGraph) -> ExecutionDTypeValidation:
         [(name, tensor.dtype, tensor.is_constant, tensor.quant is not None)
          for name, tensor in ag.tensors.items()],
         [(op.op_type, op.inputs, op.outputs) for op in ag.ops],
-        ag.model_inputs, ag.model_outputs)
+        ag.model_inputs, ag.model_outputs, state_tensor_names(ag))
     if signature_issues:
         return ExecutionDTypeValidation(dtype=None, issues=signature_issues)
 

src/tigris/capabilities.py

@@ -112,6 +112,7 @@ class KernelCapabilities:
     "Sum",
     "ReduceAll",
     "Split",
+    "Svdf",
 })
 
 _S8_REFERENCE_OPERATORS = _FLOAT_REFERENCE_OPERATORS - frozenset({
@@ -248,6 +249,7 @@ class KernelCapabilities:
     "ReduceMax": ("one axis of a rank-3 tensor; independent height or row bands; rank-4 spatial max with keepdims uses GlobalMaxPool; int8 quantization must match",),
     "ReduceMin": ("one axis of a rank-3 tensor; independent height or row bands; int8 quantization must match",),
     "ReduceAll": ("one axis of a rank-3 bool tensor; independent height or row bands",),
+    "Svdf": ("state kept between runs; int8 with int16 state and time weights; untiled",),
     "ReduceSum": ("one axis of a rank-3 tensor; independent height or row bands; int8 rejects reference arithmetic overflow",),
     "Gather": ("constant or runtime int32 indices; independent bands with constant indices; identical int8 quantization",),
     "GatherND": ("constant or runtime int32 indices; independent bands with constant indices; identical int8 quantization",),

src/tigris/dtypes.py

@@ -7,6 +7,8 @@
 
 DATA = 0
 DATA_OR_BOOL = 255
+# Data, or int16 on a tensor that holds state between runs.
+STATE = 254
 
 
 @dataclass(frozen=True)
@@ -20,9 +22,11 @@ class DTypeSignature:
 class AuxiliaryDType:
     terminal_only: bool
     requires_source: bool = False
+    state_only: bool = False
 
 
-AUXILIARY_DTYPES = {6: AuxiliaryDType(False, requires_source=True), 9: AuxiliaryDType(False)}
+AUXILIARY_DTYPES = {6: AuxiliaryDType(False, requires_source=True), 9: AuxiliaryDType(False),
+                    5: AuxiliaryDType(False, state_only=True)}
 OP_DTYPE_SIGNATURES = {kind: DTypeSignature() for kind in OP_TYPE_MAP}
 OP_DTYPE_SIGNATURES.update({
     **{kind: DTypeSignature(inputs=(DATA, 6, 6))
@@ -41,13 +45,18 @@ class AuxiliaryDType:
     "Where": DTypeSignature(inputs=(9, DATA, DATA)),
     "Cast": DTypeSignature(inputs=(9, 9, 9)),
     "ReduceAll": DTypeSignature(inputs=(9, 9, 9), outputs=(9,)),
+    # Constant operands sit between the input and the state; slots skip them.
+    # The output role admits int16 only on the next state, a state tensor.
+    "Svdf": DTypeSignature(inputs=(DATA, STATE, STATE), outputs=(STATE,)),
     **{kind: DTypeSignature(inputs=(DATA_OR_BOOL, DATA_OR_BOOL, DATA_OR_BOOL), outputs=(DATA_OR_BOOL,))
        for kind in ("Transpose", "Reshape", "Flatten")},
 })
 
 
-def check_dtype_signatures(tensors, operators, model_inputs, model_outputs):
-    """Return data tensors by dtype and errors for (name, dtype, constant, quantized) records."""
+def check_dtype_signatures(tensors, operators, model_inputs, model_outputs, state=()):
+    """Return data tensors by dtype and errors for (name, dtype, constant, quantized) records.
+    `state` names the tensors that carry state between runs."""
+    state = set(state)
     by_name = {name: (dtype, constant, quantized) for name, dtype, constant, quantized in tensors}
     data = {}
     issues = []
@@ -58,6 +67,10 @@ def check_dtype_signatures(tensors, operators, model_inputs, model_outputs):
         if policy is None:
             data.setdefault(dtype, []).append(name)
             continue
+        if policy.state_only:
+            if name not in state:
+                issues.append(f"ONNX dtype {dtype} tensor {name} must hold state between runs")
+            continue
         if quantized:
             issues.append(f"auxiliary tensor {name} cannot carry quantization")
         if policy.requires_source and sum(name in outputs for _, _, outputs in operators) != (0 if name in model_inputs else 1):
@@ -79,6 +92,11 @@ def check_dtype_signatures(tensors, operators, model_inputs, model_outputs):
                     continue
                 dtype = tensor[0]
                 expected = slots[min(position, len(slots) - 1)]
+                if expected == STATE:
+                    if dtype not in {1, 3} and not (dtype == 5 and name in state):
+                        issues.append(f"{kind} {direction} {position} ({name}) requires data "
+                                      f"or int16 state, got ONNX dtype {dtype}")
+                    continue
                 if expected == DATA_OR_BOOL:
                     if dtype not in {1, 3, 9} or dtype != by_name[inputs[0]][0]:
                         issues.append(f"{kind} must preserve its data or bool dtype")

src/tigris/emitters/binary/defs.py

@@ -70,9 +70,15 @@
 OP_ATTR_CUMSUM_OPTIONS = 11  # uint8[2], exclusive then reverse (each 0 or 1)
 
 OP_ATTR_COMPARISON_REQUANT = 13  # int32[5]: left shift, then two multiplier/shift pairs
+OP_ATTR_CONSTANTS = 14  # uint16 weight index per constant operand, in operand order;
+                        # NO_WEIGHT for an absent optional one
+OP_ATTR_SVDF = 15  # int32: rank; for int8 also the state zero point and the
+                   # input-to-state and state-to-output multiplier/shift pairs
 
 OP_ATTR_KINDS = (
     OP_ATTR_COMPARISON_REQUANT,
+    OP_ATTR_CONSTANTS,
+    OP_ATTR_SVDF,
     OP_ATTR_TRANSPOSE_PERM,
     OP_ATTR_EPSILON,
     OP_ATTR_ALPHA,
@@ -180,6 +186,7 @@
     "Cast": 82,
     "Sum": 83,
     "ReduceAll": 84,
+    "Svdf": 85,
 }
 OP_TYPE_UNKNOWN = 255