Arm backend: Fix Rescale debug names (#20787)

Fix Rescale debug names


cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

Signed-off-by: Elena Zhelezina <elena.zhelezina@arm.com>

Author: wwwind

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,

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)