[mlir][Math] Add more constant folder for Math ops.

[jacquesguan]

Summary

This revision add constant folder for abs, copysign, ctlz, cttz and
ctpop.

[ftynse]

mlir/test/Dialect/Math/canonicalize.mlir

145

Nit: a hex literal would be better here to check the correctness of the test.

[jacquesguan]

Change constant to hex format.

[jacquesguan]

mlir/test/Dialect/Math/canonicalize.mlir

145

Thanks, done.

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

let description = [{
 
     // SIMD vector element-wise absolute value.
     %f = math.abs %g : vector<4xf32>
 
     // Tensor element-wise absolute value.
     %x = math.abs %y : tensor<4x?xf8>
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // AtanOp
 //===----------------------------------------------------------------------===//
 
 def Math_AtanOp : Math_FloatUnaryOp<"atan">{
   let summary = "arcus tangent of the given value";
let description = [{
 
     // SIMD vector element-wise copysign value.
     %f = math.copysign %g, %h : vector<4xf32>
 
     // Tensor element-wise copysign value.
     %x = math.copysign %y, %z : tensor<4x?xf8>
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // CosOp
 //===----------------------------------------------------------------------===//
 
 def Math_CosOp : Math_FloatUnaryOp<"cos"> {
   let summary = "cosine of the specified value";
let description = [{
 
     // SIMD vector element-wise ctlz function value.
     %f = math.ctlz %g : vector<4xi16>
 
     // Tensor element-wise ctlz function value.
     %x = math.ctlz %y : tensor<4x?xi8>
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // CountTrailingZerosOp
 //===----------------------------------------------------------------------===//
 
 def Math_CountTrailingZerosOp : Math_IntegerUnaryOp<"cttz"> {
   let summary = "counts the trailing zeros an integer value";
   let description = [{
     The `cttz` operation computes the number of trailing zeros of an integer value.
 
     Example:
 
     ```mlir
     // Scalar cttz function value.
     %a = math.cttz %b : i32
 
     // SIMD vector element-wise cttz function value.
     %f = math.cttz %g : vector<4xi16>
 
     // Tensor element-wise cttz function value.
     %x = math.cttz %y : tensor<4x?xi8>
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // CtPopOp
 //===----------------------------------------------------------------------===//
 
 def Math_CtPopOp : Math_IntegerUnaryOp<"ctpop"> {
   let summary = "counts the number of set bits of an integer value";
   let description = [{
     The `ctpop` operation computes the number of set bits of an integer value.
 
     Example:
 
     ```mlir
     // Scalar ctpop function value.
     %a = math.ctpop %b : i32
 
     // SIMD vector element-wise ctpop function value.
     %f = math.ctpop %g : vector<4xi16>
 
     // Tensor element-wise ctpop function value.
     %x = math.ctpop %y : tensor<4x?xi8>
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // ErfOp
 //===----------------------------------------------------------------------===//
 
 def Math_ErfOp : Math_FloatUnaryOp<"erf"> {
   let summary = "error function of the specified value";

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

 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Math/IR/MathOps.cpp.inc"
 
 //===----------------------------------------------------------------------===//
+// AbsOp folder
+//===----------------------------------------------------------------------===//
+
+OpFoldResult math::AbsOp::fold(ArrayRef<Attribute> operands) {
+  auto constOperand = operands.front();
+  if (!constOperand)
+    return {};
+
+  auto attr = constOperand.dyn_cast<FloatAttr>();
+  if (!attr)
+    return {};
+
+  auto ft = getType().cast<FloatType>();
+
+  APFloat apf = attr.getValue();
+
+  if (ft.getWidth() == 64)
+    return FloatAttr::get(getType(), fabs(apf.convertToDouble()));
+
+  if (ft.getWidth() == 32)
+    return FloatAttr::get(getType(), fabsf(apf.convertToFloat()));
+
+  return {};
+}
+
+//===----------------------------------------------------------------------===//
 // CeilOp folder
 //===----------------------------------------------------------------------===//
 
 OpFoldResult math::CeilOp::fold(ArrayRef<Attribute> operands) {
   auto constOperand = operands.front();
   if (!constOperand)
     return {};
 
   auto attr = constOperand.dyn_cast<FloatAttr>();
   if (!attr)
     return {};
 
   APFloat sourceVal = attr.getValue();
   sourceVal.roundToIntegral(llvm::RoundingMode::TowardPositive);
 
   return FloatAttr::get(getType(), sourceVal);
 }
 
 //===----------------------------------------------------------------------===//
+// CopySignOp folder
+//===----------------------------------------------------------------------===//
+
+OpFoldResult math::CopySignOp::fold(ArrayRef<Attribute> operands) {
+  auto ft = getType().dyn_cast<FloatType>();
+  if (!ft)
+    return {};
+
+  APFloat vals[2]{APFloat(ft.getFloatSemantics()),
+                  APFloat(ft.getFloatSemantics())};
+  for (int i = 0; i < 2; ++i) {
+    if (!operands[i])
+      return {};
+
+    auto attr = operands[i].dyn_cast<FloatAttr>();
+    if (!attr)
+      return {};
+
+    vals[i] = attr.getValue();
+  }
+
+  vals[0].copySign(vals[1]);
+
+  return FloatAttr::get(getType(), vals[0]);
+}
+
+//===----------------------------------------------------------------------===//
+// CountLeadingZerosOp folder
+//===----------------------------------------------------------------------===//
+
+OpFoldResult math::CountLeadingZerosOp::fold(ArrayRef<Attribute> operands) {
+  auto constOperand = operands.front();
+  if (!constOperand)
+    return {};
+
+  auto attr = constOperand.dyn_cast<IntegerAttr>();
+  if (!attr)
+    return {};
+
+  return IntegerAttr::get(getType(), attr.getValue().countLeadingZeros());
+}
+
+//===----------------------------------------------------------------------===//
+// CountTrailingZerosOp folder
+//===----------------------------------------------------------------------===//
+
+OpFoldResult math::CountTrailingZerosOp::fold(ArrayRef<Attribute> operands) {
+  auto constOperand = operands.front();
+  if (!constOperand)
+    return {};
+
+  auto attr = constOperand.dyn_cast<IntegerAttr>();
+  if (!attr)
+    return {};
+
+  return IntegerAttr::get(getType(), attr.getValue().countTrailingZeros());
+}
+
+//===----------------------------------------------------------------------===//
+// CtPopOp folder
+//===----------------------------------------------------------------------===//
+
+OpFoldResult math::CtPopOp::fold(ArrayRef<Attribute> operands) {
+  auto constOperand = operands.front();
+  if (!constOperand)
+    return {};
+
+  auto attr = constOperand.dyn_cast<IntegerAttr>();
+  if (!attr)
+    return {};
+
+  return IntegerAttr::get(getType(), attr.getValue().countPopulation());
+}
+
+//===----------------------------------------------------------------------===//
 // Log2Op folder
 //===----------------------------------------------------------------------===//
 
 OpFoldResult math::Log2Op::fold(ArrayRef<Attribute> operands) {
   auto constOperand = operands.front();
   if (!constOperand)
     return {};
 

mlir/test/Dialect/Math/canonicalize.mlir

 // CHECK-LABEL: @sqrt_fold
 // CHECK: %[[cst:.+]] = arith.constant 2.000000e+00 : f32
 // CHECK: return %[[cst]]
 func @sqrt_fold() -> f32 {
   %c = arith.constant 4.0 : f32
   %r = math.sqrt %c : f32
   return %r : f32
 }
+
+// CHECK-LABEL: @abs_fold
+// CHECK: %[[cst:.+]] = arith.constant 4.000000e+00 : f32
+// CHECK: return %[[cst]]
+func @abs_fold() -> f32 {
+  %c = arith.constant -4.0 : f32
+  %r = math.abs %c : f32
+  return %r : f32
+}
+
+// CHECK-LABEL: @copysign_fold
+// CHECK: %[[cst:.+]] = arith.constant -4.000000e+00 : f32
+// CHECK: return %[[cst]]
+func @copysign_fold() -> f32 {
+  %c1 = arith.constant 4.0 : f32
+  %c2 = arith.constant -9.0 : f32
+  %r = math.copysign %c1, %c2 : f32
+  return %r : f32
+}
+
+// CHECK-LABEL: @ctlz_fold1
+// CHECK: %[[cst:.+]] = arith.constant 31 : i32
+// CHECK: return %[[cst]]
+func @ctlz_fold1() -> i32 {
+  %c = arith.constant 1 : i32
+  %r = math.ctlz %c : i32
+  return %r : i32
+}
+
+// CHECK-LABEL: @ctlz_fold2
+// CHECK: %[[cst:.+]] = arith.constant 7 : i8
+// CHECK: return %[[cst]]
+func @ctlz_fold2() -> i8 {
+  %c = arith.constant 1 : i8
+  %r = math.ctlz %c : i8
+  return %r : i8
+}
+
+// CHECK-LABEL: @cttz_fold
+// CHECK: %[[cst:.+]] = arith.constant 8 : i32
+// CHECK: return %[[cst]]
+func @cttz_fold() -> i32 {
+  %c = arith.constant 256 : i32
+  %r = math.cttz %c : i32
+  return %r : i32
+}
+
+// CHECK-LABEL: @ctpop_fold
+// CHECK: %[[cst:.+]] = arith.constant 16 : i32
+// CHECK: return %[[cst]]
+func @ctpop_fold() -> i32 {
+  %c = arith.constant 0xFF0000FF : i32
+  %r = math.ctpop %c : i32
+  return %r : i32
+}