[MLIR][TOSA] add tosa erf operator

[manupak]

Summary

This commit adds tosa erf operator and its lowering
to math lib functions.

[manupak]

Discussion on tosa spec : https://discuss.mlplatform.org/t/adding-erf-operator-to-tosa/189/9

[jpienaar]

Looks good overall.

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

471

Mathjax is supported by Hugo (well and Github markdown, but the former is more important as that's what the mlir.llvm.org website is generated with), so lets switch to that here. Some plumbing may be needed mlir.llvm.org side, but can be done in follow up.

473

Why do we need to say this as part of the op rather than this be some lowering pattern?

[eric-k256]

Agreed that this looks good, with only changes in the description that seem to be needed. As mentioned inline, part of the blame lies on the spec, and I'm going to try to correct the spec to allow legalizations to choose the best option, listing the existing version as an option rather than appearing to be required. (Let me know if there is benefit to the spec as is)

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

473

I think this is matching the existing descriptions for sigmoid and tanh, and matches what is described in the TOSA spec. On Friday I was working on how to update the spec. The specification takes what is probably too strong a statement on how the integer versions of tanh/sigmoid are defined. As you describe, it should be up to the legalizations to define the 'best' legalizations, probably based on the quantization limits. The TensorFlow Lite legalizations we've added do this already.

[manupak]

modified the equation

[manupak]

lint

[manupak]

I've changed the equation now. see if it looks good.
I'll also need help landing the patch as well :).

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

473

Yea, Im being consistent with others for now.

[eric-k256]

Looks good to me.

[manupak]

Can someone help me land this patch ?

[eric-k256]

I should be able to help land it, hopefully tomorrow.

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

let arguments = (ins
     Tosa_Tensor:$input
   );
 
   let results = (outs
     Tosa_Tensor:$output
   );
 }
 
+
+//===----------------------------------------------------------------------===//
+// Operator: erf
+//===----------------------------------------------------------------------===//
+def Tosa_ErfOp : Tosa_Op<"erf", [
+    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+                              ["inferReturnTypeComponents"]>,
+    Pure]> {
+  let summary = "Computes gauss error function of input";
+
+  let description = [{
+    Gauss error function: $ erf(x) = \frac{2}{\sqrt(\pi)} \int_{0}^{x} e^{-t^2} \,dt $
+    For quantized integer data types, the TABLE operator should be used instead
+    with the following definition.  The erf_table has 513 entries each of
+    16-bit/8-bit precision and covering the input range -4.0 to +4.0 in steps of 1/64.
+  }];
+
+  let arguments = (ins
+    Tosa_Tensor:$input
+  );
+
+  let results = (outs
+    Tosa_Tensor:$output
+  );
+}
+
 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.4
 // Operator Class: Elementwise unary/binary/ternary operators.
 // Operator Subclass: Elementwise binary ops.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: add

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

createLinalgBodyCalculationForElementwiseOp(Operation *op, ValueRange args,
   // tosa::ExpOp
   if (isa<tosa::ExpOp>(op) && isa<FloatType>(elementTy))
     return rewriter.create<mlir::math::ExpOp>(loc, resultTypes, args);
 
   // tosa::TanhOp
   if (isa<tosa::TanhOp>(op) && isa<FloatType>(elementTy))
     return rewriter.create<mlir::math::TanhOp>(loc, resultTypes, args);
 
+  // tosa::ErfOp
+  if (isa<tosa::ErfOp>(op) && elementTy.isa<FloatType>())
+    return rewriter.create<mlir::math::ErfOp>(loc, resultTypes, args);
+
   // tosa::GreaterOp
   if (isa<tosa::GreaterOp>(op) && isa<FloatType>(elementTy))
     return rewriter.create<arith::CmpFOp>(loc, arith::CmpFPredicate::OGT,
                                           args[0], args[1]);
 
   if (isa<tosa::GreaterOp>(op) && elementTy.isSignlessInteger())
     return rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sgt,
                                           args[0], args[1]);
patterns->add<
       PointwiseConverter<tosa::NegateOp>,
       PointwiseConverter<tosa::PowOp>,
       PointwiseConverter<tosa::ReciprocalOp>,
       PointwiseConverter<tosa::RsqrtOp>,
       PointwiseConverter<tosa::LogOp>,
       PointwiseConverter<tosa::ExpOp>,
       PointwiseConverter<tosa::AbsOp>,
       PointwiseConverter<tosa::TanhOp>,
