[mlir][emitc] Arith to EmitC: handle FP<->Integer conversions

Author: cferry-AMD

mlir/include/mlir/Conversion/ArithToEmitC/ArithToEmitC.h

@@ -14,7 +14,8 @@ class RewritePatternSet;
 class TypeConverter;
 
 void populateArithToEmitCPatterns(TypeConverter &typeConverter,
-                                  RewritePatternSet &patterns);
+                                  RewritePatternSet &patterns,
+                                  bool optionFloatToIntTruncates);
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_ARITHTOEMITC_ARITHTOEMITC_H

mlir/include/mlir/Conversion/Passes.td

@@ -139,7 +139,24 @@ def ArithToAMDGPUConversionPass : Pass<"convert-arith-to-amdgpu"> {
 
 def ConvertArithToEmitC : Pass<"convert-arith-to-emitc"> {
   let summary = "Convert Arith dialect to EmitC dialect";
+  let description = [{
+    This pass converts `arith` dialect operations to `emitc`.
+
+    The semantics of floating-point to integer conversions `arith.fptosi`, 
+    `arith.fptoui` require rounding towards zero. Typical C++ implementations
+    use this behavior for float-to-integer casts, but that is not mandated by 
+    C++ and there are implementation-defined means to change the default behavior.
+    
+    If casts can be guaranteed to use round-to-zero, use the 
+    `float-to-int-truncates` flag to allow conversion of `arith.fptosi` and
+    `arith.fptoui` operations.
+  }];
   let dependentDialects = ["emitc::EmitCDialect"];
+  let options = [
+    Option<"floatToIntTruncates", "float-to-int-truncates", "bool",
+           /*default=*/"false",
+           "Whether the behavior of float-to-int cast in emitc is truncation">,
+  ];
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Conversion/ArithToEmitC/ArithToEmitC.cpp

@@ -128,14 +128,87 @@ class SelectOpConversion : public OpConversionPattern<arith::SelectOp> {
   }
 };
 
+// Floating-point to integer conversions.
+template <typename CastOp>
+class FtoICastOpConversion : public OpConversionPattern<CastOp> {
+private:
+  bool floatToIntTruncates;
+
+public:
+  FtoICastOpConversion(const TypeConverter &typeConverter, MLIRContext *context,
+                       bool optionFloatToIntTruncates)
+      : OpConversionPattern<CastOp>(typeConverter, context),
+        floatToIntTruncates(optionFloatToIntTruncates) {}
+
+  LogicalResult
+  matchAndRewrite(CastOp castOp, typename CastOp::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    Type operandType = adaptor.getIn().getType();
+    if (!emitc::isSupportedFloatType(operandType))
+      return rewriter.notifyMatchFailure(castOp,
+                                         "unsupported cast source type");
+
+    if (!floatToIntTruncates)
+      return rewriter.notifyMatchFailure(
+          castOp, "conversion currently requires EmitC casts to use truncation "
+                  "as rounding mode");
+
+    Type dstType = this->getTypeConverter()->convertType(castOp.getType());
+    if (!dstType)
+      return rewriter.notifyMatchFailure(castOp, "type conversion failed");
+
+    if (!emitc::isSupportedIntegerType(dstType))
+      return rewriter.notifyMatchFailure(castOp,
+                                         "unsupported cast destination type");
+
+    rewriter.replaceOpWithNewOp<emitc::CastOp>(castOp, dstType,
+                                               adaptor.getOperands());
+
+    return success();
+  }
+};
+
+// Integer to floating-point conversions.
+template <typename CastOp>
+class ItoFCastOpConversion : public OpConversionPattern<CastOp> {
+public:
+  ItoFCastOpConversion(const TypeConverter &typeConverter, MLIRContext *context)
+      : OpConversionPattern<CastOp>(typeConverter, context) {}
+
+  LogicalResult
+  matchAndRewrite(CastOp castOp, typename CastOp::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    Type operandType = adaptor.getIn().getType();
+    if (!emitc::isSupportedIntegerType(operandType))
+      return rewriter.notifyMatchFailure(castOp,
+                                         "unsupported cast source type");
+
+    Type dstType = this->getTypeConverter()->convertType(castOp.getType());
+    if (!dstType)
+      return rewriter.notifyMatchFailure(castOp, "type conversion failed");
+
+    if (!emitc::isSupportedFloatType(dstType))
+      return rewriter.notifyMatchFailure(castOp,
+                                         "unsupported cast destination type");
+
+    rewriter.replaceOpWithNewOp<emitc::CastOp>(castOp, dstType,
+                                               adaptor.getOperands());
+
+    return success();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Pattern population
 //===----------------------------------------------------------------------===//
 
 void mlir::populateArithToEmitCPatterns(TypeConverter &typeConverter,
-                                        RewritePatternSet &patterns) {
+                                        RewritePatternSet &patterns,
+                                        bool optionFloatToIntTruncates) {
   MLIRContext *ctx = patterns.getContext();
 
   // clang-format off
@@ -148,7 +221,13 @@ void mlir::populateArithToEmitCPatterns(TypeConverter &typeConverter,
     IntegerOpConversion<arith::AddIOp, emitc::AddOp>,
     IntegerOpConversion<arith::MulIOp, emitc::MulOp>,
     IntegerOpConversion<arith::SubIOp, emitc::SubOp>,
-    SelectOpConversion
-  >(typeConverter, ctx);
+    SelectOpConversion,
+    ItoFCastOpConversion<arith::SIToFPOp>,
+    ItoFCastOpConversion<arith::UIToFPOp>
+  >(typeConverter, ctx)
+  .add<
+    FtoICastOpConversion<arith::FPToSIOp>,
+    FtoICastOpConversion<arith::FPToUIOp>
+  >(typeConverter, ctx, optionFloatToIntTruncates);
   // clang-format on
 }

mlir/lib/Conversion/ArithToEmitC/ArithToEmitCPass.cpp

@@ -29,6 +29,8 @@ using namespace mlir;
 namespace {
 struct ConvertArithToEmitC
     : public impl::ConvertArithToEmitCBase<ConvertArithToEmitC> {
+  using Base::Base;
+
   void runOnOperation() override;
 };
 } // namespace
@@ -44,7 +46,7 @@ void ConvertArithToEmitC::runOnOperation() {
   TypeConverter typeConverter;
   typeConverter.addConversion([](Type type) { return type; });
 
-  populateArithToEmitCPatterns(typeConverter, patterns);
+  populateArithToEmitCPatterns(typeConverter, patterns, floatToIntTruncates);
 
   if (failed(
           applyPartialConversion(getOperation(), target, std::move(patterns))))

mlir/test/Conversion/ArithToEmitC/arith-to-emitc-cast-truncate.mlir

@@ -0,0 +1,20 @@
+// RUN: mlir-opt -split-input-file --pass-pipeline="builtin.module(convert-arith-to-emitc{float-to-int-truncates})" %s | FileCheck %s
+
+func.func @arith_float_to_int_cast_ops(%arg0: f32, %arg1: f64) {
+  // CHECK: emitc.cast %arg0 : f32 to i32
+  %0 = arith.fptosi %arg0 : f32 to i32
+
+  // CHECK: emitc.cast %arg1 : f64 to i32
+  %1 = arith.fptosi %arg1 : f64 to i32
+
+  // CHECK: emitc.cast %arg0 : f32 to i16
+  %2 = arith.fptosi %arg0 : f32 to i16
+
+  // CHECK: emitc.cast %arg1 : f64 to i16
+  %3 = arith.fptosi %arg1 : f64 to i16
+
+  // CHECK: emitc.cast %arg0 : f32 to i32
+  %4 = arith.fptoui %arg0 : f32 to i32
+  
+  return
+}

mlir/test/Conversion/ArithToEmitC/arith-to-emitc-cast-unsupported.mlir

@@ -0,0 +1,48 @@
+// RUN: mlir-opt -split-input-file --pass-pipeline="builtin.module(convert-arith-to-emitc{float-to-int-truncates})" -verify-diagnostics %s
+
+func.func @arith_cast_tensor(%arg0: tensor<5xf32>) -> tensor<5xi32> {
+  // expected-error @+1 {{failed to legalize operation 'arith.fptosi'}}
+  %t = arith.fptosi %arg0 : tensor<5xf32> to tensor<5xi32>
+  return %t: tensor<5xi32>
+}
+
+// -----
+
+func.func @arith_cast_vector(%arg0: vector<5xf32>) -> vector<5xi32> {
+  // expected-error @+1 {{failed to legalize operation 'arith.fptosi'}}
+  %t = arith.fptosi %arg0 : vector<5xf32> to vector<5xi32>
+  return %t: vector<5xi32>
+}
+
+// -----
+
+func.func @arith_cast_bf16(%arg0: bf16) -> i32 {
+  // expected-error @+1 {{failed to legalize operation 'arith.fptosi'}}
+  %t = arith.fptosi %arg0 : bf16 to i32
+  return %t: i32
+}
+
+// -----
+
+func.func @arith_cast_f16(%arg0: f16) -> i32 {
+  // expected-error @+1 {{failed to legalize operation 'arith.fptosi'}}
+  %t = arith.fptosi %arg0 : f16 to i32
+  return %t: i32
+}
+
+
+// -----
+
+func.func @arith_cast_to_bf16(%arg0: i32) -> bf16 {
+  // expected-error @+1 {{failed to legalize operation 'arith.sitofp'}}
+  %t = arith.sitofp %arg0 : i32 to bf16
+  return %t: bf16
+}
+
+// -----
+
+func.func @arith_cast_to_f16(%arg0: i32) -> f16 {
+  // expected-error @+1 {{failed to legalize operation 'arith.sitofp'}}
+  %t = arith.sitofp %arg0 : i32 to f16
+  return %t: f16
+}

mlir/test/Conversion/ArithToEmitC/arith-to-emitc-unsupported.mlir

@@ -0,0 +1,7 @@
+// RUN: mlir-opt -split-input-file -convert-arith-to-emitc -verify-diagnostics %s
+
+func.func @arith_cast_f32(%arg0: f32) -> i32 {
+  // expected-error @+1 {{failed to legalize operation 'arith.fptosi'}}
+  %t = arith.fptosi %arg0 : f32 to i32
+  return %t: i32
+}

mlir/test/Conversion/ArithToEmitC/arith-to-emitc.mlir

@@ -93,3 +93,18 @@ func.func @arith_select(%arg0: i1, %arg1: tensor<8xi32>, %arg2: tensor<8xi32>) -
   %0 = arith.select %arg0, %arg1, %arg2 : i1, tensor<8xi32>
   return
 }
+
+// -----
+
+func.func @arith_int_to_float_cast_ops(%arg0: i8, %arg1: i64) {
+  // CHECK: emitc.cast %arg0 : i8 to f32
+  %0 = arith.sitofp %arg0 : i8 to f32
+
+  // CHECK: emitc.cast %arg1 : i64 to f32
+  %1 = arith.sitofp %arg1 : i64 to f32
+
+  // CHECK: emitc.cast %arg0 : i8 to f32
+  %2 = arith.uitofp %arg0 : i8 to f32
+
+  return
+}