Improve Intrinsics (#218)

* Improve error message

* Link error for missing builtins to issue

* Bump patch

* Check errors get thrown for missing builtin / intrinsic

* atomic_pointerset rule implementation

* Verify that _unsafe_copyto! works with abstract element types

Project.toml

@@ -1,7 +1,7 @@
 name = "Tapir"
 uuid = "07d77754-e150-4737-8c94-cd238a1fb45b"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.2.30"
+version = "0.2.31"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/rrules/builtins.jl

@@ -27,9 +27,9 @@ function rrule!!(f::CoDual{<:Core.Builtin}, args...)
  "types. In order to fix this problem, you will either need to modify your code " *
  "to avoid hitting this built-in function, or implement a method of `rrule!!` " *
  "which is specialised to this case. " *
- "Either way, please consider opening an issue at " *
- "https://github.com/compintell/Tapir.jl/ so that the issue can be fixed more " *
- "widely."
+ "Either way, please consider commenting on " *
+ "https://github.com/compintell/Tapir.jl/issues/208/ so that the issue can be " *
+ "fixed more widely."
  ))
 end
 
@@ -43,6 +43,18 @@ import ..Tapir:
  tangent_type, increment!!, @is_primitive, MinimalCtx, is_primitive, NoFData,
  zero_rdata, NoRData, tuple_map, fdata, NoRData, rdata, increment_rdata!!, zero_fcodual
 
+using Core.Intrinsics: atomic_pointerref
+
+struct MissingIntrinsicWrapperException <: Exception
+ msg::String
+end
+
+function translate(f)
+ msg = "Unable to translate the intrinsic $f into a regular Julia function. " *
+ "Please see github.com/compintell/Tapir.jl/issues/208 for more discussion."
+ throw(MissingIntrinsicWrapperException(msg))
+end
+
 # Note: performance is not considered _at_ _all_ in this implementation.
 function rrule!!(f::CoDual{<:Core.IntrinsicFunction}, args...)
  return rrule!!(CoDual(translate(Val(primal(f))), tangent(f)), args...)
@@ -105,7 +117,25 @@ end
 # atomic_pointermodify
 # atomic_pointerref
 # atomic_pointerreplace
-# atomic_pointerset
+
+@intrinsic atomic_pointerset
+function rrule!!(::CoDual{typeof(atomic_pointerset)}, p::CoDual{<:Ptr}, x::CoDual, order)
+ _p = primal(p)
+ _order = primal(order)
+ old_value = atomic_pointerref(_p, _order)
+ old_tangent = atomic_pointerref(tangent(p), _order)
+ dp = tangent(p)
+ function atomic_pointerset_pullback!!(::NoRData)
+ dx_r = atomic_pointerref(dp, _order)
+ atomic_pointerset(_p, old_value, _order)
+ atomic_pointerset(dp, old_tangent, _order)
+ return NoRData(), NoRData(), rdata(dx_r), NoRData()
+ end
+ atomic_pointerset(_p, primal(x), _order)
+ atomic_pointerset(dp, zero_tangent(primal(x)), _order)
+ return p, atomic_pointerset_pullback!!
+end
+
 # atomic_pointerswap
 
 @intrinsic bitcast
@@ -745,7 +775,10 @@ function generate_hand_written_rrule!!_test_cases(rng_ctor, ::Val{:builtins})
  # atomic_pointermodify -- NEEDS IMPLEMENTING AND TESTING
  # atomic_pointerref -- NEEDS IMPLEMENTING AND TESTING
  # atomic_pointerreplace -- NEEDS IMPLEMENTING AND TESTING
- # atomic_pointerset -- NEEDS IMPLEMENTING AND TESTING
+ (
+ true, :none, nothing,
+ IntrinsicsWrappers.atomic_pointerset, CoDual(p, dp), 1.0, :monotonic,
+ ),
  # atomic_pointerswap -- NEEDS IMPLEMENTING AND TESTING
  (false, :stability, nothing, IntrinsicsWrappers.bitcast, Int64, 5.0),
  (false, :stability, nothing, IntrinsicsWrappers.bswap_int, 5),
@@ -1030,6 +1063,10 @@ function generate_derived_rrule!!_test_cases(rng_ctor, ::Val{:builtins})
  (false, :none, nothing, getindex, randn(5), [1, 2, 2]),
  (false, :none, nothing, setindex!, randn(5), [4.0, 5.0], [1, 1]),
  (false, :none, nothing, setindex!, randn(5), [4.0, 5.0, 6.0], [1, 2, 2]),
+ (
+ false, :none, nothing,
+ Base._unsafe_copyto!, fill!(Matrix{Real}(undef, 5, 4), 1.0), 3, randn(10), 2, 4,
+ ),
  ]
  memory = Any[]
  return test_cases, memory

test/rrules/builtins.jl

@@ -24,6 +24,18 @@
 
  TestUtils.run_rrule!!_test_cases(StableRNG, Val(:builtins))
 
+ # Unhandled built-in throws an intelligible error.
+ @test_throws(
+ Tapir.MissingRuleForBuiltinException,
+ invoke(Tapir.rrule!!, Tuple{CoDual{<:Core.Builtin}}, zero_fcodual(getfield)),
+ )
+
+ # Unhandled intrinsic throws an intelligible error.
+ @test_throws(
+ Tapir.IntrinsicsWrappers.MissingIntrinsicWrapperException,
+ invoke(Tapir.IntrinsicsWrappers.translate, Tuple{Any}, Val(:foo)),
+ )
+
  @testset "Disable bitcast to differentiable type" begin
  @test_throws(
  ArgumentError,