Hybrid overloading + SCT; not pure SCT?

[oxinabox]

@ChrisRackauckas assured me that ForwardDiff2
was pure source code transformation based approach.
Rather than an overloading based approach.

My thoughts were that is was primarily an overloading based approach that uses limited source code transformation to make it possible to extend the overloads with frules.

Now its a blurry line between the two.
But the place I would cut it is that a pure source code transform approach never calls the original function passing in a special overloaded type. Not even as a fallback.
It always calls a function that it created.
This means it never has method errors,
since it created the function its going to call.

julia> using ForwardDiff2: D
[ Info: Precompiling ForwardDiff2 [994df76e-a4c1-5e1f-bd5c-23b9b5303d4f]

julia> sq(x) = x^2
sq (generic function with 1 method)

julia> D(sq)(10)*1.0
20.0

julia> sq2(x::Float64) = x^2
sq2 (generic function with 1 method)

julia> D(sq2)(10)*1.0
ERROR: MethodError: no method matching sq2(::ForwardDiff2.Dual{ForwardDiff2.Tag{Nothing},Int64,Float64})
Closest candidates are:
  sq(::Float64) at REPL[2]:1
Stacktrace:
 [1] call at /Users/oxinabox/.julia/packages/Cassette/kbN4l/src/context.jl:447 [inlined]
 [2] fallback at /Users/oxinabox/.julia/packages/Cassette/kbN4l/src/context.jl:445 [inlined]
 [3] _overdub_fallback at /Users/oxinabox/.julia/packages/Cassette/kbN4l/src/overdub.jl:486 [inlined]
 [4] _frule_overdub2 at /Users/oxinabox/JuliaEnvs/ForwardDiff2.jl/src/dual_context.jl:145 [inlined]
 [5] alternative at /Users/oxinabox/JuliaEnvs/ForwardDiff2.jl/src/dual_context.jl:186 [inlined]
 [6] #47 at /Users/oxinabox/JuliaEnvs/ForwardDiff2.jl/src/api.jl:61 [inlined]
 [7] overdub(::Cassette.Context{nametype(DualContext),Nothing,Nothing,getfield(ForwardDiff2, Symbol("##PassType#371")),Nothing,Cassette.D
isableHooks}, ::getfield(ForwardDiff2, Symbol("##47#49")){D{Int64,typeof(sq)},Float64}) at /Users/oxinabox/.julia/packages/Cassette/kbN4l
/src/overdub.jl:0
 [8] dualrun(::Function) at /Users/oxinabox/JuliaEnvs/ForwardDiff2.jl/src/dual_context.jl:192
 [9] *(::D{Int64,typeof(sq)}, ::Float64) at /Users/oxinabox/JuliaEnvs/ForwardDiff2.jl/src/api.jl:59
 [10] top-level scope at REPL[3]:1

---

To describe the rewritten function:
I am going to write it without Dual types but you can do it with Dual types, it just means allocating an extra slot in the transformed code and putting stuff into that slot.

We started with:
sq2(x::Float64) = x^2

The new code is:

function do_forward_mode(sq2, x, dsq2, dx)
    sub1_args = (Base.literal_pow, x, 2, Zero(), dx, Zero())
    sub1_res = frule(sub1_args...)
    if sub1_res === nothing
        # call the generated function to make itself
        return do_forward_mode(sub1_args...)
    else
        return sub1_res
    end
end

Doing do a more complicated one:

f(x) = g(x)*h(x)

function do_forward_mode(f, x, df, dx)
    subg_args = (g, x, Zero(), dx)
    subg_res = frule(subg_args...)
    if subg_res === nothing
        # call the generated function to make itself
        subg_res = do_forward_mode(subg_args...)
    end

    subh_args = (h, x, Zero(), dx)
    subh_res = frule(subh_args...)
    if subh_res === nothing
        # call the generated function to make itself
        subh_res = do_forward_mode(subh_args...)
    end

    # Result of `g` is being used, need to split the results
    subg_values, subg_partials = subg_res

    # Result of `h`  is being used, need to split the results
    subh_values, subh_partials = subh_res

    submul_args = (*, subg_values..., subh_values..., Zero(), subg_partials..., subh_partials...)
    submul_res = frule(sub1_args...)
    if submul_res === nothing
        # call the generated function to make itself
        return do_forward_mode(sub1_args...)
    end
    return submul_res
end

So one can see that this kind of transform is complete.
We never call the original function.
We never perform any overloaded operations.
We only call frule and out transform generating function: do_forward_mode.

---

I believe ForwardDiff2 is meant to be using this approach,
and thus bugs have slipped in and caused it to not.

I suggest to make such bugs harder to slip through tests,
that Dual stop subtyping Number and stops defining operations like + and *.
That way we know that it is correctly doing the rewrites.

Alternatively, this may be intentional.
The hybrid overloading + source code transform ForwardDiff2 is using right now
works pretty well if everything is ether a Number or an Array.

but we might need another forward mode AD package that is pure source to source to easily handle concretely typed structs, and code that is otherwise hostile to the overloading approach.

[oxinabox]

Feel free to close this if the hybrid approach is indeed intentional, and the long term plan.

[YingboMa]

Ref: JuliaLabs/Cassette.jl#157

[willtebbutt]

@YingboMa the Cassette PR doesn't address Lyndon's question. It would be helpful to understand whether the long-term goal is to take the kind of SCT approach that Lyndon discusses above, or to continue using various dual types.

[YingboMa]

That PR makes ForwardDiff2 on sq2(x::Float64) = x^2 work and Dual not being a subtype of Number.

[willtebbutt]

Sure. So is the plan to have a single Dual type that can contain arbitrary data types?

[YingboMa]

Dual can already contain arbitrary data types. The plan is to make Dual not visible.

[YingboMa]

Pure source to source is also prone to missing rule errors, so I don't think that is a good idea for the forward mode.

[willtebbutt]

Dual can already contain arbitrary data types. The plan is to make Dual not visible.

Well, yes, kind of, but it only makes sense with subtypes of Real at the minute. As I understand it, you currently have Dual for Reals and DualArray for AbstractArrays.

Pure source to source is also prone to missing rule errors, so I don't think that is a good idea for the forward mode.

Is it? Could you provide an example please?

[YingboMa]

With Dual, we know what functions to differentiate, so it is pretty simple to ignore functions like

methods, code_typed

One can solve this problem with a dependency analysis in source to source AD, but that is just doing more work. Won't the Cassette PR make FD2 as powerful as what pure source to source could be?

[shashi]

Well, yes, kind of, but it only makes sense with subtypes of Real at the minute. As I understand it, you currently have Dual for Reals and DualArray for AbstractArrays.

They could potentially become the same type, yes.

What we do is a hybrid, but we'd love to move to use Tagging in Cassette which would make it source-to-source. But tagging is pretty unusable at the moment, so it will need more work in the Cassette and compiler side.

Cassette's tagging does the same thing as DualArrays for arrays.

Won't the Cassette PR make FD2 as powerful as what pure source to source could be?

Almost. We still won't support structs that have fields restricted to subtypes of Real for example. And

it only makes sense with subtypes of Real at the minute

this problem will show up from time to time when Dual tries to fit into the type hierarchy.

[willtebbutt]

With Dual, we know what functions to differentiate, so it is pretty simple to ignore functions like

Is see. There's only a finite number of such functions though, and you can just define frules on them that are no-ops. This is essentially what Zygote does with @nograd. It seems to work fine and is really easy to fix if you find a new function that breaks stuff.

Moreover, you should generally be able to use Zero to resolve these types of problems for the cases you're considering -- if differential of all of the arguments to a function are Zero then you can just invoke the function as usual and entirely avoid doing forwards mode. The functions you're considering are on Functions, and functions generally have Zero differential.

Also, assuming that we do wind up a world where Duals can contain structs your argument about Duals avoiding these problems no longer holds. e.g.

struct Foo
    a::Float64
end
(foo::Foo)(x) = 5x * foo.a
foo = Dual(Foo(5.0), Composite{Foo}(;a=1.0.))
methods(foo)

An instance of Foo has methods, so calling methods on it is a reasonable thing to do (Zero won't save you here either, of course). In this kind of situation you really do have to think about these types of functions.

What we do is a hybrid, but we'd love to move to use Tagging in Cassette which would make it source-to-source. But tagging is pretty unusable at the moment, so it will need more work in the Cassette and compiler side.

Understood. Yeah, the current state of the Cassette's argument-local metadata propagation makes me sad.

this problem will show up from time to time when Dual tries to fit into the type hierarchy.

Agreed. But I assume that this problem will go away once you've got a generic Dual type that doesn't live in any hierarchy.

[shashi]

Cassette's argument-local metadata propagation makes me sad.

It's super nice when you don't need performance though :)

[oxinabox]

Note the source transform I propose in the to post doesn't need Tagging.
It just creates new code with 2x as many slots as what is is based on, and 2x as many arguments.
It's probably a lot easier to write with IRTools than Cassette.
And even then it's nice to fun to write.
But it is strictly easier than the code in Zygote as there is no need to mess with control flow.

I think hybrid SCT + overloading has some advantages in particular it's less likely to hit compiler edge cases, and one has to write less really annoying IR transforms.
But I think it's got disadvantages in that is harder to handle things to do with user types like arbitrary type constraint's on fields, and on functions, and indeed user provided structs in the first place.
You almost definitely can also handle these things I'm a hybrid overloading approach.
But it's easier in the pure SCT case as they basically fall out of everything you are already doing.