[MLIR][Math] Add support for f16 in the expansion of math.roundeven

Add support for f16 in the expansion of math.roundeven.
Associated GitHub issue: iree-org/iree#13522
This version addresses the build issues on Windows reported on
https://reviews.llvm.org/D157204

Test plan: ninja check-mlir check-all

Differential revision: https://reviews.llvm.org/D158234

Author: alexander-shaposhnikov

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

@@ -469,8 +469,6 @@ extern "C" MLIR_CRUNNERUTILS_EXPORT void printOpen();
 extern "C" MLIR_CRUNNERUTILS_EXPORT void printClose();
 extern "C" MLIR_CRUNNERUTILS_EXPORT void printComma();
 extern "C" MLIR_CRUNNERUTILS_EXPORT void printNewline();
-extern "C" MLIR_CRUNNERUTILS_EXPORT void printF16(uint16_t bits);  // bits!
-extern "C" MLIR_CRUNNERUTILS_EXPORT void printBF16(uint16_t bits); // bits!
 
 //===----------------------------------------------------------------------===//
 // Small runtime support library for timing execution and printing GFLOPS

mlir/include/mlir/ExecutionEngine/Float16bits.h

@@ -48,5 +48,8 @@ MLIR_FLOAT16_EXPORT std::ostream &operator<<(std::ostream &os, const f16 &f);
 // Outputs a bfloat value.
 MLIR_FLOAT16_EXPORT std::ostream &operator<<(std::ostream &os, const bf16 &d);
 
+extern "C" MLIR_FLOAT16_EXPORT void printF16(uint16_t bits);
+extern "C" MLIR_FLOAT16_EXPORT void printBF16(uint16_t bits);
+
 #undef MLIR_FLOAT16_EXPORT
 #endif // MLIR_EXECUTIONENGINE_FLOAT16BITS_H_

mlir/include/mlir/IR/BuiltinTypes.h

