[MLIR][Tosa] Fix fp canonicalization for `clamp`

[rikhuijzer]

Summary

The canonicalization for clamp removed the clamp operation when the
chosen min and max values were outside the range of what is possible to
represent with the input type. For example, if the input type is i8,
then the min and max values must be between -128 and 127. If the min and
max are, say, -200 and 200, then the clamp operation can be safely
removed.

However, as pointed out by @wyanzhao, this is wrong for floating point
types since they can represent infinity.

Fixes #62341.

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

LogicalResult matchAndRewrite(tosa::ClampOp op,
     Value input = op.getInput();
     auto inputType = llvm::dyn_cast<RankedTensorType>(op.getInput().getType());
     auto inputElementType = inputType.getElementType();
 
     if (!inputType.hasStaticShape()) {
       return failure();
     }
 
-    if (inputElementType.isF32()) {
+    if (inputElementType.isa<FloatType>()) {
+      // Unlike integer types, floating point types can represent infinity.
       auto minClamp = op.getMinFp();
       auto maxClamp = op.getMaxFp();
-      bool isMin = (minClamp.isLargest() || minClamp.isInfinity()) &&
-                   minClamp.isNegative();
-      bool isMax = (maxClamp.isLargest() || maxClamp.isInfinity()) &&
-                   !maxClamp.isNegative();
+      bool isMin = minClamp.isInfinity() && minClamp.isNegative();
+      bool isMax = maxClamp.isInfinity() && !maxClamp.isNegative();
 
       if (isMin && isMax) {
         rewriter.replaceOp(op, input);
         return success();
       }
       return failure();
     }
 

mlir/test/Dialect/Tosa/canonicalize.mlir

 
 // CHECK-LABEL: @cast_nofold
 func.func @cast_nofold(%arg0: tensor<?x1xf32>) -> tensor<?x1xi32> {
   // CHECK: "tosa.cast"
   %0 = "tosa.cast"(%arg0) : (tensor<?x1xf32>) -> tensor<?x1xi32>
   return %0 : tensor<?x1xi32>
 }
 
-// CHECK-LABEL: @clamp_not_noop
-func.func @clamp_not_noop(%arg0: tensor<4xi32>) -> tensor<4xi32> {
+// CHECK-LABEL: @clamp_i32_not_noop
+func.func @clamp_i32_not_noop(%arg0: tensor<4xi32>) -> tensor<4xi32> {
   // CHECK: "tosa.clamp"
   %0 = "tosa.clamp"(%arg0) {min_int = 1 : i64, max_int = 4 : i64, min_fp = 1.0 : f32, max_fp = 4.0 : f32} : (tensor<4xi32>) -> tensor<4xi32>
   return %0 : tensor<4xi32>
 }
 
-// CHECK-LABEL: @clamp_float_is_noop
-func.func @clamp_float_is_noop(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+// CHECK-LABEL: @clamp_f16_not_noop
+func.func @clamp_f16_not_noop(%arg0: tensor<4xf16>) -> tensor<4xf16> {
+  // CHECK: "tosa.clamp"
+  %0 = "tosa.clamp"(%arg0) {min_int = -128 : i64, max_int = 127 : i64, min_fp = -3.40282347E+38 : f32, max_fp = 3.40282347E+38 : f32} : (tensor<4xf16>) -> tensor<4xf16>
+  return %0 : tensor<4xf16>
+}
+
+// CHECK-LABEL: @clamp_f32_not_noop
+func.func @clamp_f32_not_noop(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  // CHECK: "tosa.clamp"
+  %0 = "tosa.clamp"(%arg0) {min_int = -128 : i64, max_int = 127 : i64, min_fp = -3.40282347E+38 : f32, max_fp = 3.40282347E+38 : f32} : (tensor<4xf32>) -> tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// CHECK-LABEL: @clamp_f16_is_noop
+func.func @clamp_f16_is_noop(%arg0: tensor<4xf16>) -> tensor<4xf16> {
   // CHECK: return %arg0
   // CHECK-NOT: "tosa.clamp"
-  %0 = "tosa.clamp"(%arg0) {min_int = -128 : i64, max_int = 127 : i64, min_fp = -3.40282347E+38 : f32, max_fp = 3.40282347E+38 : f32} :  (tensor<4xf32>) -> tensor<4xf32>
+  // 0xFF800000 and 0x7F800000 are respectively negative and positive F32 infinity.
+  %0 = "tosa.clamp"(%arg0) {min_int = -128 : i64, max_int = 127 : i64, min_fp = 0xFF800000 : f32, max_fp = 0x7F800000 : f32} : (tensor<4xf16>) -> tensor<4xf16>
+  return %0 : tensor<4xf16>
+}
+
+// CHECK-LABEL: @clamp_f32_is_noop
+func.func @clamp_f32_is_noop(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  // CHECK: return %arg0
+  // CHECK-NOT: "tosa.clamp"
+  // 0xFF800000 and 0x7F800000 are respectively negative and positive F32 infinity.
+  %0 = "tosa.clamp"(%arg0) {min_int = -128 : i64, max_int = 127 : i64, min_fp = 0xFF800000 : f32, max_fp = 0x7F800000 : f32} : (tensor<4xf32>) -> tensor<4xf32>
   return %0 : tensor<4xf32>
 }
 
 // CHECK-LABEL: @clamp_int8_is_noop
 func.func @clamp_int8_is_noop(%arg0: tensor<4xi8>) -> tensor<4xi8> {
   // CHECK: return %arg0
   // CHECK-NOT: "tosa.clamp"
   %0 = "tosa.clamp"(%arg0) {min_int = -128 : i64, max_int = 127 : i64, min_fp = -3.40282347E+38 : f32, max_fp = 3.40282347E+38 : f32} :  (tensor<4xi8>) -> tensor<4xi8>