[mlir][tosa] Fix tosa.cast semantics to perform rounding/clipping

[rsuderman]

Summary

Rounding to integers requires rounding (for floating points) and clipping
to the min/max values of the destination range. Added this behavior and
updated tests appropriately.

[silvas]

Code looks fine. Will let @sjarus review the rounding algorithm details.

[sjarus]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

517–518

Would be cleaner to integer clip in the integer domain here instead of the seemingly awkward expression of intMin and intMax and F32Attrs and then doing an integer range clamping in the fp Domain ?

This could then reuse the intMin/intMax further below (they'd be common for all dstTy.isa<IntegerType>() ) and swap the last two ops around - FPtoSI(rounded) followed by CmpIOp ?

[rsuderman]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

517–518

That doesn't work. To clip in the integer domain we would need to type-convert to integer, then perform the clip. F32 range exceeds i32 (and even i64), so converting type to an integer type, then clipping, would fail.

[sjarus]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

517–518

Right, realized that on second thought.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

if (isa<tosa::CastOp>(op)) {
     // Casting to boolean, floats need to only be checked as not-equal to zero.
     if (srcTy.isa<FloatType>() && dstTy.isInteger(1)) {
       Value zero =
           rewriter.create<ConstantOp>(loc, rewriter.getFloatAttr(srcTy, 0.0));
       return rewriter.create<mlir::CmpFOp>(loc, CmpFPredicate::UNE,
                                            args.front(), zero);
     }
 
-    if (mlir::FPToSIOp::areCastCompatible(srcTy, dstTy))
-      return rewriter.create<mlir::FPToSIOp>(loc, resultTypes, args,
-                                             mlir::None);
+    if (mlir::FPToSIOp::areCastCompatible(srcTy, dstTy)) {
+      auto zero =
+          rewriter.create<ConstantOp>(loc, rewriter.getF32FloatAttr(0.0f));
+      auto half =
+          rewriter.create<ConstantOp>(loc, rewriter.getF32FloatAttr(0.5f));
+
+      auto intMin = rewriter.create<ConstantOp>(
+          loc, rewriter.getF32FloatAttr(
+                   APInt::getSignedMinValue(dstTy.getIntOrFloatBitWidth())
+                       .getSExtValue()));
+
+      auto intMax = rewriter.create<ConstantOp>(
+          loc, rewriter.getF32FloatAttr(
+                   APInt::getSignedMaxValue(dstTy.getIntOrFloatBitWidth())
+                       .getSExtValue()));
+
+      auto added = rewriter.create<AddFOp>(loc, args[0], half);
+      auto subbed = rewriter.create<SubFOp>(loc, args[0], half);
+      auto negative =
+          rewriter.create<mlir::CmpFOp>(loc, CmpFPredicate::OLT, args[0], zero);
+      auto rounded =
+          rewriter.create<mlir::SelectOp>(loc, negative, subbed, added);
+
+      auto clamped = clampHelper<mlir::CmpFOp>(loc, rounded, intMin, intMax,
+                                               CmpFPredicate::OLT, rewriter);
+
+      return rewriter.create<mlir::FPToSIOp>(loc, dstTy, clamped);
+    }
 
     // Casting to boolean, integers need to only be checked as not-equal to
     // zero.
     if (srcTy.isa<IntegerType>() && dstTy.isInteger(1)) {
       Value zero =
           rewriter.create<ConstantIntOp>(loc, 0, srcTy.getIntOrFloatBitWidth());
       return rewriter.create<mlir::CmpIOp>(loc, CmpIPredicate::ne, args.front(),
                                            zero);
     }
 
     if (srcTy.isa<IntegerType>() && dstTy.isa<IntegerType>() && bitExtend)
       return rewriter.create<mlir::SignExtendIOp>(loc, resultTypes, args,
                                                   mlir::None);
 
-    if (srcTy.isa<IntegerType>() && dstTy.isa<IntegerType>() && !bitExtend)
-      return rewriter.create<mlir::TruncateIOp>(loc, resultTypes, args,
-                                                mlir::None);
+    if (srcTy.isa<IntegerType>() && dstTy.isa<IntegerType>() && !bitExtend) {
+      auto intMin = rewriter.create<ConstantIntOp>(
+          loc,
+          APInt::getSignedMinValue(dstTy.getIntOrFloatBitWidth())
+              .getSExtValue(),
+          srcTy.getIntOrFloatBitWidth());
+
+      auto intMax = rewriter.create<ConstantIntOp>(
+          loc,
+          APInt::getSignedMaxValue(dstTy.getIntOrFloatBitWidth())
+              .getSExtValue(),
+          srcTy.getIntOrFloatBitWidth());
+
+      auto clamped = clampHelper<mlir::CmpIOp>(loc, args[0], intMin, intMax,
+                                               CmpIPredicate::slt, rewriter);
+      return rewriter.create<mlir::TruncateIOp>(loc, dstTy, clamped);
+    }
   }
 
   (void)rewriter.notifyMatchFailure(
       op, "unhandled op for linalg body calculation for elementwise op");
   return nullptr;
 }
 
 static LogicalResult

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

func @test_simple_f32(%arg0: tensor<1xf32>) -> () {
   // CHECK: linalg.generic
   // CHECK: negf
   // CHECK: exp
   // CHECK: addf
   // CHECK: divf
   %20 = "tosa.sigmoid"(%0) : (tensor<1xf32>) -> tensor<1xf32>
 
   // CHECK: linalg.generic
+  // CHECK: constant 0.000000e+00
+  // CHECK: constant 5.000000e-01
+  // CHECK: constant -2.14748365E+9
+  // CHECK: constant 2.14748365E+9
+  // CHECK: addf
+  // CHECK: subf
+  // CHECK: cmpf olt
+  // CHECK: select
+  // CHECK: cmpf olt
+  // CHECK: select
+  // CHECK: cmpf olt
+  // CHECK: select
   // CHECK: fptosi
   %21 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xi32>
 
   // CHECK: linalg.generic
   // CHECK: constant 0
   // CHECK: cmpf
   %22 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xi1>
 
func @test_simple_i32(%arg0: tensor<1xi32>) -> () {
   %17 = "tosa.clamp"(%0) {min_int = 1 : i64, max_int = 5 : i64, min_fp = 1.0 : f32, max_fp = 5.0 : f32} : (tensor<1xi32>) -> tensor<1xi32>
 
   // CHECK: linalg.generic
   // CHECK: cmpi
   // CHECK: select
   %18 = "tosa.reluN"(%0) {max_int = 5 : i64, max_fp = 5.0 : f32} : (tensor<1xi32>) -> tensor<1xi32>
 
   // CHECK: linalg.generic
+  // CHECK: constant -32768
+  // CHECK: constant 32767
+  // CHECK: cmpi slt
+  // CHECK: select
+  // CHECK: cmpi slt
+  // CHECK: select
   // CHECK: trunci
   %19 = "tosa.cast"(%0) : (tensor<1xi32>) -> tensor<1xi16>
 
   // CHECK: linalg.generic
   // CHECK: yield
   %20 = "tosa.cast"(%0) : (tensor<1xi32>) -> tensor<1xi32>
 
   // CHECK: linalg.generic