Enable lowering all float types to tosa

src/Conversion/ONNXToTOSA/ConvertONNXToTOSA.cpp

@@ -120,7 +120,7 @@ void FrontendToTosaLoweringPass::runOnOperation() {
   // conversion failures. Quantized types are not supported right now.
   TypeConverter typeConverter;
   typeConverter.addConversion([](Type type) -> std::optional<Type> {
-    if (isTOSAInt(type) || isTOSAFloat(type) || type.isa<NoneType>() ||
+    if (isTOSAInt(type) || isa<FloatType>(type) || type.isa<NoneType>() ||
         isTOSABool(type))
       return type;
     return std::nullopt;

src/Conversion/ONNXToTOSA/Math/Elementwise.cpp

@@ -31,48 +31,10 @@ struct TOSADialectOp<ONNXNegOp> {
   using Op = mlir::tosa::NegateOp;
 };
 
-struct AbsIOSupportedTypes {
-  static LogicalResult checkType(
-      ConversionPatternRewriter &rewriter, Type scalarType, Operation *op) {
-    if (!mlir::isa<mlir::BFloat16Type, mlir::Float16Type, mlir::Float32Type>(
-            scalarType) &&
-        !scalarType.isSignlessInteger(/*width=*/32)) {
-      return rewriter.notifyMatchFailure(op,
-          "this operation only supports signless 32 integer or fp16, fp32"
-          " or bf16");
-    }
-    return success();
-  }
-};
-
-struct ErfIOSupportedTypes {
-  static LogicalResult checkType(
-      ConversionPatternRewriter &rewriter, Type scalarType, Operation *op) {
-    if (!mlir::isa<mlir::BFloat16Type, mlir::Float16Type, mlir::Float32Type>(
-            scalarType)) {
-      return rewriter.notifyMatchFailure(
-          op, "this operation only supports fp16, fp32 or bf16");
-    }
-    return success();
-  }
-};
-
-struct IsAnyLegalType {
-  static LogicalResult checkType(
-      ConversionPatternRewriter &rewriter, Type scalarType, Operation *op) {
-    if (!isTOSAFloat(scalarType) && !isTOSAInt(scalarType) &&
-        !isTOSABool(scalarType)) {
-      return rewriter.notifyMatchFailure(
-          op, "this operation only supports signed integer or float types");
-    }
-    return success();
-  }
-};
-
 struct IsIntOrFloat {
   static LogicalResult checkType(
       ConversionPatternRewriter &rewriter, Type scalarType, Operation *op) {
-    if (!isTOSAFloat(scalarType) && !isTOSAInt(scalarType)) {
+    if (!isa<FloatType>(scalarType) && !isTOSAInt(scalarType)) {
       return rewriter.notifyMatchFailure(
           op, "this operation only supports signed integer or float types");
     }
@@ -94,7 +56,7 @@ struct IsInt {
 struct IsFloat {
   static LogicalResult checkType(
       ConversionPatternRewriter &rewriter, Type scalarType, Operation *op) {
-    if (!isTOSAFloat(scalarType)) {
+    if (!isa<FloatType>(scalarType)) {
       return rewriter.notifyMatchFailure(
           op, "this operation only supports float types");
     }
@@ -757,9 +719,9 @@ static void populateLoweringONNXElementwiseUnaryTemplateOpToTOSAPattern(
       ONNXElementwiseUnaryOpLoweringToTOSA<ONNXNotOp, mlir::tosa::LogicalNotOp,
           IsBool, IsBool>,
       ONNXElementwiseUnaryOpLoweringToTOSA<ONNXAbsOp, mlir::tosa::AbsOp,
-          AbsIOSupportedTypes, AbsIOSupportedTypes>,
+          IsIntOrFloat, IsIntOrFloat>,
       ONNXElementwiseUnaryOpLoweringToTOSA<ONNXErfOp, mlir::tosa::ErfOp,
-          ErfIOSupportedTypes, ErfIOSupportedTypes>>(typeConverter, ctx);
+          IsFloat, IsFloat>>(typeConverter, ctx);
 
 // Tosa custom ops
 #define INSERT_ONNX_UNARY_TO_TOSA_CUSTOMOP_PATTERN(                            \

src/Conversion/ONNXToTOSA/Tensor/Resize.cpp

@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/IR/BuiltinTypes.h"
 
 #include "src/Conversion/ONNXToTOSA/DialectBuilder.hpp"
 #include "src/Conversion/ONNXToTOSA/ONNXToTOSACommon.hpp"
@@ -203,7 +204,7 @@ class ONNXResizeOpLoweringToTOSA : public ConversionPattern {
     }
 
     auto elementType = inputType.getElementType();
-    if (!(isTOSAFloat(elementType) || isTOSAInt(elementType))) {
+    if (!(isa<FloatType>(elementType) || isTOSAInt(elementType))) {
       return rewriter.notifyMatchFailure(
           resizeOp, "Element type is not supported by TOSA.");
     }

src/Conversion/ONNXToTOSA/Tensor/Transpose.cpp

@@ -53,7 +53,7 @@ class ONNXTransposeLoweringToTOSA
 
     Type inputElementType = inputType.getElementType();
 
-    if (!isTOSAFloat(inputElementType) && !isTOSAInt(inputElementType) &&
+    if (!isa<FloatType>(inputElementType) && !isTOSAInt(inputElementType) &&
         !inputElementType.isInteger(1)) {
       return rewriter.notifyMatchFailure(
           op, "input element type not supported");

test/mlir/conversion/onnx_to_tosa/Math/Elementwise.mlir

@@ -103,6 +103,15 @@ func.func @test_add_ui32(%arg0: tensor<13x21x1xui32>, %arg1: tensor<13x21x1xui32
 // CHECK:           return [[VAR_0_]] : tensor<13x21x1xui32>
 }
 
+// -----
+
+func.func @test_add_f64(%arg0: tensor<13x21x1xf64>, %arg1: tensor<13x21x1xf64>) -> tensor<13x21x1xf64> {
+  %0 = "onnx.Add"(%arg0, %arg1) : (tensor<13x21x1xf64>, tensor<13x21x1xf64>) -> tensor<13x21x1xf64>
+  "func.return"(%0) : (tensor<13x21x1xf64>) -> ()
+// CHECK-LABEL:  func.func @test_add_f64
+// CHECK-NOT:    onnx.Add
+// CHECK:        return {{.*}}: tensor<13x21x1xf64>
+}
 
 // -----
 
@@ -484,6 +493,14 @@ func.func @test_pow_broadcast(%arg0: tensor<13x21x1xf32>, %arg1: tensor<1xf32>)
 // CHECK-NEXT:      [[VAR_1_:%.+]] = tosa.pow [[PARAM_0_]], [[VAR_0_]] : (tensor<13x21x1xf32>, tensor<1x1x1xf32>) -> tensor<13x21x1xf32>
 }
 
+func.func @test_pow_f64(%arg0: tensor<13x21x1xf64>, %arg1: tensor<13x21x1xf64>) -> tensor<13x21x1xf64> {
+  %0 = "onnx.Pow"(%arg0, %arg1) : (tensor<13x21x1xf64>, tensor<13x21x1xf64>) -> tensor<13x21x1xf64>
+  "func.return"(%0) : (tensor<13x21x1xf64>) -> ()
+// CHECK-LABEL:  func @test_pow
+// CHECK-NOT:    onnx.Pow
+// CHECK:        return {{.*}}: tensor<13x21x1xf64>
+}
+
 // -----
 
 func.func @test_sqrt(%arg0: tensor<3xf32>) -> tensor<3xf32> {
@@ -519,6 +536,14 @@ func.func @test_abs_bf16(%arg0: tensor<3xbf16>) -> tensor<3xbf16> {
 // CHECK-NEXT:   }
 }
 
+func.func @test_abs_f64(%arg0: tensor<3xf64>) -> tensor<3xf64> {
+  %0 = "onnx.Abs"(%arg0) : (tensor<3xf64>) -> tensor<3xf64>
+  return %0 : tensor<3xf64>
+// CHECK-LABEL:  func @test_abs_f64
+// CHECK-NOT:    onnx.Abs
+// CHECK:        return {{.*}}: tensor<3xf64>
+}
+
 // -----
 
 func.func @test_erf_f32(%arg0: tensor<3xf32>) -> tensor<3xf32> {
@@ -541,6 +566,14 @@ func.func @test_erf_bf16(%arg0: tensor<3xbf16>) -> tensor<3xbf16> {
 // CHECK-NEXT:   }
 }
 
+func.func @test_erf_f64(%arg0: tensor<3xf64>) -> tensor<3xf64> {
+  %0 = "onnx.Erf"(%arg0) : (tensor<3xf64>) -> tensor<3xf64>
+  return %0 : tensor<3xf64>
+// CHECK-LABEL:  func @test_erf_f64
+// CHECK-NOT:    onnx.Erf
+// CHECK:        return %0 : tensor<3xf64>
+}
+
 // -----
 
 func.func @test_bitwise_not(%arg0 : tensor<10x10xi32>) -> tensor<10x10xi32> {

test/mlir/conversion/onnx_to_tosa/NN/BatchNorm.mlir

@@ -85,7 +85,6 @@ func.func @test_batchnorm_bf16_dynamic(%arg0: tensor<100x3x?x?xbf16>) -> tensor<
 }
 
 // -----
-// tosa doesn't support f64, so it should not be lowered
 
 func.func @test_batchnorm_f64(%arg0: tensor<100x3x10x10xf64>) -> tensor<100x3x10x10xf64> {
     %0 = "onnx.Constant"() {value = dense<[1.0, 2.0, 3.0]> : tensor<3xf64>} : () -> tensor<3xf64>
@@ -94,6 +93,22 @@ func.func @test_batchnorm_f64(%arg0: tensor<100x3x10x10xf64>) -> tensor<100x3x10
     %3 = "onnx.Constant"() {value = dense<[4.0, 5.0, 6.0]> : tensor<3xf64>} : () -> tensor<3xf64>
     %4 = "onnx.BatchNormalizationInferenceMode"(%arg0, %0, %1, %2, %3) {epsilon = 1.00000007E-5 : f32} : (tensor<100x3x10x10xf64>, tensor<3xf64>, tensor<3xf64>, tensor<3xf64>, tensor<3xf64>) -> tensor<100x3x10x10xf64>
     return %4 : tensor<100x3x10x10xf64>
-// CHECK:       onnx.BatchNormalizationInferenceMode
-// CHECK-NOT:   tosa
+// CHECK-LABEL: @test_batchnorm_f64
+// CHECK-SAME: ([[PARAM_0:%.*]]: tensor<100x3x10x10xf64>) -> tensor<100x3x10x10xf64> {
+// CHECK-NEXT: [[VAR_0_:%.+]] = "tosa.const"() <{value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00]> : tensor<3xf64>}> : () -> tensor<3xf64>
+// CHECK-NEXT: [[VAR_1_:%.+]]  = "tosa.const"() <{value = dense<[2.000000e+00, 3.000000e+00, 4.000000e+00]> : tensor<3xf64>}> : () -> tensor<3xf64>
+// CHECK-NEXT: [[VAR_2_:%.+]]  = "tosa.const"() <{value = dense<[3.000000e+00, 4.000000e+00, 5.000000e+00]> : tensor<3xf64>}> : () -> tensor<3xf64>
+// CHECK-NEXT: [[VAR_3_:%.+]]  = "tosa.const"() <{value = dense<[4.000000e+00, 5.000000e+00, 6.000000e+00]> : tensor<3xf64>}> : () -> tensor<3xf64>
+// CHECK-NEXT: [[VAR_4_:%.+]]  = tosa.reshape [[VAR_2_]] {new_shape = array<i64: 1, 3, 1, 1>} : (tensor<3xf64>) -> tensor<1x3x1x1xf64>
+// CHECK-NEXT: [[VAR_5_:%.+]]  = tosa.reshape [[VAR_0_]] {new_shape = array<i64: 1, 3, 1, 1>} : (tensor<3xf64>) -> tensor<1x3x1x1xf64>
+// CHECK-NEXT: [[VAR_6_:%.+]]  = tosa.reshape [[VAR_1_]] {new_shape = array<i64: 1, 3, 1, 1>} : (tensor<3xf64>) -> tensor<1x3x1x1xf64>
+// CHECK-NEXT: [[VAR_7_:%.+]]  = tosa.reshape [[VAR_3_]] {new_shape = array<i64: 1, 3, 1, 1>} : (tensor<3xf64>) -> tensor<1x3x1x1xf64>
+// CHECK-NEXT: [[VAR_8_:%.+]]  = "tosa.const"() <{value = dense<1.0000000656873453E-5> : tensor<1x1x1x1xf64>}> : () -> tensor<1x1x1x1xf64>
+// CHECK-NEXT: [[VAR_9_:%.+]]  = tosa.sub %arg0, [[VAR_4_]] : (tensor<100x3x10x10xf64>, tensor<1x3x1x1xf64>) -> tensor<100x3x10x10xf64>
+// CHECK-NEXT: [[VAR_10_:%.+]]  = tosa.add %7, [[VAR_8_]] : (tensor<1x3x1x1xf64>, tensor<1x1x1x1xf64>) -> tensor<1x3x1x1xf64>
+// CHECK-NEXT: [[VAR_11_:%.+]]  = tosa.rsqrt [[VAR_10_]] : (tensor<1x3x1x1xf64>) -> tensor<1x3x1x1xf64>
+// CHECK-NEXT: [[VAR_12_:%.+]]  = tosa.mul [[VAR_9_]], %11 {shift = 0 : i8} : (tensor<100x3x10x10xf64>, tensor<1x3x1x1xf64>) -> tensor<100x3x10x10xf64>
+// CHECK-NEXT: [[VAR_13_:%.+]]  = tosa.mul [[VAR_12_]], %5 {shift = 0 : i8} : (tensor<100x3x10x10xf64>, tensor<1x3x1x1xf64>) -> tensor<100x3x10x10xf64>
+// CHECK-NEXT: [[VAR_14_:%.+]]  = tosa.add [[VAR_13_]], [[VAR_6_]] : (tensor<100x3x10x10xf64>, tensor<1x3x1x1xf64>) -> tensor<100x3x10x10xf64>
+// CHECK-NEXT: return [[VAR_14_]] : tensor<100x3x10x10xf64>
 }

test/mlir/conversion/onnx_to_tosa/Tensor/Reshape.mlir

@@ -19,3 +19,13 @@ func.func @test_reshape_allowzero(%arg0 : tensor<12x128x1024xf32>) -> tensor<12x
 // CHECK:           [[VAR_1_:%.+]] = tosa.reshape [[PARAM_0_]] {new_shape = array<i64: 12, 128, 16, 64>} : (tensor<12x128x1024xf32>) -> tensor<12x128x16x64xf32>
 // CHECK-NEXT:      return [[VAR_1_]] : tensor<12x128x16x64xf32>
 }
+
+func.func @test_reshape_fp8(%arg0 : tensor<128x1024xf8E5M2FNUZ>) -> tensor<1x128x16x64xf8E5M2FNUZ> {
+  %0 = "onnx.Constant"() {value = dense<[-1, 128, 16, 64]> : tensor<4xi64>} : () -> tensor<4xi64>
+  %1 = "onnx.Reshape"(%arg0, %0) : (tensor<128x1024xf8E5M2FNUZ>, tensor<4xi64>) -> tensor<1x128x16x64xf8E5M2FNUZ>
+  "func.return"(%1) : (tensor<1x128x16x64xf8E5M2FNUZ>) -> ()
+// CHECK-LABEL: @test_reshape_fp8
+// CHECK-SAME: ([[PARAM_0_:%.+]] tensor<128x1024xf8E5M2FNUZ>) -> tensor<1x128x16x64xf8E5M2FNUZ> {
+// CHECK: [[VAR_1_:%.+]] = tosa.reshape %arg0 {new_shape = array<i64: 1, 128, 16, 64>} : (tensor<128x1024xf8E5M2FNUZ>) -> tensor<1x128x16x64xf8E5M2FNUZ>
+// CHECK-NEXT: return [[VAR_1_]] : tensor<1x128x16x64xf8E5M2FNUZ>
+  }

test/mlir/conversion/onnx_to_tosa/Tensor/Resize.mlir

@@ -102,6 +102,18 @@ func.func @test_resize_half_pixel_nearest_floor_downsample(%arg0: tensor<1x1x1x1
 
 // -----
 
+func.func @test_resize_f64(%arg0: tensor<1x1x1x4xf64>) -> tensor<1x1x1x12xf64> {
+    %0 = "onnx.NoValue"() {value} : () -> none
+    %1 = "onnx.Constant"() {value = dense<[1, 1, 1, 12]> : tensor<4xi64>} : () -> tensor<4xi64>
+    %2 = "onnx.Resize"(%arg0, %0, %0, %1) {coordinate_transformation_mode = "half_pixel", cubic_coeff_a = -7.500000e-01 : f32, exclude_outside = 0 : si64, extrapolation_value = 0.000000e+00 : f32, mode = "nearest", nearest_mode = "floor"} : (tensor<1x1x1x4xf64>, none, none, tensor<4xi64>) -> tensor<1x1x1x12xf64>
+    return %2 : tensor<1x1x1x12xf64>
+// CHECK-LABEL:  func.func @test_resize_f64
+// CHECK-NOT:    onnx.Resize
+// CHECK:        return {{.*}}: tensor<1x1x1x12xf64>
+}
+
+// -----
+
 func.func @test_resize_input_one(%arg0: tensor<1x1x1x1xf32>) -> tensor<1x1x4x4xf32> {
     %0 = "onnx.NoValue"() {value} : () -> none
     %1 = "onnx.Constant"() {value = dense<[1, 1, 4, 4]> : tensor<4xi64>} : () -> tensor<4xi64>

test/mlir/conversion/onnx_to_tosa/Tensor/Transpose.mlir

@@ -19,3 +19,11 @@ func.func @test_transpose(%arg0 : tensor<5x5x1x32xf32>) -> tensor<5x1x32x5xf32>
 // CHECK:           %[[VAL_2:.*]] = tosa.transpose %[[VAL_0]], %[[VAL_1]] : (tensor<5x5x1x32xf32>, tensor<4xi32>) -> tensor<5x1x32x5xf32>
 // CHECK:           return %[[VAL_2]] : tensor<5x1x32x5xf32>
 }
+
+func.func @test_transpose_f64(%arg0 : tensor<5x5x1x32xf64>) -> tensor<5x1x32x5xf64> {
+  %0 = "onnx.Transpose"(%arg0) {perm = [0, 2, 3, 1]} : (tensor<5x5x1x32xf64>) -> tensor<5x1x32x5xf64>
+  return %0 : tensor<5x1x32x5xf64>
+// CHECK-LABEL:   func.func @test_transpose
+// CHECK-NOT:     onnx.Transpose
+// CHECK:         return {{.*}}: tensor<5x1x32x5xf64>
+}

utils/clone-mlir.sh

@@ -1,3 +1,3 @@
 git clone -n https://github.com/xilinx/llvm-project.git
 # Check out a specific branch that is known to work with ONNX-MLIR.
-cd llvm-project && git checkout af893daa8d19 && cd ..
+cd llvm-project && git checkout fda272652fd65e139ed162a9c7ce521133eb34a0 && cd ..