feat(OnnxToTosa): legalization for onnx.BitwiseNot and onnx.Not.

Xilinx · Jan 3, 2024 · 482dd65 · 482dd65
1 parent 0c6a798
commit 482dd65
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 2 deletions.
diff --git a/src/Conversion/ONNXToTOSA/ConvertONNXToTOSA.cpp b/src/Conversion/ONNXToTOSA/ConvertONNXToTOSA.cpp
@@ -112,7 +112,8 @@ void FrontendToTosaLoweringPass::runOnOperation() {
   // conversion failures. Quantized types are not supported right now.
   TypeConverter typeConverter;
   typeConverter.addConversion([](Type type) -> std::optional<Type> {
-    if (isTOSASignedInt(type) || isTOSAFloat(type) || type.isa<NoneType>())
+    if (isTOSASignedInt(type) || isTOSAFloat(type) || type.isa<NoneType>() ||
+        isTOSABool(type))
       return type;
     return std::nullopt;
   });

diff --git a/src/Conversion/ONNXToTOSA/Math/Elementwise.cpp b/src/Conversion/ONNXToTOSA/Math/Elementwise.cpp
@@ -62,7 +62,29 @@ struct IsIntOrFloat {
       ConversionPatternRewriter &rewriter, Type scalarType, Operation *op) {
     if (!isTOSAFloat(scalarType) && !isTOSASignedInt(scalarType)) {
       return rewriter.notifyMatchFailure(
-          op, "this operation only support signed integer or float types");
+          op, "this operation only supports signed integer or float types");
+    }
+    return success();
+  }
+};
+
+struct IsInt {
+  static LogicalResult checkType(
+      ConversionPatternRewriter &rewriter, Type scalarType, Operation *op) {
+    if (!isTOSASignedInt(scalarType)) {
+      return rewriter.notifyMatchFailure(
+          op, "this operation only supports float types");
+    }
+    return success();
+  }
+};
+
+struct IsBool {
+  static LogicalResult checkType(
+      ConversionPatternRewriter &rewriter, Type scalarType, Operation *op) {
+    if (!isTOSABool(scalarType)) {
+      return rewriter.notifyMatchFailure(
+          op, "this operation only supports bool type");
     }
     return success();
   }
@@ -406,6 +428,10 @@ static void populateLoweringONNXElementwiseUnaryTemplateOpToTOSAPattern(
           IsIntOrFloat, IsIntOrFloat>,
       ONNXElementwiseUnaryOpLoweringToTOSA<ONNXSigmoidOp, mlir::tosa::SigmoidOp,
           IsIntOrFloat, IsIntOrFloat>,
+      ONNXElementwiseUnaryOpLoweringToTOSA<ONNXBitwiseNotOp,
+          mlir::tosa::BitwiseNotOp, IsInt, IsInt>,
+      ONNXElementwiseUnaryOpLoweringToTOSA<ONNXNotOp, mlir::tosa::LogicalNotOp,
+          IsBool, IsBool>,
       ONNXElementwiseUnaryOpLoweringToTOSA<ONNXAbsOp, mlir::tosa::AbsOp,
           AbsIOSupportedTypes, AbsIOSupportedTypes>,
       ONNXElementwiseUnaryOpLoweringToTOSA<ONNXErfOp, mlir::tosa::ErfOp,

diff --git a/src/Conversion/ONNXToTOSA/ONNXToTOSACommon.hpp b/src/Conversion/ONNXToTOSA/ONNXToTOSACommon.hpp
@@ -304,6 +304,11 @@ std::optional<mlir::Value> convertReduceOpCommon(
 // Check for valid TOSA types.
 //===----------------------------------------------------------------------===//
 
+inline bool isTOSABool(mlir::Type type) {
+  mlir::IntegerType intType = type.dyn_cast<mlir::IntegerType>();
+  return intType && intType.isSignless() && intType.getWidth() == 1;
+}
+
 inline bool isTOSASignedInt(mlir::Type type) {
   mlir::IntegerType intType = type.dyn_cast<mlir::IntegerType>();
   std::set<unsigned> intWidth{8, 16, 32, 48, 64};

diff --git a/test/mlir/conversion/onnx_to_tosa/Math/Elementwise.mlir b/test/mlir/conversion/onnx_to_tosa/Math/Elementwise.mlir
@@ -366,6 +366,30 @@ func.func @test_erf_bf16(%arg0: tensor<3xbf16>) -> tensor<3xbf16> {
 
 // -----
 
+func.func @test_bitwise_not(%arg0 : tensor<10x10xi32>) -> tensor<10x10xi32> {
+  %0 = "onnx.BitwiseNot"(%arg0) : (tensor<10x10xi32>) -> tensor<10x10xi32>
+  "func.return"(%0) : (tensor<10x10xi32>) -> ()
+// CHECK-LABEL:  func @test_bitwise_not
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<10x10xi32>) -> tensor<10x10xi32> {
+// CHECK-NEXT:      [[VAR_0_:%.+]] = "tosa.bitwise_not"([[PARAM_0_]]) : (tensor<10x10xi32>) -> tensor<10x10xi32>
+// CHECK-NEXT:      return [[VAR_0_]] : tensor<10x10xi32>
+// CHECK-NEXT:    }
+}
+
+// -----
+
+func.func @test_not(%arg0 : tensor<10x10xi1>) -> tensor<10x10xi1> {
+  %0 = "onnx.Not"(%arg0) : (tensor<10x10xi1>) -> tensor<10x10xi1>
+  "func.return"(%0) : (tensor<10x10xi1>) -> ()
+// CHECK-LABEL:  func @test_not
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<10x10xi1>) -> tensor<10x10xi1> {
+// CHECK-NEXT:      [[VAR_0_:%.+]] = "tosa.logical_not"([[PARAM_0_]]) : (tensor<10x10xi1>) -> tensor<10x10xi1>
+// CHECK-NEXT:      return [[VAR_0_]] : tensor<10x10xi1>
+// CHECK-NEXT:    }
+}
+
+// -----
+
 // Default values: alpha = 0.2, beta = 0.5
 func.func @test_hardsigmoid_default_values_f32(%arg0: tensor<3xf32>) -> tensor<3xf32> {
   %0 = "onnx.HardSigmoid"(%arg0) : (tensor<3xf32>) -> tensor<3xf32>