[mlir][tosa] Add tosa.cast to linalg lowering

[rsuderman]

Summary

Handles lowering from the tosa CastOp to the equivalent linalg lowering. It
includes support for interchange between bool, int, and floating point.

[antiagainst]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

246

So tosa::CastOp is an op from all types to all types? That's powerful.

I'm a bit puzzled by the if-else structure in the following. In my mind, it should follow the order of:

if (src == FpType && dst == FpType) {
  // Use fp to fp ops
} else if (src == FpType) {
  // Use fp to int ops and handle special case for bool
} else if (dst == FpType) {
  // use int to fp ops and handle special case for bool
} else {
  // Use int to int ops and handle special case for bool
}

or start with if (src == IntType && dst == IntType) ?

[rsuderman]

Cleaned up for clarity.

[rsuderman]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

246

I reworked most of the lowering. This should be a lot cleaner.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

if (isa<tosa::ReluNOp>(op) && elementTy.isa<IntegerType>()) {
     auto zero =
         rewriter.create<mlir::ConstantOp>(loc, IntegerAttr::get(elementTy, 0));
     auto n = createConstFromIntAttribute<int32_t>(op, "max_int", elementTy,
                                                   rewriter);
     return clampHelper<mlir::CmpIOp>(op, args, zero, n, CmpIPredicate::slt,
                                      rewriter);
   }
 
