Page MenuHomePhabricator

[mlir] Add polynomial approximation for vectorized math::Rsqrt
ClosedPublic

Authored by cota on Oct 20 2021, 4:15 PM.

Details

Summary

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.
  1. 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.

Diff Detail

Event Timeline

cota created this revision.Oct 20 2021, 4:15 PM
cota requested review of this revision.Oct 20 2021, 4:15 PM
cota updated this revision to Diff 381114.Oct 20 2021, 4:19 PM

Add missing newline's at EOFs

ezhulenev added inline comments.Oct 21 2021, 4:51 AM
mlir/include/mlir/Dialect/Math/Transforms/Passes.h
21

Document what is the purpose of this flag? Mention that some approximations are isa specific.

mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp
819

There are few variables that do not match LLVM style guide here: inf_mask, y_approx, etc...

ezhulenev added inline comments.Oct 21 2021, 5:01 AM
mlir/include/mlir/Dialect/Math/Transforms/Passes.h
21

Also maybe useAvx2/mayUseAvx2 to make it more self descriptive.

aartbik added inline comments.Oct 21 2021, 1:22 PM
mlir/include/mlir/Dialect/Math/Transforms/Passes.h
21

Perhaps not super important, but do we really want a target specific issue like "avx' to leak as high as the Math dialect? Can't we give this some math related term on what kind of approximation is used?

mlir/test/Dialect/Math/polynomial-approximation-avx2.mlir
1 ↗(On Diff #381114)

perhaps make a test with "named" CHECK that is run with flag set and not set?
Search for -check-prefix for examples

ezhulenev added inline comments.Oct 21 2021, 1:28 PM
mlir/include/mlir/Dialect/Math/Transforms/Passes.h
21

Well, that's not really "kind of approximation", but "kind of vector instructions you are allowed to use", and to me avx makes sense

cota updated this revision to Diff 381461.Oct 21 2021, 7:49 PM
cota marked 3 inline comments as done.
  • Remove stray comments from a previous iteration
  • Rename flag from avx2 to enableAvx2/enable-avx2
  • Document the enableAvx2 bool in the Options struct
  • Use FileCheck --check-prefix; remove the now-redundant _avx2.mlir file
  • Fix coding style issues
cota added inline comments.Oct 21 2021, 7:50 PM
mlir/test/Dialect/Math/polynomial-approximation-avx2.mlir
1 ↗(On Diff #381114)

Done. Didn't know about this feature, thanks!

cota updated this revision to Diff 381677.Oct 22 2021, 4:01 PM
  • Re-indented test/Dialect/Math/polynomial-approximation-avx2.mlir as suggested by pifon2a.
ezhulenev accepted this revision.Oct 23 2021, 4:47 AM
This revision is now accepted and ready to land.Oct 23 2021, 4:47 AM
This revision was automatically updated to reflect the committed changes.