Make fastmath aware of Base.literal_pow.

JuliaLang · Jan 7, 2021 · d988153 · d988153
1 parent 6c42190
commit d988153
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/base/fastmath.jl b/base/fastmath.jl
@@ -118,6 +118,9 @@ function make_fastmath(expr::Expr)
                 $arrvar[$(indvars...)] = $op($arrvar[$(indvars...)], $rhs)
             end
         end
+    elseif is_literal_int_pow(expr)
+        expr = Expr(expr.head, :(Base.literal_pow),
+            expr.args[1], expr.args[2], Val(expr.args[3]))
     end
     Base.exprarray(make_fastmath(expr.head), Base.mapany(make_fastmath, expr.args))
 end
@@ -155,6 +158,10 @@ macro fastmath(expr)
     make_fastmath(esc(expr))
 end
 
+function is_literal_int_pow(ex::Expr)
+    return (ex.head === :call && length(ex.args) == 3 &&
+        ex.args[1] === :^ && isa(ex.args[3], Integer))
+end
 
 # Basic arithmetic
 

diff --git a/test/fastmath.jl b/test/fastmath.jl
@@ -249,3 +249,10 @@ end
     @test (@fastmath "a" * "b") == "ab"
     @test (@fastmath "a" ^ 2) == "aa"
 end
+
+struct LitPowTest end
+Base.literal_pow(::typeof(Base.FastMath.pow_fast), ::LitPowTest, ::Val{p}) where {p} = 1
+Base.literal_pow(::typeof(^), ::LitPowTest, ::Val{p}) where {p} = 2
+LPT = LitPowTest()
+@test (@fastmath LPT^2) == 1
+@test LPT^2 == 2