diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py index 44e977397307..7922a2f2919e 100644 --- a/python/tvm/relax/frontend/tflite/tflite_frontend.py +++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py @@ -391,6 +391,7 @@ def __init__(self, model, subgraph, exp_tab, ctx, conversion_state=None): "STABLEHLO_REDUCE": self._convert_stablehlo_reduce, "STABLEHLO_REDUCE_WINDOW": self._convert_stablehlo_reduce_window, "STABLEHLO_REMAINDER": self._convert_stablehlo_remainder, + "STABLEHLO_RESHAPE": self._convert_stablehlo_reshape, "STABLEHLO_RNG_BIT_GENERATOR": self._convert_stablehlo_rng_bit_generator, "STABLEHLO_RSQRT": functools.partial(self._convert_stablehlo_unary, relax_op=_op.rsqrt), "STABLEHLO_SCATTER": self._convert_stablehlo_scatter, @@ -400,11 +401,13 @@ def __init__(self, model, subgraph, exp_tab, ctx, conversion_state=None): "STABLEHLO_SHIFT_LEFT": functools.partial( self._convert_stablehlo_binary, relax_op=_op.left_shift ), + "STABLEHLO_SLICE": self._convert_stablehlo_slice, "STABLEHLO_SORT": self._convert_stablehlo_sort, "STABLEHLO_SUBTRACT": functools.partial( self._convert_stablehlo_binary, relax_op=_op.subtract ), "STABLEHLO_TANH": functools.partial(self._convert_stablehlo_unary, relax_op=_op.tanh), + "STABLEHLO_TRANSPOSE": self._convert_stablehlo_transpose, "STABLEHLO_WHILE": self._convert_stablehlo_while, "SQUEEZE": self.convert_squeeze, "STRIDED_SLICE": self.convert_strided_slice, @@ -3014,6 +3017,50 @@ def _convert_stablehlo_broadcast_in_dim(self, op): reshaped = self.bb.normalize(relax.op.reshape(in_expr, intermediate_shape)) return self.bb.normalize(relax.op.broadcast_to(reshaped, output_shape)) + def _convert_stablehlo_reshape(self, op): + """Convert STABLEHLO_RESHAPE to Relax.""" + input_tensors = self.get_input_tensors(op) + assert len(input_tensors) == 1 + output_tensors = self.get_output_tensors(op) + assert len(output_tensors) == 1 + + in_expr = self.get_tensor_expr(input_tensors[0]) + output_shape = [int(d) for d in self.get_tensor_shape(output_tensors[0])] + return self.bb.normalize(relax.op.reshape(in_expr, output_shape)) + + def _convert_stablehlo_slice(self, op): + """Convert STABLEHLO_SLICE to Relax.""" + from tflite.StablehloSliceOptions import StablehloSliceOptions + + input_tensors = self.get_input_tensors(op) + assert len(input_tensors) == 1 + assert len(self.get_output_tensors(op)) == 1 + + opts = self._get_stablehlo_options(op, StablehloSliceOptions) + begin = [int(d) for d in opts.StartIndicesAsNumpy()] + end = [int(d) for d in opts.LimitIndicesAsNumpy()] + strides = [int(d) for d in opts.StridesAsNumpy()] + axes = list(range(len(begin))) + + in_expr = self.get_tensor_expr(input_tensors[0]) + return self.bb.normalize( + relax.op.strided_slice(in_expr, axes=axes, begin=begin, end=end, strides=strides) + ) + + def _convert_stablehlo_transpose(self, op): + """Convert STABLEHLO_TRANSPOSE to Relax.""" + from tflite.StablehloTransposeOptions import StablehloTransposeOptions + + input_tensors = self.get_input_tensors(op) + assert len(input_tensors) == 1 + assert len(self.get_output_tensors(op)) == 1 + + opts = self._get_stablehlo_options(op, StablehloTransposeOptions) + permutation = [int(d) for d in opts.PermutationAsNumpy()] + + in_expr = self.get_tensor_expr(input_tensors[0]) + return self.bb.normalize(relax.op.permute_dims(in_expr, axes=permutation)) + def _convert_stablehlo_iota(self, op): """Convert STABLEHLO_IOTA to Relax (arange + broadcast).""" from tflite.StablehloIotaOptions import StablehloIotaOptions diff --git a/tests/python/relax/test_frontend_tflite.py b/tests/python/relax/test_frontend_tflite.py index 590cc4ac459f..de21491df251 100644 --- a/tests/python/relax/test_frontend_tflite.py +++ b/tests/python/relax/test_frontend_tflite.py @@ -4013,7 +4013,9 @@ def _get_tflite_schema_enum(enum_name): _tfl_stablehlo_reduce_opts = _get_tflite_schema_module("StablehloReduceOptions") _tfl_stablehlo_reduce_window_opts = _get_tflite_schema_module("StablehloReduceWindowOptions") _tfl_stablehlo_scatter_opts = _get_tflite_schema_module("StablehloScatterOptions") +_tfl_stablehlo_slice_opts = _get_tflite_schema_module("StablehloSliceOptions") _tfl_stablehlo_sort_opts = _get_tflite_schema_module("StablehloSortOptions") +_tfl_stablehlo_transpose_opts = _get_tflite_schema_module("StablehloTransposeOptions") _tfl_stablehlo_while_opts = _get_tflite_schema_module("StablehloWhileOptions") _tfl_stablehlo_rng_opts = _get_tflite_schema_module("StablehloRngBitGeneratorOptions") _tfl_call_options = _get_tflite_schema_module("CallOptions") @@ -8464,6 +8466,197 @@ def main( tvm.ir.assert_structural_equal(mod, Expected) +def _build_stablehlo_reshape_model(input_shape, output_shape): + """STABLEHLO_RESHAPE with given input and output shapes.""" + builder = flatbuffers.Builder(1024) + + builtin_op = _get_stablehlo_builtin_operator("STABLEHLO_RESHAPE") + op_code = _build_operator_code(builder, builtin_op) + + tensors = [ + _build_tensor(builder, 0, input_shape), + _build_tensor(builder, 1, output_shape), + ] + op = _build_operator(builder, 0, [0], [1]) + subgraph = _build_subgraph( + builder, + tensors=tensors, + operators=[op], + inputs=[0], + outputs=[1], + ) + buffers = [_build_buffer(builder) for _ in range(2)] + return _finish_tflite_model( + builder, subgraph=subgraph, operator_codes=[op_code], buffers=buffers + ) + + +def test_stablehlo_reshape(): + """TFLite StableHLO RESHAPE lowers to Relax reshape.""" + mod = _load_model_from_buffer( + _build_stablehlo_reshape_model(input_shape=[2, 3], output_shape=[3, 2]) + ) + + @I.ir_module + class Expected: + @R.function + def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((3, 2), dtype="float32"): + R.func_attr({"num_input": 1}) + with R.dataflow(): + gv: R.Tensor((3, 2), dtype="float32") = R.reshape(x, (3, 2)) + R.output(gv) + return gv + + tvm.ir.assert_structural_equal(mod, Expected) + + +def _build_stablehlo_slice_model(input_shape, start_indices, limit_indices, strides, output_shape): + """STABLEHLO_SLICE with static start, limit, and stride attributes.""" + builder = flatbuffers.Builder(1024) + + start_vec = _tflite_int64_vector( + builder, + _tfl_stablehlo_slice_opts.StablehloSliceOptionsStartStartIndicesVector, + start_indices, + ) + limit_vec = _tflite_int64_vector( + builder, + _tfl_stablehlo_slice_opts.StablehloSliceOptionsStartLimitIndicesVector, + limit_indices, + ) + strides_vec = _tflite_int64_vector( + builder, + _tfl_stablehlo_slice_opts.StablehloSliceOptionsStartStridesVector, + strides, + ) + + _tfl_stablehlo_slice_opts.StablehloSliceOptionsStart(builder) + _tfl_stablehlo_slice_opts.StablehloSliceOptionsAddStartIndices(builder, start_vec) + _tfl_stablehlo_slice_opts.StablehloSliceOptionsAddLimitIndices(builder, limit_vec) + _tfl_stablehlo_slice_opts.StablehloSliceOptionsAddStrides(builder, strides_vec) + slice_opts = _tfl_stablehlo_slice_opts.StablehloSliceOptionsEnd(builder) + + builtin_op = _get_stablehlo_builtin_operator("STABLEHLO_SLICE") + op_code = _build_operator_code(builder, builtin_op) + + tensors = [ + _build_tensor(builder, 0, input_shape), + _build_tensor(builder, 1, output_shape), + ] + op = _build_operator( + builder, + 0, + [0], + [1], + builtin_options2_type=_tfl_builtin_options2.StablehloSliceOptions, + builtin_options2=slice_opts, + ) + subgraph = _build_subgraph( + builder, + tensors=tensors, + operators=[op], + inputs=[0], + outputs=[1], + ) + buffers = [_build_buffer(builder) for _ in range(2)] + return _finish_tflite_model( + builder, subgraph=subgraph, operator_codes=[op_code], buffers=buffers + ) + + +def test_stablehlo_slice(): + """TFLite StableHLO SLICE lowers to Relax strided_slice.""" + mod = _load_model_from_buffer( + _build_stablehlo_slice_model( + input_shape=[4, 5], + start_indices=[1, 0], + limit_indices=[4, 4], + strides=[2, 2], + output_shape=[2, 2], + ) + ) + + @I.ir_module + class Expected: + @R.function + def main(x: R.Tensor((4, 5), dtype="float32")) -> R.Tensor((2, 2), dtype="float32"): + R.func_attr({"num_input": 1}) + with R.dataflow(): + gv: R.Tensor((2, 2), dtype="float32") = R.strided_slice( + x, axes=[0, 1], begin=[1, 0], end=[4, 4], strides=[2, 2] + ) + R.output(gv) + return gv + + tvm.ir.assert_structural_equal(mod, Expected) + + +def _build_stablehlo_transpose_model(input_shape, permutation, output_shape): + """STABLEHLO_TRANSPOSE with a static permutation.""" + builder = flatbuffers.Builder(1024) + + perm_vec = _tflite_int64_vector( + builder, + _tfl_stablehlo_transpose_opts.StablehloTransposeOptionsStartPermutationVector, + permutation, + ) + _tfl_stablehlo_transpose_opts.StablehloTransposeOptionsStart(builder) + _tfl_stablehlo_transpose_opts.StablehloTransposeOptionsAddPermutation(builder, perm_vec) + transpose_opts = _tfl_stablehlo_transpose_opts.StablehloTransposeOptionsEnd(builder) + + builtin_op = _get_stablehlo_builtin_operator("STABLEHLO_TRANSPOSE") + op_code = _build_operator_code(builder, builtin_op) + + tensors = [ + _build_tensor(builder, 0, input_shape), + _build_tensor(builder, 1, output_shape), + ] + op = _build_operator( + builder, + 0, + [0], + [1], + builtin_options2_type=_tfl_builtin_options2.StablehloTransposeOptions, + builtin_options2=transpose_opts, + ) + subgraph = _build_subgraph( + builder, + tensors=tensors, + operators=[op], + inputs=[0], + outputs=[1], + ) + buffers = [_build_buffer(builder) for _ in range(2)] + return _finish_tflite_model( + builder, subgraph=subgraph, operator_codes=[op_code], buffers=buffers + ) + + +def test_stablehlo_transpose(): + """TFLite StableHLO TRANSPOSE lowers to Relax permute_dims.""" + mod = _load_model_from_buffer( + _build_stablehlo_transpose_model( + input_shape=[2, 3, 4], permutation=[1, 2, 0], output_shape=[3, 4, 2] + ) + ) + + @I.ir_module + class Expected: + @R.function + def main( + x: R.Tensor((2, 3, 4), dtype="float32") + ) -> R.Tensor((3, 4, 2), dtype="float32"): + R.func_attr({"num_input": 1}) + with R.dataflow(): + gv: R.Tensor((3, 4, 2), dtype="float32") = R.permute_dims( + x, axes=[1, 2, 0] + ) + R.output(gv) + return gv + + tvm.ir.assert_structural_equal(mod, Expected) + + def _build_stablehlo_broadcast_in_dim_model(input_shape, broadcast_dims, output_shape): """STABLEHLO_BROADCAST_IN_DIM with given broadcast dimensions.""" builder = flatbuffers.Builder(1024)