[mlir][Math] Add constant folder for sqrt.

[bondhugula]

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

109–111

These are all guaranteed by the verifier and by the op's design. We shouldn't be checking for invalid IR.

[jacquesguan]

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

109–111

Thanks, I create a new revision to remove these https://reviews.llvm.org/D121986.

[ftynse]

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

109–111

Nope, these are not guaranteed:

https://github.com/llvm/llvm-project/blob/fc8946fae7a4e9351b7775ec8407d248b8fd02b4/mlir/include/mlir/Dialect/Math/IR/MathOps.td#L41

https://github.com/llvm/llvm-project/blob/fc8946fae7a4e9351b7775ec8407d248b8fd02b4/mlir/include/mlir/IR/OpBase.td#L940

https://github.com/llvm/llvm-project/blob/fc8946fae7a4e9351b7775ec8407d248b8fd02b4/mlir/include/mlir/IR/OpBase.td#L911-L913

We just don't fold tensor or vector constants that aggressively, but it could be DenseElementsAttr.

[bondhugula]

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

109–111

Sorry, I was mistaken here. This is fine as is.

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

let description = [{
     // Scalar square root value.
     %a = math.sqrt %b : f64
     // SIMD vector element-wise square root value.
     %f = math.sqrt %g : vector<4xf32>
     // Tensor element-wise square root value.
     %x = math.sqrt %y : tensor<4x?xf32>
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TanhOp
 //===----------------------------------------------------------------------===//
 
 def Math_TanhOp : Math_FloatUnaryOp<"tanh"> {
   let summary = "hyperbolic tangent of the specified value";

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

OpFoldResult math::PowFOp::fold(ArrayRef<Attribute> operands) {
 
   if (ft.getWidth() == 32)
     return FloatAttr::get(
         getType(), powf(vals[0].convertToFloat(), vals[1].convertToFloat()));
 
   return {};
 }
 
+OpFoldResult math::SqrtOp::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 (apf.isNegative())
+    return {};
+
+  if (ft.getWidth() == 64)
+    return FloatAttr::get(getType(), sqrt(apf.convertToDouble()));
+
+  if (ft.getWidth() == 32)
+    return FloatAttr::get(getType(), sqrtf(apf.convertToFloat()));
+
+  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/Dialect/Math/canonicalize.mlir

 // CHECK-LABEL: @powf_fold
 // CHECK: %[[cst:.+]] = arith.constant 4.000000e+00 : f32
 // CHECK: return %[[cst]]
 func @powf_fold() -> f32 {
   %c = arith.constant 2.0 : f32
   %r = math.powf %c, %c : f32
   return %r : f32
 }
+
+// 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
+}