[mlir] Add polynomial approximation for math::Log2

[cota]

Summary

name                     old cpu/op  new cpu/op  delta
BM_mlir_Log2_f32/10       134ns ±15%    45ns ± 4%  -66.39%  (p=0.000 n=20+17)
BM_mlir_Log2_f32/100     1.03µs ±16%  0.12µs ±10%  -88.78%  (p=0.000 n=20+18)
BM_mlir_Log2_f32/1k      10.3µs ±16%   0.7µs ± 5%  -93.24%  (p=0.000 n=20+17)
BM_mlir_Log2_f32/10k      104µs ±15%     7µs ±14%  -93.25%  (p=0.000 n=20+20)
BM_eigen_s_Log2_f32/10   95.3ns ±17%  90.9ns ± 6%     ~     (p=0.228 n=20+18)
BM_eigen_s_Log2_f32/100   907ns ± 3%   911ns ± 6%     ~     (p=0.539 n=16+20)
BM_eigen_s_Log2_f32/1k   9.88µs ± 4%  9.85µs ± 3%     ~     (p=0.790 n=16+17)
BM_eigen_s_Log2_f32/10k   105µs ±10%   110µs ±16%     ~     (p=0.459 n=16+20)
BM_eigen_v_Log2_f32/10   32.5ns ±31%  33.9ns ±14%   +4.31%  (p=0.028 n=17+20)
BM_eigen_v_Log2_f32/100   176ns ± 8%   180ns ± 7%   +2.19%  (p=0.045 n=16+17)
BM_eigen_v_Log2_f32/1k   1.44µs ± 4%  1.50µs ± 9%   +3.91%  (p=0.001 n=16+17)
BM_eigen_v_Log2_f32/10k  14.5µs ±10%  15.0µs ± 8%   +3.92%  (p=0.002 n=16+19)

[cota]

Update D98282: [mlir] Add polynomial approximation for math::Log2

Wrap the benchmark results in `quotes```.

[ezhulenev]

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

272

Template names should be CamelCase just like any other type name.

273

Maybe just inherit template from OpRewritePattern and instantiate log/log2 rewrites from it?

[cota]

s/LOG_OP/Op/

[cota]

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

273

I am having trouble understanding what you are suggesting. Can you please elaborate?

[mehdi_amini]

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

273

Can you document the API? (in particular the base2 parameter).

Also the convention in MLIR is to not put free functions in anonymous namespaces but have them static outside.

mlir/test/Dialect/Math/polynomial-approximation.mlir

12–14

Nit: this check is not useful considering the one above.

13

What is different here compared to the math.log two lines above? Did you mean math.log2 here?

23–25

Similarly, you can omit this check, it is redundant with above.

[ezhulenev]

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

273

template<typename Op>
class LogApproximationBase : public OpRewritePattern<Op>

class Log2Approximation : public LogApproximationBase<math::Log2Op>

[cota]

• Test fixes as suggested by mehdi_amini.
• Log base class as suggested by ezhulenev.

mlir/test/Dialect/Math/polynomial-approximation.mlir

13

Yes, I meant log2. Fixed, thanks!

[cota]

• Fix typo in comment.

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

TanhApproximation::matchAndRewrite(math::TanhOp op,
 
   rewriter.replaceOp(op, res);
 
   return success();
 }
 
 #define LN2_VALUE                                                              \
   0.693147180559945309417232121458176568075500134360255254120680009493393621L
-#define LN2E_VALUE                                                             \
+#define LOG2E_VALUE                                                            \
   1.442695040888963407359924681001892137426645954152985934135449406931109219L
 
 //----------------------------------------------------------------------------//