@@ -67,6 +67,7 @@ class FloatType : public Type {
   unsigned getWidth();
 
   /// Return the width of the mantissa of this type.
+  /// The width includes the integer bit.
   unsigned getFPMantissaWidth();
 
   /// Get or create a new FloatType with bitwidth scaled by `scale`.

mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp

@@ -305,31 +305,39 @@ static LogicalResult convertRoundEvenOp(math::RoundEvenOp op,
   Type operandETy = getElementTypeOrSelf(operandTy);
   Type resultETy = getElementTypeOrSelf(resultTy);
 
-  if (!operandETy.isF32() || !resultETy.isF32()) {
-    return rewriter.notifyMatchFailure(op, "not a roundeven of f32.");
+  if (!isa<FloatType>(operandETy) || !isa<FloatType>(resultETy)) {
+    return rewriter.notifyMatchFailure(op, "not a roundeven of f16 or f32.");
   }
 
-  Type i32Ty = b.getI32Type();
-  Type f32Ty = b.getF32Type();
-  if (auto shapedTy = dyn_cast<ShapedType>(operandTy)) {
-    i32Ty = shapedTy.clone(i32Ty);
-    f32Ty = shapedTy.clone(f32Ty);
+  Type fTy = operandTy;
+  Type iTy = rewriter.getIntegerType(operandETy.getIntOrFloatBitWidth());
+  if (auto shapedTy = dyn_cast<ShapedType>(fTy)) {
+    iTy = shapedTy.clone(iTy);
   }
 
-  Value c1Float = createFloatConst(loc, f32Ty, 1.0, b);
-  Value c0 = createIntConst(loc, i32Ty, 0, b);
-  Value c1 = createIntConst(loc, i32Ty, 1, b);
-  Value cNeg1 = createIntConst(loc, i32Ty, -1, b);
-  Value c23 = createIntConst(loc, i32Ty, 23, b);
-  Value c31 = createIntConst(loc, i32Ty, 31, b);
-  Value c127 = createIntConst(loc, i32Ty, 127, b);
-  Value c2To22 = createIntConst(loc, i32Ty, 1 << 22, b);
-  Value c23Mask = createIntConst(loc, i32Ty, (1 << 23) - 1, b);
-  Value expMask = createIntConst(loc, i32Ty, (1 << 8) - 1, b);
-
-  Value operandBitcast = b.create<arith::BitcastOp>(i32Ty, operand);
+  unsigned bitWidth = operandETy.getIntOrFloatBitWidth();
+  // The width returned by getFPMantissaWidth includes the integer bit.
+  unsigned mantissaWidth =
+      llvm::cast<FloatType>(operandETy).getFPMantissaWidth() - 1;
+  unsigned exponentWidth = bitWidth - mantissaWidth - 1;
+
+  // The names of the variables correspond to f32.
+  // f32: 1 bit sign | 8 bits exponent | 23 bits mantissa.
+  // f16: 1 bit sign | 5 bits exponent | 10 bits mantissa.
+  Value c1Float = createFloatConst(loc, fTy, 1.0, b);
+  Value c0 = createIntConst(loc, iTy, 0, b);
+  Value c1 = createIntConst(loc, iTy, 1, b);
+  Value cNeg1 = createIntConst(loc, iTy, -1, b);
+  Value c23 = createIntConst(loc, iTy, mantissaWidth, b);
+  Value c31 = createIntConst(loc, iTy, bitWidth - 1, b);
+  Value c127 = createIntConst(loc, iTy, (1 << (exponentWidth - 1)) - 1, b);
+  Value c2To22 = createIntConst(loc, iTy, 1 << (mantissaWidth - 1), b);
+  Value c23Mask = createIntConst(loc, iTy, (1 << mantissaWidth) - 1, b);
+  Value expMask = createIntConst(loc, iTy, (1 << exponentWidth) - 1, b);
+
+  Value operandBitcast = b.create<arith::BitcastOp>(iTy, operand);
   Value round = b.create<math::RoundOp>(operand);
-  Value roundBitcast = b.create<arith::BitcastOp>(i32Ty, round);
+  Value roundBitcast = b.create<arith::BitcastOp>(iTy, round);
 
   // Get biased exponents for operand and round(operand)
   Value operandExp = b.create<arith::AndIOp>(
@@ -340,7 +348,7 @@ static LogicalResult convertRoundEvenOp(math::RoundEvenOp op,
   Value roundBiasedExp = b.create<arith::SubIOp>(roundExp, c127);
 
   auto safeShiftRight = [&](Value x, Value shift) -> Value {
-    // Clamp shift to valid range [0, 31] to avoid undefined behavior
+    // Clamp shift to valid range [0, bitwidth - 1] to avoid undefined behavior
     Value clampedShift = b.create<arith::MaxSIOp>(shift, c0);
     clampedShift = b.create<arith::MinSIOp>(clampedShift, c31);
     return b.create<arith::ShRUIOp>(x, clampedShift);

mlir/test/CMakeLists.txt

@@ -119,6 +119,7 @@ if(LLVM_ENABLE_PIC AND TARGET ${LLVM_NATIVE_ARCH})
     mlir-capi-execution-engine-test
     mlir_c_runner_utils
     mlir_runner_utils
+    mlir_float16_utils
   )
 endif()
 

mlir/test/Dialect/Math/expand-math.mlir

@@ -232,8 +232,8 @@ func.func @powf_func(%a: f64, %b: f64) ->f64 {
 
 // -----
 
-// CHECK-LABEL:   func.func @roundeven
-func.func @roundeven(%arg: f32) -> f32 {
+// CHECK-LABEL:   func.func @roundeven32
+func.func @roundeven32(%arg: f32) -> f32 {
   %res = math.roundeven %arg : f32
   return %res : f32
 }
@@ -331,3 +331,90 @@ func.func @roundeven(%arg: f32) -> f32 {
 // CHECK:     %[[COPYSIGN:.*]] = math.copysign %[[RESULT]], %[[VAL_0]] : f32
 
 // CHECK: return %[[COPYSIGN]] : f32
+
+// -----
+
+// CHECK-LABEL:   func.func @roundeven16
+func.func @roundeven16(%arg: f16) -> f16 {
+  %res = math.roundeven %arg : f16
+  return %res : f16
+}
+
+// CHECK-SAME:                   %[[VAL_0:.*]]: f16) -> f16 {
+// CHECK-DAG: %[[C_0:.*]] = arith.constant 0 : i16
+// CHECK-DAG: %[[C_1:.*]] = arith.constant 1 : i16
+// CHECK-DAG: %[[C_NEG_1:.*]] = arith.constant -1 : i16
+// CHECK-DAG: %[[C_1_FLOAT:.*]] = arith.constant 1.000000e+00 : f16
+// CHECK-DAG: %[[C_10:.*]] = arith.constant 10 : i16
+// CHECK-DAG: %[[C_15:.*]] = arith.constant 15 : i16
+// CHECK-DAG: %[[C_512:.*]] = arith.constant 512 : i16
+// CHECK-DAG: %[[C_1023:.*]] = arith.constant 1023 : i16
+// CHECK-DAG: %[[EXP_MASK:.*]] = arith.constant 31 : i16
+
+// CHECK:     %[[OPERAND_BITCAST:.*]] = arith.bitcast %[[VAL_0]] : f16 to i16
+// CHECK:     %[[ROUND:.*]] = math.round %[[VAL_0]] : f16
+// CHECK:     %[[ROUND_BITCAST:.*]] = arith.bitcast %[[ROUND]] : f16 to i16
+
+// Get biased exponents of `round` and `operand`
+// CHECK:     %[[SHIFTED_OPERAND_BITCAST:.*]] = arith.shrui %[[OPERAND_BITCAST]], %[[C_10]] : i16
+// CHECK:     %[[OPERAND_EXP:.*]] = arith.andi %[[SHIFTED_OPERAND_BITCAST]], %[[EXP_MASK]] : i16
+// CHECK:     %[[OPERAND_BIASED_EXP:.*]] = arith.subi %[[OPERAND_EXP]], %[[C_15]] : i16
+// CHECK:     %[[SHIFTED_ROUND_BITCAST:.*]] = arith.shrui %[[ROUND_BITCAST]], %[[C_10]] : i16
+// CHECK:     %[[ROUND_EXP:.*]] = arith.andi %[[SHIFTED_ROUND_BITCAST]], %[[EXP_MASK]] : i16
+// CHECK:     %[[ROUND_BIASED_EXP:.*]] = arith.subi %[[ROUND_EXP]], %[[C_15]] : i16
+
+// Determine if `ROUND_BITCAST` is an even whole number or a special value
+// +-inf, +-nan.
+//   Mask mantissa of `ROUND_BITCAST` with a mask shifted to the right by
+//   `ROUND_BIASED_EXP - 1`
+//   CHECK-DAG: %[[ROUND_BIASED_EXP_MINUS_1:.*]] = arith.subi %[[ROUND_BIASED_EXP]], %[[C_1]] : i16
+//   CHECK-DAG: %[[CLAMPED_SHIFT_0:.*]] = arith.maxsi %[[ROUND_BIASED_EXP_MINUS_1]], %[[C_0]] : i16
+//   CHECK-DAG: %[[CLAMPED_SHIFT_1:.*]] = arith.minsi %[[CLAMPED_SHIFT_0]], %[[C_15]] : i16
+//   CHECK-DAG: %[[SHIFTED_MANTISSA_MASK_0:.*]] = arith.shrui %[[C_1023]], %[[CLAMPED_SHIFT_1]] : i16
+//   CHECK-DAG: %[[ROUND_MASKED_MANTISSA:.*]] = arith.andi %[[ROUND_BITCAST]], %[[SHIFTED_MANTISSA_MASK_0]] : i16
+
+//   `ROUND_BITCAST` is not even whole number or special value if masked
+//   mantissa is != 0 or `ROUND_BIASED_EXP == 0`
+//   CHECK-DAG: %[[ROUND_IS_NOT_EVEN_OR_SPECIAL_0:.*]] = arith.cmpi ne, %[[ROUND_MASKED_MANTISSA]], %[[C_0]] : i16
+//   CHECK-DAG: %[[ROUND_BIASED_EXP_EQ_0:.*]] = arith.cmpi eq, %[[ROUND_BIASED_EXP]], %[[C_0]] : i16
+//   CHECK-DAG: %[[ROUND_IS_NOT_EVEN_OR_SPECIAL_1:.*]] = arith.ori %[[ROUND_IS_NOT_EVEN_OR_SPECIAL_0]], %[[ROUND_BIASED_EXP_EQ_0]] : i1
+
+// Determine if operand is halfway between two integer values
+// CHECK:     %[[OPERAND_BIASED_EXP_EQ_NEG_1:.*]] = arith.cmpi eq, %[[OPERAND_BIASED_EXP]], %[[C_NEG_1]] : i16
+// CHECK:     %[[CLAMPED_SHIFT_2:.*]] = arith.maxsi %[[OPERAND_BIASED_EXP]], %[[C_0]] : i16
+// CHECK:     %[[CLAMPED_SHIFT_3:.*]] = arith.minsi %[[CLAMPED_SHIFT_2]], %[[C_15]] : i16
+// CHECK:     %[[SHIFTED_2_TO_9:.*]] = arith.shrui %[[C_512]], %[[CLAMPED_SHIFT_3]] : i16
+
+//   A value with `0 <= BIASED_EXP < 10` is halfway between two consecutive
+//   integers if the bit at index `BIASED_EXP` starting from the left in the
+//   mantissa is 1 and all the bits to the right are zero. For the case where
+//   `BIASED_EXP == -1, the expected mantissa is all zeros.
+//   CHECK:     %[[EXPECTED_OPERAND_MASKED_MANTISSA:.*]] = arith.select %[[OPERAND_BIASED_EXP_EQ_NEG_1]], %[[C_0]], %[[SHIFTED_2_TO_9]] : i16
+
+//   Mask mantissa of `OPERAND_BITCAST` with a mask shifted to the right by
+//   `OPERAND_BIASED_EXP`
+//   CHECK:     %[[CLAMPED_SHIFT_4:.*]] = arith.maxsi %[[OPERAND_BIASED_EXP]], %[[C_0]] : i16
+//   CHECK:     %[[CLAMPED_SHIFT_5:.*]] = arith.minsi %[[CLAMPED_SHIFT_4]], %[[C_15]] : i16
+//   CHECK:     %[[SHIFTED_MANTISSA_MASK_1:.*]] = arith.shrui %[[C_1023]], %[[CLAMPED_SHIFT_5]] : i16
+//   CHECK:     %[[OPERAND_MASKED_MANTISSA:.*]] = arith.andi %[[OPERAND_BITCAST]], %[[SHIFTED_MANTISSA_MASK_1]] : i16
+
+//   The operand is halfway between two integers if the masked mantissa is equal
+//   to the expected mantissa and the biased exponent is in the range
+//   [-1,  23).
+//   CHECK-DAG: %[[OPERAND_BIASED_EXP_GE_NEG_1:.*]] = arith.cmpi sge, %[[OPERAND_BIASED_EXP]], %[[C_NEG_1]] : i16
+//   CHECK-DAG: %[[OPERAND_BIASED_EXP_LT_10:.*]] = arith.cmpi slt, %[[OPERAND_BIASED_EXP]], %[[C_10]] : i16
+//   CHECK-DAG: %[[OPERAND_IS_HALFWAY_0:.*]] = arith.cmpi eq, %[[OPERAND_MASKED_MANTISSA]], %[[EXPECTED_OPERAND_MASKED_MANTISSA]] : i16
+//   CHECK-DAG: %[[OPERAND_IS_HALFWAY_1:.*]] = arith.andi %[[OPERAND_IS_HALFWAY_0]], %[[OPERAND_BIASED_EXP_LT_10]] : i1
+//   CHECK-DAG: %[[OPERAND_IS_HALFWAY_2:.*]] = arith.andi %[[OPERAND_IS_HALFWAY_1]], %[[OPERAND_BIASED_EXP_GE_NEG_1]] : i1
+
+// Adjust rounded operand with `round(operand) - sign(operand)` to correct the
+// case where `round` rounded in the oppositve direction of `roundeven`.
+// CHECK:     %[[SIGN:.*]] = math.copysign %[[C_1_FLOAT]], %[[VAL_0]] : f16
+// CHECK:     %[[ROUND_SHIFTED:.*]] = arith.subf %[[ROUND]], %[[SIGN]] : f16
+// CHECK:     %[[NEEDS_SHIFT:.*]] = arith.andi %[[ROUND_IS_NOT_EVEN_OR_SPECIAL_1]], %[[OPERAND_IS_HALFWAY_2]] : i1
+// CHECK:     %[[RESULT:.*]] = arith.select %[[NEEDS_SHIFT]], %[[ROUND_SHIFTED]], %[[ROUND]] : f16
+
+// The `x - sign` adjustment does not preserve the sign when we are adjusting the value -1 to -0.
+// CHECK:     %[[COPYSIGN:.*]] = math.copysign %[[RESULT]], %[[VAL_0]] : f16
+
+// CHECK: return %[[COPYSIGN]] : f16

mlir/test/lit.cfg.py

@@ -98,6 +98,7 @@ def add_runtime(name):
     add_runtime("mlir_runner_utils"),
     add_runtime("mlir_c_runner_utils"),
     add_runtime("mlir_async_runtime"),
+    add_runtime("mlir_float16_utils"),
     "mlir-linalg-ods-yaml-gen",
     "mlir-reduce",
     "mlir-pdll",