From 6df544ca143c153f52e4c4f11fc5436cfa60290c Mon Sep 17 00:00:00 2001 From: Theoo1997 Date: Wed, 9 Sep 2026 18:09:11 +0300 Subject: [PATCH 1/8] [Relax][TFLite] Allow AVERAGE_POOL_2D in quantized models _SUPPORTED_QUANTIZED_OPS is checked before dispatch, and AVERAGE_POOL_2D was missing from it even though convert_pool2d already has a complete quantized branch for pool_type="average". The op was implemented and unreachable at the same time, so importing any full-integer-quantized classifier with a global average pool -- the standard head of MobileNet/ResNet/EfficientNet -- failed with OpNotImplemented. Verified by importing the MLCommons Tiny quantized CIFAR-10 ResNet, which now lowers to the expected integer pool: lv73 = R.astype(lv72, "int32") lv74 = R.nn.avg_pool2d(lv73, pool_size=[8, 8], strides=[8, 8], ...) lv75 = R.astype(lv74, "int8") Also adds TFLITE_QUANTIZED_SUPPORT.md, which documents this alongside a separate and still-open numerical bug found while validating the import: integer nn.avg_pool2d truncates the division where TFLite's quantized AveragePool rounds half away from zero, which is wrong on 537/1024 pooled values on this model and reaches the logits as 82/320 differing by up to 91. That one is described, measured and left unfixed -- the fix changes a general operator's semantics and is TVM's call to make. --- TFLITE_QUANTIZED_SUPPORT.md | 202 ++++++++++++++++++ .../relax/frontend/tflite/tflite_frontend.py | 1 + 2 files changed, 203 insertions(+) create mode 100644 TFLITE_QUANTIZED_SUPPORT.md diff --git a/TFLITE_QUANTIZED_SUPPORT.md b/TFLITE_QUANTIZED_SUPPORT.md new file mode 100644 index 000000000000..9d3f68695368 --- /dev/null +++ b/TFLITE_QUANTIZED_SUPPORT.md @@ -0,0 +1,202 @@ +# Importing a full-integer-quantized TFLite model into Relax + +Notes for TVM developers, from getting a stock `tf.lite` int8 CNN through +`tvm.relax.frontend.tflite` and running it on a CPU backend. + +There are three findings. **One is a one-line fix and is included in this +branch. One has already landed upstream. The third is a numerical correctness +bug that is still open, is not fixed here, and is the one worth your attention** +— it silently produces wrong results rather than failing to import. + +Reproducer used throughout: the MLCommons Tiny image-classification benchmark +model, a CIFAR-10 ResNet exported with `tf.lite.Optimize.DEFAULT` and an int8 +representative dataset, so every tensor is int8 and every op is quantized. + + + +```python +import tflite +from tvm.relax.frontend.tflite import from_tflite +buf = open("pretrainedResnet_quant.tflite", "rb").read() +mod = from_tflite(tflite.Model.GetRootAsModel(buf, 0)) +``` + +--- + +## 1. `AVERAGE_POOL_2D` is implemented but not listed as quantized-capable (fixed here) + +### Symptom + +``` +OpNotImplemented: The following quantized TFLite operators are not supported in +frontend TFLite yet: 'AVERAGE_POOL_2D'. +``` + +### Cause + +`OperatorConverter._SUPPORTED_QUANTIZED_OPS` is the allowlist checked *before* +dispatching, and `AVERAGE_POOL_2D` is missing from it — even though +`convert_pool2d` already has a complete quantized branch for +`pool_type="average"`. The op is implemented and unreachable at the same time. + +The set already contains `MEAN`, `REDUCE_MAX`, `RESIZE_BILINEAR` and the other +pooling-adjacent ops, so this reads as an omission rather than a decision. +`MAX_POOL_2D` is absent for what looks like the same reason; we did not need it, +so we neither added nor tested it. + +### Fix (this branch) + +```diff + "ABS", + "ADD", + "ATAN2", ++ "AVERAGE_POOL_2D", + "CEIL", +``` + +### Why it matters + +Global average pooling is the classifier head of essentially every +MobileNet / ResNet / EfficientNet variant, so this rejects most quantized image +classifiers at import. + +--- + +## 2. `relax.op.cast` does not exist — already fixed upstream + +Recorded only so the history is clear. The quantized average-pool branch used +to call `relax.op.cast`, which is the *Relay* spelling; `relax.op` exports +`astype`. Against `v0.26.dev0-198-g67bd1ea1a` this raised + +``` +AttributeError: module 'tvm.relax.op' has no attribute 'cast' +``` + +as soon as finding 1 made the branch reachable — the two bugs hid each other. +Current `main` already uses `astype`, so nothing is needed here. + +--- + +## 3. The quantized average pool computes the wrong values (OPEN — not fixed here) + +This one imports and runs cleanly and gives wrong numbers. + +### The arithmetic + +With finding 1 applied, the frontend emits — correctly, in the integer domain, +with no float round trip: + +```python +out = relax.op.astype(in_expr, "int32") +out = relax.op.nn.avg_pool2d(out, **params) +out = relax.op.astype(out, output_tensor_type_str) +``` + +`nn.avg_pool2d` on an integer tensor divides the window sum with a **truncating** +division. TFLite's quantized `AveragePool` rounds **half away from zero** +(`tensorflow/lite/kernels/internal/reference/integer_ops/pooling.h`): + +```c +acc = acc > 0 ? (acc + filter_count / 2) / filter_count + : (acc - filter_count / 2) / filter_count; +``` + +Dropping the `±filter_count/2` biases every pooled value toward zero by up to +half an LSB. It is a systematic bias, not a rounding tie: it is wrong about half +the time. + +### Measured + +Against the TFLite interpreter on this model's 8x8 -> 1x1 pool, feeding it the +interpreter's own input tensor so nothing else can contribute: + +| model of the divide | pooled values differing | +|---|---| +| TFLite round-half-away | **0 / 1024** | +| **truncate toward zero** (what relax does) | **537 / 1024 (52%)** | +| floor | 428 / 1024 | + +End to end it is much louder than that ratio suggests, because the graph ends in +an int8 `SOFTMAX`: half an LSB on a pooled feature reaches the logits as tens of +counts. Substituting only this op into an otherwise bit-exact execution of the +whole graph: + +| pool implementation | output logits differing vs the TFLite interpreter | +|---|---| +| truncating divide | **82 / 320, max error 91** | +| round half away from zero | **0 / 320** | + +The 82/320 is also what TVM produces end to end for this model, so this single +op accounts for essentially all of the divergence. + +### Why it is not fixed in this branch + +The obvious fix — make integer `nn.avg_pool2d` round half away from zero — +changes the semantics of a general operator for every integer user, which is a +decision for TVM, not for us. The alternative is to keep `nn.avg_pool2d` +untouched and have the TFLite frontend emit the window sum and the rounded +divide explicitly, which is contained but needs a sum-pooling path the frontend +does not have today. + +We took neither: our backend claims the op and implements TFLite's rounding +itself. With that in place the ResNet is **bit-exact against the TFLite +interpreter over 128 random inputs, 0/1280 logits differing**, which is what +establishes that the rounding really is the whole story. + +### Suggested regression test + +A `from_tflite` round trip on a two-op int8 graph (`CONV_2D` then +`AVERAGE_POOL_2D`) asserting the output matches the TFLite interpreter exactly. +A tolerance-based test will pass while the bug is present — the error is 1 LSB +per pooled value — so the assertion has to be exact. + +Note that TFLite's int8 average pool requires input and output to share a scale +and zero point (the frontend already asserts this), so the reference is a pure +integer average of the raw quantized bytes with no rescale. + +--- + +## What this does NOT need + +Worth stating because it is the obvious guess and it is wrong: **no layout work +is required.** ONNX forces NCHW, and `relax.transform.ConvertLayout` cannot +convert a QDQ graph anyway because `relax.quantize` / `relax.dequantize` carry +no `FRelaxInferLayout`. The TFLite frontend sidesteps that by emitting NHWC +natively, which is what a CPU backend wants. With finding 1 applied the import +produces exactly the QDQ shape a quantized conv should have: + +``` +lv = R.dequantize(x, scale_x, zp_x) # int8 -> float32 +lv1 = R.dequantize(weight, scale_w, zp_w) # int8 -> float32, per-channel +lv2 = R.nn.conv2d(lv, lv1, data_layout="NHWC", kernel_layout="HWIO") +lv3 = R.dequantize(bias, scale_b, zp_b) # int32 -> float32 +lv4 = R.add(lv2, lv3) +lv5 = R.quantize(lv4, scale_o, zp_o) # float32 -> int8 +``` + +Every scale and zero point is a compile-time constant, so a backend can fold all +three into a single per-channel requantization scale at compile time. + +--- + +## Unrelated, recorded because it costs time: `export_library` picks the wrong triple + +`Module.export_library` takes the LLVM target for its packed-imports object +(`devc.o`) from the first LLVM module it finds, and falls back to +`fcompile.get_target_triple()` when there is none. A graph fully offloaded to a +BYOC backend leaves no TIR and hence no LLVM module, so the default +`create_shared` reports the *build host's* triple and cross-compilation fails: + +``` +ld: unknown architecture of input file `.../devc.o' is incompatible with aarch64 output +``` + +Caller-side workaround: + +```python +ex.export_library(so, fcompile=_cc.cross_compiler( + CXX, options=[...], get_target_triple=_cc.get_target_by_dump_machine(CXX))) +``` + +A model that keeps *any* TIR hides this, which is why it shows up on a +single-conv test and not on a full network. diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py index 13d78c4c30a5..1edab7bf7476 100644 --- a/python/tvm/relax/frontend/tflite/tflite_frontend.py +++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py @@ -101,6 +101,7 @@ class OperatorConverter: "ABS", "ADD", "ATAN2", + "AVERAGE_POOL_2D", "CEIL", "CONCATENATION", "CONV_2D", From b5e8968044a8cafe7bb6696ded63d3dd97dc007e Mon Sep 17 00:00:00 2001 From: Theoo1997 Date: Wed, 9 Sep 2026 18:20:59 +0300 Subject: [PATCH 2/8] [Relax][TFLite] Fix rounding of the quantized AVERAGE_POOL_2D TFLite's quantized AveragePool rounds the window average HALF AWAY FROM ZERO (reference/integer_ops/pooling.h): acc = acc > 0 ? (acc + count / 2) / count : (acc - count / 2) / count relax's integer nn.avg_pool2d divides with a truncating division, so every pooled value was biased toward zero by up to half an LSB. That is a systematic bias, not a rounding tie: on the MLCommons Tiny quantized ResNet it is wrong on 510 of 1024 pooled values, and because such graphs end in an int8 softmax it reaches the logits as 312/1280 differing by up to 106. nn.avg_pool2d itself is left alone -- changing the rounding of a general operator would change semantics for every integer user. Instead the frontend takes the window sum and performs TFLite's division explicitly. Pre-scaling the input by the window size makes avg_pool2d's own division exact, so what comes back is the sum; count_include_pad=True keeps that divisor constant (the padded taps are zeros and contribute nothing); relax.op.divide on int32 truncates toward zero, which is the semantics TFLite's formula is written against; and the per-position count of non-padded taps is folded to a constant at import, since the shapes are static. Measured with the new apps/tflite_quantized/verify_quantized_tflite.py, TVM default lowering, no BYOC: isolated AVERAGE_POOL_2D whole model before 510/1024 wrong (50%) 312/1280, max 106 after 0/1024 exact 22/1280, max 23 The whole-model figure is not zero for a separate reason, noted in the doc: TVM's QDQ lowering accumulates convolutions in float32. The script therefore also slices the operator out into a standalone one-op model and compares that on its own, which is the exact test, and keys its exit status on it. --- TFLITE_QUANTIZED_SUPPORT.md | 93 ++++++--- .../verify_quantized_tflite.py | 187 ++++++++++++++++++ .../relax/frontend/tflite/tflite_frontend.py | 79 +++++++- 3 files changed, 335 insertions(+), 24 deletions(-) create mode 100755 apps/tflite_quantized/verify_quantized_tflite.py diff --git a/TFLITE_QUANTIZED_SUPPORT.md b/TFLITE_QUANTIZED_SUPPORT.md index 9d3f68695368..988ff8cd8829 100644 --- a/TFLITE_QUANTIZED_SUPPORT.md +++ b/TFLITE_QUANTIZED_SUPPORT.md @@ -3,10 +3,19 @@ Notes for TVM developers, from getting a stock `tf.lite` int8 CNN through `tvm.relax.frontend.tflite` and running it on a CPU backend. -There are three findings. **One is a one-line fix and is included in this -branch. One has already landed upstream. The third is a numerical correctness -bug that is still open, is not fixed here, and is the one worth your attention** -— it silently produces wrong results rather than failing to import. +There are three findings. Two are fixed in this branch; one had already landed +upstream. The second fix is the important one: it is a **numerical correctness +bug that silently produced wrong results** rather than failing to import. + + apps/tflite_quantized/verify_quantized_tflite.py + +reproduces all of it against the TFLite interpreter, using TVM's default +lowering with no BYOC. On the model below: + +| | isolated `AVERAGE_POOL_2D` | whole model | +|---|---|---| +| before | **510 / 1024 elements wrong (50%)** | 312 / 1280 logits, max 106 | +| after | **0 / 1024 — exact** | 22 / 1280, max 23 | Reproducer used throughout: the MLCommons Tiny image-classification benchmark model, a CIFAR-10 ResNet exported with `tf.lite.Optimize.DEFAULT` and an int8 @@ -77,9 +86,9 @@ Current `main` already uses `astype`, so nothing is needed here. --- -## 3. The quantized average pool computes the wrong values (OPEN — not fixed here) +## 3. The quantized average pool computed the wrong values (fixed here) -This one imports and runs cleanly and gives wrong numbers. +This one imported and ran cleanly and gave wrong numbers. ### The arithmetic @@ -129,26 +138,66 @@ whole graph: The 82/320 is also what TVM produces end to end for this model, so this single op accounts for essentially all of the divergence. -### Why it is not fixed in this branch +### The fix + +`nn.avg_pool2d` is left alone -- changing the rounding of a general operator +would change semantics for every integer user, which is a separate discussion. +Instead the frontend now takes the window SUM and does TFLite's division +explicitly: + +```python +window = filter_h * filter_w +acc = relax.op.astype(in_expr, "int32") +acc = relax.op.multiply(acc, relax.const(window, "int32")) +acc = relax.op.nn.avg_pool2d(acc, count_include_pad=True, **params) +acc = relax.op.astype(acc, "int32") +half = relax.const(counts // 2, "int32") +out = relax.op.where(relax.op.greater(acc, relax.const(0, "int32")), + relax.op.add(acc, half), + relax.op.subtract(acc, half)) +out = relax.op.divide(out, relax.const(counts, "int32")) +``` -The obvious fix — make integer `nn.avg_pool2d` round half away from zero — -changes the semantics of a general operator for every integer user, which is a -decision for TVM, not for us. The alternative is to keep `nn.avg_pool2d` -untouched and have the TFLite frontend emit the window sum and the rounded -divide explicitly, which is contained but needs a sum-pooling path the frontend -does not have today. +Four things make this work, each verified rather than assumed: + +* **The sum is exact.** `avg_pool2d` divides by the window size, so pre-scaling + the input by that size makes its division exact and leaves the sum behind. + `|acc| <= 255 * window^2` for 8-bit input, which the code asserts fits int32. +* **`count_include_pad=True` keeps that divisor constant.** The padded taps are + zeros, so the sum over the padded window is the sum over the valid taps. +* **`relax.op.divide` on int32 truncates toward zero**, which is the semantics + TFLite's `(acc ± count/2) / count` is written against. Confirmed by probing + it on negative operands. +* **`counts` is the number of NON-padded taps**, which varies per output + position under SAME padding. Shapes are static, so it is folded to a constant + `[1, OH, OW, 1]` array at import time instead of being computed in the graph. + +One subtlety worth knowing if you touch this: legalization widens the pooling +accumulator (TOPI uses int64 for integer pools), so the result is pinned back +to int32 with an `astype` before it meets the int32 rounding constants — +without it the import fails with a binary-op dtype mismatch. + +### Verifying it + +``` +pip install ai-edge-litert tflite +python3 apps/tflite_quantized/verify_quantized_tflite.py --model pretrainedResnet_quant.tflite +``` -We took neither: our backend claims the op and implements TFLite's rounding -itself. With that in place the ResNet is **bit-exact against the TFLite -interpreter over 128 random inputs, 0/1280 logits differing**, which is what -establishes that the rounding really is the whole story. +The script does two things. It runs the whole model through TVM's default +lowering against the interpreter, and it slices each `AVERAGE_POOL_2D` out into +a standalone one-op model, imports that, and compares it on its own — which is +the exact test, because nothing else can contribute to it. Its exit status +follows the per-operator result. -### Suggested regression test +Both comparisons are **exact, with no tolerance**, on purpose: the error this +catches is 1 LSB per pooled value, which any tolerance would hide. -A `from_tflite` round trip on a two-op int8 graph (`CONV_2D` then -`AVERAGE_POOL_2D`) asserting the output matches the TFLite interpreter exactly. -A tolerance-based test will pass while the bug is present — the error is 1 LSB -per pooled value — so the assertion has to be exact. +**The whole-model number is not zero, and that is a different issue.** TVM's +QDQ lowering dequantizes and accumulates the convolutions in float32, which +costs a count or two by itself; a backend that keeps the convolution in int32 +gets the model bit-exact. That is why the per-operator line is the one that +carries the claim here. Note that TFLite's int8 average pool requires input and output to share a scale and zero point (the frontend already asserts this), so the reference is a pure diff --git a/apps/tflite_quantized/verify_quantized_tflite.py b/apps/tflite_quantized/verify_quantized_tflite.py new file mode 100755 index 000000000000..ae72297b8397 --- /dev/null +++ b/apps/tflite_quantized/verify_quantized_tflite.py @@ -0,0 +1,187 @@ +#!/usr/bin/env python3 +"""Check a full-integer-quantized TFLite model imported into Relax against the +TFLite interpreter itself, using TVM's DEFAULT lowering -- no BYOC, no tuning. + +The comparison is EXACT on purpose. Every tensor in such a model is int8 and +every op is quantized, so an import that preserves the model's arithmetic must +reproduce the interpreter bit for bit. A tolerance would hide precisely the +class of bug this is meant to catch: a rounding rule that is off by one, which +looks negligible per layer and is not, because these graphs end in an int8 +softmax that amplifies one LSB into tens of counts. + + pip install ai-edge-litert tflite + python3 verify_quantized_tflite.py --model pretrainedResnet_quant.tflite + +Default model: the MLCommons Tiny image-classification benchmark network, + https://github.com/mlcommons/tiny/blob/master/benchmark/training/ + image_classification/trained_models/pretrainedResnet_quant.tflite +a CIFAR-10 ResNet quantized to int8 (32x32x3 in, 10 logits out) whose head is a +global AVERAGE_POOL_2D. + +Exit status is 0 only if every output element of every trial matches. +""" +import argparse +import sys + +import numpy as np + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--model", required=True, help="path to a quantized .tflite file") + ap.add_argument("--trials", type=int, default=128, help="random inputs to compare") + ap.add_argument("--seed", type=int, default=1234) + ap.add_argument("--target", default="llvm") + ap.add_argument("--op", default="AVERAGE_POOL_2D", + help="operator to check in isolation, or '' to skip") + a = ap.parse_args() + + import tvm + from tvm import relax + import tflite + from tvm.relax.frontend.tflite import from_tflite + + try: + from ai_edge_litert.interpreter import Interpreter + except ImportError: # older wheel name; numpy 1.x only + from tflite_runtime.interpreter import Interpreter + + buf = open(a.model, "rb").read() + mod = from_tflite(tflite.Model.GetRootAsModel(buf, 0)) + + dev = tvm.cpu(0) + vm = relax.VirtualMachine(tvm.compile(mod, target=tvm.target.Target(a.target)), dev) + + interp = Interpreter(model_path=a.model) + interp.allocate_tensors() + din, dout = interp.get_input_details()[0], interp.get_output_details()[0] + shape = [int(v) for v in din["shape"]] + dtype = np.dtype(din["dtype"]) + + lo, hi = (-128, 128) if dtype == np.int8 else (0, 256) + rng = np.random.default_rng(a.seed) + bad = total = worst = 0 + for _ in range(a.trials): + x = rng.integers(lo, hi, shape).astype(dtype) + interp.set_tensor(din["index"], x) + interp.invoke() + want = interp.get_tensor(dout["index"]).copy() + got = vm["main"](tvm.runtime.tensor(x, dev)).numpy() + d = np.abs(got.astype(np.int64) - want.astype(np.int64)) + bad += int((d > 0).sum()) + total += d.size + worst = max(worst, int(d.max())) + + # ---- the isolated operator check -------------------------------------- + # The whole-model number above mixes every op together, and TVM's QDQ + # lowering runs the CONVOLUTIONS in float32, which costs a count or two on + # its own. To say something exact about one operator, slice it out of the + # model into a standalone single-op .tflite, import THAT, and compare it + # against an interpreter running the same slice. Nothing else can + # contribute, so the result must be 0. + op_bad = op_total = 0 + if a.op: + try: + op_bad, op_total = _check_single_op(a, buf, Interpreter, tvm, relax, + from_tflite, tflite, dev) + except ImportError as e: + print(f"[skip] isolated {a.op} check needs ai-edge-litert's schema: {e}") + op_total = -1 + + print(f"model : {a.model}") + print(f"target : {a.target} (default lowering, no BYOC)") + print(f"trials : {a.trials} random inputs, shape {tuple(shape)} {dtype}") + print(f"result : {bad}/{total} output elements differ" + + (f", max |diff| = {worst}" if bad else "") + + f" -> {'PASS' if bad == 0 else 'FAIL'}") + if op_total > 0: + print(f"{a.op:<8}: {op_bad}/{op_total} elements differ, sliced out and run " + f"on its own -> {'PASS' if op_bad == 0 else 'FAIL'}") + if bad: + print("\nThe whole-model figure is not expected to be zero for every model:\n" + "TVM's QDQ lowering dequantizes and accumulates convolutions in float32,\n" + "which costs a count or two by itself. The per-operator line above is the\n" + "exact one -- it isolates a single op from everything else.") + return 1 if op_bad else 0 + + +def _check_single_op(a, buf, Interpreter, tvm, relax, from_tflite, tflite, dev): + """Slice every instance of `a.op` into its own one-op model and compare. + + The slice keeps the operator's constant inputs baked in with their + quantization parameters, and promotes its activations to subgraph inputs, + so it is the same computation the full graph performs. + """ + import copy + + import flatbuffers + from ai_edge_litert.schema_py_generated import Model, ModelT + + mt = ModelT.InitFromObj(Model.GetRootAsModel(buf, 0)) + sg = mt.subgraphs[0] + names = {} + from tflite.BuiltinOperator import BuiltinOperator + + for n in dir(BuiltinOperator): + if not n.startswith("_"): + names[getattr(BuiltinOperator, n)] = n + codes = [c.builtinCode for c in mt.operatorCodes] + + ref = Interpreter(model_path=a.model, experimental_preserve_all_tensors=True) + ref.allocate_tensors() + din = ref.get_input_details()[0] + shape = [int(v) for v in din["shape"]] + dtype = np.dtype(din["dtype"]) + lo, hi = (-128, 128) if dtype == np.int8 else (0, 256) + + targets = [i for i, op in enumerate(sg.operators) + if names.get(codes[op.opcodeIndex], "") == a.op] + if not targets: + return 0, 0 + + slices = [] + for i in targets: + m2 = copy.deepcopy(mt) + s2 = m2.subgraphs[0] + op = s2.operators[i] + s2.operators = [op] + s2.inputs = [t for t in op.inputs if t >= 0 and + not (m2.buffers[s2.tensors[t].buffer].data is not None + and len(m2.buffers[s2.tensors[t].buffer].data))] + s2.outputs = list(op.outputs) + b = flatbuffers.Builder(1024) + b.Finish(m2.Pack(b), b"TFL3") + content = bytes(b.Output()) + it = Interpreter(model_content=content) + it.allocate_tensors() + sub = from_tflite(tflite.Model.GetRootAsModel(content, 0)) + vm = relax.VirtualMachine( + tvm.compile(sub, target=tvm.target.Target(a.target)), dev) + slices.append((i, op, it, vm)) + + rng = np.random.default_rng(a.seed + 1) + bad = total = 0 + for _ in range(max(1, a.trials // 8)): + ref.set_tensor(din["index"], rng.integers(lo, hi, shape).astype(dtype)) + ref.invoke() + for i, op, it, vm in slices: + # feed both the interpreter and the TVM module the SAME activation + # tensors, taken from the full model's own run + args = [] + ins = it.get_input_details() + for d, t in zip(ins, [t for t in op.inputs if t >= 0][:len(ins)]): + v = np.ascontiguousarray(ref.get_tensor(t)) + it.set_tensor(d["index"], v) + args.append(tvm.runtime.tensor(v, dev)) + it.invoke() + want = it.get_tensor(it.get_output_details()[0]["index"]) + got = vm["main"](*args).numpy() + d = np.abs(got.astype(np.int64) - want.astype(np.int64)) + bad += int((d > 0).sum()) + total += d.size + return bad, total + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py index 1edab7bf7476..7a85368da523 100644 --- a/python/tvm/relax/frontend/tflite/tflite_frontend.py +++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py @@ -5641,6 +5641,34 @@ def convert_topk_v2(self, op): return out + @staticmethod + def _avg_pool2d_valid_counts(in_hw, filter_hw, stride_hw, padding, out_hw): + """Non-padded taps per output position, as an int32 [1, OH, OW, 1] array. + + TFLite's average pool divides by the number of taps that fall inside + the input (`count_include_pad=False` semantics). With no padding that + is just filter_h*filter_w everywhere; with SAME padding the windows at + the borders see fewer, so the divisor is position dependent. + """ + import numpy as _np + + in_h, in_w = in_hw + f_h, f_w = filter_hw + s_h, s_w = stride_hw + pad_top, pad_left = (padding[0], padding[1]) if len(padding) >= 2 else (0, 0) + out_h, out_w = out_hw + + def counts_1d(extent, f, s, pad_before, out): + c = _np.empty(out, dtype="int32") + for i in range(out): + start = i * s - pad_before + c[i] = min(start + f, extent) - max(start, 0) + return c + + rows = counts_1d(in_h, f_h, s_h, pad_top, out_h) + cols = counts_1d(in_w, f_w, s_w, pad_left, out_w) + return (rows[:, None] * cols[None, :]).reshape(1, out_h, out_w, 1) + def convert_pool2d(self, op, pool_type): """pool2d implementation.""" @@ -5698,8 +5726,55 @@ def convert_pool2d(self, op, pool_type): "TFLite avg_pool2dreshape requires input and output scale" "and zero points to be equal" ) - out = relax.op.astype(in_expr, "int32") - out = relax.op.nn.avg_pool2d(out, **params) + # TFLite's quantized AveragePool rounds the window average + # HALF AWAY FROM ZERO, and relax's integer avg_pool2d divides + # with a truncating division. Truncating biases every pooled + # value toward zero by up to half an LSB -- on a quantized + # ResNet that is ~52% of the pooled values, and because such + # graphs end in an int8 softmax it reaches the logits as tens + # of counts. The reference is + # tensorflow/lite/kernels/internal/reference/integer_ops/ + # pooling.h + # acc = acc > 0 ? (acc + count / 2) / count + # : (acc - count / 2) / count + # + # So take the window SUM and do that division explicitly. + # avg_pool2d divides by the window size, so pre-scaling the + # input by that size makes its division exact and what comes + # back is the sum. count_include_pad=True keeps that divisor a + # constant; the padded taps are zeros and contribute nothing, + # so the result is the sum over the VALID taps either way. + window = filter_h * filter_w + # acc holds window * sum, and |sum| <= window * 255 for any 8-bit + # input, so |acc| <= 255 * window^2 must stay inside int32. + assert 255 * window * window < (1 << 31), ( + f"pooling window {filter_h}x{filter_w} is too large for an " + "int32 exact-sum average pool" + ) + acc = relax.op.astype(in_expr, "int32") + acc = relax.op.multiply(acc, relax.const(window, "int32")) + acc = relax.op.nn.avg_pool2d(acc, count_include_pad=True, **params) + # Legalization is free to widen the pooling accumulator (TOPI + # uses int64 for integer pools), so pin the dtype back before + # mixing it with the int32 rounding constants below. + acc = relax.op.astype(acc, "int32") + + # TFLite divides by the number of NON-PADDED taps, which varies + # per output position once there is padding. The shapes are + # static, so the per-position count is folded to a constant + # here rather than computed in the graph. + counts = self._avg_pool2d_valid_counts( + (input_h, input_w), (filter_h, filter_w), + (stride_h, stride_w), params["padding"], + to_int_list(self.get_tensor_shape(output_tensor))[1:3], + ) + half = relax.const(counts // 2, "int32") + out = relax.op.where( + relax.op.greater(acc, relax.const(0, "int32")), + relax.op.add(acc, half), + relax.op.subtract(acc, half), + ) + out = relax.op.divide(out, relax.const(counts, "int32")) out = relax.op.astype(out, output_tensor_type_str) else: out = relax.op.nn.avg_pool2d(in_expr, **params) From f6ee90492d6bd6529efd3604c8f1b04474a10508 Mon Sep 17 00:00:00 2001 From: Theologis Anthimopoulos <109102287+Theoo1997@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:15:17 +0300 Subject: [PATCH 3/8] Delete TFLITE_QUANTIZED_SUPPORT.md --- TFLITE_QUANTIZED_SUPPORT.md | 251 ------------------------------------ 1 file changed, 251 deletions(-) delete mode 100644 TFLITE_QUANTIZED_SUPPORT.md diff --git a/TFLITE_QUANTIZED_SUPPORT.md b/TFLITE_QUANTIZED_SUPPORT.md deleted file mode 100644 index 988ff8cd8829..000000000000 --- a/TFLITE_QUANTIZED_SUPPORT.md +++ /dev/null @@ -1,251 +0,0 @@ -# Importing a full-integer-quantized TFLite model into Relax - -Notes for TVM developers, from getting a stock `tf.lite` int8 CNN through -`tvm.relax.frontend.tflite` and running it on a CPU backend. - -There are three findings. Two are fixed in this branch; one had already landed -upstream. The second fix is the important one: it is a **numerical correctness -bug that silently produced wrong results** rather than failing to import. - - apps/tflite_quantized/verify_quantized_tflite.py - -reproduces all of it against the TFLite interpreter, using TVM's default -lowering with no BYOC. On the model below: - -| | isolated `AVERAGE_POOL_2D` | whole model | -|---|---|---| -| before | **510 / 1024 elements wrong (50%)** | 312 / 1280 logits, max 106 | -| after | **0 / 1024 — exact** | 22 / 1280, max 23 | - -Reproducer used throughout: the MLCommons Tiny image-classification benchmark -model, a CIFAR-10 ResNet exported with `tf.lite.Optimize.DEFAULT` and an int8 -representative dataset, so every tensor is int8 and every op is quantized. - - - -```python -import tflite -from tvm.relax.frontend.tflite import from_tflite -buf = open("pretrainedResnet_quant.tflite", "rb").read() -mod = from_tflite(tflite.Model.GetRootAsModel(buf, 0)) -``` - ---- - -## 1. `AVERAGE_POOL_2D` is implemented but not listed as quantized-capable (fixed here) - -### Symptom - -``` -OpNotImplemented: The following quantized TFLite operators are not supported in -frontend TFLite yet: 'AVERAGE_POOL_2D'. -``` - -### Cause - -`OperatorConverter._SUPPORTED_QUANTIZED_OPS` is the allowlist checked *before* -dispatching, and `AVERAGE_POOL_2D` is missing from it — even though -`convert_pool2d` already has a complete quantized branch for -`pool_type="average"`. The op is implemented and unreachable at the same time. - -The set already contains `MEAN`, `REDUCE_MAX`, `RESIZE_BILINEAR` and the other -pooling-adjacent ops, so this reads as an omission rather than a decision. -`MAX_POOL_2D` is absent for what looks like the same reason; we did not need it, -so we neither added nor tested it. - -### Fix (this branch) - -```diff - "ABS", - "ADD", - "ATAN2", -+ "AVERAGE_POOL_2D", - "CEIL", -``` - -### Why it matters - -Global average pooling is the classifier head of essentially every -MobileNet / ResNet / EfficientNet variant, so this rejects most quantized image -classifiers at import. - ---- - -## 2. `relax.op.cast` does not exist — already fixed upstream - -Recorded only so the history is clear. The quantized average-pool branch used -to call `relax.op.cast`, which is the *Relay* spelling; `relax.op` exports -`astype`. Against `v0.26.dev0-198-g67bd1ea1a` this raised - -``` -AttributeError: module 'tvm.relax.op' has no attribute 'cast' -``` - -as soon as finding 1 made the branch reachable — the two bugs hid each other. -Current `main` already uses `astype`, so nothing is needed here. - ---- - -## 3. The quantized average pool computed the wrong values (fixed here) - -This one imported and ran cleanly and gave wrong numbers. - -### The arithmetic - -With finding 1 applied, the frontend emits — correctly, in the integer domain, -with no float round trip: - -```python -out = relax.op.astype(in_expr, "int32") -out = relax.op.nn.avg_pool2d(out, **params) -out = relax.op.astype(out, output_tensor_type_str) -``` - -`nn.avg_pool2d` on an integer tensor divides the window sum with a **truncating** -division. TFLite's quantized `AveragePool` rounds **half away from zero** -(`tensorflow/lite/kernels/internal/reference/integer_ops/pooling.h`): - -```c -acc = acc > 0 ? (acc + filter_count / 2) / filter_count - : (acc - filter_count / 2) / filter_count; -``` - -Dropping the `±filter_count/2` biases every pooled value toward zero by up to -half an LSB. It is a systematic bias, not a rounding tie: it is wrong about half -the time. - -### Measured - -Against the TFLite interpreter on this model's 8x8 -> 1x1 pool, feeding it the -interpreter's own input tensor so nothing else can contribute: - -| model of the divide | pooled values differing | -|---|---| -| TFLite round-half-away | **0 / 1024** | -| **truncate toward zero** (what relax does) | **537 / 1024 (52%)** | -| floor | 428 / 1024 | - -End to end it is much louder than that ratio suggests, because the graph ends in -an int8 `SOFTMAX`: half an LSB on a pooled feature reaches the logits as tens of -counts. Substituting only this op into an otherwise bit-exact execution of the -whole graph: - -| pool implementation | output logits differing vs the TFLite interpreter | -|---|---| -| truncating divide | **82 / 320, max error 91** | -| round half away from zero | **0 / 320** | - -The 82/320 is also what TVM produces end to end for this model, so this single -op accounts for essentially all of the divergence. - -### The fix - -`nn.avg_pool2d` is left alone -- changing the rounding of a general operator -would change semantics for every integer user, which is a separate discussion. -Instead the frontend now takes the window SUM and does TFLite's division -explicitly: - -```python -window = filter_h * filter_w -acc = relax.op.astype(in_expr, "int32") -acc = relax.op.multiply(acc, relax.const(window, "int32")) -acc = relax.op.nn.avg_pool2d(acc, count_include_pad=True, **params) -acc = relax.op.astype(acc, "int32") -half = relax.const(counts // 2, "int32") -out = relax.op.where(relax.op.greater(acc, relax.const(0, "int32")), - relax.op.add(acc, half), - relax.op.subtract(acc, half)) -out = relax.op.divide(out, relax.const(counts, "int32")) -``` - -Four things make this work, each verified rather than assumed: - -* **The sum is exact.** `avg_pool2d` divides by the window size, so pre-scaling - the input by that size makes its division exact and leaves the sum behind. - `|acc| <= 255 * window^2` for 8-bit input, which the code asserts fits int32. -* **`count_include_pad=True` keeps that divisor constant.** The padded taps are - zeros, so the sum over the padded window is the sum over the valid taps. -* **`relax.op.divide` on int32 truncates toward zero**, which is the semantics - TFLite's `(acc ± count/2) / count` is written against. Confirmed by probing - it on negative operands. -* **`counts` is the number of NON-padded taps**, which varies per output - position under SAME padding. Shapes are static, so it is folded to a constant - `[1, OH, OW, 1]` array at import time instead of being computed in the graph. - -One subtlety worth knowing if you touch this: legalization widens the pooling -accumulator (TOPI uses int64 for integer pools), so the result is pinned back -to int32 with an `astype` before it meets the int32 rounding constants — -without it the import fails with a binary-op dtype mismatch. - -### Verifying it - -``` -pip install ai-edge-litert tflite -python3 apps/tflite_quantized/verify_quantized_tflite.py --model pretrainedResnet_quant.tflite -``` - -The script does two things. It runs the whole model through TVM's default -lowering against the interpreter, and it slices each `AVERAGE_POOL_2D` out into -a standalone one-op model, imports that, and compares it on its own — which is -the exact test, because nothing else can contribute to it. Its exit status -follows the per-operator result. - -Both comparisons are **exact, with no tolerance**, on purpose: the error this -catches is 1 LSB per pooled value, which any tolerance would hide. - -**The whole-model number is not zero, and that is a different issue.** TVM's -QDQ lowering dequantizes and accumulates the convolutions in float32, which -costs a count or two by itself; a backend that keeps the convolution in int32 -gets the model bit-exact. That is why the per-operator line is the one that -carries the claim here. - -Note that TFLite's int8 average pool requires input and output to share a scale -and zero point (the frontend already asserts this), so the reference is a pure -integer average of the raw quantized bytes with no rescale. - ---- - -## What this does NOT need - -Worth stating because it is the obvious guess and it is wrong: **no layout work -is required.** ONNX forces NCHW, and `relax.transform.ConvertLayout` cannot -convert a QDQ graph anyway because `relax.quantize` / `relax.dequantize` carry -no `FRelaxInferLayout`. The TFLite frontend sidesteps that by emitting NHWC -natively, which is what a CPU backend wants. With finding 1 applied the import -produces exactly the QDQ shape a quantized conv should have: - -``` -lv = R.dequantize(x, scale_x, zp_x) # int8 -> float32 -lv1 = R.dequantize(weight, scale_w, zp_w) # int8 -> float32, per-channel -lv2 = R.nn.conv2d(lv, lv1, data_layout="NHWC", kernel_layout="HWIO") -lv3 = R.dequantize(bias, scale_b, zp_b) # int32 -> float32 -lv4 = R.add(lv2, lv3) -lv5 = R.quantize(lv4, scale_o, zp_o) # float32 -> int8 -``` - -Every scale and zero point is a compile-time constant, so a backend can fold all -three into a single per-channel requantization scale at compile time. - ---- - -## Unrelated, recorded because it costs time: `export_library` picks the wrong triple - -`Module.export_library` takes the LLVM target for its packed-imports object -(`devc.o`) from the first LLVM module it finds, and falls back to -`fcompile.get_target_triple()` when there is none. A graph fully offloaded to a -BYOC backend leaves no TIR and hence no LLVM module, so the default -`create_shared` reports the *build host's* triple and cross-compilation fails: - -``` -ld: unknown architecture of input file `.../devc.o' is incompatible with aarch64 output -``` - -Caller-side workaround: - -```python -ex.export_library(so, fcompile=_cc.cross_compiler( - CXX, options=[...], get_target_triple=_cc.get_target_by_dump_machine(CXX))) -``` - -A model that keeps *any* TIR hides this, which is why it shows up on a -single-conv test and not on a full network. From 569f1e65d8292d551842bcd721be75eb640bc88b Mon Sep 17 00:00:00 2001 From: Theoo1997 Date: Wed, 9 Sep 2026 19:56:10 +0300 Subject: [PATCH 4/8] [Relax][TFLite] Add regression tests for the quantized AVERAGE_POOL_2D rounding Moves the verification out of apps/ and into the suite where the rest of the TFLite frontend is tested, per convention: * test_quantized_avg_pool2d_uses_astype is renamed to ..._rounds_half_away_from_zero and updated for the new lowering, which it otherwise fails against. * A numeric test is added that pins the ANSWER rather than the shape, against the reference formula in tensorflow/lite/kernels/internal/reference/integer_ops/pooling.h, on window sums that land on the .5 boundaries where truncation and round-half-away disagree. This is the test that would have caught the original bug: a truncating divide is wrong on about half of all inputs while looking perfectly reasonable in the IR. Both tests fail against the truncating lowering and pass with the fix. Also corrects the comment on the int32 pin after the pool. The widening is done by relax's own type inference for integer avg_pool2d when the window is large enough that an int32 accumulator could overflow -- an 8x8 pool over int32 comes back as int64, a 2x2 one stays int32 -- not by legalization as previously stated. --- .../verify_quantized_tflite.py | 187 ------------------ .../relax/frontend/tflite/tflite_frontend.py | 15 +- tests/python/relax/test_frontend_tflite.py | 110 ++++++++++- 3 files changed, 114 insertions(+), 198 deletions(-) delete mode 100755 apps/tflite_quantized/verify_quantized_tflite.py diff --git a/apps/tflite_quantized/verify_quantized_tflite.py b/apps/tflite_quantized/verify_quantized_tflite.py deleted file mode 100755 index ae72297b8397..000000000000 --- a/apps/tflite_quantized/verify_quantized_tflite.py +++ /dev/null @@ -1,187 +0,0 @@ -#!/usr/bin/env python3 -"""Check a full-integer-quantized TFLite model imported into Relax against the -TFLite interpreter itself, using TVM's DEFAULT lowering -- no BYOC, no tuning. - -The comparison is EXACT on purpose. Every tensor in such a model is int8 and -every op is quantized, so an import that preserves the model's arithmetic must -reproduce the interpreter bit for bit. A tolerance would hide precisely the -class of bug this is meant to catch: a rounding rule that is off by one, which -looks negligible per layer and is not, because these graphs end in an int8 -softmax that amplifies one LSB into tens of counts. - - pip install ai-edge-litert tflite - python3 verify_quantized_tflite.py --model pretrainedResnet_quant.tflite - -Default model: the MLCommons Tiny image-classification benchmark network, - https://github.com/mlcommons/tiny/blob/master/benchmark/training/ - image_classification/trained_models/pretrainedResnet_quant.tflite -a CIFAR-10 ResNet quantized to int8 (32x32x3 in, 10 logits out) whose head is a -global AVERAGE_POOL_2D. - -Exit status is 0 only if every output element of every trial matches. -""" -import argparse -import sys - -import numpy as np - - -def main(): - ap = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - ap.add_argument("--model", required=True, help="path to a quantized .tflite file") - ap.add_argument("--trials", type=int, default=128, help="random inputs to compare") - ap.add_argument("--seed", type=int, default=1234) - ap.add_argument("--target", default="llvm") - ap.add_argument("--op", default="AVERAGE_POOL_2D", - help="operator to check in isolation, or '' to skip") - a = ap.parse_args() - - import tvm - from tvm import relax - import tflite - from tvm.relax.frontend.tflite import from_tflite - - try: - from ai_edge_litert.interpreter import Interpreter - except ImportError: # older wheel name; numpy 1.x only - from tflite_runtime.interpreter import Interpreter - - buf = open(a.model, "rb").read() - mod = from_tflite(tflite.Model.GetRootAsModel(buf, 0)) - - dev = tvm.cpu(0) - vm = relax.VirtualMachine(tvm.compile(mod, target=tvm.target.Target(a.target)), dev) - - interp = Interpreter(model_path=a.model) - interp.allocate_tensors() - din, dout = interp.get_input_details()[0], interp.get_output_details()[0] - shape = [int(v) for v in din["shape"]] - dtype = np.dtype(din["dtype"]) - - lo, hi = (-128, 128) if dtype == np.int8 else (0, 256) - rng = np.random.default_rng(a.seed) - bad = total = worst = 0 - for _ in range(a.trials): - x = rng.integers(lo, hi, shape).astype(dtype) - interp.set_tensor(din["index"], x) - interp.invoke() - want = interp.get_tensor(dout["index"]).copy() - got = vm["main"](tvm.runtime.tensor(x, dev)).numpy() - d = np.abs(got.astype(np.int64) - want.astype(np.int64)) - bad += int((d > 0).sum()) - total += d.size - worst = max(worst, int(d.max())) - - # ---- the isolated operator check -------------------------------------- - # The whole-model number above mixes every op together, and TVM's QDQ - # lowering runs the CONVOLUTIONS in float32, which costs a count or two on - # its own. To say something exact about one operator, slice it out of the - # model into a standalone single-op .tflite, import THAT, and compare it - # against an interpreter running the same slice. Nothing else can - # contribute, so the result must be 0. - op_bad = op_total = 0 - if a.op: - try: - op_bad, op_total = _check_single_op(a, buf, Interpreter, tvm, relax, - from_tflite, tflite, dev) - except ImportError as e: - print(f"[skip] isolated {a.op} check needs ai-edge-litert's schema: {e}") - op_total = -1 - - print(f"model : {a.model}") - print(f"target : {a.target} (default lowering, no BYOC)") - print(f"trials : {a.trials} random inputs, shape {tuple(shape)} {dtype}") - print(f"result : {bad}/{total} output elements differ" - + (f", max |diff| = {worst}" if bad else "") - + f" -> {'PASS' if bad == 0 else 'FAIL'}") - if op_total > 0: - print(f"{a.op:<8}: {op_bad}/{op_total} elements differ, sliced out and run " - f"on its own -> {'PASS' if op_bad == 0 else 'FAIL'}") - if bad: - print("\nThe whole-model figure is not expected to be zero for every model:\n" - "TVM's QDQ lowering dequantizes and accumulates convolutions in float32,\n" - "which costs a count or two by itself. The per-operator line above is the\n" - "exact one -- it isolates a single op from everything else.") - return 1 if op_bad else 0 - - -def _check_single_op(a, buf, Interpreter, tvm, relax, from_tflite, tflite, dev): - """Slice every instance of `a.op` into its own one-op model and compare. - - The slice keeps the operator's constant inputs baked in with their - quantization parameters, and promotes its activations to subgraph inputs, - so it is the same computation the full graph performs. - """ - import copy - - import flatbuffers - from ai_edge_litert.schema_py_generated import Model, ModelT - - mt = ModelT.InitFromObj(Model.GetRootAsModel(buf, 0)) - sg = mt.subgraphs[0] - names = {} - from tflite.BuiltinOperator import BuiltinOperator - - for n in dir(BuiltinOperator): - if not n.startswith("_"): - names[getattr(BuiltinOperator, n)] = n - codes = [c.builtinCode for c in mt.operatorCodes] - - ref = Interpreter(model_path=a.model, experimental_preserve_all_tensors=True) - ref.allocate_tensors() - din = ref.get_input_details()[0] - shape = [int(v) for v in din["shape"]] - dtype = np.dtype(din["dtype"]) - lo, hi = (-128, 128) if dtype == np.int8 else (0, 256) - - targets = [i for i, op in enumerate(sg.operators) - if names.get(codes[op.opcodeIndex], "") == a.op] - if not targets: - return 0, 0 - - slices = [] - for i in targets: - m2 = copy.deepcopy(mt) - s2 = m2.subgraphs[0] - op = s2.operators[i] - s2.operators = [op] - s2.inputs = [t for t in op.inputs if t >= 0 and - not (m2.buffers[s2.tensors[t].buffer].data is not None - and len(m2.buffers[s2.tensors[t].buffer].data))] - s2.outputs = list(op.outputs) - b = flatbuffers.Builder(1024) - b.Finish(m2.Pack(b), b"TFL3") - content = bytes(b.Output()) - it = Interpreter(model_content=content) - it.allocate_tensors() - sub = from_tflite(tflite.Model.GetRootAsModel(content, 0)) - vm = relax.VirtualMachine( - tvm.compile(sub, target=tvm.target.Target(a.target)), dev) - slices.append((i, op, it, vm)) - - rng = np.random.default_rng(a.seed + 1) - bad = total = 0 - for _ in range(max(1, a.trials // 8)): - ref.set_tensor(din["index"], rng.integers(lo, hi, shape).astype(dtype)) - ref.invoke() - for i, op, it, vm in slices: - # feed both the interpreter and the TVM module the SAME activation - # tensors, taken from the full model's own run - args = [] - ins = it.get_input_details() - for d, t in zip(ins, [t for t in op.inputs if t >= 0][:len(ins)]): - v = np.ascontiguousarray(ref.get_tensor(t)) - it.set_tensor(d["index"], v) - args.append(tvm.runtime.tensor(v, dev)) - it.invoke() - want = it.get_tensor(it.get_output_details()[0]["index"]) - got = vm["main"](*args).numpy() - d = np.abs(got.astype(np.int64) - want.astype(np.int64)) - bad += int((d > 0).sum()) - total += d.size - return bad, total - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py index 7a85368da523..1a603762b37b 100644 --- a/python/tvm/relax/frontend/tflite/tflite_frontend.py +++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py @@ -5754,9 +5754,12 @@ def convert_pool2d(self, op, pool_type): acc = relax.op.astype(in_expr, "int32") acc = relax.op.multiply(acc, relax.const(window, "int32")) acc = relax.op.nn.avg_pool2d(acc, count_include_pad=True, **params) - # Legalization is free to widen the pooling accumulator (TOPI - # uses int64 for integer pools), so pin the dtype back before - # mixing it with the int32 rounding constants below. + # relax widens the output dtype of an INTEGER avg_pool2d when the + # window is large enough that an int32 accumulator could + # overflow -- an 8x8 pool over int32 comes back as int64, a 2x2 + # one stays int32. Pin it so the rounding arithmetic below always + # meets int32 constants; without this an 8x8 pool fails to import + # with "Binary operators must have the same datatype". acc = relax.op.astype(acc, "int32") # TFLite divides by the number of NON-PADDED taps, which varies @@ -5764,8 +5767,10 @@ def convert_pool2d(self, op, pool_type): # static, so the per-position count is folded to a constant # here rather than computed in the graph. counts = self._avg_pool2d_valid_counts( - (input_h, input_w), (filter_h, filter_w), - (stride_h, stride_w), params["padding"], + (input_h, input_w), + (filter_h, filter_w), + (stride_h, stride_w), + params["padding"], to_int_list(self.get_tensor_shape(output_tensor))[1:3], ) half = relax.const(counts // 2, "int32") diff --git a/tests/python/relax/test_frontend_tflite.py b/tests/python/relax/test_frontend_tflite.py index e7ebbeaf92ce..f57bca7245d6 100644 --- a/tests/python/relax/test_frontend_tflite.py +++ b/tests/python/relax/test_frontend_tflite.py @@ -11293,8 +11293,8 @@ def main() -> R.Tensor((2,), dtype="float32"): tvm.ir.assert_structural_equal(mod, Expected) -def test_quantized_avg_pool2d_uses_astype(): - """Quantized AVERAGE_POOL_2D casts through int32 with R.astype.""" +def test_quantized_avg_pool2d_rounds_half_away_from_zero(): + """Quantized AVERAGE_POOL_2D takes the window sum and rounds like TFLite.""" builder = flatbuffers.Builder(1024) qparams = _build_quantization_parameters( @@ -11373,24 +11373,122 @@ def main(tvmgen_tensor_0: R.Tensor((1, 2, 2, 1), dtype="int8")) -> R.Tensor( ): with R.dataflow(): lv: R.Tensor((1, 2, 2, 1), dtype="int32") = R.astype(tvmgen_tensor_0, dtype="int32") - lv1: R.Tensor((1, 1, 1, 1), dtype="int32") = R.nn.avg_pool2d( - lv, + # scaling by the window size makes avg_pool2d's own division + # exact, so what comes back is the window SUM + lv1: R.Tensor((1, 2, 2, 1), dtype="int32") = R.multiply(lv, R.const(4, "int32")) + lv2: R.Tensor((1, 1, 1, 1), dtype="int32") = R.nn.avg_pool2d( + lv1, pool_size=[2, 2], strides=[1, 1], dilation=[1, 1], padding=[0, 0, 0, 0], ceil_mode=False, - count_include_pad=False, + count_include_pad=True, layout="NHWC", out_layout="NHWC", ) - gv: R.Tensor((1, 1, 1, 1), dtype="int8") = R.astype(lv1, dtype="int8") + lv3: R.Tensor((1, 1, 1, 1), dtype="int32") = R.astype(lv2, dtype="int32") + # acc > 0 ? (acc + count/2) / count : (acc - count/2) / count + lv4: R.Tensor((1, 1, 1, 1), dtype="bool") = R.greater(lv3, R.const(0, "int32")) + lv5: R.Tensor((1, 1, 1, 1), dtype="int32") = R.add( + lv3, R.const(np.full((1, 1, 1, 1), 2, dtype="int32")) + ) + lv6: R.Tensor((1, 1, 1, 1), dtype="int32") = R.subtract( + lv3, R.const(np.full((1, 1, 1, 1), 2, dtype="int32")) + ) + lv7: R.Tensor((1, 1, 1, 1), dtype="int32") = R.where(lv4, lv5, lv6) + lv8: R.Tensor((1, 1, 1, 1), dtype="int32") = R.divide( + lv7, R.const(np.full((1, 1, 1, 1), 4, dtype="int32")) + ) + gv: R.Tensor((1, 1, 1, 1), dtype="int8") = R.astype(lv8, dtype="int8") R.output(gv) return gv tvm.ir.assert_structural_equal(mod, Expected) +def test_quantized_avg_pool2d_matches_tflite_rounding_numerically(): + """The pooled values must match TFLite's reference integer average pool. + + The structural test above pins the shape of the lowering; this one pins the + ANSWER, which is what actually regressed: a truncating division looks + reasonable and is wrong on about half of all inputs. The reference is + tensorflow/lite/kernels/internal/reference/integer_ops/pooling.h, + + acc = acc > 0 ? (acc + count / 2) / count : (acc - count / 2) / count + + and the inputs below are chosen so the window sums land on and around the + .5 boundaries where truncation and round-half-away disagree. + """ + builder = flatbuffers.Builder(1024) + + qparams = _build_quantization_parameters( + builder, scale=[0.5], zero_point=[0], quantized_dimension=0 + ) + input_tensor = _build_tensor( + builder, 0, [1, 2, 2, 8], tensor_type=_tfl_tensor_type.INT8, quantization=qparams + ) + output_tensor = _build_tensor( + builder, 1, [1, 1, 1, 8], tensor_type=_tfl_tensor_type.INT8, quantization=qparams + ) + + _tfl_pool2d_options.Pool2DOptionsStart(builder) + _tfl_pool2d_options.Pool2DOptionsAddPadding(builder, _tfl_padding.VALID) + _tfl_pool2d_options.Pool2DOptionsAddStrideH(builder, 2) + _tfl_pool2d_options.Pool2DOptionsAddStrideW(builder, 2) + _tfl_pool2d_options.Pool2DOptionsAddFilterHeight(builder, 2) + _tfl_pool2d_options.Pool2DOptionsAddFilterWidth(builder, 2) + _tfl_pool2d_options.Pool2DOptionsAddFusedActivationFunction(builder, _tfl_activation_fn.NONE) + pool_opts = _tfl_pool2d_options.Pool2DOptionsEnd(builder) + + avg_pool_op = _build_operator( + builder, + 0, + [0], + [1], + builtin_options_type=_tfl_builtin_options.Pool2DOptions, + builtin_options=pool_opts, + ) + subgraph = _build_subgraph( + builder, + tensors=[input_tensor, output_tensor], + operators=[avg_pool_op], + inputs=[0], + outputs=[1], + ) + operator_codes = [_build_operator_code(builder, _tfl_builtin_operator.AVERAGE_POOL_2D)] + buf = _finish_tflite_model( + builder, + subgraph=subgraph, + operator_codes=operator_codes, + buffers=[_build_buffer(builder), _build_buffer(builder)], + ) + + if hasattr(tflite.Model, "Model"): + tflite_model = tflite.Model.Model.GetRootAsModel(buf, 0) + else: + tflite_model = tflite.Model.GetRootAsModel(buf, 0) + mod = from_tflite(tflite_model) + + dev = tvm.cpu(0) + vm = relax.VirtualMachine(tvm.compile(mod, target=tvm.target.Target("llvm")), dev) + + # one channel per window sum in [-7, 7]: +-1.5, +-1.75 and the exact halves + # are where a truncating divide diverges from round-half-away + sums = [5, 6, 7, -5, -6, -7, 2, -2] + x = np.zeros((1, 2, 2, 8), dtype="int8") + for c, total in enumerate(sums): + x[0, 0, 0, c] = total + got = vm["main"](tvm.runtime.tensor(x, dev)).numpy().reshape(-1) + + count = 4 + want = np.array( + [(t + count // 2) // count if t > 0 else -((-t + count // 2) // count) for t in sums], + dtype="int8", + ) + np.testing.assert_array_equal(got, want) + + def test_quantized_conv2d_per_tensor_uses_qdq(): """Quantized Conv2D with per-tensor quantization uses DQ -> conv2d -> Q.""" builder = flatbuffers.Builder(2048) From 67058e87eabbd4cd67fb30b23f11b83ff36dba59 Mon Sep 17 00:00:00 2001 From: Theologis Anthimopoulos <109102287+Theoo1997@users.noreply.github.com> Date: Wed, 9 Sep 2026 22:44:36 +0300 Subject: [PATCH 5/8] Update test_frontend_tflite.py --- tests/python/relax/test_frontend_tflite.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/python/relax/test_frontend_tflite.py b/tests/python/relax/test_frontend_tflite.py index f57bca7245d6..0a34785c0af0 100644 --- a/tests/python/relax/test_frontend_tflite.py +++ b/tests/python/relax/test_frontend_tflite.py @@ -57,7 +57,6 @@ def _get_mod_from_cfunc(cfunc): mod["main"] = mod["main"].without_attr("params") return mod - def verify(TestClass, expected=None): if isinstance(TestClass, type): cf = TestClass().func.get_concrete_function() From 68cd8699fe891b23399078dd3c70ad77209e7834 Mon Sep 17 00:00:00 2001 From: Theologis Anthimopoulos <109102287+Theoo1997@users.noreply.github.com> Date: Wed, 9 Sep 2026 22:48:37 +0300 Subject: [PATCH 6/8] Update test_frontend_tflite.py --- tests/python/relax/test_frontend_tflite.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/python/relax/test_frontend_tflite.py b/tests/python/relax/test_frontend_tflite.py index 0a34785c0af0..dcc60f450049 100644 --- a/tests/python/relax/test_frontend_tflite.py +++ b/tests/python/relax/test_frontend_tflite.py @@ -11418,6 +11418,8 @@ def test_quantized_avg_pool2d_matches_tflite_rounding_numerically(): and the inputs below are chosen so the window sums land on and around the .5 boundaries where truncation and round-half-away disagree. + + The activation was also tested with resnet. """ builder = flatbuffers.Builder(1024) From a2adbb252de5c6725d3a40066be2bce5ea5e752c Mon Sep 17 00:00:00 2001 From: Theoo1997 Date: Fri, 11 Sep 2026 07:58:46 +0300 Subject: [PATCH 7/8] [Relax][TFLite] Size the quantized AVERAGE_POOL_2D accumulator from the input type The exact-sum lowering pre-scales the input by the window size, so its accumulator must hold max|x| * window^2. The overflow guard assumed 8-bit input (255 * window^2), but the branch also accepts int16: a VALID 17x17 pool of int16 32767 needs 2.7e9 and returned -18657 instead of 32767. Derive the bound from the input dtype's range (128 int8, 255 uint8, 32768 int16) and use int32 when it fits, int64 otherwise; the rounding and the divide run in the same type. int8 lowering is unchanged. Adds test_quantized_avg_pool2d_large_window_does_not_overflow (int16 +-extremes, int8, uint8, 17x17 window). --- .../relax/frontend/tflite/tflite_frontend.py | 45 +++++++---- tests/python/relax/test_frontend_tflite.py | 78 +++++++++++++++++++ 2 files changed, 108 insertions(+), 15 deletions(-) diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py index 1a603762b37b..33795c7dcd9a 100644 --- a/python/tvm/relax/frontend/tflite/tflite_frontend.py +++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py @@ -5745,22 +5745,37 @@ def convert_pool2d(self, op, pool_type): # constant; the padded taps are zeros and contribute nothing, # so the result is the sum over the VALID taps either way. window = filter_h * filter_w - # acc holds window * sum, and |sum| <= window * 255 for any 8-bit - # input, so |acc| <= 255 * window^2 must stay inside int32. - assert 255 * window * window < (1 << 31), ( - f"pooling window {filter_h}x{filter_w} is too large for an " - "int32 exact-sum average pool" - ) - acc = relax.op.astype(in_expr, "int32") - acc = relax.op.multiply(acc, relax.const(window, "int32")) + # The accumulator must hold the PRE-SCALED window sum. Every + # element is scaled by `window` and then `window` of them are + # summed, so |acc| <= max|x| * window^2 -- and max|x| is set by + # the INPUT type, which is not always 8-bit: this branch also + # takes int16 (a 17x17 pool of int16 32767 needs 2.7e9, past + # int32). Size the accumulator from the actual bound: int32 + # whenever it fits, which keeps the usual int8 case as cheap as + # before, else int64. + in_info = np.iinfo(self.get_tensor_type_str(input_tensor.tensor.Type())) + max_abs = max(-int(in_info.min), int(in_info.max)) # 128 int8, 255 uint8 + bound = max_abs * window * window + if bound < (1 << 31): + acc_dtype = "int32" + elif bound < (1 << 63): + acc_dtype = "int64" + else: + raise tvm.error.OpAttributeUnImplemented( + f"pooling window {filter_h}x{filter_w} is too large for an " + "exact-sum average pool even with an int64 accumulator" + ) + acc = relax.op.astype(in_expr, acc_dtype) + acc = relax.op.multiply(acc, relax.const(window, acc_dtype)) acc = relax.op.nn.avg_pool2d(acc, count_include_pad=True, **params) # relax widens the output dtype of an INTEGER avg_pool2d when the # window is large enough that an int32 accumulator could # overflow -- an 8x8 pool over int32 comes back as int64, a 2x2 - # one stays int32. Pin it so the rounding arithmetic below always - # meets int32 constants; without this an 8x8 pool fails to import - # with "Binary operators must have the same datatype". - acc = relax.op.astype(acc, "int32") + # one stays int32. Pin it back to the accumulator type so the + # rounding arithmetic below always meets constants of the same + # dtype; without this an 8x8 pool fails to import with "Binary + # operators must have the same datatype". + acc = relax.op.astype(acc, acc_dtype) # TFLite divides by the number of NON-PADDED taps, which varies # per output position once there is padding. The shapes are @@ -5773,13 +5788,13 @@ def convert_pool2d(self, op, pool_type): params["padding"], to_int_list(self.get_tensor_shape(output_tensor))[1:3], ) - half = relax.const(counts // 2, "int32") + half = relax.const(counts // 2, acc_dtype) out = relax.op.where( - relax.op.greater(acc, relax.const(0, "int32")), + relax.op.greater(acc, relax.const(0, acc_dtype)), relax.op.add(acc, half), relax.op.subtract(acc, half), ) - out = relax.op.divide(out, relax.const(counts, "int32")) + out = relax.op.divide(out, relax.const(counts, acc_dtype)) out = relax.op.astype(out, output_tensor_type_str) else: out = relax.op.nn.avg_pool2d(in_expr, **params) diff --git a/tests/python/relax/test_frontend_tflite.py b/tests/python/relax/test_frontend_tflite.py index dcc60f450049..121ae0045955 100644 --- a/tests/python/relax/test_frontend_tflite.py +++ b/tests/python/relax/test_frontend_tflite.py @@ -11490,6 +11490,84 @@ def test_quantized_avg_pool2d_matches_tflite_rounding_numerically(): np.testing.assert_array_equal(got, want) +@pytest.mark.parametrize( + "tensor_type, dtype, value", + [ + # The pre-scaled window sum is max|x| * window^2: for int16 17x17 that is + # 2.7e9, past int32, and an int32 accumulator returned -18657 / 18656. + (_tfl_tensor_type.INT16, "int16", 32767), + (_tfl_tensor_type.INT16, "int16", -32768), + (_tfl_tensor_type.INT8, "int8", 127), + (_tfl_tensor_type.INT8, "int8", -128), + (_tfl_tensor_type.UINT8, "uint8", 255), + ], +) +def test_quantized_avg_pool2d_large_window_does_not_overflow(tensor_type, dtype, value): + """A 17x17 VALID average pool of a constant input must return that constant. + + The lowering computes the exact window sum by pre-scaling the input by the + window size, so its accumulator has to hold max|x| * window^2 -- which + depends on the INPUT type, not just the window. A constant input makes the + answer obvious (the average of N copies of v is v) and puts the sum at its + extreme. + """ + size = 17 + builder = flatbuffers.Builder(1024) + qparams = _build_quantization_parameters( + builder, scale=[0.5], zero_point=[0], quantized_dimension=0 + ) + input_tensor = _build_tensor( + builder, 0, [1, size, size, 1], tensor_type=tensor_type, quantization=qparams + ) + output_tensor = _build_tensor( + builder, 1, [1, 1, 1, 1], tensor_type=tensor_type, quantization=qparams + ) + + _tfl_pool2d_options.Pool2DOptionsStart(builder) + _tfl_pool2d_options.Pool2DOptionsAddPadding(builder, _tfl_padding.VALID) + _tfl_pool2d_options.Pool2DOptionsAddStrideH(builder, 1) + _tfl_pool2d_options.Pool2DOptionsAddStrideW(builder, 1) + _tfl_pool2d_options.Pool2DOptionsAddFilterHeight(builder, size) + _tfl_pool2d_options.Pool2DOptionsAddFilterWidth(builder, size) + _tfl_pool2d_options.Pool2DOptionsAddFusedActivationFunction(builder, _tfl_activation_fn.NONE) + pool_opts = _tfl_pool2d_options.Pool2DOptionsEnd(builder) + + avg_pool_op = _build_operator( + builder, + 0, + [0], + [1], + builtin_options_type=_tfl_builtin_options.Pool2DOptions, + builtin_options=pool_opts, + ) + subgraph = _build_subgraph( + builder, + tensors=[input_tensor, output_tensor], + operators=[avg_pool_op], + inputs=[0], + outputs=[1], + ) + operator_codes = [_build_operator_code(builder, _tfl_builtin_operator.AVERAGE_POOL_2D)] + buf = _finish_tflite_model( + builder, + subgraph=subgraph, + operator_codes=operator_codes, + buffers=[_build_buffer(builder), _build_buffer(builder)], + ) + + if hasattr(tflite.Model, "Model"): + tflite_model = tflite.Model.Model.GetRootAsModel(buf, 0) + else: + tflite_model = tflite.Model.GetRootAsModel(buf, 0) + mod = from_tflite(tflite_model) + + dev = tvm.cpu(0) + vm = relax.VirtualMachine(tvm.compile(mod, target=tvm.target.Target("llvm")), dev) + x = np.full((1, size, size, 1), value, dtype=dtype) + got = vm["main"](tvm.runtime.tensor(x, dev)).numpy().reshape(-1) + np.testing.assert_array_equal(got, np.array([value], dtype=dtype)) + + def test_quantized_conv2d_per_tensor_uses_qdq(): """Quantized Conv2D with per-tensor quantization uses DQ -> conv2d -> Q.""" builder = flatbuffers.Builder(2048) From 26f71feb2816eb069690e1102cb67a778a108e70 Mon Sep 17 00:00:00 2001 From: Theoo1997 Date: Tue, 15 Sep 2026 17:17:47 +0300 Subject: [PATCH 8/8] [Relax][TFLite] Derive the quantized AVERAGE_POOL_2D divisor from the Relax shapes The per-position divisor (and the SAME padding) were folded to constants from the SERIALIZED TFLite tensor shapes, which are stale once from_tflite(..., shape_dict=...) overrides the input dimensions. For an int8 model with input (1,4,4,1), a 2x2 VALID pool and stride 2, overriding the input to (1,2,2,1) silently broadcast the result back to (1,2,2,1) instead of (1,1,1,1), and (1,6,6,1) failed to import because the actual 3x3 pooled output could not broadcast with the 2x2 counts tensor. Read the input H/W from the Relax input expression and the output H/W from the pooled Relax tensor. The input extents also drive the SAME padding, so that is now correct for overridden shapes on every Pool2D path (average, max, L2). Non-static shapes raise OpAttributeUnImplemented. Adds test_quantized_avg_pool2d_follows_shape_dict_override: the two VALID overrides from the review plus two SAME overrides whose border windows see fewer taps, checked for shape and against a NumPy reference of TFLite's integer average pool. --- .../relax/frontend/tflite/tflite_frontend.py | 31 ++++- tests/python/relax/test_frontend_tflite.py | 112 ++++++++++++++++++ 2 files changed, 139 insertions(+), 4 deletions(-) diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py index 33795c7dcd9a..fe6c420d5bae 100644 --- a/python/tvm/relax/frontend/tflite/tflite_frontend.py +++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py @@ -5641,6 +5641,22 @@ def convert_topk_v2(self, op): return out + def _pool2d_static_hw(self, expr, what): + """(H, W) of an NHWC Relax tensor, which must be static. + + Pool2D's SAME padding and the quantized average pool's divisor are + folded to constants, so they need concrete extents -- read from the + Relax expression rather than the serialized TFLite tensor, which is + stale once shape_dict overrides the input shape.""" + shape = self._infer_shape(expr) + try: + _, h, w, _ = [int(v) for v in shape] + except (TypeError, ValueError) as err: + raise tvm.error.OpAttributeUnImplemented( + f"Pool2D requires a static NHWC {what} shape, got {shape}" + ) from err + return h, w + @staticmethod def _avg_pool2d_valid_counts(in_hw, filter_hw, stride_hw, padding, out_hw): """Non-padded taps per output position, as an int32 [1, OH, OW, 1] array. @@ -5707,7 +5723,11 @@ def convert_pool2d(self, op, pool_type): in_expr = self.get_expr(input_tensor_idx) - _, input_h, input_w, _ = to_int_list(self.get_tensor_shape(input_tensor)) + # Take H and W from the Relax input, not from the serialized TFLite + # shape: from_tflite(..., shape_dict=...) can override the input + # dimensions, and then the flatbuffer's shapes are stale. SAME padding + # and the quantized divisor below both depend on the real extents. + input_h, input_w = self._pool2d_static_hw(in_expr, "input") if padding == Padding.VALID: pass @@ -5775,18 +5795,21 @@ def convert_pool2d(self, op, pool_type): # rounding arithmetic below always meets constants of the same # dtype; without this an 8x8 pool fails to import with "Binary # operators must have the same datatype". - acc = relax.op.astype(acc, acc_dtype) + acc = self.bb.normalize(relax.op.astype(acc, acc_dtype)) # TFLite divides by the number of NON-PADDED taps, which varies # per output position once there is padding. The shapes are # static, so the per-position count is folded to a constant - # here rather than computed in the graph. + # here rather than computed in the graph. Both extents come from + # the Relax graph -- the input above and the pooled output here + # -- so the constant matches the tensor it divides even when + # shape_dict overrides the serialized input shape. counts = self._avg_pool2d_valid_counts( (input_h, input_w), (filter_h, filter_w), (stride_h, stride_w), params["padding"], - to_int_list(self.get_tensor_shape(output_tensor))[1:3], + self._pool2d_static_hw(acc, "pooled output"), ) half = relax.const(counts // 2, acc_dtype) out = relax.op.where( diff --git a/tests/python/relax/test_frontend_tflite.py b/tests/python/relax/test_frontend_tflite.py index 121ae0045955..14c14edeb5d7 100644 --- a/tests/python/relax/test_frontend_tflite.py +++ b/tests/python/relax/test_frontend_tflite.py @@ -11568,6 +11568,118 @@ def test_quantized_avg_pool2d_large_window_does_not_overflow(tensor_type, dtype, np.testing.assert_array_equal(got, np.array([value], dtype=dtype)) +def _tflite_int_avg_pool2d_reference(x, filter_hw, stride_hw, same_padding): + """TFLite's reference integer AveragePool on an NHWC array: divide the sum over + the NON-PADDED taps, rounding half away from zero (reference/integer_ops/pooling.h).""" + n, in_h, in_w, c = x.shape + (f_h, f_w), (s_h, s_w) = filter_hw, stride_hw + + def axis(extent, f, s): + if not same_padding: + return (extent - f) // s + 1, 0 + out = -(-extent // s) + return out, max((out - 1) * s + f - extent, 0) // 2 + + out_h, pad_h = axis(in_h, f_h, s_h) + out_w, pad_w = axis(in_w, f_w, s_w) + y = np.zeros((n, out_h, out_w, c), dtype=x.dtype) + for oy in range(out_h): + for ox in range(out_w): + y0, x0 = oy * s_h - pad_h, ox * s_w - pad_w + win = x[:, max(y0, 0) : min(y0 + f_h, in_h), max(x0, 0) : min(x0 + f_w, in_w), :] + count = win.shape[1] * win.shape[2] + acc = win.astype("int64").sum(axis=(1, 2)) + y[:, oy, ox, :] = np.where( + acc > 0, (acc + count // 2) // count, -((-acc + count // 2) // count) + ) + return y + + +@pytest.mark.parametrize( + "padding, override, want_hw", + [ + # the review's cases: the serialized model is 4x4 -> 2x2 + (_tfl_padding.VALID, (1, 2, 2, 1), (1, 1)), # counts tensor used to broadcast to 2x2 + (_tfl_padding.VALID, (1, 6, 6, 1), (3, 3)), # used to fail: 3x3 vs a 2x2 counts tensor + # SAME: the border windows see fewer taps, and the padding itself must be + # derived from the overridden extent, not the serialized 4 + (_tfl_padding.SAME, (1, 5, 5, 1), (3, 3)), + (_tfl_padding.SAME, (1, 7, 3, 1), (4, 2)), + ], +) +def test_quantized_avg_pool2d_follows_shape_dict_override(padding, override, want_hw): + """The quantized AVERAGE_POOL_2D divisor must come from the Relax shapes. + + The per-position counts (and SAME padding) are folded to constants. They used + to be built from the SERIALIZED TFLite shapes, which are stale once + from_tflite(..., shape_dict=...) overrides the input: a smaller input got a + counts tensor that silently broadcast the result back to the old output + shape, and a larger one failed to import. + """ + builder = flatbuffers.Builder(1024) + qparams = _build_quantization_parameters( + builder, scale=[0.5], zero_point=[0], quantized_dimension=0 + ) + input_tensor = _build_tensor( + builder, 0, [1, 4, 4, 1], tensor_type=_tfl_tensor_type.INT8, quantization=qparams + ) + output_tensor = _build_tensor( + builder, 1, [1, 2, 2, 1], tensor_type=_tfl_tensor_type.INT8, quantization=qparams + ) + + _tfl_pool2d_options.Pool2DOptionsStart(builder) + _tfl_pool2d_options.Pool2DOptionsAddPadding(builder, padding) + _tfl_pool2d_options.Pool2DOptionsAddStrideH(builder, 2) + _tfl_pool2d_options.Pool2DOptionsAddStrideW(builder, 2) + _tfl_pool2d_options.Pool2DOptionsAddFilterHeight(builder, 2) + _tfl_pool2d_options.Pool2DOptionsAddFilterWidth(builder, 2) + _tfl_pool2d_options.Pool2DOptionsAddFusedActivationFunction(builder, _tfl_activation_fn.NONE) + pool_opts = _tfl_pool2d_options.Pool2DOptionsEnd(builder) + + avg_pool_op = _build_operator( + builder, + 0, + [0], + [1], + builtin_options_type=_tfl_builtin_options.Pool2DOptions, + builtin_options=pool_opts, + ) + subgraph = _build_subgraph( + builder, + tensors=[input_tensor, output_tensor], + operators=[avg_pool_op], + inputs=[0], + outputs=[1], + ) + operator_codes = [_build_operator_code(builder, _tfl_builtin_operator.AVERAGE_POOL_2D)] + buf = _finish_tflite_model( + builder, + subgraph=subgraph, + operator_codes=operator_codes, + buffers=[_build_buffer(builder), _build_buffer(builder)], + ) + + if hasattr(tflite.Model, "Model"): + tflite_model = tflite.Model.Model.GetRootAsModel(buf, 0) + else: + tflite_model = tflite.Model.GetRootAsModel(buf, 0) + from tvm.relax.frontend.tflite.tflite_frontend import _input_type + + input_name = next(iter(_input_type(tflite_model)[0])) + mod = from_tflite(tflite_model, shape_dict={input_name: override}) + + dev = tvm.cpu(0) + vm = relax.VirtualMachine(tvm.compile(mod, target=tvm.target.Target("llvm")), dev) + x = np.random.default_rng(0).integers(-128, 128, size=override, dtype=np.int64).astype("int8") + got = vm["main"](tvm.runtime.tensor(x, dev)).numpy() + + assert got.shape == (1, *want_hw, 1) + want = _tflite_int_avg_pool2d_reference( + x, (2, 2), (2, 2), same_padding=padding == _tfl_padding.SAME + ) + np.testing.assert_array_equal(got, want) + + def test_quantized_conv2d_per_tensor_uses_qdq(): """Quantized Conv2D with per-tensor quantization uses DQ -> conv2d -> Q.""" builder = flatbuffers.Builder(2048)