Faster math, primarily in broadcasts.

Author: chriselrod

src/LoopVectorization.jl

@@ -33,8 +33,8 @@ using Base.Meta: isexpr
 using DocStringExtensions
 import LinearAlgebra # for check_args
 
-using Base.FastMath: add_fast, sub_fast, mul_fast, div_fast, inv_fast, abs2_fast, rem_fast, max_fast, min_fast, pow_fast
-using SLEEFPirates: log_fast, log2_fast, log10_fast, pow
+using Base.FastMath: add_fast, sub_fast, mul_fast, div_fast, inv_fast, abs2_fast, rem_fast, max_fast, min_fast, pow_fast, sqrt_fast
+using SLEEFPirates: log_fast, log2_fast, log10_fast, pow, sin_fast, cos_fast, sincos_fast
 
 using ArrayInterface
 using ArrayInterface: OptionallyStaticUnitRange, OptionallyStaticRange, Zero, One, StaticBool, True, False, reduce_tup, indices, UpTri, LoTri

src/modeling/costs.jl

@@ -207,23 +207,6 @@ const COST = Dict{Symbol,InstructionCost}(
     :vfmsub231 => InstructionCost(4,0.5), # - and * will fuse into this, so much of the time they're not twice as expensive
     :vfnmadd231 => InstructionCost(4,0.5), # + and -* will fuse into this, so much of the time they're not twice as expensive
     :vfnmsub231 => InstructionCost(4,0.5), # - and -* will fuse into this, so much of the time they're not twice as expensive
-    # :vfmadd! => InstructionCost(4,0.5), # + and * will fuse into this, so much of the time they're not twice as expensive
-    # :vfnmadd! => InstructionCost(4,0.5), # + and -* will fuse into this, so much of the time they're not twice as expensive
-    # :vfmsub! => InstructionCost(4,0.5), # + and * will fuse into this, so much of the time they're not twice as expensive
-    # :vfnmsub! => InstructionCost(4,0.5), # + and -* will fuse into this, so much of the time they're not twice as expensive
-    # :vfmadd_fast => InstructionCost(4,0.5), # + and * will fuse into this, so much of the time they're not twice as expensive
-    # :vfmsub_fast => InstructionCost(4,0.5), # - and * will fuse into this, so much of the time they're not twice as expensive
-    # :vfnmadd_fast => InstructionCost(4,0.5), # + and -* will fuse into this, so much of the time they're not twice as expensive
-    # :vfnmsub_fast => InstructionCost(4,0.5), # - and -* will fuse into this, so much of the time they're not twice as expensive
-    # :vfmaddaddone => InstructionCost(4,0.5), # - and -* will fuse into this, so much of the time they're not twice as expensive
-    # :vmullog2 => InstructionCost(4,0.5),
-    # :vmullog2add! => InstructionCost(4,0.5),
-    # :vmullog10 => InstructionCost(4,0.5),
-    # :vmullog10add! => InstructionCost(4,0.5),
-    # :vdivlog2 => InstructionCost(13,4.0,-2.0),
-    # :vdivlog2add! =>InstructionCost(13,4.0,-2.0),
-    # :vdivlog10 => InstructionCost(13,4.0,-2.0),
-    # :vdivlog10add! =>InstructionCost(13,4.0,-2.0),
     :sqrt => InstructionCost(15,4.0,-2.0),
     :sqrt_fast => InstructionCost(15,4.0,-2.0),
     :log => InstructionCost(-3.0, 15, 30, 11),
@@ -242,14 +225,14 @@ const COST = Dict{Symbol,InstructionCost}(
     :sin => InstructionCost(-3,30.0,60.0,23),
     :cos => InstructionCost(-3,27.0,60.0,26),
     :sincos => InstructionCost(-3,37.0,85.0,26),
+    :sincos_fast => InstructionCost(-3,37.0,85.0,26),
     :sinpi => InstructionCost(18,15.0,68.0,23),
     :cospi => InstructionCost(18,15.0,68.0,26),
     :sincospi => InstructionCost(25,37.0,70.0,26),
     :log_fast => InstructionCost(20,20.0,40.0,20),
     :exp_fast => InstructionCost(20,20.0,20.0,18),
     :sin_fast => InstructionCost(18,15.0,68.0,23),
     :cos_fast => InstructionCost(18,15.0,68.0,26),
-    :sincos => InstructionCost(-3,37.0,85.0,26),
     :sinpi_fast => InstructionCost(18,15.0,68.0,23),
     :cospi_fast => InstructionCost(18,15.0,68.0,26),
     :sincospi_fast => InstructionCost(25,22.0,70.0,26),
@@ -502,35 +485,29 @@ end
 isreductcombineinstr(instr::Instruction) = isreductcombineinstr(instr.instr)
 
 const FUNCTIONSYMBOLS = IdDict{Type{<:Function},Instruction}(
-    typeof(+) => :(+),
-    typeof(VectorizationBase.vadd) => :(+),
-    # typeof(VectorizationBase.vadd!) => :(+),
-    typeof(Base.FastMath.add_fast) => :add_fast,
-    typeof(-) => :(-),
-    typeof(VectorizationBase.vsub) => :(-),
-    # typeof(VectorizationBase.vsub!) => :(-),
-    typeof(Base.FastMath.sub_fast) => :sub_fast,
-    typeof(*) => :(*),
-    typeof(VectorizationBase.vmul) => :(*),
-    # typeof(VectorizationBase.vmul!) => :(*),
-    typeof(Base.FastMath.mul_fast) => :mul_fast,
-    typeof(/) => :(/),
-    typeof(^) => :(^),
-    # typeof(VectorizationBase.vfdiv) => :(/),
-    # typeof(VectorizationBase.vfdiv!) => :(/),
-    typeof(VectorizationBase.vdiv) => :(/),
+    typeof(+) => :add_fast,
+    typeof(VectorizationBase.vadd) => :add_fast,
+    typeof(add_fast) => :add_fast,
+    typeof(-) => :sub_fast,
+    typeof(VectorizationBase.vsub) => :sub_fast,
+    typeof(sub_fast) => :sub_fast,
+    typeof(*) => :mul_fast,
+    typeof(VectorizationBase.vmul) => :mul_fast,
+    typeof(mul_fast) => :mul_fast,
+    typeof(/) => :div_fast,
+    typeof(^) => :pow_fast,
+    typeof(VectorizationBase.vdiv) => :div_fast,
     typeof(÷) => :(÷),
-    typeof(Base.FastMath.div_fast) => :div_fast,
-    typeof(Base.FastMath.div_fast) => :div_fast,
-    typeof(Base.FastMath.rem_fast) => :rem_fast,
+    typeof(div_fast) => :div_fast,
+    typeof(rem_fast) => :rem_fast,
     typeof(==) => :(==),
     typeof(!=) => :(!=),
     typeof(isequal) => :isequal,
     typeof(isnan) => :isnan,
     typeof(isinf) => :isinf,
     typeof(isfinite) => :isfinite,
     typeof(abs) => :abs,
-    typeof(abs2) => :abs2,
+    typeof(abs2) => :abs2_fast,
     typeof(abs2_fast) => :abs2_fast,
     typeof(~) => :(~),
     typeof(!) => :(!),
@@ -541,15 +518,15 @@ const FUNCTIONSYMBOLS = IdDict{Type{<:Function},Instruction}(
     typeof(<) => :(<),
     typeof(>=) => :(>=),
     typeof(<=) => :(<=),
-    typeof(inv) => :inv,
+    typeof(inv) => :inv_fast,
     typeof(inv_fast) => :inv_fast,
-    typeof(muladd) => :muladd,
-    typeof(fma) => :fma,
-    typeof(VectorizationBase.vfma) => :vfma,
-    typeof(VectorizationBase.vmuladd) => :vmuladd,
-    typeof(VectorizationBase.vfmsub) => :vfmsub,
-    typeof(VectorizationBase.vfnmadd) => :vfnmadd,
-    typeof(VectorizationBase.vfnmsub) => :vfnmsub,
+    typeof(muladd) => :vmuladd_fast,
+    typeof(fma) => :vfma_fast,
+    typeof(VectorizationBase.vfma) => :vfma_fast,
+    typeof(VectorizationBase.vmuladd) => :vmuladd_fast,
+    typeof(VectorizationBase.vfmsub) => :vfmsub_fast,
+    typeof(VectorizationBase.vfnmadd) => :vfnmadd_fast,
+    typeof(VectorizationBase.vfnmsub) => :vfnmsub_fast,
     typeof(VectorizationBase.vfma_fast) => :vfma_fast,
     typeof(VectorizationBase.vmuladd_fast) => :vmuladd_fast,
     typeof(VectorizationBase.vfmsub_fast) => :vfmsub_fast,
@@ -559,30 +536,12 @@ const FUNCTIONSYMBOLS = IdDict{Type{<:Function},Instruction}(
     typeof(VectorizationBase.vfmsub231) => :vfmsub231,
     typeof(VectorizationBase.vfnmadd231) => :vfnmadd231,
     typeof(VectorizationBase.vfnmsub231) => :vfnmsub231,
-    # typeof(VectorizationBase.vfmadd!) => :vfmadd!,
-    # typeof(VectorizationBase.vfnmadd!) => :vfnmadd!,
-    # typeof(VectorizationBase.vfmsub!) => :vfmsub!,
-    # typeof(VectorizationBase.vfnmsub!) => :vfnmsub!,
-    # typeof(VectorizationBase.vfmadd_fast) => :vfmadd_fast,
-    # typeof(VectorizationBase.vfmsub_fast) => :vfmsub_fast,
-    # typeof(VectorizationBase.vfnmadd_fast) => :vfnmadd_fast,
-    # typeof(VectorizationBase.vfnmsub_fast) => :vfnmsub_fast,
-    # typeof(vfmaddaddone) => :vfmaddaddone,
-    # typeof(vmullog2) => :vmullog2,
-    # typeof(vmullog2add!) => :vmullog2add!,
-    # typeof(vmullog10) => :vmullog10,
-    # typeof(vmullog10add!) => :vmullog10add!,
-    # typeof(vdivlog2) => :vdivlog2,
-    # typeof(vdivlog2add!) => :vdivlog2add!,
-    # typeof(vdivlog10) => :vdivlog10,
-    # typeof(vdivlog10add!) => :vdivlog10add!,
-    typeof(sqrt) => :sqrt,
-    typeof(Base.FastMath.sqrt_fast) => :sqrt,
-    # typeof(VectorizationBase.vsqrt) => :sqrt,
-    typeof(log) => :log,
-    typeof(log2) => :log2,
-    typeof(log10) => :log10,
-    typeof(Base.FastMath.log_fast) => :log,
+    typeof(sqrt) => :sqrt_fast,
+    typeof(sqrt_fast) => :sqrt_fast,
+    typeof(log) => :log_fast,
+    typeof(log2) => :log2_fast,
+    typeof(log10) => :log10_fast,
+    typeof(log_fast) => :log_fast,
     typeof(log1p) => :log1p,
     # typeof(VectorizationBase.vlog) => :log,
     typeof(SLEEFPirates.log) => :log,
@@ -593,20 +552,23 @@ const FUNCTIONSYMBOLS = IdDict{Type{<:Function},Instruction}(
     typeof(expm1) => :expm1,
     # typeof(VectorizationBase.vexp) => :exp,
     typeof(SLEEFPirates.exp) => :exp,
-    typeof(sin) => :sin,
-    typeof(Base.FastMath.sin_fast) => :sin,
-    typeof(SLEEFPirates.sin) => :sin,
-    typeof(cos) => :cos,
-    typeof(Base.FastMath.cos_fast) => :cos,
-    typeof(SLEEFPirates.cos) => :cos,
-    typeof(sincos) => :sincos,
-    typeof(Base.FastMath.sincos_fast) => :sincos,
-    typeof(SLEEFPirates.sincos) => :sincos,
+    typeof(sin) => :sin_fast,
+    typeof(Base.FastMath.sin_fast) => :sin_fast,
+    typeof(SLEEFPirates.sin) => :sin_fast,
+    typeof(cos) => :cos_fast,
+    typeof(Base.FastMath.cos_fast) => :cos_fast,
+    typeof(SLEEFPirates.cos) => :cos_fast,
+    typeof(sincos) => :sincos_fast,
+    typeof(Base.FastMath.sincos_fast) => :sincos_fast,
+    typeof(SLEEFPirates.sincos) => :sincos_fast,
+    typeof(tan) => :tan_fast,
+    typeof(Base.FastMath.sincos_fast) => :sincos_fast,
+    typeof(SLEEFPirates.sincos) => :sincos_fast,
     typeof(Base.tanh) => :tanh,
     typeof(tanh_fast) => :tanh_fast,
     typeof(sigmoid_fast) => :sigmoid_fast,
-    typeof(max) => :max,
-    typeof(min) => :min,
+    typeof(max) => :max_fast,
+    typeof(min) => :min_fast,
     typeof(max_fast) => :max_fast,
     typeof(min_fast) => :min_fast,
     typeof(relu) => :relu,
@@ -618,7 +580,7 @@ const FUNCTIONSYMBOLS = IdDict{Type{<:Function},Instruction}(
     typeof(Base.ifelse) => :ifelse,
     typeof(ifelse) => :ifelse,
     typeof(identity) => :identity,
-    typeof(conj) => :conj
+    typeof(conj) => :identity#conj
     # typeof(zero) => :zero,
     # typeof(one) => :one,
     # typeof(axes) => :axes,

src/vectorizationbase_compat/contract_pass.jl

@@ -50,7 +50,14 @@ function append_args_skip!(call, args, i, mod)
     call
 end
 
-fastfunc(f) = get(VectorizationBase.FASTDICT, f, f)
+function fastfunc(f)
+  i = findfirst(Base.Fix2(===,f), (:sin,:cos,:sincos))
+  if i === nothing
+    get(VectorizationBase.FASTDICT, f, f)
+  else
+    (:sin_fast,:cos_fast,:sincos_fast)[i]
+  end
+end
 function muladd_arguments!(argv, mod, f = first(argv))
     if f === :*
         argv[1] = :mul_fast