Handle xlogy limit (EnzymeAD#1615)

* Handle xlogy limit * with test * fixup
MilesCranmer · Jul 11, 2024 · ff9d320 · ff9d320
1 parent c83fcf8
commit ff9d320
Show file tree

Hide file tree

Showing 7 changed files with 86 additions and 111 deletions.
diff --git a/Project.toml b/Project.toml
@@ -18,11 +18,13 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [extensions]
 EnzymeChainRulesCoreExt = "ChainRulesCore"
+EnzymeLogExpFunctionsExt = "LogExpFunctions"
 EnzymeSpecialFunctionsExt = "SpecialFunctions"
 EnzymeStaticArraysExt = "StaticArrays"
 
@@ -33,6 +35,7 @@ EnzymeCore = "0.7.5"
 Enzyme_jll = "0.0.133"
 GPUCompiler = "0.21, 0.22, 0.23, 0.24, 0.25, 0.26"
 LLVM = "6.1, 7, 8"
+LogExpFunctions = "0.3"
 ObjectFile = "0.4"
 Preferences = "1.4"
 SpecialFunctions = "1, 2"
@@ -41,5 +44,6 @@ julia = "1.6"
 
 [extras]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
diff --git a/ext/EnzymeLogExpFunctionsExt.jl b/ext/EnzymeLogExpFunctionsExt.jl
@@ -0,0 +1,10 @@
+module EnzymeLogExpFunctionsExt
+
+using LogExpFunctions
+using Enzyme
+
+function __init__()
+ Enzyme.Compiler.known_ops[typeof(LogExpFunctions.xlogy)] = (:xlogy_jl, 2, nothing)
+end
+
+end
diff --git a/src/compiler.jl b/src/compiler.jl
@@ -101,6 +101,58 @@ Dict{DataType, Tuple{Symbol, Int, Union{Nothing, Tuple{Symbol, DataType}}}}(
 @static if VERSION >= v"1.8.0"
  known_ops[typeof(Base.fma_emulated)] = (:fma, 3, nothing)
 end
+@inline function find_math_method(@nospecialize(func), sparam_vals)
+ if func ∈ keys(known_ops)
+ name, arity, toinject = known_ops[func]
+ Tys = (Float32, Float64)
+
+ if length(sparam_vals) == arity
+ T = first(sparam_vals)
+ legal = T ∈ Tys
+
+ if legal
+ if name == :ldexp
+ if !(sparam_vals[2] <: Integer)
+ legal = false
+ end
+ elseif name == :pow
+ if sparam_vals[2] <: Integer
+ name = :powi
+ elseif sparam_vals[2] != T
+ legal = false
+ end
+ elseif name == :jl_rem2pi
+ else
+ if !all(==(T), sparam_vals)
+ legal = false
+ end
+ end
+ end
+ if legal
+ return name, toinject, T
+ end
+ end
+ end 
+
+ if func ∈ keys(cmplx_known_ops)
+ name, arity, toinject = cmplx_known_ops[func]
+ Tys = (Complex{Float32}, Complex{Float64})
+ if length(sparam_vals) == arity
+ T = first(sparam_vals)
+ legal = T ∈ Tys
+
+ if legal
+ if !all(==(T), sparam_vals)
+ legal = false
+ end
+ end
+ if legal
+ return name, toinject, T
+ end
+ end
+ end
+ return nothing, nothing, nothing
+end
 
 const nofreefns = Set{String}((
  "ijl_f_isdefined", "jl_f_isdefined",
@@ -5621,61 +5673,8 @@ function GPUCompiler.codegen(output::Symbol, job::CompilerJob{<:EnzymeTarget};
  end
  continue
  end
-
- @inline function find_math_method()
- if func ∈ keys(known_ops)
- name, arity, toinject = known_ops[func]
- Tys = (Float32, Float64)
-
- if length(sparam_vals) == arity
- T = first(sparam_vals)
- legal = T ∈ Tys
-
- if legal
- if name == :ldexp
- if !(sparam_vals[2] <: Integer)
- legal = false
- end
- elseif name == :pow
- if sparam_vals[2] <: Integer
- name = :powi
- elseif sparam_vals[2] != T
- legal = false
- end
- elseif name == :jl_rem2pi
- else
- if !all(==(T), sparam_vals)
- legal = false
- end
- end
- end
- if legal
- return name, toinject, T
- end
- end
- end 
-
- if func ∈ keys(cmplx_known_ops)
- name, arity, toinject = cmplx_known_ops[func]
- Tys = (Complex{Float32}, Complex{Float64})
- if length(sparam_vals) == arity
- T = first(sparam_vals)
- legal = T ∈ Tys
-
- if legal
- if !all(==(T), sparam_vals)
- legal = false
- end
- end
- if legal
- return name, toinject, T
- end
- end
- end
- return nothing, nothing, nothing
- end
 
- name, toinject, T = find_math_method()
+ name, toinject, T = find_math_method(func, sparam_vals)
  if name === nothing
  continue
  end

diff --git a/src/compiler/interpreter.jl b/src/compiler/interpreter.jl
@@ -108,65 +108,11 @@ function is_primitive_func(@nospecialize(TT))
  if ft == typeof(Enzyme.pmap)
  return true
  end
- if ft === typeof(Base.rem2pi)
- if TT <: Tuple{ft, Float32, <:Any} || TT <: Tuple{ft, Float64, <:Any} || TT <: Tuple{ft, Float16, <:Any}
- return true
- end
- end
-
- if ft == typeof(Base.inv) || ft == typeof(Base.sqrt)
- if TT <: Tuple{ft, Complex{Float32}} || TT <: Tuple{ft, Complex{Float64}}
- return true
- end
- end
-
- @static if VERSION >= v"1.9-"
- if ft === typeof(Base.rem)
- if TT <: Tuple{ft, Float32, Float32} || TT <: Tuple{ft, Float64, Float64}
- return true
- end
- end
+ match = Enzyme.Compiler.find_math_method(ft, TT.parameters[2:end])[1]
+ if match !== nothing
+ return true
  end
 
- if ft === typeof(Base.cbrt) || ft === typeof(Base.sin) || ft === typeof(Base.cos) ||
- ft === typeof(Base.sinc) ||
- ft === typeof(Base.tan) || ft === typeof(Base.exp) || ft === typeof(Base.FastMath.exp_fast) ||
- ft === typeof(Base.exp10) ||
- ft === typeof(Base.exp2) ||
- ft === typeof(Base.expm1) ||
- ft === typeof(Base.log) || ft === typeof(Base.FastMath.log) ||
- ft === typeof(Base.log1p) ||
- ft === typeof(Base.log2) ||
- ft === typeof(Base.log10) ||
- ft === typeof(Base.asin) ||
- ft === typeof(Base.acos) ||
- ft === typeof(Base.atan) ||
- ft === typeof(Base.sinpi) ||
- ft === typeof(Base.cospi) ||
- ft === typeof(Base.sinh) || ft === typeof(Base.FastMath.sinh_fast) ||
- ft === typeof(Base.cosh) || ft === typeof(Base.FastMath.cosh_fast) ||
- ft === typeof(Base.tanh) || ft === typeof(Base.FastMath.tanh_fast) ||
- ft === typeof(Base.sqrt) || ft === typeof(Base.sincos) || ft === typeof(Base.sincospi)
- if TT <: Tuple{ft, Float32} || TT <: Tuple{ft, Float64} || TT <: Tuple{ft, Float16}
- return true
- end
- end
-@static if VERSION < v"1.8.0"
-else
- if ft === typeof(Base.fma_emulated)
- if TT <: Tuple{ft, Float32, Float32, Float32} || TT <: Tuple{ft, Float64, Float64, Float64}
- return true
- end
- end
-end
- if ft === typeof(Base.:^) || ft === typeof(Base.atan)
- if TT <: Tuple{ft, Float32, Float32} || TT <: Tuple{ft, Float64, Float64}
- return true
- end
- if TT <: Tuple{ft, Float32, <:Integer} || TT <: Tuple{ft, Float64, <:Integer}
- return true
- end
- end
  # FIXME(@wsmoses): For which types should we not inline?
  if ft === typeof(Base.wait) || ft === typeof(Base._wait) || ft === typeof(Base.enq_work) ||
  ft === typeof(Base.Threads.threadid) || ft == typeof(Base.Threads.nthreads) ||

diff --git a/test/Project.toml b/test/Project.toml
@@ -11,6 +11,7 @@ InlineStrings = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 LLVM_jll = "86de99a1-58d6-5da7-8064-bd56ce2e322c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

diff --git a/test/ext/logexpfunctions.jl b/test/ext/logexpfunctions.jl
@@ -0,0 +1,14 @@
+using LogExpFunctions
+
+
+xlogydiff(x) = xlogy(x[1], 23.0)
+@testset "LogExpFunctions" begin
+
+ x = [0.0]
+
+ grad_forward = Enzyme.gradient(Enzyme.Forward, xlogydiff, x)
+ grad_reverse = Enzyme.gradient(Enzyme.Reverse, xlogydiff, x)
+
+ @test grad_forward[1] ≈ log(23.0) 
+ @test grad_reverse[1] ≈ log(23.0) 
+end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -3511,6 +3511,7 @@ end
  @testset "ChainRulesCore ext" begin
  include("ext/chainrulescore.jl")
  end
+ include("ext/logexpfunctions.jl")
 end