TosaToLinalg: Fixed tosa.cast for small unsigned integers

Author: mgehre-amd

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

@@ -469,7 +469,8 @@ createLinalgBodyCalculationForElementwiseOp(Operation *op, ValueRange args,
                                             args.front(), zero);
     }
 
-    if (arith::FPToSIOp::areCastCompatible(srcTy, dstTy)) {
+    if (dstTy.isSignlessInteger() &&
+        arith::FPToSIOp::areCastCompatible(srcTy, dstTy)) {
       auto intMin = rewriter.create<arith::ConstantOp>(
           loc, rewriter.getF32FloatAttr(
                    APInt::getSignedMinValue(dstTy.getIntOrFloatBitWidth())
@@ -484,16 +485,33 @@ createLinalgBodyCalculationForElementwiseOp(Operation *op, ValueRange args,
 
       auto clamped = clampFloatHelper(loc, rounded, intMin, intMax, rewriter);
 
-      if (dstTy.isUnsignedInteger()) {
-        auto cast = rewriter.create<arith::FPToUIOp>(
-          loc, rewriter.getIntegerType(dstTy.getIntOrFloatBitWidth()), clamped);
-        return rewriter.create<UnrealizedConversionCastOp>(
-          loc, dstTy, cast->getResult(0)).getResult(0);
-      }
-
       return rewriter.create<arith::FPToSIOp>(loc, dstTy, clamped);
     }
 
+    if (dstTy.isUnsignedInteger() &&
+        arith::FPToUIOp::areCastCompatible(srcTy, dstTy)) {
+      auto intMin = rewriter.create<arith::ConstantOp>(
+          loc, rewriter.getF32FloatAttr(
+                   APInt::getMinValue(dstTy.getIntOrFloatBitWidth())
+                       .getZExtValue()));
+
+      auto intMax = rewriter.create<arith::ConstantOp>(
+          loc, rewriter.getF32FloatAttr(
+                   APInt::getMaxValue(dstTy.getIntOrFloatBitWidth())
+                       .getZExtValue()));
+
+      auto rounded = rewriter.create<math::RoundEvenOp>(loc, args[0]);
+
+      auto clamped = clampFloatHelper(loc, rounded, intMin, intMax, rewriter);
+
+      auto cast = rewriter.create<arith::FPToUIOp>(
+          loc, rewriter.getIntegerType(dstTy.getIntOrFloatBitWidth()), clamped);
+      // arith is signless, so temporarily cast back to being unsigned.
+      return rewriter
+          .create<UnrealizedConversionCastOp>(loc, dstTy, cast->getResult(0))
+          .getResult(0);
+    }
+
     // Casting to boolean, integers need to only be checked as not-equal to
     // zero.
     if (srcTy.isa<IntegerType>() && dstTy.isInteger(1)) {

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

@@ -2,12 +2,12 @@
 
 func.func @test_cast(%arg0: tensor<1xf32>) -> tensor<1xui3> {
   // CHECK: linalg.generic
-  // CHECK: arith.constant -4.000000e+00
-  // CHECK: arith.constant 3.000000e+00
+  // CHECK: arith.constant 0.000000e+00
+  // CHECK: arith.constant 7.000000e+00
   // CHECK: math.roundeven
   // CHECK: arith.minf
   // CHECK: arith.maxf
-  // CHECK: arith.fptoui
+  // CHECK: arith.fptoui {{.*}} : f32 to i3
   // CHECK: builtin.unrealized_conversion_cast
   %1 = "tosa.cast"(%arg0) : (tensor<1xf32>) -> tensor<1xui3>