Arm backend: Prevent illegal fusion in FuseEqualPlaceholdersPass (#11781)

- Constant placeholders with same values but different data types, such
as int32 and fp32, shouldn't be fused into a single placeholder.
Otherwise, some operators will have operands with mismatched dtypes.
- Fix the bug by adding a dtype check to fuse only constants with
matching types and same values.

Signed-off-by: Yufeng Shi <[email protected]>

Author: YufengShi-dudu

backends/arm/_passes/fuse_equal_placeholders_pass.py

@@ -49,7 +49,11 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
                 if tensor2 is None:
                     continue
 
-                if torch.equal(tensor1, tensor2):
+                if (
+                    tensor1.dtype == tensor2.dtype
+                    and tensor1.shape == tensor2.shape
+                    and torch.allclose(tensor1, tensor2, atol=1e-08)
+                ):
                     eq_nodes.append(node2)
 
             if len(eq_nodes) > 1:

backends/arm/test/passes/test_fuse_equal_placeholders_ops_pass.py

@@ -10,7 +10,10 @@
 from executorch.backends.arm._passes.fuse_equal_placeholders_pass import (
     FuseEqualPlaceholdersPass,
 )
-from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
+from executorch.backends.arm.test.tester.test_pipeline import (
+    PassPipeline,
+    TosaPipelineMI,
+)
 
 input_t = Tuple[torch.Tensor]  # Input x
 
@@ -54,6 +57,25 @@ def forward(self, x):
         return self.fc1(x) + self.fc2(x)
 
 
+class NotFuseTensorWithDifferentType(torch.nn.Module):
+
+    ops_before_pass = {}
+    ops_after_pass = {}
+    ops_not_after_pass = []
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        """
+        Args:
+            x: A float tensor (dtype=torch.float32)
+            y: An int tensor (dtype=torch.int32)
+        """
+        a = torch.tensor(1.0, dtype=torch.float32)
+        b = torch.tensor(1, dtype=torch.int32)
+        m = x < a
+        n = y > b
+        return m, n
+
+
 def test_fuse_equal_placeholders_constants_tosa_MI():
     module = FuseWeightsConstants()
     data = (torch.rand(1, 2, 8),)
@@ -94,3 +116,24 @@ def test_fuse_equal_placeholders_state_dict_tosa_MI():
     assert len(state_dict_keys) == 2, "FuseEqualPlaceholders state_dict failed"
     assert "_common" in state_dict_keys[0], "FuseEqualPlaceholders state_dict failed"
     assert "_common" in state_dict_keys[1], "FuseEqualPlaceholders state_dict failed"
+
+
+def test_not_fuse_tensor_with_different_type_MI():
+    module = NotFuseTensorWithDifferentType()
+    data = (
+        torch.rand(
+            1,
+        ),
+        torch.randint(
+            0,
+            10,
+            (1,),
+            dtype=torch.int,
+        ),
+    )
+    pipeline = TosaPipelineMI[input_t](
+        module,
+        data,
+        aten_op=[],
+    )
+    pipeline.run()