[MLIR][TOSA] Fix converting tosa.clamp and tosa.relu to linalg

[jungpark-mlir]

Summary

Tosa to Linalg conversion crashes when input tensor is a float type other than fp32.
Because tosa.clamp and tosa.reluN have fp32 min/max attribute which is converted as arith.constant with the attribute type.
This commit fixes the crash by correctly setting the float constant type from the input tensor.

[krzysz00]

In "things I'm just now noticing now that we're trying to send this upstream", should we add a test?

[eric-k256]

This doesn't directly impact this change, but in the TOSA specification, we've removed the ReluN operator, since it's trivially implementable with Clamp. I will look at removing ReluN from the dialect to line up with the specification.

[jpienaar]

In D128630#3613348, @krzysz00 wrote:

In "things I'm just now noticing now that we're trying to send this upstream", should we add a test?

That SGTM

[jungpark-mlir]

Add test clamp f16

[krzysz00]

@jpienaar Would you be willing to approve this, or do you not make sense as a reviewer?

[krzysz00]

@eric-k256 Would you be willing to approve here?

[eric-k256]

It looks good to me, although I don't have commit rights, so it might not be all the approvals you need.

As another side note, we recently released a new version of the TOSA specification, where we call out non-fp32 data types directly for operators and their attributes: https://www.mlplatform.org/tosa/tosa_spec.html. Our goal is to do some fine tuning to make sure that the specification part looks good, and then try to align the dialect. This should end up with better fp16 support in the long run.

[krzysz00]

And I do have the bits, so landed. Thanks for the external pair of eyes!

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

if (isa<tosa::CeilOp>(op) && elementTy.isa<FloatType>())
     return rewriter.create<math::CeilOp>(loc, resultTypes, args);
 
   // tosa::FloorOp
   if (isa<tosa::FloorOp>(op) && elementTy.isa<FloatType>())
     return rewriter.create<math::FloorOp>(loc, resultTypes, args);
 
   // tosa::ClampOp
   if (isa<tosa::ClampOp>(op) && elementTy.isa<FloatType>()) {
-    auto min = rewriter.create<arith::ConstantOp>(loc, elementTy,
-                                                  op->getAttr("min_fp"));
-    auto max = rewriter.create<arith::ConstantOp>(loc, elementTy,
-                                                  op->getAttr("max_fp"));
+    bool losesInfo = false;
+    APFloat min_apf = op->getAttr("min_fp").cast<FloatAttr>().getValue();
+    APFloat max_apf = op->getAttr("max_fp").cast<FloatAttr>().getValue();
+    min_apf.convert(elementTy.cast<FloatType>().getFloatSemantics(),
+                    APFloat::rmNearestTiesToEven, &losesInfo);
+    max_apf.convert(elementTy.cast<FloatType>().getFloatSemantics(),
+                    APFloat::rmNearestTiesToEven, &losesInfo);
+    auto min = rewriter.create<arith::ConstantOp>(
+        loc, elementTy, rewriter.getFloatAttr(elementTy, min_apf));
+    auto max = rewriter.create<arith::ConstantOp>(
+        loc, elementTy, rewriter.getFloatAttr(elementTy, max_apf));
     return clampHelper<arith::CmpFOp>(loc, args[0], min, max,
                                       arith::CmpFPredicate::OLT, rewriter);
   }
 
   if (isa<tosa::ClampOp>(op) && elementTy.isa<IntegerType>()) {
     auto intTy = elementTy.cast<IntegerType>();
     int32_t min = static_cast<int32_t>(
         op->getAttr("min_int").cast<IntegerAttr>().getValue().getSExtValue());
if (isa<tosa::ClampOp>(op) && elementTy.isa<IntegerType>()) {
     return clampHelper<arith::CmpIOp>(loc, args[0], minVal, maxVal,
                                       arith::CmpIPredicate::slt, rewriter);
   }
 
   // tosa::ReluNOp
   if (isa<tosa::ReluNOp>(op) && elementTy.isa<FloatType>()) {
     auto zero =
         rewriter.create<arith::ConstantOp>(loc, FloatAttr::get(elementTy, 0));
-    auto n = rewriter.create<arith::ConstantOp>(loc, elementTy,
-                                                op->getAttr("max_fp"));
+    bool losesInfo = false;
+    APFloat max_apf = op->getAttr("max_fp").cast<FloatAttr>().getValue();
+    max_apf.convert(elementTy.cast<FloatType>().getFloatSemantics(),
+                    APFloat::rmNearestTiesToEven, &losesInfo);
+    auto n = rewriter.create<arith::ConstantOp>(
+        loc, elementTy, rewriter.getFloatAttr(elementTy, max_apf));
     return clampHelper<arith::CmpFOp>(loc, args[0], zero, n,
                                       arith::CmpFPredicate::OLT, rewriter);
   }
 
   if (isa<tosa::ReluNOp>(op) && elementTy.isa<IntegerType>()) {
     auto zero =
         rewriter.create<arith::ConstantOp>(loc, IntegerAttr::get(elementTy, 0));
     auto n = createConstFromIntAttribute<int32_t>(op, "max_int", elementTy,

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

func.func @test_i8(%arg0: tensor<1xi8>) -> () {
   // CHECK: %[[SEL2:.+]] = arith.select %[[CMP2]], %[[C127]], %[[SEL1]]
   %1 = "tosa.clamp"(%arg0) {min_int = -130 : i64, max_int = 130 : i64, min_fp = 0.0 : f32, max_fp = 0.0 : f32} : (tensor<1xi8>) -> tensor<1xi8>
 
   return
 }
 
 // -----
 
+// CHECK-LABEL: @test_clamp_f16
+func.func @test_clamp_f16(%arg0: tensor<1xf16>) -> () {
+  // CHECK: linalg.generic
+  // CHECK-DAG: %[[C0:.+]] = arith.constant 0.0
+  // CHECK-DAG: %[[C6:.+]] = arith.constant 6.0
+  // CHECK-DAG: %[[CMP1:.+]] = arith.cmpf olt, %arg1, %[[C0]]
+  // CHECK-DAG: %[[SEL1:.+]] = arith.select %[[CMP1]], %[[C0]]
+  // CHECK-DAG: %[[CMP2:.+]] = arith.cmpf olt, %[[C6]], %arg1
+  // CHECK: %[[SEL2:.+]] = arith.select %[[CMP2]], %[[C6]], %[[SEL1]]
+  %0 = "tosa.clamp"(%arg0) {min_int = 0 : i64, max_int = 0 : i64, min_fp = 0.0 : f32, max_fp = 6.0 : f32} : (tensor<1xf16>) -> tensor<1xf16>
+
+  return
+}
+
+// -----
+
 // CHECK-LABEL: @test_bool
 func.func @test_bool(%arg0: tensor<1xi1>, %arg1: tensor<1xi1>) -> () {
   // CHECK: linalg.generic
   // CHECK: and
   %0 = "tosa.logical_and"(%arg0, %arg1) : (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1>
 
   // CHECK: linalg.generic
   // CHECK: or