Arm backend: Support bool inputs for aten.bitwise_not via logical_not (pytorch#16417)

TOSA bitwise_not does not accept boolean inputs. For boolean tensors,
aten.bitwise_not is equivalent to aten.logical_not, so rewrite to enable
lowering.


Signed-off-by: Yufeng Shi <yufeng.shi@arm.com>

Author: YufengShi-dudu

backends/arm/_passes/__init__.py

@@ -113,6 +113,9 @@
 from .replace_scalar_with_tensor_pass import (  # noqa
     ReplaceScalarWithTensorByProfilePass,
 )
+from .rewrite_bool_bitwise_not_to_logical_not_pass import (  # noqa
+    RewriteBoolBitwiseNotToLogicalNotPass,
+)
 from .rewrite_bool_to_fp32_cast_via_int8_pass import (  # noqa
     RewriteBoolToFp32CastViaInt8Pass,
 )

backends/arm/_passes/arm_pass_manager.py

@@ -103,6 +103,7 @@
     RemoveNoopPass,
     ReplaceInfAndLimitValuesPass,
     ReplaceScalarWithTensorByProfilePass,
+    RewriteBoolBitwiseNotToLogicalNotPass,
     RewriteBoolToFp32CastViaInt8Pass,
     RewriteConvPass,
     RewriteMatmulPass,
@@ -222,6 +223,7 @@ def _tosa_pipeline(
         self.add_passes(
             [
                 FuseQuantizedActivationPass(),
+                RewriteBoolBitwiseNotToLogicalNotPass(),
                 RewriteBoolToFp32CastViaInt8Pass(),
                 ConvertToClampPass(),
                 DecomposeTOSAUnsupportedClampPass(),

backends/arm/_passes/rewrite_bool_bitwise_not_to_logical_not_pass.py

@@ -0,0 +1,43 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+from typing import Set, Type
+
+import torch
+from executorch.backends.arm._passes import ArmPass
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass
+
+
+class RewriteBoolBitwiseNotToLogicalNotPass(ArmPass):
+    """
+    Rewrites ``aten.bitwise_not`` on boolean tensors to ``aten.logical_not``.
+
+    TOSA ``bitwise_not`` does not support boolean inputs. On boolean tensors,
+    ``bitwise_not`` is equivalent to ``logical_not``, so this rewrite preserves
+    semantics while enabling lowering.
+    """
+
+    _passes_required_after: Set[Type[ExportPass]] = set()
+
+    _TARGET_OPS = {
+        exir_ops.edge.aten.bitwise_not.default,
+    }
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op not in self._TARGET_OPS:
+            return super().call_operator(op, args, kwargs, meta)
+
+        if meta["val"].dtype == torch.bool:
+            x = args[0]
+            return super().call_operator(
+                exir_ops.edge.aten.logical_not.default,
+                (x,),
+                {},
+                meta,
+            )
+
+        return super().call_operator(op, args, kwargs, meta)

backends/arm/test/ops/test_bitwise_not.py

@@ -29,6 +29,7 @@
     "pattern2_int8": 0xCC * torch.ones(1, 2, 2, 2, dtype=torch.int8),
     "pattern2_int16": 0xCCCC * torch.ones(1, 2, 2, 2, dtype=torch.int16),
     "pattern2_int32": 0xCCCCCCCC * torch.ones(1, 2, 2, 2, dtype=torch.int32),
+    "pattern_bool": torch.tensor([True, False, True], dtype=torch.bool),
     "rand_rank2": torch.randint(-128, 127, (10, 10), dtype=torch.int8),
     "rand_rank4": torch.randint(-128, 127, (1, 10, 10, 10), dtype=torch.int8),
 }