35553d452b
This patch adds a polynomial approximation that matches the approximation in Eigen. Note that the approximation only applies to vectorized inputs; the scalar rsqrt is left unmodified. The approximation is protected with a flag since it emits an AVX2 intrinsic (generated via the X86Vector). This is the only reasonably clean way that I could find to generate the exact approximation that I wanted (i.e. an identical one to Eigen's). I considered two alternatives: 1. Introduce a Rsqrt intrinsic in LLVM, which doesn't exist yet. I believe this is because there is no definition of Rsqrt that all backends could agree on, since hardware instructions that implement it have widely varying degrees of precision. This is something that the standard could mandate, but Rsqrt is not part of IEEE754, so I don't think this option is feasible. 2. Emit fdiv(1.0, sqrt) with fast math flags to allow reciprocal transformations. Although portable, this doesn't allow us to generate exactly the code we want; it is the LLVM backend, and not MLIR, who controls what code is generated based on the target CPU. Reviewed By: ezhulenev Differential Revision: https://reviews.llvm.org/D112192
68 lines
2.4 KiB
C++
68 lines
2.4 KiB
C++
//===- TestPolynomialApproximation.cpp - Test math ops approximations -----===//
|
|
//
|
|
// 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 contains test passes for expanding math operations into
|
|
// polynomial approximations.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
|
|
#include "mlir/Dialect/Math/IR/Math.h"
|
|
#include "mlir/Dialect/Math/Transforms/Passes.h"
|
|
#include "mlir/Dialect/Vector/VectorOps.h"
|
|
#include "mlir/Dialect/X86Vector/X86VectorDialect.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace {
|
|
struct TestMathPolynomialApproximationPass
|
|
: public PassWrapper<TestMathPolynomialApproximationPass, FunctionPass> {
|
|
TestMathPolynomialApproximationPass() = default;
|
|
TestMathPolynomialApproximationPass(
|
|
const TestMathPolynomialApproximationPass &pass) {}
|
|
|
|
void runOnFunction() override;
|
|
void getDependentDialects(DialectRegistry ®istry) const override {
|
|
registry.insert<arith::ArithmeticDialect, math::MathDialect,
|
|
vector::VectorDialect>();
|
|
if (enableAvx2)
|
|
registry.insert<x86vector::X86VectorDialect>();
|
|
}
|
|
StringRef getArgument() const final {
|
|
return "test-math-polynomial-approximation";
|
|
}
|
|
StringRef getDescription() const final {
|
|
return "Test math polynomial approximations";
|
|
}
|
|
|
|
Option<bool> enableAvx2{
|
|
*this, "enable-avx2",
|
|
llvm::cl::desc("Enable approximations that emit AVX2 intrinsics via the "
|
|
"X86Vector dialect"),
|
|
llvm::cl::init(false)};
|
|
};
|
|
} // end anonymous namespace
|
|
|
|
void TestMathPolynomialApproximationPass::runOnFunction() {
|
|
RewritePatternSet patterns(&getContext());
|
|
MathPolynomialApproximationOptions approx_options;
|
|
approx_options.enableAvx2 = enableAvx2;
|
|
populateMathPolynomialApproximationPatterns(patterns, approx_options);
|
|
(void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
|
|
}
|
|
|
|
namespace mlir {
|
|
namespace test {
|
|
void registerTestMathPolynomialApproximationPass() {
|
|
PassRegistration<TestMathPolynomialApproximationPass>();
|
|
}
|
|
} // namespace test
|
|
} // namespace mlir
|