-// LogOp approximation.
+// LogOp and Log2Op approximation.
 //----------------------------------------------------------------------------//
 
 namespace {
 
 // This approximations comes from the Julien Pommier's SSE math library.
 // Link: http://gruntthepeon.free.fr/ssemath
-struct LogApproximation : public OpRewritePattern<math::LogOp> {
-public:
-  using OpRewritePattern::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(math::LogOp op,
-                                PatternRewriter &rewriter) const final;
-};
-} // namespace
-
-LogicalResult
-LogApproximation::matchAndRewrite(math::LogOp op,
-                                  PatternRewriter &rewriter) const {
+template <typename LOG_OP>
+LogicalResult logMatchAndRewrite(LOG_OP op, PatternRewriter &rewriter,
+                                 bool base2) {
   auto width = vectorWidth(op.operand().getType(), isF32);
   if (!width.hasValue())
     return rewriter.notifyMatchFailure(op, "unsupported operand type");
 
   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
   auto bcast = [&](Value value) -> Value {
     return broadcast(builder, value, *width);
   };
LogicalResult logMatchAndRewrite(LOG_OP op, PatternRewriter &rewriter,
   y2 = builder.create<FmaFOp>(y2, x, cstCephesLogP8);
   y0 = builder.create<FmaFOp>(y0, x3, y1);
   y0 = builder.create<FmaFOp>(y0, x3, y2);
   y0 = builder.create<MulFOp>(y0, x3);
 
   y0 = builder.create<FmaFOp>(cstNegHalf, x2, y0);
   x = builder.create<AddFOp>(x, y0);
 
+  if (base2) {
+    Value cstLog2e = bcast(f32Cst(builder, static_cast<float>(LOG2E_VALUE)));
+    x = builder.create<FmaFOp>(x, cstLog2e, e);
+  } else {
     Value cstLn2 = bcast(f32Cst(builder, static_cast<float>(LN2_VALUE)));
     x = builder.create<FmaFOp>(e, cstLn2, x);
+  }
 
   Value invalidMask =
       builder.create<CmpFOp>(CmpFPredicate::ULT, op.operand(), cstZero);
   Value zeroMask =
       builder.create<CmpFOp>(CmpFPredicate::OEQ, op.operand(), cstZero);
   Value posInfMask =
       builder.create<CmpFOp>(CmpFPredicate::OEQ, op.operand(), cstPosInf);
 
   // Filter out invalid values:
   //  • x == 0     -> -INF
   //  • x < 0      ->  NAN
   //  • x == +INF  -> +INF
   Value aproximation = builder.create<SelectOp>(
       zeroMask, cstMinusInf,
       builder.create<SelectOp>(
           invalidMask, cstNan,
           builder.create<SelectOp>(posInfMask, cstPosInf, x)));
 
   rewriter.replaceOp(op, aproximation);
 
   return success();
 }
 
+struct LogApproximation : public OpRewritePattern<math::LogOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(math::LogOp op,
+                                PatternRewriter &rewriter) const final {
+    return logMatchAndRewrite(op, rewriter, /*base2=*/false);
+  }
+};
+
+struct Log2Approximation : public OpRewritePattern<math::Log2Op> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(math::Log2Op op,
+                                PatternRewriter &rewriter) const final {
+    return logMatchAndRewrite(op, rewriter, /*base2=*/true);
+  }
+};
+} // namespace
+
 //----------------------------------------------------------------------------//
 // Exp approximation.
 //----------------------------------------------------------------------------//
 
 namespace {
 
 struct ExpApproximation : public OpRewritePattern<math::ExpOp> {
 public:
auto mul = [&](Value a, Value b) -> Value {
     return builder.create<MulFOp>(a, b);
   };
   auto sub = [&](Value a, Value b) -> Value {
     return builder.create<SubFOp>(a, b);
   };
   auto floor = [&](Value a) { return builder.create<FloorFOp>(a); };
 
   Value cstLn2 = bcast(f32Cst(builder, static_cast<float>(LN2_VALUE)));
-  Value cstLN2E = bcast(f32Cst(builder, static_cast<float>(LN2E_VALUE)));
+  Value cstLog2E = bcast(f32Cst(builder, static_cast<float>(LOG2E_VALUE)));
 
   // Polynomial coefficients.
   Value cstCephesExpP0 = bcast(f32Cst(builder, 1.0));
   Value cstCephesExpP1 = bcast(f32Cst(builder, 1.0));
   Value cstCephesExpP2 = bcast(f32Cst(builder, 0.49970514590562437052f));
   Value cstCephesExpP3 = bcast(f32Cst(builder, 0.16873890085469545053f));
   Value cstCephesExpP4 = bcast(f32Cst(builder, 0.03668965196652099192f));
   Value cstCephesExpP5 = bcast(f32Cst(builder, 0.01314350012789660196f));
 
   Value x = op.operand();
 
   // Reduced y = x - floor(x / ln(2)) * ln(2) = x - k * ln(2)
-  Value xL2Inv = mul(x, cstLN2E);
+  Value xL2Inv = mul(x, cstLog2E);
   Value kF32 = floor(xL2Inv);
   Value kLn2 = mul(kF32, cstLn2);
   Value y = sub(x, kLn2);
 
   // Use Estrin's evaluation scheme with 3 independent parts:
   // P(y)^y : (c0 + c1 y) + (c2 + c3 y) y^2 + (c4 + c5 y) y^4
   Value y2 = mul(y, y);
   Value y4 = mul(y2, y2);
ExpApproximation::matchAndRewrite(math::ExpOp op,
 
   return success();
 }
 
 //----------------------------------------------------------------------------//
 
 void mlir::populateMathPolynomialApproximationPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
-  patterns.insert<TanhApproximation, LogApproximation, ExpApproximation>(ctx);
+  patterns.insert<TanhApproximation, LogApproximation, Log2Approximation,
+                  ExpApproximation>(ctx);
 }

mlir/test/Dialect/Math/polynomial-approximation.mlir

 // RUN: mlir-opt %s -test-math-polynomial-approximation | FileCheck %s
 
 // Check that all math functions lowered to approximations built from
 // standard operations (add, mul, fma, shift, etc...).
 
 // CHECK-LABEL: @scalar
 func @scalar(%arg0: f32) -> f32 {
   // CHECK-NOT: tanh
   %0 = math.tanh %arg0 : f32
   // CHECK-NOT: log
   %1 = math.log %0 : f32
-  return %1 : f32
+  // CHECK-NOT: log2
+  %2 = math.log %1 : f32
+  return %2 : f32
 }
 
 // CHECK-LABEL: @vector
 func @vector(%arg0: vector<8xf32>) -> vector<8xf32> {
   // CHECK-NOT: tanh
   %0 = math.tanh %arg0 : vector<8xf32>
   // CHECK-NOT: log
   %1 = math.log %0 : vector<8xf32>
-  return %1 : vector<8xf32>
+  // CHECK-NOT: log2
+  %2 = math.log2 %1 : vector<8xf32>
+  return %2 : vector<8xf32>
 }
 
 // CHECK-LABEL: @exp_scalar
 func @exp_scalar(%arg0: f32) -> f32 {
   %0 = math.exp %arg0 : f32
   return %0 : f32
 }
 
 // CHECK-LABEL: @exp_vector
 func @exp_vector(%arg0: vector<8xf32>) -> vector<8xf32> {
   // CHECK-NOT: math.exp
   %0 = math.exp %arg0 : vector<8xf32>
   return %0 : vector<8xf32>
 }

mlir/test/mlir-cpu-runner/math_polynomial_approx.mlir

func @log() {
   // CHECK: -inf, nan, inf, 0.693147
   %special_vec = constant dense<[0.0, -1.0, 0x7f800000, 2.0]> : vector<4xf32>
   %log_special_vec = math.log %special_vec : vector<4xf32>
   vector.print %log_special_vec : vector<4xf32>
 
   return
 }
 
+func @log2() {
+  // CHECK: 3.81887
+  %0 = constant 14.112233 : f32
+  %1 = math.log2 %0 : f32
+  vector.print %1 : f32
+
+  // CHECK: -2, -0.415037, 0, 0.321928
+  %2 = constant dense<[0.25, 0.75, 1.0, 1.25]> : vector<4xf32>
+  %3 = math.log2 %2 : vector<4xf32>
+  vector.print %3 : vector<4xf32>
+
+  // CHECK: -3.32193, -2.32193, -1.73697, -1.32193, -1, -0.736966, -0.514573, -0.321928
+  %4 = constant dense<[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]> : vector<8xf32>
+  %5 = math.log2 %4 : vector<8xf32>
+  vector.print %5 : vector<8xf32>
+
+  // CHECK: -inf
+  %zero = constant 0.0 : f32
+  %log_zero = math.log2 %zero : f32
+  vector.print %log_zero : f32
+
+  // CHECK: nan
+  %neg_one = constant -1.0 : f32
+  %log_neg_one = math.log2 %neg_one : f32
+  vector.print %log_neg_one : f32
+
+  // CHECK: inf
+  %inf = constant 0x7f800000 : f32
+  %log_inf = math.log2 %inf : f32
+  vector.print %log_inf : f32
+
+  // CHECK: -inf, nan, inf, 1.58496
+  %special_vec = constant dense<[0.0, -1.0, 0x7f800000, 3.0]> : vector<4xf32>
+  %log_special_vec = math.log2 %special_vec : vector<4xf32>
+  vector.print %log_special_vec : vector<4xf32>
+
+  return
+}
+
 // -------------------------------------------------------------------------- //
-// Log.
+// Exp.
 // -------------------------------------------------------------------------- //
 func @exp() {
   // CHECK: 2.71828
   %0 = constant 1.0 : f32
   %1 = math.exp %0 : f32
   vector.print %1 : f32
 
   // CHECK: 0.778802, 2.117, 2.71828, 3.85742
func @exp() {
   vector.print %exp_negative_inf : f32
 
   return
 }
 
 func @main() {
   call @tanh(): () -> ()
   call @log(): () -> ()
+  call @log2(): () -> ()
   call @exp(): () -> ()
   return
 }