Arm backend: Keep integer scalars exact in the scalar-to-tensor pass

Author: abm-77

backends/arm/_passes/replace_scalar_with_tensor_pass.py

@@ -17,6 +17,7 @@
 
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.pass_base import ExportPass
+from torch._ops import OpOverload
 
 from .arm_pass import ArmPass
 
@@ -103,6 +104,35 @@ def __init__(self, tfa_pass=False, *args, **kwargs):
         # Actual selection is done per-call in call_operator.
         super().__init__(tfa_pass, _all_ops, *args, **kwargs)
 
+    def get_replacement(self, op, args, kwargs, meta):
+        # The base pass casts the scalar through float(), which rounds an
+        # integer above 2**53 before it reaches an int64 tensor.
+        if isinstance(args[1], int) and not isinstance(args[1], bool):
+            full_op = (
+                torch.ops.aten.full.default
+                if isinstance(op, OpOverload)
+                else exir_ops.edge.aten.full.default
+            )
+            tensor = args[0].to_tensor()
+            return ExportPass.call_operator(
+                self,
+                self.scalar_to_tensor_ops[op],
+                (
+                    args[0],
+                    ExportPass.call_operator(
+                        self,
+                        full_op,
+                        ((1,), args[1]),
+                        {"dtype": tensor.dtype, "device": tensor.device},
+                        meta,
+                    ),
+                    *args[2:],
+                ),
+                kwargs,
+                meta,
+            )
+        return super().get_replacement(op, args, kwargs, meta)
+
     def should_run_pass(self, graph_module: torch.fx.GraphModule) -> bool:
         for node in graph_module.graph.nodes:
             if node.op == "call_function" and node.target in _all_ops: