[TOSA] Lower boolean aten.bitwise_not to logical_not (#4364)

- Fix TorchToTosa's shared unary pattern so AtenBitwiseNotOp with i1
outputs emits tosa.logical_not instead of the tosa.bitwise_not.
- Add a regression in test/Conversion/TorchToTosa/basic.mlir that checks
the lowering path for a bool tensor.
- Add a regression end-to-end test for bitwise_not with boolean.

Author: Lallapallooza

lib/Conversion/TorchToTosa/TorchToTosa.cpp

@@ -30,6 +30,7 @@
 #include <numeric>
 #include <optional>
 #include <random>
+#include <type_traits>
 
 #include "mlir/Dialect/Tosa/Utils/QuantUtils.h"
 
@@ -97,6 +98,16 @@ class ConvertAtenUnaryOp : public OpConversionPattern<AtenOpT> {
 
     self = tosa::tosaCastTensorToType(rewriter, self, outType).value();
 
+    if constexpr (std::is_same_v<AtenOpT, AtenBitwiseNotOp>) {
+      if (auto intTy = dyn_cast<IntegerType>(outType.getElementType())) {
+        if (intTy.getWidth() == 1) {
+          rewriter.replaceOpWithNewOp<tosa::LogicalNotOp>(op, outType, self);
+          return success();
+        }
+      }
+      // otherwise fall through to standard emission
+    }
+
     rewriter.replaceOpWithNewOp<TosaOpT>(op, outType, self);
 
     return success();

projects/pt1/python/torch_mlir_e2e_test/test_suite/elementwise.py

@@ -5175,6 +5175,29 @@ def ElementwiseBitwiseNotInt32Module_basic(module, tu: TestUtils):
 # ==============================================================================
 
 
+class ElementwiseBitwiseNotBoolModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([-1, -1], torch.bool, True),
+        ]
+    )
+    def forward(self, x):
+        return torch.bitwise_not(x)
+
+
+@register_test_case(module_factory=lambda: ElementwiseBitwiseNotBoolModule())
+def ElementwiseBitwiseNotBoolModule_basic(module, tu: TestUtils):
+    module.forward(tu.randint(3, 4, low=0, high=2).to(torch.bool))
+
+
+# ==============================================================================
+
+
 class ElementwiseSubTensorInt8Module(torch.nn.Module):
     def __init__(self):
         super().__init__()

test/Conversion/TorchToTosa/basic.mlir

@@ -135,6 +135,20 @@ func.func @torch.aten.bitwise_not$basic(%arg0: !torch.vtensor<[?,?],f32>) -> !to
 
 // -----
 
+// CHECK-LABEL:   func.func @torch.aten.bitwise_not$bool(
+// CHECK-SAME:                               %[[ARG:.*]]: !torch.vtensor<[2,3],i1>) -> !torch.vtensor<[2,3],i1> {
+// CHECK:           %[[ARG_BUILTIN:.*]] = torch_c.to_builtin_tensor %[[ARG]] : !torch.vtensor<[2,3],i1> -> tensor<2x3xi1>
+// CHECK:           %[[RESULT_BUILTIN:.*]] = tosa.logical_not %[[ARG_BUILTIN]] : (tensor<2x3xi1>) -> tensor<2x3xi1>
+// CHECK:           %[[RESULT:.*]] = torch_c.from_builtin_tensor %[[RESULT_BUILTIN]] : tensor<2x3xi1> -> !torch.vtensor<[2,3],i1>
+// CHECK:           return %[[RESULT]] : !torch.vtensor<[2,3],i1>
+// CHECK:         }
+func.func @torch.aten.bitwise_not$bool(%arg0: !torch.vtensor<[2,3],i1>) -> !torch.vtensor<[2,3],i1> {
+  %0 = torch.aten.bitwise_not %arg0 : !torch.vtensor<[2,3],i1> -> !torch.vtensor<[2,3],i1>
+  return %0 : !torch.vtensor<[2,3],i1>
+}
+
+// -----
+
 // CHECK-LABEL:   func.func @torch.aten.ceil$basic(
 // CHECK-SAME:                                %[[VAL_0:.*]]: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
 // CHECK:           %[[VAL_1:.*]] = torch_c.to_builtin_tensor %[[VAL_0]] : !torch.vtensor<[?,?],f32> -> tensor<?x?xf32>