1//===- TestPolynomialApproximation.cpp - Test math ops approximations -----===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file contains test passes for expanding math operations into
10// polynomial approximations.
11//
12//===----------------------------------------------------------------------===//
13
14#include "mlir/Dialect/Arith/IR/Arith.h"
15#include "mlir/Dialect/Math/IR/Math.h"
16#include "mlir/Dialect/Math/Transforms/Passes.h"
17#include "mlir/Dialect/Vector/IR/VectorOps.h"
18#include "mlir/Dialect/X86Vector/X86VectorDialect.h"
19#include "mlir/Pass/Pass.h"
20#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
21
22using namespace mlir;
23
24namespace {
25struct TestMathPolynomialApproximationPass
26 : public PassWrapper<TestMathPolynomialApproximationPass, OperationPass<>> {
27 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
28 TestMathPolynomialApproximationPass)
29
30 TestMathPolynomialApproximationPass() = default;
31 TestMathPolynomialApproximationPass(
32 const TestMathPolynomialApproximationPass &pass)
33 : PassWrapper(pass) {}
34
35 void runOnOperation() override;
36 void getDependentDialects(DialectRegistry &registry) const override {
37 registry.insert<arith::ArithDialect, math::MathDialect,
38 vector::VectorDialect>();
39 if (enableAvx2)
40 registry.insert<x86vector::X86VectorDialect>();
41 }
42 StringRef getArgument() const final {
43 return "test-math-polynomial-approximation";
44 }
45 StringRef getDescription() const final {
46 return "Test math polynomial approximations";
47 }
48
49 Option<bool> enableAvx2{
50 *this, "enable-avx2",
51 llvm::cl::desc("Enable approximations that emit AVX2 intrinsics via the "
52 "X86Vector dialect"),
53 llvm::cl::init(Val: false)};
54};
55} // namespace
56
57void TestMathPolynomialApproximationPass::runOnOperation() {
58 RewritePatternSet patterns(&getContext());
59 MathPolynomialApproximationOptions approxOptions;
60 approxOptions.enableAvx2 = enableAvx2;
61 populateMathPolynomialApproximationPatterns(patterns, options: approxOptions);
62 (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
63}
64
65namespace mlir {
66namespace test {
67void registerTestMathPolynomialApproximationPass() {
68 PassRegistration<TestMathPolynomialApproximationPass>();
69}
70} // namespace test
71} // namespace mlir
72

source code of mlir/test/lib/Dialect/Math/TestPolynomialApproximation.cpp