[Mlir-commits] [mlir] 279a659 - [mlir][math] lower rsqrt to sqrt + fdiv (#91344)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon May 13 01:15:43 PDT 2024
Author: Corentin Ferry
Date: 2024-05-13T10:15:39+02:00
New Revision: 279a659e9772e48d95ad7d81f6deb00ee31e35e1
URL: https://github.com/llvm/llvm-project/commit/279a659e9772e48d95ad7d81f6deb00ee31e35e1
DIFF: https://github.com/llvm/llvm-project/commit/279a659e9772e48d95ad7d81f6deb00ee31e35e1.diff
LOG: [mlir][math] lower rsqrt to sqrt + fdiv (#91344)
This commit creates an expansion pattern to lower math.rsqrt(x) into
fdiv(1, sqrt(x)).
Added:
Modified:
mlir/include/mlir/Dialect/Math/Transforms/Passes.h
mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp
mlir/test/Dialect/Math/expand-math.mlir
mlir/test/lib/Dialect/Math/TestExpandMath.cpp
mlir/test/mlir-cpu-runner/test-expand-math-approx.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Math/Transforms/Passes.h b/mlir/include/mlir/Dialect/Math/Transforms/Passes.h
index 24e6d9a8d98e0..ba69772515647 100644
--- a/mlir/include/mlir/Dialect/Math/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Math/Transforms/Passes.h
@@ -42,6 +42,7 @@ void populateExpandPowFPattern(RewritePatternSet &patterns);
void populateExpandFPowIPattern(RewritePatternSet &patterns);
void populateExpandRoundFPattern(RewritePatternSet &patterns);
void populateExpandRoundEvenPattern(RewritePatternSet &patterns);
+void populateExpandRsqrtPattern(RewritePatternSet &patterns);
void populateMathAlgebraicSimplificationPatterns(RewritePatternSet &patterns);
struct MathPolynomialApproximationOptions {
diff --git a/mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp b/mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp
index 5ccf3b6d72a2c..80569d95137c3 100644
--- a/mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp
@@ -615,6 +615,23 @@ static LogicalResult convertRoundEvenOp(math::RoundEvenOp op,
return success();
}
+// Convert `math.rsqrt` into `arith.divf` + `math.sqrt`
+static LogicalResult convertRsqrtOp(math::RsqrtOp op,
+ PatternRewriter &rewriter) {
+
+ auto operand = op.getOperand();
+ auto operandTy = operand.getType();
+ auto eTy = getElementTypeOrSelf(operandTy);
+ if (!isa<FloatType>(eTy))
+ return failure();
+
+ Location loc = op->getLoc();
+ auto constOneFloat = createFloatConst(loc, operandTy, 1.0, rewriter);
+ auto sqrtOp = rewriter.create<math::SqrtOp>(loc, operand);
+ rewriter.replaceOpWithNewOp<arith::DivFOp>(op, constOneFloat, sqrtOp);
+ return success();
+}
+
void mlir::populateExpandCtlzPattern(RewritePatternSet &patterns) {
patterns.add(convertCtlzOp);
}
@@ -678,3 +695,7 @@ void mlir::populateExpandFloorFPattern(RewritePatternSet &patterns) {
void mlir::populateExpandRoundEvenPattern(RewritePatternSet &patterns) {
patterns.add(convertRoundEvenOp);
}
+
+void mlir::populateExpandRsqrtPattern(RewritePatternSet &patterns) {
+ patterns.add(convertRsqrtOp);
+}
diff --git a/mlir/test/Dialect/Math/expand-math.mlir b/mlir/test/Dialect/Math/expand-math.mlir
index 3d94b55126d09..016a7bbdeb569 100644
--- a/mlir/test/Dialect/Math/expand-math.mlir
+++ b/mlir/test/Dialect/Math/expand-math.mlir
@@ -658,3 +658,73 @@ func.func @math_fpowi_to_powf_scalar(%0 : f32, %1: i64) -> f32 {
// CHECK: %[[AND:.*]] = arith.andi %[[CMPF1]], %[[CMPF]] : i1
// CHECK: %[[SEL:.*]] = arith.select %[[AND]], %[[MUL1]], %[[EXP]] : f32
// CHECK: return %[[SEL]] : f32
+
+// -----
+
+// CHECK-LABEL: func.func @rsqrt
+// CHECK-SAME: (%[[ARG:.*]]: f16)
+// CHECK-SAME: -> f16
+// CHECK-DAG: %[[CST:.*]] = arith.constant 1.000000e+00 : f16
+// CHECK-DAG: %[[SQRT:.*]] = math.sqrt %[[ARG]] : f16
+// CHECK-DAG: %[[DIV:.*]] = arith.divf %[[CST]], %[[SQRT]] : f16
+// CHECK: return %[[DIV]] : f16
+func.func @rsqrt16(%float: f16) -> (f16) {
+ %float_result = math.rsqrt %float : f16
+ return %float_result : f16
+}
+
+// -----
+
+// CHECK-LABEL: func.func @rsqrt
+// CHECK-SAME: (%[[ARG:.*]]: f32)
+// CHECK-SAME: -> f32
+// CHECK-DAG: %[[CST:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK-DAG: %[[SQRT:.*]] = math.sqrt %[[ARG]] : f32
+// CHECK-DAG: %[[DIV:.*]] = arith.divf %[[CST]], %[[SQRT]] : f32
+// CHECK: return %[[DIV]] : f32
+func.func @rsqrt32(%float: f32) -> (f32) {
+ %float_result = math.rsqrt %float : f32
+ return %float_result : f32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @rsqrt
+// CHECK-SAME: (%[[ARG:.*]]: f64)
+// CHECK-SAME: -> f64
+// CHECK-DAG: %[[CST:.*]] = arith.constant 1.000000e+00 : f64
+// CHECK-DAG: %[[SQRT:.*]] = math.sqrt %[[ARG]] : f64
+// CHECK-DAG: %[[DIV:.*]] = arith.divf %[[CST]], %[[SQRT]] : f64
+// CHECK: return %[[DIV]] : f64
+func.func @rsqrt64(%float: f64) -> (f64) {
+ %float_result = math.rsqrt %float : f64
+ return %float_result : f64
+}
+
+// -----
+
+// CHECK-LABEL: func.func @rsqrt_vec
+// CHECK-SAME: (%[[ARG:.*]]: vector<5xf32>)
+// CHECK-SAME: -> vector<5xf32>
+// CHECK-DAG: %[[CST:.*]] = arith.constant dense<1.000000e+00> : vector<5xf32>
+// CHECK-DAG: %[[SQRT:.*]] = math.sqrt %[[ARG]] : vector<5xf32>
+// CHECK-DAG: %[[DIV:.*]] = arith.divf %[[CST]], %[[SQRT]] : vector<5xf32>
+// CHECK: return %[[DIV]] : vector<5xf32>
+func.func @rsqrt_vec(%float: vector<5xf32>) -> (vector<5xf32>) {
+ %float_result = math.rsqrt %float : vector<5xf32>
+ return %float_result : vector<5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @rsqrt_tns
+// CHECK-SAME: (%[[ARG:.*]]: tensor<5x8xf32>)
+// CHECK-SAME: -> tensor<5x8xf32>
+// CHECK-DAG: %[[CST:.*]] = arith.constant dense<1.000000e+00> : tensor<5x8xf32>
+// CHECK-DAG: %[[SQRT:.*]] = math.sqrt %[[ARG]] : tensor<5x8xf32>
+// CHECK-DAG: %[[DIV:.*]] = arith.divf %[[CST]], %[[SQRT]] : tensor<5x8xf32>
+// CHECK: return %[[DIV]] : tensor<5x8xf32>
+func.func @rsqrt_tns(%float: tensor<5x8xf32>) -> (tensor<5x8xf32>) {
+ %float_result = math.rsqrt %float : tensor<5x8xf32>
+ return %float_result : tensor<5x8xf32>
+}
diff --git a/mlir/test/lib/Dialect/Math/TestExpandMath.cpp b/mlir/test/lib/Dialect/Math/TestExpandMath.cpp
index da48ccb6e5e08..69af2a08b97bd 100644
--- a/mlir/test/lib/Dialect/Math/TestExpandMath.cpp
+++ b/mlir/test/lib/Dialect/Math/TestExpandMath.cpp
@@ -52,6 +52,7 @@ void TestExpandMathPass::runOnOperation() {
populateExpandFPowIPattern(patterns);
populateExpandRoundFPattern(patterns);
populateExpandRoundEvenPattern(patterns);
+ populateExpandRsqrtPattern(patterns);
(void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
}
diff --git a/mlir/test/mlir-cpu-runner/test-expand-math-approx.mlir b/mlir/test/mlir-cpu-runner/test-expand-math-approx.mlir
index 2b72acde6a3bb..9b929b3c864dc 100644
--- a/mlir/test/mlir-cpu-runner/test-expand-math-approx.mlir
+++ b/mlir/test/mlir-cpu-runner/test-expand-math-approx.mlir
@@ -833,6 +833,46 @@ func.func @atanh() {
return
}
+// -------------------------------------------------------------------------- //
+// Rsqrt.
+// -------------------------------------------------------------------------- //
+
+func.func @rsqrt_f32(%a : f32) {
+ %r = math.rsqrt %a : f32
+ vector.print %r : f32
+ return
+}
+
+func.func @rsqrt_3xf32(%a : vector<3xf32>) {
+ %r = math.rsqrt %a : vector<3xf32>
+ vector.print %r : vector<3xf32>
+ return
+}
+
+func.func @rsqrt() {
+ // CHECK: 1
+ %zero = arith.constant 1.0 : f32
+ call @rsqrt_f32(%zero) : (f32) -> ()
+
+ // CHECK: 0.707107
+ %cst1 = arith.constant 2.0 : f32
+ call @rsqrt_f32(%cst1) : (f32) -> ()
+
+ // CHECK: inf
+ %cst2 = arith.constant 0.0 : f32
+ call @rsqrt_f32(%cst2) : (f32) -> ()
+
+ // CHECK: -nan
+ %cst3 = arith.constant -1.0 : f32
+ call @rsqrt_f32(%cst3) : (f32) -> ()
+
+ // CHECK: 0.5, 1.41421, 0.57735
+ %vec_x = arith.constant dense<[4.0, 0.5, 3.0]> : vector<3xf32>
+ call @rsqrt_3xf32(%vec_x) : (vector<3xf32>) -> ()
+
+ return
+}
+
func.func @main() {
call @exp2f() : () -> ()
call @roundf() : () -> ()
@@ -844,5 +884,6 @@ func.func @main() {
call @asinh() : () -> ()
call @acosh() : () -> ()
call @atanh() : () -> ()
+ call @rsqrt() : () -> ()
return
}
More information about the Mlir-commits
mailing list