diff --git a/mlir/include/mlir/Dialect/Math/IR/MathOps.td b/mlir/include/mlir/Dialect/Math/IR/MathOps.td --- a/mlir/include/mlir/Dialect/Math/IR/MathOps.td +++ b/mlir/include/mlir/Dialect/Math/IR/MathOps.td @@ -724,6 +724,7 @@ %x = math.sqrt %y : tensor<4x?xf32> ``` }]; + let hasFolder = 1; } //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/Math/IR/MathOps.cpp b/mlir/lib/Dialect/Math/IR/MathOps.cpp --- a/mlir/lib/Dialect/Math/IR/MathOps.cpp +++ b/mlir/lib/Dialect/Math/IR/MathOps.cpp @@ -101,6 +101,31 @@ return {}; } +OpFoldResult math::SqrtOp::fold(ArrayRef operands) { + auto constOperand = operands.front(); + if (!constOperand) + return {}; + + auto attr = constOperand.dyn_cast(); + if (!attr) + return {}; + + auto ft = getType().cast(); + + 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, diff --git a/mlir/test/Dialect/Math/canonicalize.mlir b/mlir/test/Dialect/Math/canonicalize.mlir --- a/mlir/test/Dialect/Math/canonicalize.mlir +++ b/mlir/test/Dialect/Math/canonicalize.mlir @@ -82,3 +82,12 @@ %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 +}