[Mlir-commits] [mlir] 9093d6e - [mlir][x86] Move AVX rsqrt approximation from Math to X86 transforms (#236807)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 29 08:58:06 PDT 2026
Author: Markus Böck
Date: 2026-10-29T16:58:01+01:00
New Revision: 9093d6e8dbdc453680c1b4b8adea8fce78005b4a
URL: https://github.com/llvm/llvm-project/commit/9093d6e8dbdc453680c1b4b8adea8fce78005b4a
DIFF: https://github.com/llvm/llvm-project/commit/9093d6e8dbdc453680c1b4b8adea8fce78005b4a.diff
LOG: [mlir][x86] Move AVX rsqrt approximation from Math to X86 transforms (#236807)
The `math.rsqrt` polynomial approximation emitted `x86.avx.rsqrt`, which
made `MLIRMathTransforms` depend on `MLIRX86Dialect` even though the
pattern was only enabled behind an AVX2 opt-in.
Move the pattern into `MLIRX86Transforms` and expose it as
`x86::populateMathRsqrtApproximationPatterns` and the transform op
`transform.apply_patterns.x86.math_rsqrt_approximation`.
Users who relied on the AVX2 rsqrt approximation should call
`x86::populateMathRsqrtApproximationPatterns` (C++) or use
`transform.apply_patterns.x86.math_rsqrt_approximation`.
Added:
mlir/lib/Dialect/X86/Transforms/MathRsqrtApproximation.cpp
mlir/test/Dialect/X86/math-rsqrt-approximation.mlir
Modified:
mlir/include/mlir/Dialect/Math/TransformOps/MathTransformOps.td
mlir/include/mlir/Dialect/Math/Transforms/Passes.h
mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
mlir/include/mlir/Dialect/X86/Transforms.h
mlir/lib/Dialect/Math/TransformOps/CMakeLists.txt
mlir/lib/Dialect/Math/TransformOps/MathTransformOps.cpp
mlir/lib/Dialect/Math/Transforms/CMakeLists.txt
mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp
mlir/lib/Dialect/X86/TransformOps/CMakeLists.txt
mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp
mlir/lib/Dialect/X86/Transforms/CMakeLists.txt
mlir/test/Dialect/Math/polynomial-approximation.mlir
mlir/test/mlir-runner/X86/math-polynomial-approx-avx2.mlir
mlir/test/python/dialects/transform_math_ext.py
mlir/test/python/dialects/transform_x86_ext.py
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Math/TransformOps/MathTransformOps.td b/mlir/include/mlir/Dialect/Math/TransformOps/MathTransformOps.td
index 32ed7416eb6046..e71b88c973df88 100644
--- a/mlir/include/mlir/Dialect/Math/TransformOps/MathTransformOps.td
+++ b/mlir/include/mlir/Dialect/Math/TransformOps/MathTransformOps.td
@@ -29,11 +29,9 @@ def ApplyPolynomialApproximationPatternsOp : Op<Transform_Dialect,
[DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
let description = [{
Enables polynomial approximations for math functions.
- Optionally, allows usage of x86 AVX2 intrinsics.
}];
- let arguments = (ins UnitAttr:$enable_avx2);
- let assemblyFormat = "(`enable_avx2` $enable_avx2^)? attr-dict";
+ let assemblyFormat = "attr-dict";
}
#endif // MATH_TRANSFORM_OPS
diff --git a/mlir/include/mlir/Dialect/Math/Transforms/Passes.h b/mlir/include/mlir/Dialect/Math/Transforms/Passes.h
index b3abbf728a3c68..d9abde0a7f2288 100644
--- a/mlir/include/mlir/Dialect/Math/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Math/Transforms/Passes.h
@@ -35,11 +35,6 @@ void populateExpansionPatterns(RewritePatternSet &patterns,
void populateMathAlgebraicSimplificationPatterns(RewritePatternSet &patterns);
-struct MathPolynomialApproximationOptions {
- // Enables the use of AVX2 intrinsics in some of the approximations.
- bool enableAvx2 = false;
-};
-
void populatePolynomialApproximateTanhPattern(RewritePatternSet &patterns);
void populatePolynomialApproximateErfPattern(RewritePatternSet &patterns);
void populatePolynomialApproximateErfcPattern(RewritePatternSet &patterns);
@@ -65,9 +60,7 @@ void populateMathPolynomialApproximationPatterns(
// Prefer calling these functions directly:
// * populateMathF32ExpansionPatterns(patterns, predicate)
// * populateMathPolynomialApproximationPatterns(patterns, predicate)
-void populateMathPolynomialApproximationPatterns(
- RewritePatternSet &patterns,
- const MathPolynomialApproximationOptions &options = {});
+void populateMathPolynomialApproximationPatterns(RewritePatternSet &patterns);
void populateUpliftToFMAPatterns(RewritePatternSet &patterns);
diff --git a/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td b/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
index b5859f6d27adea..5976a38b273188 100644
--- a/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
+++ b/mlir/include/mlir/Dialect/X86/TransformOps/X86TransformOps.td
@@ -114,4 +114,16 @@ def ApplyVectorContractMultiLevelUnrollForAMXPatternsOp : Op<Transform_Dialect,
let assemblyFormat = "attr-dict";
}
+def ApplyMathRsqrtApproximationPatternsOp : Op<Transform_Dialect,
+ "apply_patterns.x86.math_rsqrt_approximation",
+ [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+ let description = [{
+ Collect patterns that approximate vectorized f32 `math.rsqrt` operations
+ using the AVX `vrsqrtps` instruction (`x86.avx.rsqrt`) followed by one
+ Newton-Raphson step. The inner vector dimension must be a multiple of 8.
+ }];
+
+ let assemblyFormat = "attr-dict";
+}
+
#endif // X86_TRANSFORM_OPS
diff --git a/mlir/include/mlir/Dialect/X86/Transforms.h b/mlir/include/mlir/Dialect/X86/Transforms.h
index 6d28cf5b081d6b..a4a154cd351e9f 100644
--- a/mlir/include/mlir/Dialect/X86/Transforms.h
+++ b/mlir/include/mlir/Dialect/X86/Transforms.h
@@ -122,6 +122,11 @@ void populatePrepareVectorContractLoopsPatterns(RewritePatternSet &patterns);
void populateVectorContractMultiLevelUnrollForAMXPatterns(
RewritePatternSet &patterns);
+// Approximates vectorized f32 `math.rsqrt` with the AVX `vrsqrtps`
+// instruction followed by one Newton-Raphson step. Requires the inner vector
+// dimension to be a multiple of 8.
+void populateMathRsqrtApproximationPatterns(RewritePatternSet &patterns);
+
//===----------------------------------------------------------------------===//
/// Helpers extracted from:
/// - clang/lib/Headers/avxintrin.h
diff --git a/mlir/lib/Dialect/Math/TransformOps/CMakeLists.txt b/mlir/lib/Dialect/Math/TransformOps/CMakeLists.txt
index c10e0ce55801ec..bfad63e1da352a 100644
--- a/mlir/lib/Dialect/Math/TransformOps/CMakeLists.txt
+++ b/mlir/lib/Dialect/Math/TransformOps/CMakeLists.txt
@@ -14,5 +14,4 @@ add_mlir_dialect_library(MLIRMathTransformOps
MLIRMathTransforms
MLIRTransformDialect
MLIRVectorDialect
- MLIRX86Dialect
)
diff --git a/mlir/lib/Dialect/Math/TransformOps/MathTransformOps.cpp b/mlir/lib/Dialect/Math/TransformOps/MathTransformOps.cpp
index 37fc18a6b5da95..6720c962616d55 100644
--- a/mlir/lib/Dialect/Math/TransformOps/MathTransformOps.cpp
+++ b/mlir/lib/Dialect/Math/TransformOps/MathTransformOps.cpp
@@ -13,7 +13,6 @@
#include "mlir/Dialect/Math/Transforms/Passes.h"
#include "mlir/Dialect/Transform/IR/TransformDialect.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
-#include "mlir/Dialect/X86/X86Dialect.h"
using namespace mlir;
@@ -28,11 +27,8 @@ void transform::ApplyF32ExpansionPatternsOp::populatePatterns(
void transform::ApplyPolynomialApproximationPatternsOp::populatePatterns(
RewritePatternSet &patterns) {
- bool enableAvx2 = getEnableAvx2();
- populateMathPolynomialApproximationPatterns(patterns, [&](StringRef name) {
- // `rsqrt` requires AVX2 so always gate it.
- return enableAvx2 || name != math::RsqrtOp::getOperationName();
- });
+ populateMathPolynomialApproximationPatterns(patterns,
+ [](StringRef) { return true; });
}
//===----------------------------------------------------------------------===//
@@ -50,7 +46,6 @@ class MathTransformDialectExtension
declareGeneratedDialect<arith::ArithDialect>();
declareGeneratedDialect<math::MathDialect>();
declareGeneratedDialect<vector::VectorDialect>();
- declareGeneratedDialect<x86::X86Dialect>();
registerTransformOps<
#define GET_OP_LIST
#include "mlir/Dialect/Math/TransformOps/MathTransformOps.cpp.inc"
diff --git a/mlir/lib/Dialect/Math/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Math/Transforms/CMakeLists.txt
index 3cbfe03d784ab1..c85cc8e12377a5 100644
--- a/mlir/lib/Dialect/Math/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Math/Transforms/CMakeLists.txt
@@ -22,6 +22,5 @@ add_mlir_dialect_library(MLIRMathTransforms
MLIRPass
MLIRTensorDialect
MLIRTransforms
- MLIRX86Dialect
MLIRVectorDialect
)
diff --git a/mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp b/mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp
index 96b5a25da113ab..0e1517aa2596fe 100644
--- a/mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp
@@ -19,10 +19,8 @@
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/Math/Transforms/Approximation.h"
#include "mlir/Dialect/Math/Transforms/Passes.h"
-#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
-#include "mlir/Dialect/X86/X86Dialect.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
@@ -74,99 +72,6 @@ static Value broadcast(ImplicitLocOpBuilder &builder, Value value,
return shape ? BroadcastOp::create(builder, type, value) : value;
}
-//----------------------------------------------------------------------------//
-// Helper function to handle n-D vectors with 1-D operations.
-//----------------------------------------------------------------------------//
-
-// Expands and unrolls n-D vector operands into multiple fixed size 1-D vectors
-// and calls the compute function with 1-D vector operands. Stitches back all
-// results into the original n-D vector result.
-//
-// Examples: vectorWidth = 8
-// - vector<4x8xf32> unrolled 4 times
-// - vector<16xf32> expanded to vector<2x8xf32> and unrolled 2 times
-// - vector<4x16xf32> expanded to vector<4x2x8xf32> and unrolled 4*2 times
-//
-// Some math approximations rely on ISA-specific operations that only accept
-// fixed size 1-D vectors (e.g. AVX expects vectors of width 8).
-//
-// It is the caller's responsibility to verify that the inner dimension is
-// divisible by the vectorWidth, and that all operands have the same vector
-// shape.
-static Value
-handleMultidimensionalVectors(ImplicitLocOpBuilder &builder,
- ValueRange operands, int64_t vectorWidth,
- llvm::function_ref<Value(ValueRange)> compute) {
- assert(!operands.empty() && "operands must be not empty");
- assert(vectorWidth > 0 && "vector width must be larger than 0");
-
- VectorType inputType = cast<VectorType>(operands[0].getType());
- ArrayRef<int64_t> inputShape = inputType.getShape();
-
- // If input shape matches target vector width, we can just call the
- // user-provided compute function with the operands.
- if (inputShape == llvm::ArrayRef(vectorWidth))
- return compute(operands);
-
- // Check if the inner dimension has to be expanded, or we can directly iterate
- // over the outer dimensions of the vector.
- int64_t innerDim = inputShape.back();
- int64_t expansionDim = innerDim / vectorWidth;
- assert((innerDim % vectorWidth == 0) && "invalid inner dimension size");
-
- // Maybe expand operands to the higher rank vector shape that we'll use to
- // iterate over and extract one dimensional vectors.
- SmallVector<int64_t> expandedShape(inputShape);
- SmallVector<Value> expandedOperands(operands);
-
- if (expansionDim > 1) {
- // Expand shape from [..., innerDim] to [..., expansionDim, vectorWidth].
- expandedShape.insert(expandedShape.end() - 1, expansionDim);
- expandedShape.back() = vectorWidth;
-
- for (unsigned i = 0; i < operands.size(); ++i) {
- auto operand = operands[i];
- auto eltType = cast<VectorType>(operand.getType()).getElementType();
- auto expandedType = VectorType::get(expandedShape, eltType);
- expandedOperands[i] =
- vector::ShapeCastOp::create(builder, expandedType, operand);
- }
- }
-
- // Iterate over all outer dimensions of the compute shape vector type.
- auto iterationDims = ArrayRef<int64_t>(expandedShape).drop_back();
- int64_t maxIndex = computeMaxLinearIndex(iterationDims);
- auto strides = computeStrides(iterationDims);
-
- // Compute results for each one dimensional vector.
- SmallVector<Value> results(maxIndex);
-
- for (int64_t i = 0; i < maxIndex; ++i) {
- auto offsets = delinearize(i, strides);
-
- SmallVector<Value> extracted(expandedOperands.size());
- for (const auto &tuple : llvm::enumerate(expandedOperands))
- extracted[tuple.index()] =
- vector::ExtractOp::create(builder, tuple.value(), offsets);
-
- results[i] = compute(extracted);
- }
-
- // Stitch results together into one large vector.
- Type resultEltType = cast<VectorType>(results[0].getType()).getElementType();
- Type resultExpandedType = VectorType::get(expandedShape, resultEltType);
- Value result = arith::ConstantOp::create(
- builder, resultExpandedType, builder.getZeroAttr(resultExpandedType));
-
- for (int64_t i = 0; i < maxIndex; ++i)
- result = vector::InsertOp::create(builder, results[i], result,
- delinearize(i, strides));
-
- // Reshape back to the original vector shape.
- return vector::ShapeCastOp::create(
- builder, VectorType::get(inputShape, resultEltType), result);
-}
-
//----------------------------------------------------------------------------//
// Helper functions to create constants.
//----------------------------------------------------------------------------//
@@ -1695,77 +1600,6 @@ CbrtApproximation::matchAndRewrite(math::CbrtOp op,
return success();
}
-//----------------------------------------------------------------------------//
-// Rsqrt approximation.
-//----------------------------------------------------------------------------//
-
-namespace {
-struct RsqrtApproximation : public OpRewritePattern<math::RsqrtOp> {
- using Base::Base;
-
- LogicalResult matchAndRewrite(math::RsqrtOp op,
- PatternRewriter &rewriter) const final;
-};
-} // namespace
-
-LogicalResult
-RsqrtApproximation::matchAndRewrite(math::RsqrtOp op,
- PatternRewriter &rewriter) const {
- if (!getElementTypeOrSelf(op.getOperand()).isF32())
- return rewriter.notifyMatchFailure(op, "unsupported operand type");
-
- std::optional<VectorShape> shape = vectorShape(op.getOperand());
-
- // Only support already-vectorized rsqrt's.
- if (!shape || shape->sizes.empty() || shape->sizes.back() % 8 != 0)
- return rewriter.notifyMatchFailure(op, "unsupported operand type");
-
- ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
- auto bcast = [&](Value value) -> Value {
- return broadcast(builder, value, shape);
- };
-
- Value cstPosInf = bcast(f32FromBits(builder, 0x7f800000u));
- Value cstOnePointFive = bcast(f32Cst(builder, 1.5f));
- Value cstNegHalf = bcast(f32Cst(builder, -0.5f));
- Value cstMinNormPos = bcast(f32FromBits(builder, 0x00800000u));
-
- Value negHalf = arith::MulFOp::create(builder, op.getOperand(), cstNegHalf);
-
- // Select only the inverse sqrt of positive normals (denormals are
- // flushed to zero).
- Value ltMinMask = arith::CmpFOp::create(builder, arith::CmpFPredicate::OLT,
- op.getOperand(), cstMinNormPos);
- Value infMask = arith::CmpFOp::create(builder, arith::CmpFPredicate::OEQ,
- op.getOperand(), cstPosInf);
- Value notNormalFiniteMask = arith::OrIOp::create(builder, ltMinMask, infMask);
-
- // Compute an approximate result.
- Value yApprox = handleMultidimensionalVectors(
- builder, op->getOperands(), 8, [&builder](ValueRange operands) -> Value {
- return x86::avx::RsqrtOp::create(builder, operands);
- });
-
- // Do a single step of Newton-Raphson iteration to improve the approximation.
- // This uses the formula y_{n+1} = y_n * (1.5 - y_n * (0.5 * x) * y_n).
- // It is essential to evaluate the inner term like this because forming
- // y_n^2 may over- or underflow.
- Value inner = arith::MulFOp::create(builder, negHalf, yApprox);
- Value fma = math::FmaOp::create(builder, yApprox, inner, cstOnePointFive);
- Value yNewton = arith::MulFOp::create(builder, yApprox, fma);
-
- // Select the result of the Newton-Raphson step for positive normal arguments.
- // For other arguments, choose the output of the intrinsic. This will
- // return rsqrt(+inf) = 0, rsqrt(x) = NaN if x < 0, and rsqrt(x) = +inf if
- // x is zero or a positive denormalized float (equivalent to flushing positive
- // denormalized inputs to zero).
- Value res =
- arith::SelectOp::create(builder, notNormalFiniteMask, yApprox, yNewton);
- rewriter.replaceOp(op, res);
-
- return success();
-}
-
//----------------------------------------------------------------------------//
void mlir::populatePolynomialApproximateTanhPattern(
@@ -1866,8 +1700,6 @@ void mlir::populateMathPolynomialApproximationPatterns(
patterns, predicate, benefit);
populateMathPolynomialApproximationPattern<Log1pOp, Log1pApproximation>(
patterns, predicate, benefit);
- populateMathPolynomialApproximationPattern<RsqrtOp, RsqrtApproximation>(
- patterns, predicate, benefit);
populateMathPolynomialApproximationPattern<
SinOp, SinAndCosApproximation<true, math::SinOp>>(patterns, predicate,
benefit);
@@ -1876,8 +1708,7 @@ void mlir::populateMathPolynomialApproximationPatterns(
}
void mlir::populateMathPolynomialApproximationPatterns(
- RewritePatternSet &patterns,
- const MathPolynomialApproximationOptions &options) {
+ RewritePatternSet &patterns) {
mlir::populateMathF32ExpansionPatterns(patterns, [](StringRef name) -> bool {
return llvm::is_contained(
{math::AtanOp::getOperationName(), math::Atan2Op::getOperationName(),
@@ -1905,12 +1736,4 @@ void mlir::populateMathPolynomialApproximationPatterns(
math::CosOp::getOperationName()},
name);
});
-
- if (options.enableAvx2) {
- auto predicateRsqrt = [](StringRef name) {
- return name == math::RsqrtOp::getOperationName();
- };
- mlir::populateMathF32ExpansionPatterns(patterns, predicateRsqrt);
- mlir::populateMathPolynomialApproximationPatterns(patterns, predicateRsqrt);
- }
}
diff --git a/mlir/lib/Dialect/X86/TransformOps/CMakeLists.txt b/mlir/lib/Dialect/X86/TransformOps/CMakeLists.txt
index b9036b78c3ac0a..c7eb32d001350f 100644
--- a/mlir/lib/Dialect/X86/TransformOps/CMakeLists.txt
+++ b/mlir/lib/Dialect/X86/TransformOps/CMakeLists.txt
@@ -10,6 +10,7 @@ add_mlir_dialect_library(MLIRX86TransformOps
MLIRIR
MLIRLLVMCommonConversion
MLIRLLVMDialect
+ MLIRMathDialect
MLIRMemRefDialect
MLIRVectorDialect
MLIRSideEffectInterfaces
diff --git a/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp b/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp
index 610c3dcff63348..d3d196d9d0889f 100644
--- a/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp
+++ b/mlir/lib/Dialect/X86/TransformOps/X86TransformOps.cpp
@@ -11,6 +11,7 @@
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Transform/IR/TransformDialect.h"
#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
@@ -70,6 +71,11 @@ void mlir::transform::ApplyVectorContractMultiLevelUnrollForAMXPatternsOp::
x86::populateVectorContractMultiLevelUnrollForAMXPatterns(patterns);
}
+void mlir::transform::ApplyMathRsqrtApproximationPatternsOp::populatePatterns(
+ RewritePatternSet &patterns) {
+ x86::populateMathRsqrtApproximationPatterns(patterns);
+}
+
//===----------------------------------------------------------------------===//
// Transform op registration
//===----------------------------------------------------------------------===//
@@ -84,7 +90,9 @@ class X86TransformDialectExtension
X86TransformDialectExtension() {
declareGeneratedDialect<affine::AffineDialect>();
declareGeneratedDialect<arith::ArithDialect>();
+ declareGeneratedDialect<math::MathDialect>();
declareGeneratedDialect<memref::MemRefDialect>();
+ declareGeneratedDialect<vector::VectorDialect>();
declareGeneratedDialect<x86::X86Dialect>();
declareGeneratedDialect<LLVM::LLVMDialect>();
registerTransformOps<
diff --git a/mlir/lib/Dialect/X86/Transforms/CMakeLists.txt b/mlir/lib/Dialect/X86/Transforms/CMakeLists.txt
index 31fe12a05cc119..9f9bdd989edba2 100644
--- a/mlir/lib/Dialect/X86/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/X86/Transforms/CMakeLists.txt
@@ -1,6 +1,7 @@
add_mlir_dialect_library(MLIRX86Transforms
AVXTranspose.cpp
LegalizeForLLVMExport.cpp
+ MathRsqrtApproximation.cpp
VectorContractToFMA.cpp
VectorContractToPackedTypeDotProduct.cpp
VectorContractBF16ToFMA.cpp
@@ -19,6 +20,7 @@ add_mlir_dialect_library(MLIRX86Transforms
MLIRLinalgDialect
MLIRLLVMCommonConversion
MLIRLLVMDialect
+ MLIRMathDialect
MLIRMemRefDialect
MLIRSCFDialect
MLIRUBDialect
diff --git a/mlir/lib/Dialect/X86/Transforms/MathRsqrtApproximation.cpp b/mlir/lib/Dialect/X86/Transforms/MathRsqrtApproximation.cpp
new file mode 100644
index 00000000000000..da31f17f2273a5
--- /dev/null
+++ b/mlir/lib/Dialect/X86/Transforms/MathRsqrtApproximation.cpp
@@ -0,0 +1,191 @@
+//===- MathRsqrtApproximation.cpp - Approximate math.rsqrt with AVX -------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements an approximation of `math.rsqrt` using the AVX
+// `vrsqrtps` instruction refined by one Newton-Raphson step.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/X86/Transforms.h"
+#include "mlir/Dialect/X86/X86Dialect.h"
+#include "mlir/IR/ImplicitLocOpBuilder.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/TypeUtilities.h"
+
+using namespace mlir;
+
+// Expands and unrolls n-D vector operands into multiple fixed size 1-D vectors
+// and calls the compute function with 1-D vector operands. Stitches back all
+// results into the original n-D vector result.
+//
+// Examples: vectorWidth = 8
+// - vector<4x8xf32> unrolled 4 times
+// - vector<16xf32> expanded to vector<2x8xf32> and unrolled 2 times
+// - vector<4x16xf32> expanded to vector<4x2x8xf32> and unrolled 4*2 times
+//
+// It is the caller's responsibility to verify that the inner dimension is
+// divisible by the vectorWidth, and that all operands have the same vector
+// shape.
+static Value
+handleMultidimensionalVectors(ImplicitLocOpBuilder &builder,
+ ValueRange operands, int64_t vectorWidth,
+ llvm::function_ref<Value(ValueRange)> compute) {
+ assert(!operands.empty() && "operands must be not empty");
+ assert(vectorWidth > 0 && "vector width must be larger than 0");
+
+ VectorType inputType = cast<VectorType>(operands[0].getType());
+ ArrayRef<int64_t> inputShape = inputType.getShape();
+
+ // If input shape matches target vector width, we can just call the
+ // user-provided compute function with the operands.
+ if (inputShape == llvm::ArrayRef(vectorWidth))
+ return compute(operands);
+
+ // Check if the inner dimension has to be expanded, or we can directly iterate
+ // over the outer dimensions of the vector.
+ int64_t innerDim = inputShape.back();
+ int64_t expansionDim = innerDim / vectorWidth;
+ assert((innerDim % vectorWidth == 0) && "invalid inner dimension size");
+
+ // Maybe expand operands to the higher rank vector shape that we'll use to
+ // iterate over and extract one dimensional vectors.
+ SmallVector<int64_t> expandedShape(inputShape);
+ SmallVector<Value> expandedOperands(operands);
+
+ if (expansionDim > 1) {
+ // Expand shape from [..., innerDim] to [..., expansionDim, vectorWidth].
+ expandedShape.insert(expandedShape.end() - 1, expansionDim);
+ expandedShape.back() = vectorWidth;
+
+ for (unsigned i = 0; i < operands.size(); ++i) {
+ auto operand = operands[i];
+ auto eltType = cast<VectorType>(operand.getType()).getElementType();
+ auto expandedType = VectorType::get(expandedShape, eltType);
+ expandedOperands[i] =
+ vector::ShapeCastOp::create(builder, expandedType, operand);
+ }
+ }
+
+ // Iterate over all outer dimensions of the compute shape vector type.
+ auto iterationDims = ArrayRef<int64_t>(expandedShape).drop_back();
+ int64_t maxIndex = computeMaxLinearIndex(iterationDims);
+ auto strides = computeStrides(iterationDims);
+
+ // Compute results for each one dimensional vector.
+ SmallVector<Value> results(maxIndex);
+
+ for (int64_t i = 0; i < maxIndex; ++i) {
+ auto offsets = delinearize(i, strides);
+
+ SmallVector<Value> extracted(expandedOperands.size());
+ for (const auto &tuple : llvm::enumerate(expandedOperands))
+ extracted[tuple.index()] =
+ vector::ExtractOp::create(builder, tuple.value(), offsets);
+
+ results[i] = compute(extracted);
+ }
+
+ // Stitch results together into one large vector.
+ Type resultEltType = cast<VectorType>(results[0].getType()).getElementType();
+ Type resultExpandedType = VectorType::get(expandedShape, resultEltType);
+ Value result = arith::ConstantOp::create(
+ builder, resultExpandedType, builder.getZeroAttr(resultExpandedType));
+
+ for (int64_t i = 0; i < maxIndex; ++i)
+ result = vector::InsertOp::create(builder, results[i], result,
+ delinearize(i, strides));
+
+ // Reshape back to the original vector shape.
+ return vector::ShapeCastOp::create(
+ builder, VectorType::get(inputShape, resultEltType), result);
+}
+
+namespace {
+struct RsqrtApproximation : public OpRewritePattern<math::RsqrtOp> {
+ using Base::Base;
+
+ LogicalResult matchAndRewrite(math::RsqrtOp op,
+ PatternRewriter &rewriter) const final;
+};
+} // namespace
+
+LogicalResult
+RsqrtApproximation::matchAndRewrite(math::RsqrtOp op,
+ PatternRewriter &rewriter) const {
+ if (!getElementTypeOrSelf(op.getOperand()).isF32())
+ return rewriter.notifyMatchFailure(op, "unsupported operand type");
+
+ // Only support already-vectorized rsqrt's.
+ auto vectorType = dyn_cast<VectorType>(op.getOperand().getType());
+ if (!vectorType || vectorType.getRank() == 0 ||
+ vectorType.getShape().back() % 8 != 0)
+ return rewriter.notifyMatchFailure(op, "unsupported operand type");
+
+ ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
+ auto bcastF32FromBits = [&](uint32_t bits) -> Value {
+ Value i32Value = arith::ConstantOp::create(
+ builder, builder.getI32IntegerAttr(static_cast<int32_t>(bits)));
+ Value f32Value =
+ arith::BitcastOp::create(builder, builder.getF32Type(), i32Value);
+ return vector::BroadcastOp::create(builder, vectorType, f32Value);
+ };
+ auto bcastF32 = [&](float value) -> Value {
+ Value f32Value =
+ arith::ConstantOp::create(builder, builder.getF32FloatAttr(value));
+ return vector::BroadcastOp::create(builder, vectorType, f32Value);
+ };
+
+ Value cstPosInf = bcastF32FromBits(0x7f800000u);
+ Value cstOnePointFive = bcastF32(1.5f);
+ Value cstNegHalf = bcastF32(-0.5f);
+ Value cstMinNormPos = bcastF32FromBits(0x00800000u);
+
+ Value negHalf = arith::MulFOp::create(builder, op.getOperand(), cstNegHalf);
+
+ // Select only the inverse sqrt of positive normals (denormals are
+ // flushed to zero).
+ Value ltMinMask = arith::CmpFOp::create(builder, arith::CmpFPredicate::OLT,
+ op.getOperand(), cstMinNormPos);
+ Value infMask = arith::CmpFOp::create(builder, arith::CmpFPredicate::OEQ,
+ op.getOperand(), cstPosInf);
+ Value notNormalFiniteMask = arith::OrIOp::create(builder, ltMinMask, infMask);
+
+ // Compute an approximate result.
+ Value yApprox = handleMultidimensionalVectors(
+ builder, op->getOperands(), 8, [&builder](ValueRange operands) -> Value {
+ return x86::avx::RsqrtOp::create(builder, operands);
+ });
+
+ // Do a single step of Newton-Raphson iteration to improve the approximation.
+ // This uses the formula y_{n+1} = y_n * (1.5 - y_n * (0.5 * x) * y_n).
+ // It is essential to evaluate the inner term like this because forming
+ // y_n^2 may over- or underflow.
+ Value inner = arith::MulFOp::create(builder, negHalf, yApprox);
+ Value fma = math::FmaOp::create(builder, yApprox, inner, cstOnePointFive);
+ Value yNewton = arith::MulFOp::create(builder, yApprox, fma);
+
+ // Select the result of the Newton-Raphson step for positive normal arguments.
+ // For other arguments, choose the output of the intrinsic. This will
+ // return rsqrt(+inf) = 0, rsqrt(x) = NaN if x < 0, and rsqrt(x) = +inf if
+ // x is zero or a positive denormalized float (equivalent to flushing positive
+ // denormalized inputs to zero).
+ Value res =
+ arith::SelectOp::create(builder, notNormalFiniteMask, yApprox, yNewton);
+ rewriter.replaceOp(op, res);
+
+ return success();
+}
+
+void mlir::x86::populateMathRsqrtApproximationPatterns(
+ RewritePatternSet &patterns) {
+ patterns.add<RsqrtApproximation>(patterns.getContext());
+}
diff --git a/mlir/test/Dialect/Math/polynomial-approximation.mlir b/mlir/test/Dialect/Math/polynomial-approximation.mlir
index 41c7a2819e6e85..42d56be288cd41 100644
--- a/mlir/test/Dialect/Math/polynomial-approximation.mlir
+++ b/mlir/test/Dialect/Math/polynomial-approximation.mlir
@@ -1,6 +1,4 @@
// RUN: mlir-opt %s -transform-interpreter | FileCheck %s
-// RUN: mlir-opt %s -transform-interpreter=entry-point=avx2 \
-// RUN: | FileCheck --check-prefix=AVX2 %s
// Check that all math functions lowered to approximations built from
// standard operations (add, mul, fma, shift, etc...).
@@ -669,11 +667,9 @@ func.func @tanh_scalable_vector(%arg0: vector<[8]xf32>) -> vector<[8]xf32> {
return %0 : vector<[8]xf32>
}
-// We only approximate rsqrt for vectors and when the AVX2 option is enabled.
+// rsqrt is only approximated by target-specific patterns.
// CHECK-LABEL: func @rsqrt_scalar
-// AVX2-LABEL: func @rsqrt_scalar
// CHECK: math.rsqrt
-// AVX2: math.rsqrt
func.func @rsqrt_scalar(%arg0: f32) -> f32 {
%0 = math.rsqrt %arg0 : f32
return %0 : f32
@@ -681,23 +677,6 @@ func.func @rsqrt_scalar(%arg0: f32) -> f32 {
// CHECK-LABEL: func @rsqrt_vector_8xf32
// CHECK: math.rsqrt
-// AVX2-LABEL: func @rsqrt_vector_8xf32(
-// AVX2-SAME: %[[VAL_0:.*]]: vector<8xf32>) -> vector<8xf32> {
-// AVX2: %[[VAL_1:.*]] = arith.constant dense<0x7F800000> : vector<8xf32>
-// AVX2: %[[VAL_2:.*]] = arith.constant dense<1.500000e+00> : vector<8xf32>
-// AVX2: %[[VAL_3:.*]] = arith.constant dense<-5.000000e-01> : vector<8xf32>
-// AVX2: %[[VAL_4:.*]] = arith.constant dense<1.17549435E-38> : vector<8xf32>
-// AVX2: %[[VAL_5:.*]] = arith.mulf %[[VAL_0]], %[[VAL_3]] : vector<8xf32>
-// AVX2: %[[VAL_6:.*]] = arith.cmpf olt, %[[VAL_0]], %[[VAL_4]] : vector<8xf32>
-// AVX2: %[[VAL_7:.*]] = arith.cmpf oeq, %[[VAL_0]], %[[VAL_1]] : vector<8xf32>
-// AVX2: %[[VAL_8:.*]] = arith.ori %[[VAL_6]], %[[VAL_7]] : vector<8xi1>
-// AVX2: %[[VAL_9:.*]] = x86.avx.rsqrt %[[VAL_0]] : vector<8xf32>
-// AVX2: %[[VAL_10:.*]] = arith.mulf %[[VAL_5]], %[[VAL_9]] : vector<8xf32>
-// AVX2: %[[VAL_11:.*]] = math.fma %[[VAL_9]], %[[VAL_10]], %[[VAL_2]] : vector<8xf32>
-// AVX2: %[[VAL_12:.*]] = arith.mulf %[[VAL_9]], %[[VAL_11]] : vector<8xf32>
-// AVX2: %[[VAL_13:.*]] = arith.select %[[VAL_8]], %[[VAL_9]], %[[VAL_12]] : vector<8xi1>, vector<8xf32>
-// AVX2: return %[[VAL_13]] : vector<8xf32>
-// AVX2: }
func.func @rsqrt_vector_8xf32(%arg0: vector<8xf32>) -> vector<8xf32> {
%0 = math.rsqrt %arg0 : vector<8xf32>
return %0 : vector<8xf32>
@@ -707,8 +686,6 @@ func.func @rsqrt_vector_8xf32(%arg0: vector<8xf32>) -> vector<8xf32> {
//
// CHECK-LABEL: func @rsqrt_vector_5xf32
// CHECK: math.rsqrt
-// AVX2-LABEL: func @rsqrt_vector_5xf32
-// AVX2: math.rsqrt
func.func @rsqrt_vector_5xf32(%arg0: vector<5xf32>) -> vector<5xf32> {
%0 = math.rsqrt %arg0 : vector<5xf32>
return %0 : vector<5xf32>
@@ -719,18 +696,6 @@ func.func @rsqrt_vector_5xf32(%arg0: vector<5xf32>) -> vector<5xf32> {
//
// CHECK-LABEL: func @rsqrt_vector_16xf32
// CHECK: math.rsqrt
-// AVX2-LABEL: func @rsqrt_vector_16xf32(
-// AVX2-SAME: %[[ARG:.*]]: vector<16xf32>
-// AVX2-SAME: ) -> vector<16xf32>
-// AVX2: %[[INIT:.*]] = arith.constant dense<0.000000e+00> : vector<2x8xf32>
-// AVX2: %[[EXPAND:.*]] = vector.shape_cast %[[ARG]] : vector<16xf32> to vector<2x8xf32>
-// AVX2: %[[VEC0:.*]] = vector.extract %[[EXPAND]][0]
-// AVX2: %[[RSQRT0:.*]] = x86.avx.rsqrt %[[VEC0]]
-// AVX2: %[[VEC1:.*]] = vector.extract %[[EXPAND]][1]
-// AVX2: %[[RSQRT1:.*]] = x86.avx.rsqrt %[[VEC1]]
-// AVX2: %[[RESULT0:.*]] = vector.insert %[[RSQRT0]], %[[INIT]] [0]
-// AVX2: %[[RESULT1:.*]] = vector.insert %[[RSQRT1]], %[[RESULT0]] [1]
-// AVX2: %[[RSQRT:.*]] = vector.shape_cast %[[RESULT1]] : vector<2x8xf32> to vector<16xf32>
func.func @rsqrt_vector_16xf32(%arg0: vector<16xf32>) -> vector<16xf32> {
%0 = math.rsqrt %arg0 : vector<16xf32>
return %0 : vector<16xf32>
@@ -740,18 +705,6 @@ func.func @rsqrt_vector_16xf32(%arg0: vector<16xf32>) -> vector<16xf32> {
//
// CHECK-LABEL: func @rsqrt_vector_2x8xf32
// CHECK: math.rsqrt
-// AVX2-LABEL: func @rsqrt_vector_2x8xf32(
-// AVX2-SAME: %[[ARG:.*]]: vector<2x8xf32>
-// AVX2-SAME: ) -> vector<2x8xf32>
-// AVX2: %[[INIT:.*]] = arith.constant dense<0.000000e+00> : vector<2x8xf32>
-// AVX2-NOT: vector.shape_cast
-// AVX2: %[[VEC0:.*]] = vector.extract %[[ARG]][0]
-// AVX2: %[[RSQRT0:.*]] = x86.avx.rsqrt %[[VEC0]]
-// AVX2: %[[VEC1:.*]] = vector.extract %[[ARG]][1]
-// AVX2: %[[RSQRT1:.*]] = x86.avx.rsqrt %[[VEC1]]
-// AVX2: %[[RESULT0:.*]] = vector.insert %[[RSQRT0]], %[[INIT]] [0]
-// AVX2: %[[RESULT1:.*]] = vector.insert %[[RSQRT1]], %[[RESULT0]] [1]
-// AVX2-NOT: vector.shape_cast
func.func @rsqrt_vector_2x8xf32(%arg0: vector<2x8xf32>) -> vector<2x8xf32> {
%0 = math.rsqrt %arg0 : vector<2x8xf32>
return %0 : vector<2x8xf32>
@@ -762,24 +715,6 @@ func.func @rsqrt_vector_2x8xf32(%arg0: vector<2x8xf32>) -> vector<2x8xf32> {
//
// CHECK-LABEL: func @rsqrt_vector_2x16xf32
// CHECK: math.rsqrt
-// AVX2-LABEL: func @rsqrt_vector_2x16xf32(
-// AVX2-SAME: %[[ARG:.*]]: vector<2x16xf32>
-// AVX2-SAME: ) -> vector<2x16xf32>
-// AVX2: %[[INIT:.*]] = arith.constant dense<0.000000e+00> : vector<2x2x8xf32>
-// AVX2: %[[EXPAND:.*]] = vector.shape_cast %[[ARG]] : vector<2x16xf32> to vector<2x2x8xf32>
-// AVX2: %[[VEC00:.*]] = vector.extract %[[EXPAND]][0, 0]
-// AVX2: %[[RSQRT00:.*]] = x86.avx.rsqrt %[[VEC00]]
-// AVX2: %[[VEC01:.*]] = vector.extract %[[EXPAND]][0, 1]
-// AVX2: %[[RSQRT01:.*]] = x86.avx.rsqrt %[[VEC01]]
-// AVX2: %[[VEC10:.*]] = vector.extract %[[EXPAND]][1, 0]
-// AVX2: %[[RSQRT10:.*]] = x86.avx.rsqrt %[[VEC10]]
-// AVX2: %[[VEC11:.*]] = vector.extract %[[EXPAND]][1, 1]
-// AVX2: %[[RSQRT11:.*]] = x86.avx.rsqrt %[[VEC11]]
-// AVX2: %[[RESULT0:.*]] = vector.insert %[[RSQRT00]], %[[INIT]] [0, 0]
-// AVX2: %[[RESULT1:.*]] = vector.insert %[[RSQRT01]], %[[RESULT0]] [0, 1]
-// AVX2: %[[RESULT2:.*]] = vector.insert %[[RSQRT10]], %[[RESULT1]] [1, 0]
-// AVX2: %[[RESULT3:.*]] = vector.insert %[[RSQRT11]], %[[RESULT2]] [1, 1]
-// AVX2: %[[RSQRT:.*]] = vector.shape_cast %[[RESULT3]] : vector<2x2x8xf32> to vector<2x16xf32>
func.func @rsqrt_vector_2x16xf32(%arg0: vector<2x16xf32>) -> vector<2x16xf32> {
%0 = math.rsqrt %arg0 : vector<2x16xf32>
return %0 : vector<2x16xf32>
@@ -1046,13 +981,6 @@ func.func @math_zero_rank(%arg0 : vector<f16>) -> vector<f16> {
return %11 : vector<f16>
}
-// AVX2-LABEL: @rsqrt_f16
-func.func @rsqrt_f16(%arg0 : vector<2x8xf16>) -> vector<2x8xf16> {
- // AVX2-NOT: math.rsqrt
- %0 = "math.rsqrt"(%arg0) : (vector<2x8xf16>) -> vector<2x8xf16>
- return %0 : vector<2x8xf16>
-}
-
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
@@ -1062,13 +990,4 @@ module attributes {transform.with_named_sequence} {
} : !transform.any_op
transform.yield
}
-
- transform.named_sequence @avx2(%root: !transform.any_op {transform.readonly}) {
- %func = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
- transform.apply_patterns to %func {
- transform.apply_patterns.math.f32_expansion
- transform.apply_patterns.math.polynomial_approximation enable_avx2
- } : !transform.any_op
- transform.yield
- }
}
diff --git a/mlir/test/Dialect/X86/math-rsqrt-approximation.mlir b/mlir/test/Dialect/X86/math-rsqrt-approximation.mlir
new file mode 100644
index 00000000000000..5a8b052cf645b5
--- /dev/null
+++ b/mlir/test/Dialect/X86/math-rsqrt-approximation.mlir
@@ -0,0 +1,125 @@
+// RUN: mlir-opt %s -transform-interpreter | FileCheck %s
+
+// Check the AVX-based approximation of math.rsqrt.
+
+// We only approximate rsqrt for vectors.
+// CHECK-LABEL: func @rsqrt_scalar
+// CHECK: math.rsqrt
+func.func @rsqrt_scalar(%arg0: f32) -> f32 {
+ %0 = math.rsqrt %arg0 : f32
+ return %0 : f32
+}
+
+// CHECK-LABEL: func @rsqrt_vector_8xf32(
+// CHECK-SAME: %[[VAL_0:.*]]: vector<8xf32>) -> vector<8xf32> {
+// CHECK: %[[VAL_1:.*]] = arith.constant dense<0x7F800000> : vector<8xf32>
+// CHECK: %[[VAL_2:.*]] = arith.constant dense<1.500000e+00> : vector<8xf32>
+// CHECK: %[[VAL_3:.*]] = arith.constant dense<-5.000000e-01> : vector<8xf32>
+// CHECK: %[[VAL_4:.*]] = arith.constant dense<1.17549435E-38> : vector<8xf32>
+// CHECK: %[[VAL_5:.*]] = arith.mulf %[[VAL_0]], %[[VAL_3]] : vector<8xf32>
+// CHECK: %[[VAL_6:.*]] = arith.cmpf olt, %[[VAL_0]], %[[VAL_4]] : vector<8xf32>
+// CHECK: %[[VAL_7:.*]] = arith.cmpf oeq, %[[VAL_0]], %[[VAL_1]] : vector<8xf32>
+// CHECK: %[[VAL_8:.*]] = arith.ori %[[VAL_6]], %[[VAL_7]] : vector<8xi1>
+// CHECK: %[[VAL_9:.*]] = x86.avx.rsqrt %[[VAL_0]] : vector<8xf32>
+// CHECK: %[[VAL_10:.*]] = arith.mulf %[[VAL_5]], %[[VAL_9]] : vector<8xf32>
+// CHECK: %[[VAL_11:.*]] = math.fma %[[VAL_9]], %[[VAL_10]], %[[VAL_2]] : vector<8xf32>
+// CHECK: %[[VAL_12:.*]] = arith.mulf %[[VAL_9]], %[[VAL_11]] : vector<8xf32>
+// CHECK: %[[VAL_13:.*]] = arith.select %[[VAL_8]], %[[VAL_9]], %[[VAL_12]] : vector<8xi1>, vector<8xf32>
+// CHECK: return %[[VAL_13]] : vector<8xf32>
+// CHECK: }
+func.func @rsqrt_vector_8xf32(%arg0: vector<8xf32>) -> vector<8xf32> {
+ %0 = math.rsqrt %arg0 : vector<8xf32>
+ return %0 : vector<8xf32>
+}
+
+// Virtual vector width is not a multiple of an AVX2 vector width.
+//
+// CHECK-LABEL: func @rsqrt_vector_5xf32
+// CHECK: math.rsqrt
+func.func @rsqrt_vector_5xf32(%arg0: vector<5xf32>) -> vector<5xf32> {
+ %0 = math.rsqrt %arg0 : vector<5xf32>
+ return %0 : vector<5xf32>
+}
+
+// One dimensional virtual vector expanded and unrolled into multiple AVX2-sized
+// vectors.
+//
+// CHECK-LABEL: func @rsqrt_vector_16xf32(
+// CHECK-SAME: %[[ARG:.*]]: vector<16xf32>
+// CHECK-SAME: ) -> vector<16xf32>
+// CHECK: %[[INIT:.*]] = arith.constant dense<0.000000e+00> : vector<2x8xf32>
+// CHECK: %[[EXPAND:.*]] = vector.shape_cast %[[ARG]] : vector<16xf32> to vector<2x8xf32>
+// CHECK: %[[VEC0:.*]] = vector.extract %[[EXPAND]][0]
+// CHECK: %[[RSQRT0:.*]] = x86.avx.rsqrt %[[VEC0]]
+// CHECK: %[[VEC1:.*]] = vector.extract %[[EXPAND]][1]
+// CHECK: %[[RSQRT1:.*]] = x86.avx.rsqrt %[[VEC1]]
+// CHECK: %[[RESULT0:.*]] = vector.insert %[[RSQRT0]], %[[INIT]] [0]
+// CHECK: %[[RESULT1:.*]] = vector.insert %[[RSQRT1]], %[[RESULT0]] [1]
+// CHECK: %[[RSQRT:.*]] = vector.shape_cast %[[RESULT1]] : vector<2x8xf32> to vector<16xf32>
+func.func @rsqrt_vector_16xf32(%arg0: vector<16xf32>) -> vector<16xf32> {
+ %0 = math.rsqrt %arg0 : vector<16xf32>
+ return %0 : vector<16xf32>
+}
+
+// Two dimensional virtual vector unrolled into multiple AVX2-sized vectors.
+//
+// CHECK-LABEL: func @rsqrt_vector_2x8xf32(
+// CHECK-SAME: %[[ARG:.*]]: vector<2x8xf32>
+// CHECK-SAME: ) -> vector<2x8xf32>
+// CHECK: %[[INIT:.*]] = arith.constant dense<0.000000e+00> : vector<2x8xf32>
+// CHECK-NOT: vector.shape_cast
+// CHECK: %[[VEC0:.*]] = vector.extract %[[ARG]][0]
+// CHECK: %[[RSQRT0:.*]] = x86.avx.rsqrt %[[VEC0]]
+// CHECK: %[[VEC1:.*]] = vector.extract %[[ARG]][1]
+// CHECK: %[[RSQRT1:.*]] = x86.avx.rsqrt %[[VEC1]]
+// CHECK: %[[RESULT0:.*]] = vector.insert %[[RSQRT0]], %[[INIT]] [0]
+// CHECK: %[[RESULT1:.*]] = vector.insert %[[RSQRT1]], %[[RESULT0]] [1]
+// CHECK-NOT: vector.shape_cast
+func.func @rsqrt_vector_2x8xf32(%arg0: vector<2x8xf32>) -> vector<2x8xf32> {
+ %0 = math.rsqrt %arg0 : vector<2x8xf32>
+ return %0 : vector<2x8xf32>
+}
+
+// Two dimensional virtual vector expanded and unrolled into multiple AVX2-sized
+// vectors.
+//
+// CHECK-LABEL: func @rsqrt_vector_2x16xf32(
+// CHECK-SAME: %[[ARG:.*]]: vector<2x16xf32>
+// CHECK-SAME: ) -> vector<2x16xf32>
+// CHECK: %[[INIT:.*]] = arith.constant dense<0.000000e+00> : vector<2x2x8xf32>
+// CHECK: %[[EXPAND:.*]] = vector.shape_cast %[[ARG]] : vector<2x16xf32> to vector<2x2x8xf32>
+// CHECK: %[[VEC00:.*]] = vector.extract %[[EXPAND]][0, 0]
+// CHECK: %[[RSQRT00:.*]] = x86.avx.rsqrt %[[VEC00]]
+// CHECK: %[[VEC01:.*]] = vector.extract %[[EXPAND]][0, 1]
+// CHECK: %[[RSQRT01:.*]] = x86.avx.rsqrt %[[VEC01]]
+// CHECK: %[[VEC10:.*]] = vector.extract %[[EXPAND]][1, 0]
+// CHECK: %[[RSQRT10:.*]] = x86.avx.rsqrt %[[VEC10]]
+// CHECK: %[[VEC11:.*]] = vector.extract %[[EXPAND]][1, 1]
+// CHECK: %[[RSQRT11:.*]] = x86.avx.rsqrt %[[VEC11]]
+// CHECK: %[[RESULT0:.*]] = vector.insert %[[RSQRT00]], %[[INIT]] [0, 0]
+// CHECK: %[[RESULT1:.*]] = vector.insert %[[RSQRT01]], %[[RESULT0]] [0, 1]
+// CHECK: %[[RESULT2:.*]] = vector.insert %[[RSQRT10]], %[[RESULT1]] [1, 0]
+// CHECK: %[[RESULT3:.*]] = vector.insert %[[RSQRT11]], %[[RESULT2]] [1, 1]
+// CHECK: %[[RSQRT:.*]] = vector.shape_cast %[[RESULT3]] : vector<2x2x8xf32> to vector<2x16xf32>
+func.func @rsqrt_vector_2x16xf32(%arg0: vector<2x16xf32>) -> vector<2x16xf32> {
+ %0 = math.rsqrt %arg0 : vector<2x16xf32>
+ return %0 : vector<2x16xf32>
+}
+
+// CHECK-LABEL: @rsqrt_f16
+func.func @rsqrt_f16(%arg0 : vector<2x8xf16>) -> vector<2x8xf16> {
+ // CHECK-NOT: math.rsqrt
+ %0 = "math.rsqrt"(%arg0) : (vector<2x8xf16>) -> vector<2x8xf16>
+ return %0 : vector<2x8xf16>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.math.f32_expansion
+ transform.apply_patterns.x86.math_rsqrt_approximation
+ } : !transform.any_op
+ transform.yield
+ }
+}
diff --git a/mlir/test/mlir-runner/X86/math-polynomial-approx-avx2.mlir b/mlir/test/mlir-runner/X86/math-polynomial-approx-avx2.mlir
index 6a1bc237c00a89..0a522dedadf62d 100644
--- a/mlir/test/mlir-runner/X86/math-polynomial-approx-avx2.mlir
+++ b/mlir/test/mlir-runner/X86/math-polynomial-approx-avx2.mlir
@@ -47,7 +47,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func {
- transform.apply_patterns.math.polynomial_approximation enable_avx2
+ transform.apply_patterns.x86.math_rsqrt_approximation
} : !transform.any_op
transform.yield
}
diff --git a/mlir/test/python/dialects/transform_math_ext.py b/mlir/test/python/dialects/transform_math_ext.py
index 418c837332ead2..8c5f7f7d00598f 100644
--- a/mlir/test/python/dialects/transform_math_ext.py
+++ b/mlir/test/python/dialects/transform_math_ext.py
@@ -38,11 +38,3 @@ def testPolynomialApproximation():
# CHECK-LABEL: TEST: testPolynomialApproximation
# CHECK: apply_patterns
# CHECK: transform.apply_patterns.math.polynomial_approximation{{$}}
-
-
- at run_apply_patterns
-def testPolynomialApproximationAvx2():
- math.ApplyPolynomialApproximationPatternsOp(enable_avx2=True)
- # CHECK-LABEL: TEST: testPolynomialApproximationAvx2
- # CHECK: apply_patterns
- # CHECK: transform.apply_patterns.math.polynomial_approximation enable_avx2
diff --git a/mlir/test/python/dialects/transform_x86_ext.py b/mlir/test/python/dialects/transform_x86_ext.py
index b4034f02dc0b9e..fdd1afc4b8f486 100644
--- a/mlir/test/python/dialects/transform_x86_ext.py
+++ b/mlir/test/python/dialects/transform_x86_ext.py
@@ -38,3 +38,5 @@ def non_configurable_patterns():
x86.ApplySinkVectorProducerOpsPatternsOp()
# CHECK: transform.apply_patterns.x86.shuffle_vector_fma_ops
x86.ApplyShuffleVectorFMAOpsPatternsOp()
+ # CHECK: transform.apply_patterns.x86.math_rsqrt_approximation
+ x86.ApplyMathRsqrtApproximationPatternsOp()
More information about the Mlir-commits
mailing list