[mlir][Math] Replace some constant folder functions with common folder functions.

Differential Revision: https://reviews.llvm.org/D123485

Author: jacquesguan

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

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/CommonFolders.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/IR/Builders.h"
 
@@ -25,119 +26,65 @@ using namespace mlir::math;
 //===----------------------------------------------------------------------===//
 
 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 {};
+  return constFoldUnaryOp<FloatAttr>(operands, [](const APFloat &a) {
+    APFloat result(a);
+    return abs(result);
+  });
 }
 
 //===----------------------------------------------------------------------===//
 // 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);
+  return constFoldUnaryOp<FloatAttr>(operands, [](const APFloat &a) {
+    APFloat result(a);
+    result.roundToIntegral(llvm::RoundingMode::TowardPositive);
+    return result;
+  });
 }
 
 //===----------------------------------------------------------------------===//
 // 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]);
+  return constFoldBinaryOp<FloatAttr>(operands,
+                                      [](const APFloat &a, const APFloat &b) {
+                                        APFloat result(a);
+                                        result.copySign(b);
+                                        return result;
+                                      });
 }
 
 //===----------------------------------------------------------------------===//
 // 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());
+  return constFoldUnaryOp<IntegerAttr>(operands, [](const APInt &a) {
+    return APInt(a.getBitWidth(), a.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());
+  return constFoldUnaryOp<IntegerAttr>(operands, [](const APInt &a) {
+    return APInt(a.getBitWidth(), a.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());
+  return constFoldUnaryOp<IntegerAttr>(operands, [](const APInt &a) {
+    return APInt(a.getBitWidth(), a.countPopulation());
+  });
 }
 
 //===----------------------------------------------------------------------===//