+      PointwiseConverter<tosa::ErfOp>,
       PointwiseConverter<tosa::BitwiseAndOp>,
       PointwiseConverter<tosa::BitwiseOrOp>,
       PointwiseConverter<tosa::BitwiseNotOp>,
       PointwiseConverter<tosa::BitwiseXorOp>,
       PointwiseConverter<tosa::LogicalAndOp>,
       PointwiseConverter<tosa::LogicalNotOp>,
       PointwiseConverter<tosa::LogicalOrOp>,
       PointwiseConverter<tosa::LogicalXorOp>,

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

 NARY_SHAPE_INFER(tosa::PowOp)
 NARY_SHAPE_INFER(tosa::ReciprocalOp)
 NARY_SHAPE_INFER(tosa::RescaleOp)
 NARY_SHAPE_INFER(tosa::ReverseOp)
 NARY_SHAPE_INFER(tosa::RsqrtOp)
 NARY_SHAPE_INFER(tosa::SelectOp)
 NARY_SHAPE_INFER(tosa::SubOp)
 NARY_SHAPE_INFER(tosa::TanhOp)
+NARY_SHAPE_INFER(tosa::ErfOp)
 NARY_SHAPE_INFER(tosa::SigmoidOp)
 #undef PRED_SHAPE_INFER
 
 static LogicalResult poolingInferReturnTypes(
     const ValueShapeRange &operands, DictionaryAttr attributes,
     ArrayRef<int64_t> kernel, ArrayRef<int64_t> stride, ArrayRef<int64_t> pad,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   ShapeAdaptor inputShape = operands.getShape(0);

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

func.func @test_simple_f32(%arg0: tensor<1xf32>) -> () {
   // CHECK: linalg.generic
   // CHECK: arith.truncf
   %22 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xf16>
 
   // CHECK: linalg.generic
   // CHECK: arith.divf
   %23 = "tosa.reciprocal"(%0) : (tensor<1xf32>) -> tensor<1xf32>
 
+  // CHECK: linalg.generic
+  // CHECK: math.erf
+  %24 = "tosa.erf"(%0) : (tensor<1xf32>) -> tensor<1xf32>
+
   return
 }
 
 // -----
 
 // CHECK-LABEL: @test_simple_f16
 func.func @test_simple_f16(%arg0: tensor<1xf16>) -> () {
 

mlir/test/Dialect/Tosa/ops.mlir

 // -----
 // CHECK-LABEL: tanh
 func.func @test_tanh(%arg0: tensor<13x21x3xf32>) -> tensor<13x21x3xf32> {
   %0 = "tosa.tanh"(%arg0) : (tensor<13x21x3xf32>) -> tensor<13x21x3xf32>
   return %0 : tensor<13x21x3xf32>
 }
 
 // -----
+// CHECK-LABEL: erf
+func.func @test_erf(%arg0: tensor<13x21x3xf32>) -> tensor<13x21x3xf32> {
+  %0 = "tosa.erf"(%arg0) : (tensor<13x21x3xf32>) -> tensor<13x21x3xf32>
+  return %0 : tensor<13x21x3xf32>
+}
+
+// -----
 // CHECK-LABEL: add
 func.func @test_add(%arg0: tensor<13x21x1xf32>, %arg1: tensor<13x21x3xf32>) -> tensor<13x21x3xf32> {
   %0 = "tosa.add"(%arg0, %arg1) : (tensor<13x21x1xf32>, tensor<13x21x3xf32>) -> tensor<13x21x3xf32>
   return %0 : tensor<13x21x3xf32>
 }
 
 // -----
 // CHECK-LABEL: arithmetic_right_shift

mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir

func.func @test_unary_f32(%arg0 : tensor<4xf32>) -> () {
   // CHECK: "tosa.tanh"(%arg0) : (tensor<4xf32>) -> tensor<4xf32>
   %10 = "tosa.tanh"(%arg0) : (tensor<4xf32>) -> tensor<*xf32>
 
   // CHECK: "tosa.sigmoid"(%arg0) : (tensor<4xf32>) -> tensor<4xf32>
   %11 = "tosa.sigmoid"(%arg0) : (tensor<4xf32>) -> tensor<*xf32>
 
   // CHECK: "tosa.cast"(%arg0) : (tensor<4xf32>) -> tensor<4xi32>
   %12 = "tosa.cast"(%arg0) : (tensor<4xf32>) -> tensor<*xi32>
+
+  // CHECK: "tosa.erf"(%arg0) : (tensor<4xf32>) -> tensor<4xf32>
+  %13 = "tosa.erf"(%arg0) : (tensor<4xf32>) -> tensor<*xf32>
   return
 }
 
 // -----
 
 // CHECK-LABEL: @test_unary_i32
 func.func @test_unary_i32(%arg0 : tensor<4xi32>) -> () {
   // CHECK: "tosa.abs"(%arg0) : (tensor<4xi32>) -> tensor<4xi32>