use rationalize for literal powers, test

Author: chriselrod

src/parse/add_compute.jl

@@ -472,42 +472,57 @@ function add_pow!(
   elseif x isa Number
     return add_constant!(ls, x ^ p, elementbytes, var)::Operation
   end
-  pint = round(Int, p)
+  local pnum::Int, pden::Int
+  if p isa Integer
+    pnum = Int(p)::Int
+    pden = 1
+  else
+    prational = rationalize(p)
+    @unpack num, den = prational
+    pnum = convert(Int,num)::Int
+    pden = convert(Int,den)::Int
+  end
+  if pden == 1
+    nothing
+  elseif pden == 2
+    if pnum == 1
+      return add_compute!(ls, var, :sqrt, [xop], elementbytes)
+    else
+      xop = add_compute!(ls, gensym!(ls,"root"), :sqrt, [xop], elementbytes)
+    end
+  elseif pden == 3
+    if pnum == 1
+      return add_compute!(ls, var, :cbrt, [xop], elementbytes)
+    else
+      xop = add_compute!(ls, gensym!(ls,"cbroot"), :cbrt, [xop], elementbytes)
+    end
+  elseif pden == 4
+    xop = add_compute!(ls, gensym!(ls,"root"), :sqrt, [xop], elementbytes)
+    if pnum == 1
+      return add_compute!(ls, var, :sqrt, [xop], elementbytes)
+    else
+      xop = add_compute!(ls, gensym!(ls,"root"), :sqrt, [xop], elementbytes)
+    end
+  else
+    pop = add_constant!(ls, p, elementbytes)
+    return add_compute!(ls, var, :^, [xop, pop], elementbytes)
+  end
+  pint = pnum
   if pint == -1
     return add_compute!(ls, var, :inv, [xop], elementbytes)
   elseif pint < 0
     xop = add_compute!(ls, gensym!(ls, "inverse"), :inv, [xop], elementbytes)
-    p = -p
-    pint = - pint
-  end
-  if p == 0.5
-    return add_compute!(ls, var, :sqrt, [xop], elementbytes)
-  elseif p == 1/3
-    return add_compute!(ls, var, :cbrt, [xop], elementbytes)
-  elseif p == 2/3
-    xop = add_compute!(ls, gensym!(ls, "cbrt"), :cbrt, [xop], elementbytes)
-    return add_compute!(ls, var, :abs2_fast, [xop], elementbytes)
-  elseif p == 0.75
-    xop = add_compute!(ls, gensym!(ls, "root1"), :sqrt, [xop], elementbytes)
-    xop = add_compute!(ls, gensym!(ls, "root2"), :sqrt, [xop], elementbytes)
-    pint = 3
-  elseif p == 0.25
-    xop = add_compute!(ls, gensym!(ls, "root1"), :sqrt, [xop], elementbytes)
-    return add_compute!(ls, var, :sqrt, [xop], elementbytes)
-  elseif p != pint
-    pop = add_constant!(ls, p, elementbytes)
-    return add_compute!(ls, var, :^, [xop, pop], elementbytes)
+    pint = -pint
   end
   if pint == 0
     op = Operation(length(operations(ls)), var, elementbytes, LOOPCONSTANT, constant, NODEPENDENCY, Symbol[], NOPARENTS)
     push!(ls.preamble_funcofeltypes, (identifier(op),MULTIPLICATIVE_IN_REDUCTIONS))
     return pushop!(ls, op)
-  elseif pint == 1
+  elseif pint == 1#requires `pden ≠ 1`.
     return add_compute!(ls, var, :identity, [xop], elementbytes)
   elseif pint == 2
     return add_compute!(ls, var, :abs2_fast, [xop], elementbytes)
   end
-
   # Implementation from https://github.com/JuliaLang/julia/blob/a965580ba7fd0e8314001521df254e30d686afbf/base/intfuncs.jl#L216
   t = trailing_zeros(pint) + 1
   pint >>= t

test/special.jl

@@ -263,6 +263,14 @@
             y[i] = x[i] ^ p[i]
         end; y
     end
+  @generated function vpow!(y, x, ::Val{p}) where {p}
+    quote
+      @turbo for i ∈ eachindex(y,x)
+        y[i] = x[i] ^ $p
+      end
+      return y
+    end
+  end
 
     function csetanh!(y, z, x)
         for j in axes(x, 2)
@@ -427,7 +435,12 @@
         @test vpowf!(r1, x) ≈ (r2 .= x .^ 2.3)
         @test vpowf!(r1, x, -1.7) ≈ (r2 .= x .^ -1.7)
         p = randn(length(x));
-        @test vpowf!(r1, x, x) ≈ (r2 .= x .^ x)
+      @test vpowf!(r1, x, x) ≈ (r2 .= x .^ x)
+      @test vpow!(r1, x, Val(0.75)) ≈ (r2 .= x .^ 0.75)
+      @test vpow!(r1, x, Val(2/3)) ≈ (r2 .= x .^ (2/3))
+      @test vpow!(r1, x, Val(0.5)) == (r2 .= sqrt.(x))
+      @test vpow!(r1, x, Val(1/4)) ≈ (r2 .= x .^ (1/4))
+      @test vpow!(r1, x, Val(4.5)) ≈ (r2 .= x .^ 4.5)
 
         X = rand(T, N, M); Z = rand(T, N, M);
         Y1 = similar(X); Y2 = similar(Y1);