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 22 using namespace mlir; 23 24 namespace { 25 struct 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 ®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 } // namespace 56 57 void 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 65 namespace mlir { 66 namespace test { 67 void registerTestMathPolynomialApproximationPass() { 68 PassRegistration<TestMathPolynomialApproximationPass>(); 69 } 70 } // namespace test 71 } // namespace mlir 72