mlir/test/lib/Dialect/Math/TestPolynomialApproximation.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-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 &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(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)applyPatternsAndFoldGreedily(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