[mlir][tosa] Add tosa.reciprocal and tosa.sigmoid lowerings

[NatashaKnk]

Summary

Lowering reciprocal and sigmoid elementwise operations to linalg dialect.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

if (isa<tosa::MulOp>(op) && elementTy.isa<FloatType>()) {
     if (dyn_cast<tosa::MulOp>(op).shift() != 0) {
       (void)rewriter.notifyMatchFailure(op,
                                         "Cannot have shift value for float");
       return nullptr;
     }
     return rewriter.create<mlir::MulFOp>(loc, resultTypes, args);
   }
 
+  // tosa::ReciprocalOp
+  if (isa<tosa::ReciprocalOp>(op) && elementTy.isa<FloatType>()) {
+    auto one =
+        rewriter.create<mlir::ConstantOp>(loc, FloatAttr::get(elementTy, 1));
+    return rewriter.create<mlir::DivFOp>(loc, resultTypes, one, args[0]);
+  }
+
   if (isa<tosa::MulOp>(op) && elementTy.isa<IntegerType>()) {
     Value a = args[0];
     Value b = args[1];
     auto shift =
         op->getAttr("shift").cast<IntegerAttr>().getValue().getSExtValue();
     if (shift > 0) {
       auto shiftConst =
           rewriter.create<ConstantIntOp>(loc, shift, /*bitwidth=*/8);
if (isa<tosa::ReluNOp>(op) && elementTy.isa<IntegerType>()) {
     auto zero =
         rewriter.create<mlir::ConstantOp>(loc, IntegerAttr::get(elementTy, 0));
     auto n = createConstFromIntAttribute<int32_t>(op, "max_int", elementTy,
                                                   rewriter);
     return clampHelper<mlir::CmpIOp>(loc, args[0], zero, n, CmpIPredicate::slt,
                                      rewriter);
   }
 
+  // tosa::SigmoidOp
+  if (isa<tosa::SigmoidOp>(op) && elementTy.isa<FloatType>()) {
+    auto one =
+        rewriter.create<mlir::ConstantOp>(loc, FloatAttr::get(elementTy, 1));
+    auto negate = rewriter.create<mlir::NegFOp>(loc, resultTypes, args[0]);
+    auto exp = rewriter.create<mlir::math::ExpOp>(loc, resultTypes, negate);
+    auto added = rewriter.create<mlir::AddFOp>(loc, resultTypes, exp, one);
+    return rewriter.create<mlir::DivFOp>(loc, resultTypes, one, added);
+  }
+
   // tosa::CastOp
   if (isa<tosa::CastOp>(op)) {
     Type srcTy = elementTy;
     Type dstTy = resultTypes.front();
     bool bitExtend =
         srcTy.getIntOrFloatBitWidth() < dstTy.getIntOrFloatBitWidth();
 
     if (srcTy == dstTy)

 };
 
 } // namespace
 
 void mlir::tosa::populateTosaToLinalgOnTensorsConversionPatterns(
     RewritePatternSet *patterns) {
   patterns->add<
       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::MulOp>, PointwiseConverter<tosa::ReciprocalOp>,
+      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::BitwiseNotOp>,
       PointwiseConverter<tosa::BitwiseXorOp>,
       PointwiseConverter<tosa::LogicalAndOp>,
       PointwiseConverter<tosa::LogicalNotOp>,
       PointwiseConverter<tosa::LogicalOrOp>,
       PointwiseConverter<tosa::LogicalXorOp>, PointwiseConverter<tosa::CastOp>,
       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>,
+      PointwiseConverter<tosa::SigmoidOp>, IdentityNConverter<tosa::IdentityOp>,
       IdentityNConverter<tosa::IdentityNOp>, ReduceConverter<tosa::ReduceMinOp>,
       ReduceConverter<tosa::ReduceMaxOp>, ReduceConverter<tosa::ReduceSumOp>,
-      ReduceConverter<tosa::ReduceProdOp>, ArgMaxConverter, ConcatConverter, PadConverter,
-      ReshapeConverter, RescaleConverter, ReverseConverter, TileConverter,
-      TransposeConverter, MatMulConverter, FullyConnectedConverter>(
-        patterns->getContext());
+      ReduceConverter<tosa::ReduceProdOp>, ArgMaxConverter, ConcatConverter,
+      PadConverter, ReshapeConverter, RescaleConverter, ReverseConverter,
+      TileConverter, TransposeConverter, MatMulConverter,
+      FullyConnectedConverter>(patterns->getContext());
 }

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

func @test_simple_f32(%arg0: tensor<1xf32>) -> () {
   %17 = "tosa.clamp"(%0) {min_int = 1 : i64, max_int = 5 : i64, min_fp = 1.0 : f32, max_fp = 5.0 : f32} : (tensor<1xf32>) -> tensor<1xf32>
 
   // CHECK: linalg.generic
   // CHECK: cmpf
   // CHECK: select
   %18 = "tosa.reluN"(%0) {max_int = 5 : i64, max_fp = 5.0 : f32} : (tensor<1xf32>) -> tensor<1xf32>
 
   // CHECK: linalg.generic
+  // CHECK: negf
+  // CHECK: exp
+  // CHECK: addf
+  // CHECK: divf
+  %19 = "tosa.sigmoid"(%0) : (tensor<1xf32>) -> tensor<1xf32>
+
+  // CHECK: linalg.generic
   // CHECK: fptosi
-  %19 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xi32>
+  %20 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xi32>
 
   // CHECK: linalg.generic
   // CHECK: constant 0
   // CHECK: cmpf
-  %20 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xi1>
+  %21 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xi1>
 
   // CHECK: linalg.generic
   // CHECK: fptrunc
-  %21 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xf16>
+  %22 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xf16>
 
   // CHECK: linalg.generic
   // CHECK: yield
-  %22 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xf32>
+  %23 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xf32>
+
+  // CHECK: linalg.generic
+  // CHECK: divf
+  %24 = "tosa.reciprocal"(%0) : (tensor<1xf32>) -> tensor<1xf32>
 
   return
 }
 
 // -----
 
 // CHECK-LABEL: @test_simple_f16
 func @test_simple_f16(%arg0: tensor<1xf16>) -> () {