72 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9// This file contains test passes for expanding math operations into10// 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 TestMathPolynomialApproximationPass26 : 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 ®istry) 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(false)};54};55} // namespace56 57void TestMathPolynomialApproximationPass::runOnOperation() {58 RewritePatternSet patterns(&getContext());59 MathPolynomialApproximationOptions approxOptions;60 approxOptions.enableAvx2 = enableAvx2;61 populateMathPolynomialApproximationPatterns(patterns, approxOptions);62 (void)applyPatternsGreedily(getOperation(), std::move(patterns));63}64 65namespace mlir {66namespace test {67void registerTestMathPolynomialApproximationPass() {68 PassRegistration<TestMathPolynomialApproximationPass>();69}70} // namespace test71} // namespace mlir72