[mlir][Math] Add TruncOp.

[jacquesguan]

Summary

This patch adds TruncOp for Math, it returns the operand rounded to the nearest integer not larger in magnitude than the operand. And this patch also adds the correspond llvm intrinsic op.

[Mogball]

What's the difference between this and arith.truncf?

[jacquesguan]

In D133342#3772301, @Mogball wrote:

What's the difference between this and arith.truncf?

arith.truncf would use the default rounding mode and has the similar semantic with fptrunc in llvm ir; math.trunc should always use strict rounding to the nearest integer not larger in magnitude than the operand and has the similar semantic with llvm.trunc in llvm ir or trunc in libm.

[mehdi_amini]

In D133342#3773604, @jacquesguan wrote:

In D133342#3772301, @Mogball wrote:

What's the difference between this and arith.truncf?

arith.truncf would use the default rounding mode and has the similar semantic with fptrunc in llvm ir; math.trunc should always use strict rounding to the nearest integer not larger in magnitude than the operand and has the similar semantic with llvm.trunc in llvm ir or trunc in libm.

Actually more fundamentally:

• arith.truncf is really a "truncate" in the sense of "reducing bit size", for example "convert from f32 to f16 and round". But it does not produce an integer value.
• llvm.trunc is operating on the same bit-size (f32->f32, f64->f64) and is actually just rounding to the nearest integer. It is coming alongside floor and ceil.

[Mogball]

mlir/include/mlir/Dialect/Math/IR/MathOps.td

818–819

This seems like a repeat of the first sentence

820

mlir/lib/Conversion/MathToLibm/MathToLibm.cpp

170

Should this be PromoteOpToF32?

[jacquesguan]

Address comment and rebase main.

mlir/include/mlir/Dialect/Math/IR/MathOps.td

818–819

Done.

820

Done.

mlir/lib/Conversion/MathToLibm/MathToLibm.cpp

170

Done.

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

 def LLVM_LogOp : LLVM_UnaryIntrinsicOp<"log">;
 def LLVM_Prefetch : LLVM_ZeroResultIntrOp<"prefetch", [0]> {
   let arguments = (ins LLVM_Type:$addr, LLVM_Type:$rw, LLVM_Type:$hint,
                    LLVM_Type:$cache);
 }
 def LLVM_SinOp : LLVM_UnaryIntrinsicOp<"sin">;
 def LLVM_RoundEvenOp : LLVM_UnaryIntrinsicOp<"roundeven">;
 def LLVM_RoundOp : LLVM_UnaryIntrinsicOp<"round">;
+def LLVM_FTruncOp : LLVM_UnaryIntrinsicOp<"trunc">;
 def LLVM_SqrtOp : LLVM_UnaryIntrinsicOp<"sqrt">;
 def LLVM_PowOp : LLVM_BinarySameArgsIntrinsicOp<"pow">;
 def LLVM_PowIOp : LLVM_BinaryIntrinsicOp<"powi">;
 def LLVM_BitReverseOp : LLVM_UnaryIntrinsicOp<"bitreverse">;
 def LLVM_CountLeadingZerosOp : LLVM_CountZerosIntrinsicOp<"ctlz">;
 def LLVM_CountTrailingZerosOp : LLVM_CountZerosIntrinsicOp<"cttz">;
 def LLVM_CtPopOp : LLVM_UnaryIntrinsicOp<"ctpop">;
 def LLVM_MaxNumOp : LLVM_BinarySameArgsIntrinsicOp<"maxnum">;

mlir/include/mlir/Dialect/Math/IR/MathOps.td

