[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

293

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

293

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>(loc, args[0], zero, n, CmpIPredicate::slt,
                                      rewriter);
   }
 
+  // tosa::CastOp
+  if (isa<tosa::CastOp>(op)) {
+    Type srcTy = elementTy;
+    Type dstTy = resultTypes.front();
+    bool bitExtend =
+        srcTy.getIntOrFloatBitWidth() < dstTy.getIntOrFloatBitWidth();
+
+    if (srcTy == dstTy)
+      return args.front();
+
+    if (srcTy.isa<FloatType>() && dstTy.isa<FloatType>() && bitExtend)
+      return rewriter.create<mlir::FPExtOp>(loc, resultTypes, args, mlir::None);
+
+    if (srcTy.isa<FloatType>() && dstTy.isa<FloatType>() && !bitExtend)
+      return rewriter.create<mlir::FPTruncOp>(loc, resultTypes, args,
+                                              mlir::None);
+
+    // 1-bit integers need to be treated as signless.
+    if (srcTy.isInteger(1) && mlir::UIToFPOp::areCastCompatible(srcTy, dstTy))
+      return rewriter.create<mlir::UIToFPOp>(loc, resultTypes, args,
+                                             mlir::None);
+
+    if (srcTy.isInteger(1) && dstTy.isa<IntegerType>() && bitExtend)
+      return rewriter.create<mlir::ZeroExtendIOp>(loc, resultTypes, args,
+                                                  mlir::None);
+
+    // All other si-to-fp conversions should be handled by SIToFP.
+    if (mlir::SIToFPOp::areCastCompatible(srcTy, dstTy))
+      return rewriter.create<mlir::SIToFPOp>(loc, resultTypes, args,
+                                             mlir::None);
+
+    // 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);
+
+    // 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);
+  }
+
   (void)rewriter.notifyMatchFailure(
       op, "unhandled op for linalg body calculation for elementwise op");
   return nullptr;
 }
 
 static LogicalResult
 elementwiseMatchAndRewriteHelper(Operation *operation,
                                  PatternRewriter &rewriter) {
patterns->insert<
       PointwiseConverter<tosa::LogOp>, PointwiseConverter<tosa::ExpOp>,
       PointwiseConverter<tosa::AbsOp>, PointwiseConverter<tosa::TanhOp>,
       PointwiseConverter<tosa::BitwiseAndOp>,
       PointwiseConverter<tosa::BitwiseOrOp>,
       PointwiseConverter<tosa::BitwiseXorOp>,
       PointwiseConverter<tosa::LogicalAndOp>,
       PointwiseConverter<tosa::LogicalNotOp>,
       PointwiseConverter<tosa::LogicalOrOp>,
-      PointwiseConverter<tosa::LogicalXorOp>,
+      PointwiseConverter<tosa::LogicalXorOp>, 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>, ConcatOpConversion,
       ReshapeOpConverter, TransposeConverter, RescaleOpConverter>(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-LABEL: @test_bool
 func @test_bool(%arg0: tensor<1xi1>, %arg1: tensor<1xi1>) -> () {
   // CHECK: linalg.generic