[mlir][tosa] Add tosa.logical_* to linalg lowerings

[rsuderman]

Summary

Adds lowerings for logical_* boolean operations. Each of these ops only operate
on booleans allowing simple lowerings.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

createLinalgBodyCalculationForElementwiseOp(Operation *op, ValueRange args,
   // tosa::BitwiseXOrOp
   if (isa<tosa::BitwiseXorOp>(op) && elementTy.isa<IntegerType>())
     return rewriter.create<mlir::XOrOp>(loc, resultTypes, args);
 
   // tosa::LogicalLeftShiftOp
   if (isa<tosa::LogicalLeftShiftOp>(op) && elementTy.isa<IntegerType>())
     return rewriter.create<mlir::ShiftLeftOp>(loc, resultTypes, args);
 
-  // tosa::LogicalrightShiftOp
+  // tosa::LogicalRightShiftOp
   if (isa<tosa::LogicalRightShiftOp>(op) && elementTy.isa<IntegerType>())
     return rewriter.create<mlir::UnsignedShiftRightOp>(loc, resultTypes, args);
 
+  // tosa::LogicalAnd
+  if (isa<tosa::LogicalAndOp>(op) && elementTy.isInteger(1))
+    return rewriter.create<mlir::AndOp>(loc, resultTypes, args);
+
+  // tosa::LogicalNot
+  if (isa<tosa::LogicalNotOp>(op) && elementTy.isInteger(1)) {
+    auto one = rewriter.create<mlir::ConstantOp>(
+        loc, rewriter.getIntegerAttr(elementTy, 1));
+    return rewriter.create<mlir::XOrOp>(loc, resultTypes, args[0], one);
+  }
+
+  // tosa::LogicalOr
+  if (isa<tosa::LogicalOrOp>(op) && elementTy.isInteger(1))
+    return rewriter.create<mlir::OrOp>(loc, resultTypes, args);
+
+  // tosa::LogicalXor
+  if (isa<tosa::LogicalXorOp>(op) && elementTy.isInteger(1))
+    return rewriter.create<mlir::XOrOp>(loc, resultTypes, args);
+
   // tosa::PowOp
   if (isa<tosa::PowOp>(op) && elementTy.isa<FloatType>())
     return rewriter.create<mlir::math::PowFOp>(loc, resultTypes, args);
 
   // tosa::RsqrtOp
   if (isa<tosa::RsqrtOp>(op) && elementTy.isa<FloatType>())
     return rewriter.create<mlir::math::RsqrtOp>(loc, resultTypes, args);
 
patterns->insert<
       PointwiseConverter<tosa::AddOp>, PointwiseConverter<tosa::SubOp>,
       PointwiseConverter<tosa::MulOp>, PointwiseConverter<tosa::NegateOp>,
       PointwiseConverter<tosa::PowOp>, PointwiseConverter<tosa::RsqrtOp>,
       PointwiseConverter<tosa::LogOp>, PointwiseConverter<tosa::ExpOp>,
       PointwiseConverter<tosa::AbsOp>, PointwiseConverter<tosa::TanhOp>,
       PointwiseConverter<tosa::BitwiseAndOp>,
       PointwiseConverter<tosa::BitwiseOrOp>,
       PointwiseConverter<tosa::BitwiseXorOp>,
+      PointwiseConverter<tosa::LogicalAndOp>,
+      PointwiseConverter<tosa::LogicalNotOp>,
+      PointwiseConverter<tosa::LogicalOrOp>,
+      PointwiseConverter<tosa::LogicalXorOp>,
       PointwiseConverter<tosa::LogicalLeftShiftOp>,
       PointwiseConverter<tosa::LogicalRightShiftOp>,
       PointwiseConverter<tosa::SelectOp>, PointwiseConverter<tosa::GreaterOp>,
       PointwiseConverter<tosa::GreaterEqualOp>,
       PointwiseConverter<tosa::MaximumOp>, PointwiseConverter<tosa::MinimumOp>,
       PointwiseConverter<tosa::CeilOp>, PointwiseConverter<tosa::FloorOp>,
       PointwiseConverter<tosa::ClampOp>, PointwiseConverter<tosa::ReluNOp>,
       IdentityNConverter<tosa::IdentityOp>,
       IdentityNConverter<tosa::IdentityNOp>, ReduceConverter<tosa::ReduceMinOp>,
       ReduceConverter<tosa::ReduceMaxOp>, ReduceConverter<tosa::ReduceSumOp>,
       ReduceConverter<tosa::ReduceProdOp>, ConcatOpConversion,
       ReshapeOpConverter, TransposeConverter, RescaleOpConverter>(context);
 }

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

func @test_simple_i32(%arg0: tensor<1xi32>) -> () {
   // CHECK: select
   %15 = "tosa.reluN"(%0) {max_int = 5 : i64, max_fp = 5.0 : f32} : (tensor<1xi32>) -> tensor<1xi32>
 
   return
 }
 
 // -----
 
+// CHECK-LABEL: @test_bool
+func @test_bool(%arg0: tensor<1xi1>, %arg1: tensor<1xi1>) -> () {
+  // CHECK: linalg.generic
+  // CHECK: and
+  %0 = "tosa.logical_and"(%arg0, %arg1) : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1>
+
+  // CHECK: linalg.generic
+  // CHECK: or
+  %1 = "tosa.logical_or"(%arg0, %arg1) : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1>
+
+  // CHECK: linalg.generic
+  // CHECK: xor
+  %2 = "tosa.logical_xor"(%arg0, %arg1) : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1>
+
+  // CHECK: linalg.generic
+  // CHECK: constant true
+  // CHECK: xor
+  %3 = "tosa.logical_not"(%arg0) : (tensor<1xi1>) -> tensor<1xi1>
+
+  return
+}
+
+// -----
+
 // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 // CHECK-LABEL: @test_reshape_downrank
 func @test_reshape_downrank(%arg0: tensor<2x3xf32>) -> tensor<6xf32> {
   // CHECK: [[RESHAPE:%.+]] = linalg.tensor_reshape %arg0 [#[[$MAP0]]]
   %0 = "tosa.reshape"(%arg0) {new_shape = [6]} : (tensor<2x3xf32>) -> tensor<6xf32>
   // CHECK: return [[RESHAPE]]
   return %0 : tensor<6xf32>
 }