From 40d0cffe8fa815db317e2c3149242a0b0340a99c Mon Sep 17 00:00:00 2001 From: THINKER-ONLY <3193304954@qq.com> Date: Wed, 13 May 2026 14:43:00 +0800 Subject: [PATCH] [Relax][Frontend][TFLite] Add REDUCE_WINDOW support This commit adds Relax TFLite frontend support for the builtin REDUCE_WINDOW operator. The converter parses ReduceWindowOptions from BuiltinOptions2, validates the static window attributes, and lowers numeric and boolean reductions through topi.sliding_window plus Relax reductions. The implementation covers ADD, MUL, MINIMUM, MAXIMUM, ALL, and ANY reduce functions. Empty output shapes are handled directly with Relax zeros, while quantized REDUCE_WINDOW and dynamic window attributes are left unsupported with explicit errors. Tests add minimal hand-built TFLite flatbuffer fixtures and structural-equal coverage for all supported reduce functions, empty output dimensions, unsupported reduce functions, rank mismatch, and invalid stride values. --- .../relax/frontend/tflite/tflite_frontend.py | 164 ++++++++ tests/python/relax/test_frontend_tflite.py | 379 ++++++++++++++++++ 2 files changed, 543 insertions(+) diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py index 145e953394cd..6420283875a0 100644 --- a/python/tvm/relax/frontend/tflite/tflite_frontend.py +++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py @@ -212,6 +212,7 @@ def __init__(self, model, subgraph, exp_tab, ctx): "REDUCE_MAX": functools.partial(self._convert_reduce, relax_op=_op.max), "REDUCE_MIN": functools.partial(self._convert_reduce, relax_op=_op.min), "REDUCE_PROD": functools.partial(self._convert_reduce, relax_op=_op.prod), + "REDUCE_WINDOW": self.convert_reduce_window, "RELU": self.convert_relu, "RELU6": self.convert_relu6, "RELU_N1_TO_1": self.convert_relu_n1_to_1, @@ -2462,6 +2463,169 @@ def _convert_reduce(self, relax_op, op): return out + def convert_reduce_window(self, op): + """Convert TFLite REDUCE_WINDOW.""" + + from tflite.BuiltinOptions2 import BuiltinOptions2 + from tflite.ReduceWindowFunction import ReduceWindowFunction + from tflite.ReduceWindowOptions import ReduceWindowOptions + + input_tensors = self.get_input_tensors(op) + output_tensors = self.get_output_tensors(op) + assert len(input_tensors) == 5, "input tensors length should be 5" + assert len(output_tensors) == 1, "output tensors length should be 1" + + if op.BuiltinOptions2Type() != BuiltinOptions2.ReduceWindowOptions: + raise tvm.error.OpAttributeUnImplemented( + "TFLite REDUCE_WINDOW requires ReduceWindowOptions." + ) + + ( + input_tensor, + init_tensor, + window_shape_tensor, + window_strides_tensor, + window_dilations_tensor, + ) = input_tensors + output_tensor = output_tensors[0] + + if any( + self.has_expr(tensor.tensor_idx) + for tensor in [window_shape_tensor, window_strides_tensor, window_dilations_tensor] + ): + raise tvm.error.OpNotImplemented( + "TFLite REDUCE_WINDOW requires constant window_shape, " + "window_strides, and window_dilations." + ) + + input_shape = to_int_list(self.get_tensor_shape(input_tensor)) + output_shape = to_int_list(self.get_tensor_shape(output_tensor)) + input_dtype = self.get_tensor_type_str(input_tensor.tensor.Type()) + output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type()) + + if input_tensor.qnn_params or output_tensor.qnn_params: + raise tvm.error.OpNotImplemented( + "Quantized TFLite REDUCE_WINDOW is not yet supported in the Relax frontend." + ) + + if input_dtype != output_dtype: + raise tvm.error.OpAttributeUnImplemented( + "TFLite REDUCE_WINDOW requires input and output dtypes to match." + ) + + if len(to_int_list(self.get_tensor_shape(init_tensor))) != 0: + raise tvm.error.OpNotImplemented( + "TFLite REDUCE_WINDOW only supports scalar init_value." + ) + + options = ReduceWindowOptions() + op_options = op.BuiltinOptions2() + options.Init(op_options.Bytes, op_options.Pos) + reduce_function = options.ReduceFunction() + + if reduce_function == ReduceWindowFunction.UNSUPPORTED: + raise tvm.error.OpNotImplemented( + "TFLite REDUCE_WINDOW with UNSUPPORTED reduce_function is not supported." + ) + + window_shape = to_int_list(self.get_tensor_value(window_shape_tensor)) + window_strides = to_int_list(self.get_tensor_value(window_strides_tensor)) + window_dilations = to_int_list(self.get_tensor_value(window_dilations_tensor)) + rank = len(input_shape) + + if not (len(window_shape) == len(window_strides) == len(window_dilations) == rank): + raise tvm.error.OpAttributeUnImplemented( + "TFLite REDUCE_WINDOW window_shape, window_strides, and window_dilations " + "must match input rank." + ) + + if any(value <= 0 for value in window_shape + window_strides + window_dilations): + raise tvm.error.OpAttributeUnImplemented( + "TFLite REDUCE_WINDOW window dimensions, strides, and dilations must be positive." + ) + + dilated_window_shape = [ + (window_dim - 1) * dilation + 1 + for window_dim, dilation in zip(window_shape, window_dilations) + ] + expected_output_shape = [ + 0 if input_dim < dilated_dim else (input_dim - dilated_dim) // stride + 1 + for input_dim, dilated_dim, stride in zip( + input_shape, dilated_window_shape, window_strides + ) + ] + + numeric_reduce_functions = ( + ReduceWindowFunction.ADD, + ReduceWindowFunction.MUL, + ReduceWindowFunction.MINIMUM, + ReduceWindowFunction.MAXIMUM, + ) + bool_reduce_functions = ( + ReduceWindowFunction.ALL, + ReduceWindowFunction.ANY, + ) + + if reduce_function in numeric_reduce_functions and input_dtype == "bool": + raise tvm.error.OpAttributeUnImplemented( + "TFLite REDUCE_WINDOW numeric reductions expect numeric input." + ) + if reduce_function in bool_reduce_functions and input_dtype != "bool": + raise tvm.error.OpAttributeUnImplemented( + "TFLite REDUCE_WINDOW boolean reductions expect bool input." + ) + + if output_shape != expected_output_shape: + raise tvm.error.OpAttributeUnImplemented( + "TFLite REDUCE_WINDOW output shape does not match input/window parameters." + ) + + if any(output_dim == 0 for output_dim in output_shape): + return relax.op.zeros(output_shape, output_dtype) + + data = self.get_tensor_expr(input_tensor) + init_value = self.get_tensor_expr(init_tensor) + + windowed = relax.op.call_dps_packed( + "topi.sliding_window", + ( + data, + 0, + relax.ShapeExpr(dilated_window_shape), + relax.ShapeExpr(window_strides), + ), + out_sinfo=relax.TensorStructInfo(output_shape + dilated_window_shape, input_dtype), + ) + + if any(dilation != 1 for dilation in window_dilations): + windowed = relax.op.strided_slice( + windowed, + axes=list(range(rank, 2 * rank)), + begin=[0] * rank, + end=dilated_window_shape, + strides=window_dilations, + ) + + reduce_axes = list(range(rank, 2 * rank)) + if reduce_function == ReduceWindowFunction.ADD: + return relax.op.add(relax.op.sum(windowed, axis=reduce_axes), init_value) + if reduce_function == ReduceWindowFunction.MUL: + return relax.op.multiply(relax.op.prod(windowed, axis=reduce_axes), init_value) + if reduce_function == ReduceWindowFunction.MINIMUM: + return relax.op.minimum(relax.op.min(windowed, axis=reduce_axes), init_value) + if reduce_function == ReduceWindowFunction.MAXIMUM: + return relax.op.maximum(relax.op.max(windowed, axis=reduce_axes), init_value) + if reduce_function == ReduceWindowFunction.ALL: + reduced = relax.op.min(relax.op.astype(windowed, "int8"), axis=reduce_axes) + return relax.op.logical_and(relax.op.astype(reduced, "bool"), init_value) + if reduce_function == ReduceWindowFunction.ANY: + reduced = relax.op.max(relax.op.astype(windowed, "int8"), axis=reduce_axes) + return relax.op.logical_or(relax.op.astype(reduced, "bool"), init_value) + + raise tvm.error.OpNotImplemented( + f"TFLite REDUCE_WINDOW reduce_function {reduce_function} is not supported." + ) + def _convert_reduce_bool(self, relax_op, op): """Convert TFLite REDUCE_ANY / REDUCE_ALL (bool-only ops). diff --git a/tests/python/relax/test_frontend_tflite.py b/tests/python/relax/test_frontend_tflite.py index a53906d2f147..0a55271b7c6d 100644 --- a/tests/python/relax/test_frontend_tflite.py +++ b/tests/python/relax/test_frontend_tflite.py @@ -3673,6 +3673,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_reduce_window_options = _get_tflite_schema_module("ReduceWindowOptions") _tfl_sparsity_parameters = _get_tflite_schema_module("SparsityParameters") _tfl_subgraph = _get_tflite_schema_module("SubGraph") _tfl_tensor = _get_tflite_schema_module("Tensor") @@ -3683,6 +3684,7 @@ def _get_tflite_schema_enum(enum_name): _tfl_dimension_type = _get_tflite_schema_enum("DimensionType") _tfl_fc_weights_format = _get_tflite_schema_enum("FullyConnectedOptionsWeightsFormat") _tfl_padding = _get_tflite_schema_enum("Padding") +_tfl_reduce_window_function = _get_tflite_schema_enum("ReduceWindowFunction") _tfl_sparse_index_vector = _get_tflite_schema_enum("SparseIndexVector") _tfl_tensor_type = _get_tflite_schema_enum("TensorType") @@ -3704,6 +3706,13 @@ def _tflite_int32_vector(builder, start_vector_fn, values): return builder.EndVector() +def _tflite_int64_vector(builder, start_vector_fn, values): + start_vector_fn(builder, len(values)) + for value in reversed(values): + builder.PrependInt64(value) + return builder.EndVector() + + def _tflite_offset_vector(builder, start_vector_fn, offsets): start_vector_fn(builder, len(offsets)) for offset in reversed(offsets): @@ -3845,6 +3854,376 @@ def _load_model_from_buffer(model_bytes): return mod +def _build_reduce_window_options(builder, reduce_function): + _tfl_reduce_window_options.ReduceWindowOptionsStart(builder) + _tfl_reduce_window_options.ReduceWindowOptionsAddReduceFunction(builder, reduce_function) + return _tfl_reduce_window_options.ReduceWindowOptionsEnd(builder) + + +def _reduce_window_output_shape(input_shape, window_shape, window_strides, window_dilations): + output_shape = [] + for input_dim, window_dim, stride, dilation in zip( + input_shape, window_shape, window_strides, window_dilations + ): + dilated_window = (window_dim - 1) * dilation + 1 + if stride <= 0: + output_shape.append(0) + elif input_dim < dilated_window: + output_shape.append(0) + else: + output_shape.append((input_dim - dilated_window) // stride + 1) + return tuple(output_shape) + + +def _build_reduce_window_model( + *, + input_shape, + init_value, + window_shape, + window_strides, + window_dilations, + output_shape=None, + reduce_function, + tensor_type=None, + value_dtype=np.float32, +): + builder = flatbuffers.Builder(1024) + if tensor_type is None: + tensor_type = _tfl_tensor_type.FLOAT32 + + input_tensor_idx = 0 + init_tensor_idx = 1 + window_shape_tensor_idx = 2 + window_strides_tensor_idx = 3 + window_dilations_tensor_idx = 4 + output_tensor_idx = 5 + + if output_shape is None: + output_shape = _reduce_window_output_shape( + input_shape, window_shape, window_strides, window_dilations + ) + + input_tensor = _build_tensor(builder, 1, input_shape, tensor_type=tensor_type) + init_tensor = _build_tensor(builder, 2, [], tensor_type=tensor_type) + window_shape_tensor = _build_tensor( + builder, 3, [len(window_shape)], tensor_type=_tfl_tensor_type.INT64 + ) + window_strides_tensor = _build_tensor( + builder, 4, [len(window_strides)], tensor_type=_tfl_tensor_type.INT64 + ) + window_dilations_tensor = _build_tensor( + builder, 5, [len(window_dilations)], tensor_type=_tfl_tensor_type.INT64 + ) + output_tensor = _build_tensor(builder, 6, output_shape, tensor_type=tensor_type) + + reduce_window_opts = _build_reduce_window_options(builder, reduce_function) + reduce_window_op = _build_operator( + builder, + 0, + [ + input_tensor_idx, + init_tensor_idx, + window_shape_tensor_idx, + window_strides_tensor_idx, + window_dilations_tensor_idx, + ], + [output_tensor_idx], + builtin_options2_type=_tfl_builtin_options2.ReduceWindowOptions, + builtin_options2=reduce_window_opts, + ) + + subgraph = _build_subgraph( + builder, + tensors=[ + input_tensor, + init_tensor, + window_shape_tensor, + window_strides_tensor, + window_dilations_tensor, + output_tensor, + ], + operators=[reduce_window_op], + inputs=[input_tensor_idx], + outputs=[output_tensor_idx], + ) + operator_codes = [_build_operator_code(builder, _tfl_builtin_operator.REDUCE_WINDOW)] + + buffers = [ + _build_buffer(builder), + _build_buffer(builder), + _build_buffer(builder, np.asarray([init_value], dtype=value_dtype).tobytes()), + _build_buffer(builder, np.asarray(window_shape, dtype=np.int64).tobytes()), + _build_buffer(builder, np.asarray(window_strides, dtype=np.int64).tobytes()), + _build_buffer(builder, np.asarray(window_dilations, dtype=np.int64).tobytes()), + _build_buffer(builder), + ] + + return _finish_tflite_model( + builder, subgraph=subgraph, operator_codes=operator_codes, buffers=buffers + ) + + +def _from_reduce_window_model(**kwargs): + return _load_model_from_buffer(_build_reduce_window_model(**kwargs)) + + +def _reduce_window_dilated_shape(window_shape, window_dilations): + return [ + (window_dim - 1) * dilation + 1 + for window_dim, dilation in zip(window_shape, window_dilations) + ] + + +def _make_reduce_window_numeric_expected( + *, + input_shape, + init_value, + window_shape, + window_strides, + window_dilations, + reduce_op, + combine_op, + dtype="float32", +): + output_shape = _reduce_window_output_shape( + input_shape, window_shape, window_strides, window_dilations + ) + dilated_window_shape = _reduce_window_dilated_shape(window_shape, window_dilations) + rank = len(input_shape) + + bb = relax.BlockBuilder() + x = relax.Var("tvmgen_tensor_0", relax.TensorStructInfo(input_shape, dtype)) + with bb.function("main", [x]): + with bb.dataflow(): + windowed = bb.emit( + relax.op.call_dps_packed( + "topi.sliding_window", + ( + x, + 0, + relax.ShapeExpr(dilated_window_shape), + relax.ShapeExpr(window_strides), + ), + out_sinfo=relax.TensorStructInfo( + output_shape + tuple(dilated_window_shape), dtype + ), + ) + ) + if any(dilation != 1 for dilation in window_dilations): + windowed = bb.emit( + relax.op.strided_slice( + windowed, + axes=list(range(rank, 2 * rank)), + begin=[0] * rank, + end=dilated_window_shape, + strides=window_dilations, + ) + ) + reduced = bb.emit(reduce_op(windowed, axis=list(range(rank, 2 * rank)))) + gv = bb.emit_output(combine_op(reduced, relax.const(init_value, dtype))) + bb.emit_func_output(gv) + + mod = bb.get() + mod["main"] = mod["main"].with_attr("num_input", 1) + return mod + + +def _make_reduce_window_bool_expected( + *, + input_shape, + init_value, + window_shape, + window_strides, + window_dilations, + reduce_op, + combine_op, +): + output_shape = _reduce_window_output_shape( + input_shape, window_shape, window_strides, window_dilations + ) + dilated_window_shape = _reduce_window_dilated_shape(window_shape, window_dilations) + rank = len(input_shape) + + bb = relax.BlockBuilder() + x = relax.Var("tvmgen_tensor_0", relax.TensorStructInfo(input_shape, "bool")) + with bb.function("main", [x]): + with bb.dataflow(): + windowed = bb.emit( + relax.op.call_dps_packed( + "topi.sliding_window", + ( + x, + 0, + relax.ShapeExpr(dilated_window_shape), + relax.ShapeExpr(window_strides), + ), + out_sinfo=relax.TensorStructInfo( + output_shape + tuple(dilated_window_shape), "bool" + ), + ) + ) + cast_windowed = bb.emit(relax.op.astype(windowed, "int8")) + reduced = bb.emit(reduce_op(cast_windowed, axis=list(range(rank, 2 * rank)))) + reduced_bool = bb.emit(relax.op.astype(reduced, "bool")) + gv = bb.emit_output(combine_op(reduced_bool, relax.const(init_value, "bool"))) + bb.emit_func_output(gv) + + mod = bb.get() + mod["main"] = mod["main"].with_attr("num_input", 1) + return mod + + +def _make_reduce_window_empty_expected(*, input_shape, output_shape, dtype="float32"): + bb = relax.BlockBuilder() + x = relax.Var("tvmgen_tensor_0", relax.TensorStructInfo(input_shape, dtype)) + with bb.function("main", [x]): + with bb.dataflow(): + gv = bb.emit_output(relax.op.zeros(output_shape, dtype)) + bb.emit_func_output(gv) + + mod = bb.get() + mod["main"] = mod["main"].with_attr("num_input", 1) + return mod + + +def test_reduce_window_unsupported_function(): + with pytest.raises(tvm.error.OpNotImplemented, match="UNSUPPORTED reduce_function"): + _from_reduce_window_model( + input_shape=(4,), + init_value=0.0, + window_shape=[2], + window_strides=[1], + window_dilations=[1], + reduce_function=_tfl_reduce_window_function.UNSUPPORTED, + ) + + +@pytest.mark.parametrize( + "reduce_function, reduce_op, combine_op", + [ + (_tfl_reduce_window_function.ADD, relax.op.sum, relax.op.add), + (_tfl_reduce_window_function.MUL, relax.op.prod, relax.op.multiply), + (_tfl_reduce_window_function.MINIMUM, relax.op.min, relax.op.minimum), + (_tfl_reduce_window_function.MAXIMUM, relax.op.max, relax.op.maximum), + ], +) +def test_reduce_window_numeric_modes(reduce_function, reduce_op, combine_op): + input_shape = (4, 5) + init_value = 1.0 + window_shape = [2, 2] + window_strides = [1, 2] + window_dilations = [2, 1] + mod = _from_reduce_window_model( + input_shape=input_shape, + init_value=init_value, + window_shape=window_shape, + window_strides=window_strides, + window_dilations=window_dilations, + reduce_function=reduce_function, + ) + expected = _make_reduce_window_numeric_expected( + input_shape=input_shape, + init_value=init_value, + window_shape=window_shape, + window_strides=window_strides, + window_dilations=window_dilations, + reduce_op=reduce_op, + combine_op=combine_op, + ) + tvm.ir.assert_structural_equal(mod, expected) + + +@pytest.mark.parametrize( + "reduce_function, reduce_op, combine_op, init_value", + [ + (_tfl_reduce_window_function.ALL, relax.op.min, relax.op.logical_and, True), + (_tfl_reduce_window_function.ANY, relax.op.max, relax.op.logical_or, False), + ], +) +def test_reduce_window_bool_modes(reduce_function, reduce_op, combine_op, init_value): + input_shape = (5,) + window_shape = [3] + window_strides = [2] + window_dilations = [1] + mod = _from_reduce_window_model( + input_shape=input_shape, + init_value=init_value, + window_shape=window_shape, + window_strides=window_strides, + window_dilations=window_dilations, + reduce_function=reduce_function, + tensor_type=_tfl_tensor_type.BOOL, + value_dtype=np.bool_, + ) + expected = _make_reduce_window_bool_expected( + input_shape=input_shape, + init_value=init_value, + window_shape=window_shape, + window_strides=window_strides, + window_dilations=window_dilations, + reduce_op=reduce_op, + combine_op=combine_op, + ) + tvm.ir.assert_structural_equal(mod, expected) + + +def test_reduce_window_empty_output_dimension(): + input_shape = (2,) + window_shape = [3] + window_strides = [1] + window_dilations = [1] + mod = _from_reduce_window_model( + input_shape=input_shape, + init_value=0.0, + window_shape=window_shape, + window_strides=window_strides, + window_dilations=window_dilations, + reduce_function=_tfl_reduce_window_function.ADD, + ) + expected = _make_reduce_window_empty_expected( + input_shape=input_shape, + output_shape=(0,), + ) + tvm.ir.assert_structural_equal(mod, expected) + + +def test_reduce_window_mismatched_window_rank(): + with pytest.raises(tvm.error.OpAttributeUnImplemented, match="must match input rank"): + _from_reduce_window_model( + input_shape=(4, 5), + init_value=0.0, + window_shape=[2], + window_strides=[1], + window_dilations=[1], + reduce_function=_tfl_reduce_window_function.ADD, + ) + + +def test_reduce_window_non_positive_stride(): + with pytest.raises(tvm.error.OpAttributeUnImplemented, match="must be positive"): + _from_reduce_window_model( + input_shape=(4,), + init_value=0.0, + window_shape=[2], + window_strides=[0], + window_dilations=[1], + reduce_function=_tfl_reduce_window_function.ADD, + ) + + +def test_reduce_window_inconsistent_output_shape(): + with pytest.raises(tvm.error.OpAttributeUnImplemented, match="output shape"): + _from_reduce_window_model( + input_shape=(5,), + init_value=0.0, + window_shape=[2], + window_strides=[1], + window_dilations=[1], + output_shape=(3,), + reduce_function=_tfl_reduce_window_function.ADD, + ) + + def _get_stablehlo_builtin_operator(builtin_name): if not hasattr(_tfl_builtin_operator, builtin_name): pytest.skip(f"TFLite schema does not provide BuiltinOperator.{builtin_name}")