let description = [{
     ```mlir
     // Scalar round operation.
     %a = math.round %b : f64
     ```
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// TruncOp
+//===----------------------------------------------------------------------===//
+
+def Math_TruncOp : Math_FloatUnaryOp<"trunc"> {
+  let summary = "trunc of the specified value";
+  let description = [{
+    Syntax:
+
+    ```
+    operation ::= ssa-id `=` `math.trunc` ssa-use `:` type
+    ```
+
+    The `trunc` operation returns the operand rounded to the nearest integer
+    value in floating-point format. It takes one operand of floating point type
+    (i.e., scalar, tensor or vector) and produces one result of the same type.
+    The operation rounds the argument to the nearest integer value in
+    floating-point format, rounding to the nearest integer not larger in
+    magnitude than the operand., regardless of the current rounding direction.
+
+    Example:
+
+    ```mlir
+    // Scalar trunc operation.
+    %a = math.trunc %b : f64
+    ```
+  }];
+  let hasFolder = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // FPowIOp
 //===----------------------------------------------------------------------===//
 
 def Math_FPowIOp : Math_Op<"fpowi",
     [SameOperandsAndResultShape, AllTypesMatch<["lhs", "result"]>]> {
   let summary = "floating point raised to the signed integer power";
   let description = [{
     Syntax:

mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp

 using LogOpLowering = VectorConvertToLLVMPattern<math::LogOp, LLVM::LogOp>;
 using PowFOpLowering = VectorConvertToLLVMPattern<math::PowFOp, LLVM::PowOp>;
 using RoundEvenOpLowering =
     VectorConvertToLLVMPattern<math::RoundEvenOp, LLVM::RoundEvenOp>;
 using RoundOpLowering =
     VectorConvertToLLVMPattern<math::RoundOp, LLVM::RoundOp>;
 using SinOpLowering = VectorConvertToLLVMPattern<math::SinOp, LLVM::SinOp>;
 using SqrtOpLowering = VectorConvertToLLVMPattern<math::SqrtOp, LLVM::SqrtOp>;
+using FTruncOpLowering =
+    VectorConvertToLLVMPattern<math::TruncOp, LLVM::FTruncOp>;
 
 // A `CtLz/CtTz/absi(a)` is converted into `CtLz/CtTz/absi(a, false)`.
 template <typename MathOp, typename LLVMOp>
 struct IntOpWithFlagLowering : public ConvertOpToLLVMPattern<MathOp> {
   using ConvertOpToLLVMPattern<MathOp>::ConvertOpToLLVMPattern;
   using Super = IntOpWithFlagLowering<MathOp, LLVMOp>;
 
   LogicalResult
patterns.add<
     Log1pOpLowering,
     Log2OpLowering,
     LogOpLowering,
     PowFOpLowering,
     RoundEvenOpLowering,
     RoundOpLowering,
     RsqrtOpLowering,
     SinOpLowering,
-    SqrtOpLowering
+    SqrtOpLowering,
+    FTruncOpLowering
   >(converter);
   // clang-format on
 }
 
 std::unique_ptr<Pass> mlir::createConvertMathToLLVMPass() {
   return std::make_unique<ConvertMathToLLVMPass>();
 }

mlir/lib/Conversion/MathToLibm/MathToLibm.cpp

rewriter.replaceOpWithNewOp<func::CallOp>(op, name, op.getType(),
                                             op->getOperands());
 
   return success();
 }
 
 void mlir::populateMathToLibmConversionPatterns(
     RewritePatternSet &patterns, PatternBenefit benefit,
     llvm::Optional<PatternBenefit> log1pBenefit) {
-  patterns
-      .add<VecOpToScalarOp<math::Atan2Op>, VecOpToScalarOp<math::ExpM1Op>,
+  patterns.add<VecOpToScalarOp<math::Atan2Op>, VecOpToScalarOp<math::ExpM1Op>,
                VecOpToScalarOp<math::TanhOp>, VecOpToScalarOp<math::CosOp>,
                VecOpToScalarOp<math::SinOp>, VecOpToScalarOp<math::ErfOp>,
-           VecOpToScalarOp<math::RoundEvenOp>, VecOpToScalarOp<math::RoundOp>,
-           VecOpToScalarOp<math::AtanOp>, VecOpToScalarOp<math::TanOp>>(
-          patterns.getContext(), benefit);
+               VecOpToScalarOp<math::RoundEvenOp>,
+               VecOpToScalarOp<math::RoundOp>, VecOpToScalarOp<math::AtanOp>,
+               VecOpToScalarOp<math::TanOp>, VecOpToScalarOp<math::TruncOp>>(
   patterns.add<PromoteOpToF32<math::Atan2Op>, PromoteOpToF32<math::ExpM1Op>,
                PromoteOpToF32<math::TanhOp>, PromoteOpToF32<math::CosOp>,
                PromoteOpToF32<math::SinOp>, PromoteOpToF32<math::ErfOp>,
                PromoteOpToF32<math::RoundEvenOp>, PromoteOpToF32<math::RoundOp>,
-               PromoteOpToF32<math::AtanOp>, PromoteOpToF32<math::TanOp>>(
-      patterns.getContext(), benefit);
+               PromoteOpToF32<math::AtanOp>, PromoteOpToF32<math::TanOp>,
+               VecOpToScalarOp<math::TruncOp>>(patterns.getContext(), benefit);
   patterns.add<ScalarOpToLibmCall<math::AtanOp>>(patterns.getContext(), "atanf",
                                                  "atan", benefit);
   patterns.add<ScalarOpToLibmCall<math::Atan2Op>>(patterns.getContext(),
                                                   "atan2f", "atan2", benefit);
   patterns.add<ScalarOpToLibmCall<math::ErfOp>>(patterns.getContext(), "erff",
                                                 "erf", benefit);
   patterns.add<ScalarOpToLibmCall<math::ExpM1Op>>(patterns.getContext(),
                                                   "expm1f", "expm1", benefit);
void mlir::populateMathToLibmConversionPatterns(
   patterns.add<ScalarOpToLibmCall<math::SinOp>>(patterns.getContext(), "sinf",
                                                 "sin", benefit);
   patterns.add<ScalarOpToLibmCall<math::Log1pOp>>(
       patterns.getContext(), "log1pf", "log1p", log1pBenefit.value_or(benefit));
   patterns.add<ScalarOpToLibmCall<math::FloorOp>>(patterns.getContext(),
                                                   "floorf", "floor", benefit);
   patterns.add<ScalarOpToLibmCall<math::CeilOp>>(patterns.getContext(), "ceilf",
                                                  "ceil", benefit);
+  patterns.add<ScalarOpToLibmCall<math::TruncOp>>(patterns.getContext(),
+                                                  "truncf", "trunc", benefit);
 }
 
 namespace {
 struct ConvertMathToLibmPass
     : public impl::ConvertMathToLibmBase<ConvertMathToLibmPass> {
   void runOnOperation() override;
 };
 } // namespace

mlir/lib/Dialect/Math/IR/MathOps.cpp

return constFoldUnaryOpConditional<FloatAttr>(
         case 32:
           return APFloat(tanhf(a.convertToFloat()));
         default:
           return {};
         }
       });
 }
 
+//===----------------------------------------------------------------------===//
+// TruncOp folder
+//===----------------------------------------------------------------------===//
+
+OpFoldResult math::TruncOp::fold(ArrayRef<Attribute> operands) {
+  return constFoldUnaryOpConditional<FloatAttr>(
+      operands, [](const APFloat &a) -> Optional<APFloat> {
+        switch (a.getSizeInBits(a.getSemantics())) {
+        case 64:
+          return APFloat(trunc(a.convertToDouble()));
+        case 32:
+          return APFloat(truncf(a.convertToFloat()));
+        default:
+          return {};
+        }
+      });
+}
+
 /// Materialize an integer or floating point constant.
 Operation *math::MathDialect::materializeConstant(OpBuilder &builder,
                                                   Attribute value, Type type,
                                                   Location loc) {
   return builder.create<arith::ConstantOp>(loc, value, type);
 }

mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir

 
 // CHECK-LABEL: func @roundeven(
 // CHECK-SAME: f32
 func.func @roundeven(%arg0 : f32) {
   // CHECK: "llvm.intr.roundeven"(%arg0) : (f32) -> f32
   %0 = math.roundeven %arg0 : f32
   func.return
 }
+
+// -----
+
+// CHECK-LABEL: func @trunc(
+// CHECK-SAME: f32
+func.func @trunc(%arg0 : f32) {
+  // CHECK: "llvm.intr.trunc"(%arg0) : (f32) -> f32
+  %0 = math.trunc %arg0 : f32
+  func.return
+}

mlir/test/Conversion/MathToLibm/convert-to-libm.mlir

 // CHECK-DAG: @tan(f64) -> f64 attributes {llvm.readnone}
 // CHECK-DAG: @tanf(f32) -> f32 attributes {llvm.readnone}
 // CHECK-DAG: @tanh(f64) -> f64 attributes {llvm.readnone}
 // CHECK-DAG: @tanhf(f32) -> f32 attributes {llvm.readnone}
 // CHECK-DAG: @round(f64) -> f64 attributes {llvm.readnone}
 // CHECK-DAG: @roundf(f32) -> f32 attributes {llvm.readnone}
 // CHECK-DAG: @roundeven(f64) -> f64 attributes {llvm.readnone}
 // CHECK-DAG: @roundevenf(f32) -> f32 attributes {llvm.readnone}
+// CHECK-DAG: @trunc(f64) -> f64 attributes {llvm.readnone}
+// CHECK-DAG: @truncf(f32) -> f32 attributes {llvm.readnone}
 // CHECK-DAG: @cos(f64) -> f64 attributes {llvm.readnone}
 // CHECK-DAG: @cosf(f32) -> f32 attributes {llvm.readnone}
 // CHECK-DAG: @sin(f64) -> f64 attributes {llvm.readnone}
 // CHECK-DAG: @sinf(f32) -> f32 attributes {llvm.readnone}
 // CHECK-DAG: @floor(f64) -> f64 attributes {llvm.readnone}
 // CHECK-DAG: @floorf(f32) -> f32 attributes {llvm.readnone}
 // CHECK-DAG: @ceil(f64) -> f64 attributes {llvm.readnone}
 // CHECK-DAG: @ceilf(f32) -> f32 attributes {llvm.readnone}
func.func @roundeven_caller(%float: f32, %double: f64) -> (f32, f64) {
   // CHECK-DAG: %[[FLOAT_RESULT:.*]] = call @roundevenf(%[[FLOAT]]) : (f32) -> f32
   %float_result = math.roundeven %float : f32
   // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = call @roundeven(%[[DOUBLE]]) : (f64) -> f64
   %double_result = math.roundeven %double : f64
   // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
   return %float_result, %double_result : f32, f64
 }
 
+// CHECK-LABEL: func @trunc_caller
+// CHECK-SAME: %[[FLOAT:.*]]: f32
+// CHECK-SAME: %[[DOUBLE:.*]]: f64
+func.func @trunc_caller(%float: f32, %double: f64) -> (f32, f64) {
+  // CHECK-DAG: %[[FLOAT_RESULT:.*]] = call @truncf(%[[FLOAT]]) : (f32) -> f32
+  %float_result = math.trunc %float : f32
+  // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = call @trunc(%[[DOUBLE]]) : (f64) -> f64
+  %double_result = math.trunc %double : f64
+  // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
+  return %float_result, %double_result : f32, f64
+}
 
 // CHECK-LABEL: func @cos_caller
 // CHECK-SAME: %[[FLOAT:.*]]: f32
 // CHECK-SAME: %[[DOUBLE:.*]]: f64
 func.func @cos_caller(%float: f32, %double: f64) -> (f32, f64)  {
   // CHECK-DAG: %[[FLOAT_RESULT:.*]] = call @cosf(%[[FLOAT]]) : (f32) -> f32
   %float_result = math.cos %float : f32
   // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = call @cos(%[[DOUBLE]]) : (f64) -> f64
func.func @roundeven_vec_caller(%float: vector<2xf32>, %double: vector<2xf64>) -> (vector<2xf32>, vector<2xf64>) {
   // CHECK:           %[[IN1_F64:.*]] = vector.extract %[[VAL_1]][1] : vector<2xf64>
   // CHECK:           %[[OUT1_F64:.*]] = call @roundeven(%[[IN1_F64]]) : (f64) -> f64
   // CHECK:           %[[VAL_17:.*]] = vector.insert %[[OUT1_F64]], %[[VAL_14]] [1] : f64 into vector<2xf64>
   %double_result = math.roundeven %double : vector<2xf64>
   // CHECK:           return %[[VAL_11]], %[[VAL_17]] : vector<2xf32>, vector<2xf64>
   return %float_result, %double_result : vector<2xf32>, vector<2xf64>
 }
 
+// CHECK-LABEL:   func @trunc_vec_caller(
+// CHECK-SAME:                           %[[VAL_0:.*]]: vector<2xf32>,
+// CHECK-SAME:                           %[[VAL_1:.*]]: vector<2xf64>) -> (vector<2xf32>, vector<2xf64>) {
+func.func @trunc_vec_caller(%float: vector<2xf32>, %double: vector<2xf64>) -> (vector<2xf32>, vector<2xf64>) {
+  // CHECK-DAG:       %[[CVF:.*]] = arith.constant dense<0.000000e+00> : vector<2xf32>
+  // CHECK-DAG:       %[[CVD:.*]] = arith.constant dense<0.000000e+00> : vector<2xf64>
+  // CHECK:           %[[IN0_F32:.*]] = vector.extract %[[VAL_0]][0] : vector<2xf32>
+  // CHECK:           %[[OUT0_F32:.*]] = call @truncf(%[[IN0_F32]]) : (f32) -> f32
+  // CHECK:           %[[VAL_8:.*]] = vector.insert %[[OUT0_F32]], %[[CVF]] [0] : f32 into vector<2xf32>
+  // CHECK:           %[[IN1_F32:.*]] = vector.extract %[[VAL_0]][1] : vector<2xf32>
+  // CHECK:           %[[OUT1_F32:.*]] = call @truncf(%[[IN1_F32]]) : (f32) -> f32
+  // CHECK:           %[[VAL_11:.*]] = vector.insert %[[OUT1_F32]], %[[VAL_8]] [1] : f32 into vector<2xf32>
+  %float_result = math.trunc %float : vector<2xf32>
+  // CHECK:           %[[IN0_F64:.*]] = vector.extract %[[VAL_1]][0] : vector<2xf64>
+  // CHECK:           %[[OUT0_F64:.*]] = call @trunc(%[[IN0_F64]]) : (f64) -> f64
+  // CHECK:           %[[VAL_14:.*]] = vector.insert %[[OUT0_F64]], %[[CVD]] [0] : f64 into vector<2xf64>
+  // CHECK:           %[[IN1_F64:.*]] = vector.extract %[[VAL_1]][1] : vector<2xf64>
+  // CHECK:           %[[OUT1_F64:.*]] = call @trunc(%[[IN1_F64]]) : (f64) -> f64
+  // CHECK:           %[[VAL_17:.*]] = vector.insert %[[OUT1_F64]], %[[VAL_14]] [1] : f64 into vector<2xf64>
+  %double_result = math.trunc %double : vector<2xf64>
+  // CHECK:           return %[[VAL_11]], %[[VAL_17]] : vector<2xf32>, vector<2xf64>
+  return %float_result, %double_result : vector<2xf32>, vector<2xf64>
+}
 
 // CHECK-LABEL: func @tan_caller
 // CHECK-SAME: %[[FLOAT:.*]]: f32
 // CHECK-SAME: %[[DOUBLE:.*]]: f64
 // CHECK-SAME: %[[HALF:.*]]: f16
 // CHECK-SAME: %[[BFLOAT:.*]]: bf16
 func.func @tan_caller(%float: f32, %double: f64, %half: f16, %bfloat: bf16) -> (f32, f64, f16, bf16)  {
   // CHECK: %[[FLOAT_RESULT:.*]] = call @tanf(%[[FLOAT]]) : (f32) -> f32

mlir/test/Dialect/Math/canonicalize.mlir

 // CHECK-NEXT:   return %[[cst]]
 func.func @atan2_fold_vec() -> (vector<4xf32>) {
   %v1 = arith.constant dense<[0.0, 0.0, 1.0, 1.0]> : vector<4xf32>
   %v2 = arith.constant dense<[1.0, 1.0, 2.0, 2.0]> : vector<4xf32>
   %0 = math.atan2 %v1, %v2 : vector<4xf32>
   return %0 : vector<4xf32>
 }
 
+// CHECK-LABEL: @trunc_fold
+// CHECK-NEXT: %[[cst:.+]] = arith.constant 1.000000e+00 : f32
+// CHECK-NEXT:   return %[[cst]]
+func.func @trunc_fold() -> f32 {
+  %c = arith.constant 1.1 : f32
+  %r = math.trunc %c : f32
+  return %r : f32
+}
+
+// CHECK-LABEL: @trunc_fold_vec
+// CHECK-NEXT: %[[cst:.+]] = arith.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, -1.000000e+00]> : vector<4xf32>
+// CHECK-NEXT:   return %[[cst]]
+func.func @trunc_fold_vec() -> (vector<4xf32>) {
+  %v = arith.constant dense<[0.5, -0.5, 1.5, -1.5]> : vector<4xf32>
+  %0 = math.trunc %v : vector<4xf32>
+  return %0 : vector<4xf32>
+}

mlir/test/Dialect/Math/ops.mlir

func.func @ipowi(%i: i32, %v: vector<4xi32>, %t: tensor<4x4x?xi32>) {
   // CHECK: %{{.*}} = math.ipowi %[[I]], %[[I]] : i32
   %0 = math.ipowi %i, %i : i32
   // CHECK: %{{.*}} = math.ipowi %[[V]], %[[V]] : vector<4xi32>
   %1 = math.ipowi %v, %v : vector<4xi32>
   // CHECK: %{{.*}} = math.ipowi %[[T]], %[[T]] : tensor<4x4x?xi32>
   %2 = math.ipowi %t, %t : tensor<4x4x?xi32>
   return
 }
+
+// CHECK-LABEL: func @trunc(
+// CHECK-SAME:             %[[F:.*]]: f32, %[[V:.*]]: vector<4xf32>, %[[T:.*]]: tensor<4x4x?xf32>)
+func.func @trunc(%f: f32, %v: vector<4xf32>, %t: tensor<4x4x?xf32>) {
+  // CHECK: %{{.*}} = math.trunc %[[F]] : f32
+  %0 = math.trunc %f : f32
+  // CHECK: %{{.*}} = math.trunc %[[V]] : vector<4xf32>
+  %1 = math.trunc %v : vector<4xf32>
+  // CHECK: %{{.*}} = math.trunc %[[T]] : tensor<4x4x?xf32>
+  %2 = math.trunc %t : tensor<4x4x?xf32>
+  return
+}