+  // tosa::CastOp
+  if (isa<tosa::CastOp>(op)) {
+    Type sourceType = elementTy;
+    Type targetType = resultTypes.front();
+
+    // A boolean value is considered to be unsigned when converting to
+    // floating-point. Otherwise, it will become `-1`.
+    if (sourceType.isInteger(/*width=*/1) &&
+        mlir::UIToFPOp::areCastCompatible(sourceType, targetType)) {
+      return rewriter.create<mlir::UIToFPOp>(loc, resultTypes, args,
+                                             mlir::None);
+    } else if (mlir::SIToFPOp::areCastCompatible(sourceType, targetType)) {
+      return rewriter.create<mlir::SIToFPOp>(loc, resultTypes, args,
+                                             mlir::None);
+    } else if (sourceType.isa<FloatType>() && targetType.isa<FloatType>()) {
+      FloatType src = sourceType.cast<FloatType>();
+      FloatType res = targetType.cast<FloatType>();
+      if (src.getWidth() > res.getWidth()) {
+        return rewriter.create<mlir::FPTruncOp>(loc, resultTypes, args,
+                                                mlir::None);
+      } else if (src.getWidth() < res.getWidth()) {
+        return rewriter.create<mlir::FPExtOp>(loc, resultTypes, args,
+                                              mlir::None);
+      }
+      // No conversion is needed for the same width floats
+      return args.front();
+    }
+    if (targetType.isInteger(/*width=*/1)) {
+      // When casting to bool, we need to compare whether the value is equal to
+      // zero.
+      if (sourceType.isSignlessInteger()) {
+        Value zeroIntVal = rewriter.create<::mlir::ConstantIntOp>(
+            loc, 0, sourceType.cast<IntegerType>().getWidth());
+        return rewriter.create<mlir::CmpIOp>(loc, CmpIPredicate::ne,
+                                             args.front(), zeroIntVal);
+      } else if (sourceType.isa<FloatType>()) {
+        Value zero = rewriter.create<ConstantOp>(
+            loc, rewriter.getFloatAttr(sourceType, 0.0));
+        return rewriter.create<mlir::CmpFOp>(loc, CmpFPredicate::UNE,
+                                             args.front(), zero);
+      }
+    }
+    if (sourceType.isSignlessInteger() && targetType.isSignlessInteger()) {
+      IntegerType src = sourceType.cast<IntegerType>();
+      IntegerType res = targetType.cast<IntegerType>();
+      if (src.getWidth() > res.getWidth()) {
+        return rewriter.create<mlir::TruncateIOp>(loc, resultTypes, args,
+                                                  mlir::None);
+      } else if (src.getWidth() == 1) {
+        // Special case boolean values, so they get casted to `1` instead of
+        // `-1`.
+        return rewriter.create<mlir::ZeroExtendIOp>(loc, resultTypes, args,
+                                                    mlir::None);
+      } else if (src.getWidth() < res.getWidth()) {
+        return rewriter.create<mlir::SignExtendIOp>(loc, resultTypes, args,
+                                                    mlir::None);
+      }
+      // No conversion is needed for the same width integers
+      return args.front();
+    }
+    if (mlir::FPToSIOp::areCastCompatible(sourceType, targetType)) {
+      return rewriter.create<mlir::FPToSIOp>(loc, resultTypes, args,
+                                             mlir::None);
+    }
+    return nullptr;
+  }
+
   (void)rewriter.notifyMatchFailure(
       op, "unhandled op for linalg body calculation for elementwise op");
   return nullptr;
 }
 
 static LogicalResult
 elementwiseMatchAndRewriteHelper(Operation *operation,
                                  PatternRewriter &rewriter) {
void mlir::tosa::populateTosaToLinalgOnTensorsConversionPatterns(
   patterns->insert<
       PointwiseConverter<tosa::AddOp>, PointwiseConverter<tosa::SubOp>,
       PointwiseConverter<tosa::MulOp>, PointwiseConverter<tosa::NegateOp>,
       PointwiseConverter<tosa::PowOp>, PointwiseConverter<tosa::RsqrtOp>,
       PointwiseConverter<tosa::LogOp>, PointwiseConverter<tosa::ExpOp>,
       PointwiseConverter<tosa::AbsOp>, PointwiseConverter<tosa::TanhOp>,
       PointwiseConverter<tosa::BitwiseAndOp>,
       PointwiseConverter<tosa::BitwiseOrOp>,
-      PointwiseConverter<tosa::BitwiseXorOp>,
+      PointwiseConverter<tosa::BitwiseXorOp>, PointwiseConverter<tosa::CastOp>,
       PointwiseConverter<tosa::LogicalLeftShiftOp>,
       PointwiseConverter<tosa::LogicalRightShiftOp>,
       PointwiseConverter<tosa::SelectOp>, PointwiseConverter<tosa::GreaterOp>,
       PointwiseConverter<tosa::GreaterEqualOp>,
       PointwiseConverter<tosa::MaximumOp>, PointwiseConverter<tosa::MinimumOp>,
       PointwiseConverter<tosa::CeilOp>, PointwiseConverter<tosa::FloorOp>,
       PointwiseConverter<tosa::ClampOp>, PointwiseConverter<tosa::ReluNOp>,
       IdentityNConverter<tosa::IdentityOp>,
       IdentityNConverter<tosa::IdentityNOp>, ReduceConverter<tosa::ReduceMinOp>,
       ReduceConverter<tosa::ReduceMaxOp>, ReduceConverter<tosa::ReduceSumOp>,
       ReduceConverter<tosa::ReduceProdOp>, ReshapeOpConverter,
       TransposeConverter>(context);
 }

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

func @test_simple_f32(%arg0: tensor<1xf32>) -> () {
   // CHECK: select
   %17 = "tosa.clamp"(%0) {min_int = 1 : i64, max_int = 5 : i64, min_fp = 1.0 : f32, max_fp = 5.0 : f32} : (tensor<1xf32>) -> tensor<1xf32>
 
   // CHECK: linalg.generic
   // CHECK: cmpf
   // CHECK: select
   %18 = "tosa.reluN"(%0) {max_int = 5 : i64, max_fp = 5.0 : f32} : (tensor<1xf32>) -> tensor<1xf32>
 
+  // CHECK: linalg.generic
+  // CHECK: fptosi
+  %19 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xi32>
+
+  // CHECK: linalg.generic
+  // CHECK: constant 0
+  // CHECK: cmpf
+  %20 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xi1>
+
+  // CHECK: linalg.generic
+  // CHECK: fptrunc
+  %21 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xf16>
+
+  // CHECK: linalg.generic
+  // CHECK: yield
+  %22 = "tosa.cast"(%0) : (tensor<1xf32>) -> tensor<1xf32>
+
+  return
+}
+
+// -----
+
+// CHECK-LABEL: @test_simple_f16
+func @test_simple_f16(%arg0: tensor<1xf16>) -> () {
+
+  // CHECK: linalg.generic
+  // CHECK: fpext
+  %0 = "tosa.cast"(%arg0) : (tensor<1xf16>) -> tensor<1xf32>
+
   return
 }
 
 // -----
 
 // CHECK-LABEL: @test_simple_i32
 func @test_simple_i32(%arg0: tensor<1xi32>) -> () {
   // CHECK: linalg.generic
func @test_simple_i32(%arg0: tensor<1xi32>) -> () {
   // CHECK: select
   %14 = "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
   %15 = "tosa.reluN"(%0) {max_int = 5 : i64, max_fp = 5.0 : f32} : (tensor<1xi32>) -> tensor<1xi32>
 
+  // CHECK: linalg.generic
+  // CHECK: trunci
+  %16 = "tosa.cast"(%0) : (tensor<1xi32>) -> tensor<1xi16>
+
+  // CHECK: linalg.generic
+  // CHECK: yield
+  %17 = "tosa.cast"(%0) : (tensor<1xi32>) -> tensor<1xi32>
+
+  // CHECK: linalg.generic
+  // CHECK: sexti
+  %18 = "tosa.cast"(%0) : (tensor<1xi32>) -> tensor<1xi64>
+
+  // CHECK: linalg.generic
+  // CHECK: constant 0
+  // CHECK: cmpi
+  %19 = "tosa.cast"(%0) : (tensor<1xi32>) -> tensor<1xi1>
+
+  // CHECK: linalg.generic
+  // CHECK: sitofp
+  %20 = "tosa.cast"(%0) : (tensor<1xi32>) -> tensor<1xf32>
+
   return
 }
 
 // -----
 
 // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 // CHECK-LABEL: @test_reshape_downrank
 func @test_reshape_downrank(%arg0: tensor<2x3xf32>) -> tensor<6xf32> {