Get pullback without running primal pass for @scalar_rules

[GiggleLiu]

julia> @benchmark cos(x) setup=(x=0.5)
BenchmarkTools.Trial: 
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     4.409 ns (0.00% GC)
  median time:      4.421 ns (0.00% GC)
  mean time:        4.530 ns (0.00% GC)
  maximum time:     101.587 ns (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     1000

julia> @benchmark rrule(sin, x)[2](1.0) setup=(x=0.5)
BenchmarkTools.Trial: 
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     7.840 ns (0.00% GC)
  median time:      8.102 ns (0.00% GC)
  mean time:        8.157 ns (0.00% GC)
  maximum time:     35.264 ns (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     999

Because the macro @scalar_rule generates the following code.

julia> MacroTools.prettify(@macroexpand @scalar_rule sinc(x) cosc(x))
quote
    if !(sinc isa ChainRulesCore.Type) && ChainRulesCore.fieldcount(ChainRulesCore.typeof(sinc)) > 0
        ChainRulesCore.throw(ChainRulesCore.ArgumentError("@scalar_rule cannot be used on closures/functors (such as $(sinc))"))
    end
    function (ChainRulesCore.ChainRulesCore).frule((ChainRulesCore._, Δ1), ::ChainRulesCore.typeof(sinc), x::Number)
        Ω = sinc(x)
        nothing
        return (Ω, cosc(x) * Δ1)
    end
    function (ChainRulesCore.ChainRulesCore).rrule(::ChainRulesCore.typeof(sinc), x::Number)
        Ω = sinc(x)
        nothing
        return (Ω, function sinc_pullback(gull)
                    return (ChainRulesCore.NO_FIELDS, ChainRulesCore.conj(cosc(x)) * gull)
                end)
    end
end

In order to make the code more friendly to packages that want to make use of these scalar rules. The following generated code might be better?

julia> MacroTools.prettify(@macroexpand @scalar_rule sinc(x) cosc(x))
quote
    if !(sinc isa ChainRulesCore.Type) && ChainRulesCore.fieldcount(ChainRulesCore.typeof(sinc)) > 0
        ChainRulesCore.throw(ChainRulesCore.ArgumentError("@scalar_rule cannot be used on closures/functors (such as $(sinc))"))
    end
    function (ChainRulesCore.ChainRulesCore).frule((ChainRulesCore._, Δ1), ::ChainRulesCore.typeof(sinc), x::Number)
        Ω = sinc(x)
        nothing
        return (Ω, cosc(x) * Δ1)
    end

    function scalar_pullback(::ChainRulesCore.typeof(sinc), x::Number)
           function sinc_pullback(gull)
                return (ChainRulesCore.NO_FIELDS, ChainRulesCore.conj(cosc(x)) * gull)
           end
    end

    function (ChainRulesCore.ChainRulesCore).rrule(f::ChainRulesCore.typeof(sinc), x::Number)
        Ω = sinc(x)
        nothing
        return (Ω, scalar_pullback(f, x))
    end
end

Note: there are limited number of functions like + that do not need to know the value of x. We can define them separately.

[oxinabox]

julia> @benchmark cos(x) setup=(x=0.5)
...
julia> @benchmark rrule(sin, x)[2](1.0) setup=(x=0.5)
...

Those are very different benchmarks.
rrule(sin,x)[2](1.0) needs to peform not just cos(x) but also a getindex and 1.0*x

A more correct benchmark is:

julia> const x = 0.5
0.5

 julia> @benchmark $(rrule(sin, x)[2])(d) setup=(d=2)
BenchmarkTools.Trial:
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     4.705 ns (0.00% GC)
  median time:      5.174 ns (0.00% GC)
  mean time:        5.299 ns (0.00% GC)
  maximum time:     30.412 ns (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     1000

julia> @benchmark $(dd->cos(x)*dd)(d) setup=(d=2)
BenchmarkTools.Trial:
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     1.977 ns (0.00% GC)
  median time:      2.030 ns (0.00% GC)
  mean time:        2.140 ns (0.00% GC)
  maximum time:     16.229 ns (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     1000

Which does still show some overhead.
I am not sure what exactly is up with that,
looking at the code it generates with @code_warntype and everything looks great.

---

I tried benchmarking you code for sinc against what ChainRules currently has.
I don't see why your code would be faster, but perhaps i am missing something?
They look functionally identical to me.
The benchmarks also show they perform identically.

Code from https://github.com/JuliaDiff/ChainRules.jl/issues/289#issue-723498328

julia> @benchmark $(rrule(sinc, x)[2])(d) setup=(d=2)
BenchmarkTools.Trial:
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     22.660 ns (0.00% GC)
  median time:      24.313 ns (0.00% GC)
  mean time:        24.998 ns (0.00% GC)
  maximum time:     109.339 ns (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     996

Current master

julia> @benchmark $(rrule(sinc, x)[2])(d) setup=(d=2)
BenchmarkTools.Trial:
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     20.062 ns (0.00% GC)
  median time:      21.683 ns (0.00% GC)
  mean time:        22.289 ns (0.00% GC)
  maximum time:     110.750 ns (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     996

Looking at the LLVM code they generate they are identical (not posting both since identical)

julia> @code_llvm (rrule(sinc, x)[2])(0.2)
;  @ REPL[13]:2 within `sinc_pullback'
define { double } @julia_sinc_pullback_2316([1 x double]* nocapture nonnull readonly align 8 dereferenceable(8) %0, double %1) {
top:
; └
;  @ REPL[13]:3 within `sinc_pullback'
  %2 = getelementptr inbounds [1 x double], [1 x double]* %0, i64 0, i64 0
; ┌ @ special/trig.jl:1093 within `cosc'
   %3 = load double, double* %2, align 8
   %4 = call double @j__cosc_2318(double %3)
; └
; ┌ @ float.jl:405 within `*'
   %5 = fmul double %4, %1
; └
  %.fca.0.insert = insertvalue { double } undef, double %5, 0
  ret { double } %.fca.0.insert
}

[GiggleLiu]

In some case, people might want the gradient function df/dx, without running the forward code.

It is impossible to extract the gradient function without running forward if the code is generated like this.

NOTE: I find this issue when I try to make use of ChainRules in NiLang.

[oxinabox]

Oh that is a very different question to what I thought you were asking.

Its not possible in general to this.
It is possible in many cases to do this, including most (maybe all) current scalar rules.

It could be handy, it is a thing DiffRules.jl used to do.

It would prevent some optimizations for sure, and it also might be weird to have some functions offer this and others not.

I guess the bigger worry is that people writting traditional ADs might end up relying on it a lot and write themselves into a corner.
Nabla.jl relies on that kind of pattern a lot and it causes problems for it.
NiLang is not a traditional AD though; so that fear probably does not apply.

[GiggleLiu]

Yeah, I can imagine that sometimes it will cause performance issues. Maybe a better solution is creating a new macro?

e.g.

@naive_scalar_rule sin(x) cos(x)

[oxinabox]

Maybe a better solution is creating a new macro?

I think for anything that can be defined using @scalar_rule it is possible to do this.
Because @scalar_rule doesn't let you change how the primal is computed.

[GiggleLiu]

Maybe a better solution is creating a new macro?

I think for anything that can be defined using @scalar_rule it is possible to do this.
Because @scalar_rule doesn't let you change how the primal is computed.

That's fair. Then it is a matter of how people use this macro, a kind of convention like

naive_grad(::typeof(sin), x) = cos(x)

for f in [:sin, :cos, ...]
    @eval @scalar_rule $f(x) naive_grad($f, x)
end

[shashi]

What is the proposed solution? I think just being able to pass in Omega is simple and workable?

[oxinabox]

@shashi proposed a very similar need in #245
for ModelingToolkit.
ModellingToolkit, like NiLang is different to a standard AD.

@shashi proposed in #245 a different API, that was adding a Ω keyword argument to frule that defaults to computing the primal result.
A few problems with that is that it would:

• Make it impossible to write frules for functions that have kwargs called Ω.
• Cause ModelingToolkit to deal with a inconsistent API for frule since some rules will accept that, and some won't. Since some can't.

We note importantly that it is not possible to write all rules in a form that it is OK to only have access to the final value off the primal pass.
Sometimes you want to run a different primal computation, to get key extra output for computing the derivative at same time (examples: (computing sincos to do sin and cos; or the extended ODE thing));
and sometimes you want to capture intermediate state (example: handling getindex where the user may pass in a mask array or a index object, its much easier (and slightly more efficient) to do the translation to plain numerical indicies once and use that to compute primal and reuse it in derivative computation)

However, there is a large subset of rules that can be written in such away as to not need to touch how the primal is computed.
and a subset of them (and in particular the subset that ModellingToolkit is using already via DiffRules.jl) are the functions defined via @scalar_rule.
since @scalar_rule explicitly doesn't allow the computation of the primal to be changed, only used.

I like @GiggleLui's idea that make @scalar_rule create a get_propagator function of somekind, that is used inside the created frule/rrule.
Probably a separate get_pushforward and get_pullback

It was basically where we expose propagator function that is internally called with-in the rule.
It would look more like:

function get_pushforward(Ω, ::typeof(hypot), z::Complex)
    (_, Δz) -> _realconjtimes(z, Δz) / ifelse(iszero(Ω), one(Ω), Ω)
end

function frule((_, Δz), ::typeof(hypot), z::Complex)
    ∂Ω = _realconjtimes(z, Δz) / ifelse(iszero(Ω), one(Ω), Ω)
    return Ω, get_pushforward(Ω, hypot, z)
end

Not sure on the exact API and where the Ω should be inserted.
Possibly even get_pushforward(hypot, z)(Ω)((_, Δz)) ?

For some things that have more efficient frules that directly benifit e.g. from the fusing of the primal computation with the pushforward computation, (like sin via sincos)
we would still be able to manually write the get_pushforward/get_pullback that would be less efficient but useful to NiLang and ModellingToolkit.

Would also allow for the thing that I think is a bad idea that @YingboMa wanted to do with trying to generate higher-order rules using AD.
But it might also allow for trying generate higher order rules using symbolic differentiation, which I am less dubious about.
(Still dubious, but hey researchers exist to prove me wrong. Both could be great.)