From 1c3e129c787f3eb078ce31c326b103a8cef1389b Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Thu, 2 Jul 2026 18:25:22 +0800 Subject: [PATCH] [Fix][Relax][TFLite] Use astype for frontend casts --- .../relax/frontend/tflite/tflite_frontend.py | 6 +- tests/python/relax/test_frontend_tflite.py | 148 ++++++++++++++++++ 2 files changed, 151 insertions(+), 3 deletions(-) diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py index 6918c600f65e..a2db13546689 100644 --- a/python/tvm/relax/frontend/tflite/tflite_frontend.py +++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py @@ -5627,9 +5627,9 @@ def convert_pool2d(self, op, pool_type): "TFLite avg_pool2dreshape requires input and output scale" "and zero points to be equal" ) - out = relax.op.cast(in_expr, dtype="int32") + out = relax.op.astype(in_expr, "int32") out = relax.op.nn.avg_pool2d(out, **params) - out = relax.op.cast(out, dtype=output_tensor_type_str) + out = relax.op.astype(out, output_tensor_type_str) else: out = relax.op.nn.avg_pool2d(in_expr, **params) elif pool_type == "max": @@ -7327,7 +7327,7 @@ def convert_dequantize(self, op): in_expr = self.exp_tab.new_const( input_value, dtype=dtype, source_name=input_tensor.tensor.Name() ) - out = relax.cast(in_expr, dtype="float32") + out = relax.op.astype(in_expr, "float32") return out in_expr = self.get_expr(input_tensor.tensor_idx) diff --git a/tests/python/relax/test_frontend_tflite.py b/tests/python/relax/test_frontend_tflite.py index 3397d08abb57..716fe1222ba8 100644 --- a/tests/python/relax/test_frontend_tflite.py +++ b/tests/python/relax/test_frontend_tflite.py @@ -4901,6 +4901,7 @@ def _get_tflite_schema_enum(enum_name): _tfl_model = _get_tflite_schema_module("Model") _tfl_operator = _get_tflite_schema_module("Operator") _tfl_operator_code = _get_tflite_schema_module("OperatorCode") +_tfl_pool2d_options = _get_tflite_schema_module("Pool2DOptions") _tfl_quantization_parameters = _get_tflite_schema_module("QuantizationParameters") _tfl_sparsity_parameters = _get_tflite_schema_module("SparsityParameters") _tfl_subgraph = _get_tflite_schema_module("SubGraph") @@ -11182,6 +11183,153 @@ def main(x: R.Tensor((2,), dtype="int8")) -> R.Tensor((2,), dtype="float32"): tvm.ir.assert_structural_equal(mod, Expected) +def test_dequantize_float16_uses_astype(): + """TFLite DEQUANTIZE float16 -> float32 uses R.astype.""" + builder = flatbuffers.Builder(1024) + + input_data = np.array([1.5, -2.0], dtype=np.float16) + + input_tensor = _build_tensor(builder, 0, [2], tensor_type=_tfl_tensor_type.FLOAT16) + output_tensor = _build_tensor(builder, 1, [2], tensor_type=_tfl_tensor_type.FLOAT32) + + dequantize_op = _build_operator(builder, 0, [0], [1]) + subgraph = _build_subgraph( + builder, + tensors=[input_tensor, output_tensor], + operators=[dequantize_op], + inputs=[], + outputs=[1], + ) + operator_codes = [_build_operator_code(builder, _tfl_builtin_operator.DEQUANTIZE)] + input_buffer = _build_buffer(builder, input_data.tobytes()) + output_buffer = _build_buffer(builder) + buf = _finish_tflite_model( + builder, + subgraph=subgraph, + operator_codes=operator_codes, + buffers=[input_buffer, output_buffer], + ) + + 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) + mod["main"] = mod["main"].without_attr("params") + + @I.ir_module + class Expected: + @R.function + def main() -> R.Tensor((2,), dtype="float32"): + R.func_attr({"num_input": 0}) + with R.dataflow(): + gv: R.Tensor((2,), dtype="float32") = R.astype( + R.const(np.array([1.5, -2.0], dtype=np.float16)), dtype="float32" + ) + R.output(gv) + return gv + + tvm.ir.assert_structural_equal(mod, Expected) + + +def test_quantized_avg_pool2d_uses_astype(): + """Quantized AVERAGE_POOL_2D casts through int32 with R.astype.""" + builder = flatbuffers.Builder(1024) + + qparams = _build_quantization_parameters( + builder, scale=[0.5], zero_point=[3], quantized_dimension=0 + ) + input_tensor = _build_tensor( + builder, + 0, + [1, 2, 2, 1], + tensor_type=_tfl_tensor_type.INT8, + quantization=qparams, + ) + output_tensor = _build_tensor( + builder, + 1, + [1, 1, 1, 1], + 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, 1) + _tfl_pool2d_options.Pool2DOptionsAddStrideW(builder, 1) + _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) + + subgraph = tflite_model.Subgraphs(0) + bb = relax.BlockBuilder() + exp_tab = tflite_frontend.ExprTable() + input_var = relax.Var("tvmgen_tensor_0", relax.TensorType((1, 2, 2, 1), dtype="int8")) + exp_tab.set_expr("tvmgen_tensor_0", input_var) + converter = tflite_frontend.OperatorConverter(tflite_model, subgraph, exp_tab, bb) + with bb.function("main", [input_var]): + with bb.dataflow(): + output = converter.convert_pool2d(subgraph.Operators(0), "average") + gv = bb.emit_output(output) + bb.emit_func_output(gv) + mod = bb.get() + + @I.ir_module + class Expected: + @R.function + def main(tvmgen_tensor_0: R.Tensor((1, 2, 2, 1), dtype="int8")) -> R.Tensor( + (1, 1, 1, 1), dtype="int8" + ): + 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, + pool_size=[2, 2], + strides=[1, 1], + dilation=[1, 1], + padding=[0, 0, 0, 0], + ceil_mode=False, + count_include_pad=False, + layout="NHWC", + out_layout="NHWC", + ) + gv: R.Tensor((1, 1, 1, 1), dtype="int8") = R.astype(lv1, dtype="int8") + R.output(gv) + return gv + + tvm.ir.assert_structural_equal(mod, Expected) + + def test_quantized_conv2d_per_tensor_uses_qdq(): """Quantized Conv2D with per-tensor quantization uses DQ -> conv2d -> Q.""" builder = flatbuffers.Builder(2048)