diff --git a/backends/arm/operators/op_tosa_rescale.py b/backends/arm/operators/op_tosa_rescale.py index dfaabecb41b..48622b62172 100644 --- a/backends/arm/operators/op_tosa_rescale.py +++ b/backends/arm/operators/op_tosa_rescale.py @@ -150,71 +150,67 @@ def _create_const_ops_for_rescale( return [multipliers.name, shifts.name, input_zp.name, output_zp.name] -def _build_rescale( - tosa_fb: Any, - scale: list[float], - input_node: Any, - output_name: str, - output_type: Any, - input_zp: list[int], - output_zp: list[int], - rounding_mode: ts.RoundingMode, - per_channel: bool = False, - is_scale32: bool = True, - input_unsigned: bool = False, - output_unsigned: bool = False, -): - """Insert a TOSA RESCALE operator configured for the quantized path. +@register_node_visitor +class RescaleVisitor(NodeVisitor): + target = "tosa.RESCALE.default" - Args: - tosa_fb (Any): Graph builder receiving the RESCALE operator. - scale (list[float]): Scale factors applied during rescaling. - input_node (Any): Input tensor node feeding the operator. - output_name (str): Name assigned to the RESCALE output tensor. - output_type (ts.DType): Data type of the output tensor. - input_zp (list[int]): Quantization zero points for the input tensor. - output_zp (list[int]): Quantization zero points for the output tensor. - rounding_mode (ts.RoundingMode): Rounding policy for the RESCALE op. - per_channel (bool): Whether scales are applied per output channel. - is_scale32 (bool): Declared scale width; ignored when the input type is - ``ts.DType.INT48``. + def _build_rescale( + self, + node: Node, + tosa_graph: Any, + scale: list[float], + input_node: TosaArg, + output: TosaArg, + input_zp: list[int], + output_zp: list[int], + rounding_mode: ts.RoundingMode, + per_channel: bool = False, + input_unsigned: bool = False, + output_unsigned: bool = False, + ) -> None: + """Insert a TOSA RESCALE operator configured for the quantized path. - """ - scaleWidth = 16 if input_node.dtype == ts.DType.INT48 else 32 - is_scale32 = False if input_node.dtype == ts.DType.INT48 else True - multipliers, shifts = _compute_multiplier_and_shift(scale, scaleWidth) - rescale_inputs = _create_const_ops_for_rescale( - tosa_fb, - is_scale32, - input_node.dtype, - output_name, - multipliers, - shifts, - input_zp, - output_zp, - output_type, - ts, - ) - attr_rescale = ts.TosaSerializerAttribute() - attr_rescale.RescaleAttribute( - scale32=is_scale32, - rounding_mode=rounding_mode, - per_channel=per_channel, - input_unsigned=input_unsigned, - output_unsigned=output_unsigned, - ) + RESCALE is serialized through NodeVisitor._serialize_operator so that + debug-location metadata is attached consistently with other TOSA ops. + The scale width is derived from the input dtype: INT48 uses 16-bit + multipliers, otherwise 32-bit multipliers are used. - tosa_fb.addOperator( - ts.Op.RESCALE, - [input_node.name, *rescale_inputs], - [output_name], - attr_rescale, - ) + """ + scale_width = 16 if input_node.dtype == ts.DType.INT48 else 32 + is_scale32 = input_node.dtype != ts.DType.INT48 + multipliers, shifts = _compute_multiplier_and_shift(scale, scale_width) -@register_node_visitor -class RescaleVisitor(NodeVisitor): - target = "tosa.RESCALE.default" + rescale_inputs = _create_const_ops_for_rescale( + tosa_graph, + is_scale32, + input_node.dtype, + output.name, + multipliers, + shifts, + input_zp, + output_zp, + output.dtype, + ts, + ) + + attr_rescale = ts.TosaSerializerAttribute() + attr_rescale.RescaleAttribute( + scale32=is_scale32, + rounding_mode=rounding_mode, + per_channel=per_channel, + input_unsigned=input_unsigned, + output_unsigned=output_unsigned, + ) + + self._serialize_operator( + node, + tosa_graph, + ts.Op.RESCALE, + [input_node.name, *rescale_inputs], + [output.name], + attr_rescale, + ) def define_node( self, @@ -254,12 +250,12 @@ def define_node( raise ValueError( f"If output dtype is not int8 or int16, output_zp must be 0. Got {ts.DTypeNames[output_dtype]}, {output_zp=}" ) - _build_rescale( - tosa_graph, + self._build_rescale( + node=node, + tosa_graph=tosa_graph, scale=scales, input_node=inputs[0], - output_name=output.name, - output_type=output.dtype, + output=output, input_zp=[input_zp], output_zp=[output_zp], rounding_mode=ts.RoundingMode.SINGLE_ROUND, diff --git a/backends/arm/test/misc/test_debug_feats.py b/backends/arm/test/misc/test_debug_feats.py index 774dfd41f98..246a64dcda6 100644 --- a/backends/arm/test/misc/test_debug_feats.py +++ b/backends/arm/test/misc/test_debug_feats.py @@ -21,10 +21,17 @@ EthosU55PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) from executorch.backends.test.harness.stages import StageType input_t1 = Tuple[torch.Tensor] # Input x +input_t2 = Tuple[torch.Tensor, torch.Tensor] + + +class AddModel(torch.nn.Module): + def forward(self, x, y): + return x + y class Linear(torch.nn.Module): @@ -363,3 +370,47 @@ def test_fail_dump_ops_u55_INT(capsys, test_data: input_t1): error_msg = "Can not get operator distribution for Vela command stream." with pytest.raises(NotImplementedError, match=error_msg): pipeline.run() + + +def _is_rescale_op(op: dict) -> bool: + return op.get("op") == "RESCALE" or op.get("opcode") == "RESCALE" + + +@common.SkipIfNoModelConverter +def test_vgf_rescale_tosa_debug_location_is_not_empty(): + # Repro test for a bug with debug loc for RESCALE + with tempfile.TemporaryDirectory() as tmpdir: + pipeline = VgfPipeline[input_t2]( + module=AddModel(), + test_data=(torch.ones(5), 2 * torch.ones(5)), + aten_op="torch.ops.aten.add.Tensor", + exir_op="executorch_exir_dialects_edge__ops_aten_add_Tensor", + run_on_vulkan_runtime=False, + vgf_compiler_flags="--emit-debug-info", + symmetric_io_quantization=True, + custom_path=tmpdir, + tosa_debug_mode=ArmCompileSpec.DebugMode.TOSA, + tosa_spec="TOSA-1.0+INT", + ) + + pipeline.run() + + tosa_files = list(Path(tmpdir).glob("*.tosa")) + assert tosa_files, "Expected VGF lowering to dump a TOSA artifact" + + rescale_ops = [] + for tosa_file in tosa_files: + with tosa_file.open("rb") as f: + tosa_json = dbg_tosa_fb_to_json(f.read()) + + ops = tosa_json["regions"][0]["blocks"][0]["operators"] + rescale_ops.extend([op for op in ops if _is_rescale_op(op)]) + + assert ( + rescale_ops + ), "Expected the quantized AddModel repro to emit TOSA RESCALE ops" + + for op in rescale_ops: + location_text = op["location"]["text"] + assert location_text, f"RESCALE op has empty debug location: {op}" + json.loads(location_text)