JuliaDiff / JuliaDiff/ReverseDiff.jl

Record `Broadcast.broadcasted` instead of `Broadcast.broadcast`

Open
#215 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Julia
Stars
393
Forks
60
Avg merge
18h 24m
Merged PRs (30d)
8

Description

IIUC ReverseDiff records broadcast and uses ForwardDiff to specialize further on broadcasted statements whenever possible, leading to much better performance than if one were to trace through all the operations using ReverseDiff.TrackedReal.

Unfortunately this means that once one tries to make use of Broadcast.broadcasted, i.e. lazy broadcasting, this is not recorded and we end up taking the less desirable path of tracing through the broadcast using ReverseDiff.TrackedReal:

julia> using ReverseDiff

julia> f(x) = sum(exp.(x))
f (generic function with 1 method)

julia> f_tape = ReverseDiff.GradientTape(f, (rand(10, ),))
typename(ReverseDiff.GradientTape)(f)

julia> g(x) = sum(Broadcast.instantiate(Broadcast.broadcasted(exp, x)))
g (generic function with 1 method)

julia> g_tape = ReverseDiff.GradientTape(g, (rand(10, ),))
typename(ReverseDiff.GradientTape)(g)

julia> length(g_tape.tape)
19

julia> g_tape.tape
19-element Vector{ReverseDiff.AbstractInstruction}:
 ScalarInstruction(exp):
  input:  TrackedReal<76u>(0.5057895559423533, 0.0, 6vX, 1, 5d7)
  output: TrackedReal<AW2>(1.6582943198420965, 0.0, 6vX, ---)
  cache:  Base.RefValue{Float64}(1.6582943198420965)
 ScalarInstruction(exp):
  input:  TrackedReal<KUi>(0.13213345349262395, 0.0, 6vX, 2, 5d7)
  output: TrackedReal<JmS>(1.141260614319831, 0.0, 6vX, ---)
  cache:  Base.RefValue{Float64}(1.141260614319831)
 ScalarInstruction(+):
  input:  (TrackedReal<AW2>(1.6582943198420965, 0.0, 6vX, ---),
           TrackedReal<JmS>(1.141260614319831, 0.0, 6vX, ---))
  output: TrackedReal<Cng>(2.7995549341619275, 0.0, 6vX, ---)
  cache:  Base.RefValue{StaticArraysCore.SVector{2, Float64}}([1.0, 1.0])
 ScalarInstruction(exp):
  input:  TrackedReal<GdN>(0.034478177830953305, 0.0, 6vX, 3, 5d7)
  output: TrackedReal<DNm>(1.03507944045113, 0.0, 6vX, ---)
  cache:  Base.RefValue{Float64}(1.03507944045113)
 ScalarInstruction(+):
  input:  (TrackedReal<Cng>(2.7995549341619275, 0.0, 6vX, ---),
           TrackedReal<DNm>(1.03507944045113, 0.0, 6vX, ---))
  output: TrackedReal<DEh>(3.8346343746130573, 0.0, 6vX, ---)
  cache:  Base.RefValue{StaticArraysCore.SVector{2, Float64}}([1.0, 1.0])
 ScalarInstruction(exp):
  input:  TrackedReal<CEJ>(0.04867133616730335, 0.0, 6vX, 4, 5d7)
  output: TrackedReal<GVI>(1.0498752380105207, 0.0, 6vX, ---)
  cache:  Base.RefValue{Float64}(1.0498752380105207)
 ScalarInstruction(+):
  input:  (TrackedReal<DEh>(3.8346343746130573, 0.0, 6vX, ---),
           TrackedReal<GVI>(1.0498752380105207, 0.0, 6vX, ---))
  output: TrackedReal<7HD>(4.884509612623578, 0.0, 6vX, ---)
  cache:  Base.RefValue{StaticArraysCore.SVector{2, Float64}}([1.0, 1.0])
 ScalarInstruction(exp):
  input:  TrackedReal<Ft8>(0.8637862831888328, 0.0, 6vX, 5, 5d7)
  output: TrackedReal<5lH>(2.3721252497772487, 0.0, 6vX, ---)
  cache:  Base.RefValue{Float64}(2.3721252497772487)
 ScalarInstruction(+):
  input:  (TrackedReal<7HD>(4.884509612623578, 0.0, 6vX, ---),
           TrackedReal<5lH>(2.3721252497772487, 0.0, 6vX, ---))
  output: TrackedReal<HpK>(7.256634862400826, 0.0, 6vX, ---)
  cache:  Base.RefValue{StaticArraysCore.SVector{2, Float64}}([1.0, 1.0])
 ScalarInstruction(exp):
  input:  TrackedReal<Gmo>(0.0039196786165185404, 0.0, 6vX, 6, 5d7)
  output: TrackedReal<1R0>(1.0039273706035023, 0.0, 6vX, ---)
  cache:  Base.RefValue{Float64}(1.0039273706035023)
 ScalarInstruction(+):
  input:  (TrackedReal<HpK>(7.256634862400826, 0.0, 6vX, ---),
           TrackedReal<1R0>(1.0039273706035023, 0.0, 6vX, ---))
  output: TrackedReal<8hX>(8.260562233004329, 0.0, 6vX, ---)
  cache:  Base.RefValue{StaticArraysCore.SVector{2, Float64}}([1.0, 1.0])
 ScalarInstruction(exp):
  input:  TrackedReal<L8R>(0.9153223594101434, 0.0, 6vX, 7, 5d7)
  output: TrackedReal<4qL>(2.4975802406432295, 0.0, 6vX, ---)
  cache:  Base.RefValue{Float64}(2.4975802406432295)
 ScalarInstruction(+):
  input:  (TrackedReal<8hX>(8.260562233004329, 0.0, 6vX, ---),
           TrackedReal<4qL>(2.4975802406432295, 0.0, 6vX, ---))
  output: TrackedReal<17T>(10.758142473647558, 0.0, 6vX, ---)
  cache:  Base.RefValue{StaticArraysCore.SVector{2, Float64}}([1.0, 1.0])
 ScalarInstruction(exp):
  input:  TrackedReal<JYS>(0.15063946146751517, 0.0, 6vX, 8, 5d7)
  output: TrackedReal<6gL>(1.1625774285521715, 0.0, 6vX, ---)
  cache:  Base.RefValue{Float64}(1.1625774285521715)
 ScalarInstruction(+):
  input:  (TrackedReal<17T>(10.758142473647558, 0.0, 6vX, ---),
           TrackedReal<6gL>(1.1625774285521715, 0.0, 6vX, ---))
  output: TrackedReal<Dvu>(11.92071990219973, 0.0, 6vX, ---)
  cache:  Base.RefValue{StaticArraysCore.SVector{2, Float64}}([1.0, 1.0])
 ScalarInstruction(exp):
  input:  TrackedReal<B28>(0.3010502862135006, 0.0, 6vX, 9, 5d7)
  output: TrackedReal<5Ku>(1.3512772904478805, 0.0, 6vX, ---)
  cache:  Base.RefValue{Float64}(1.3512772904478805)
 ScalarInstruction(+):
  input:  (TrackedReal<Dvu>(11.92071990219973, 0.0, 6vX, ---),
           TrackedReal<5Ku>(1.3512772904478805, 0.0, 6vX, ---))
  output: TrackedReal<2ZC>(13.27199719264761, 0.0, 6vX, ---)
  cache:  Base.RefValue{StaticArraysCore.SVector{2, Float64}}([1.0, 1.0])
 ScalarInstruction(exp):
  input:  TrackedReal<Cpk>(0.02748173107946794, 0.0, 6vX, 10, 5d7)
  output: TrackedReal<9uF>(1.0278628369912384, 0.0, 6vX, ---)
  cache:  Base.RefValue{Float64}(1.0278628369912384)
 ScalarInstruction(+):
  input:  (TrackedReal<2ZC>(13.27199719264761, 0.0, 6vX, ---),
           TrackedReal<9uF>(1.0278628369912384, 0.0, 6vX, ---))
  output: TrackedReal<J8N>(14.299860029638849, 0.0, 6vX, ---)
  cache:  Base.RefValue{StaticArraysCore.SVector{2, Float64}}([1.0, 1.0])

julia> @benchmark ReverseDiff.gradient($f, $(randn(1000)))
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (min … max):   9.308 μs …  1.218 ms  ┊ GC (min … max): 0.00% … 95.95%
 Time  (median):     10.105 μs              ┊ GC (median):    0.00%
 Time  (mean ± σ):   12.275 μs ± 36.408 μs  ┊ GC (mean ± σ):  9.10% ±  3.05%

  ▇██▇▆▅▅▄▃▁                                       ▁▁▁▁▁▁     ▂
  ██████████▇▅▄▅▄▄▃▃▄▅▅▅▄▅▄▄▁▃▄▁▄▃▃▁▁▁▁▁▁▃▄▃▁▃▄▅▅▆████████▇▇▇ █
  9.31 μs      Histogram: log(frequency) by time      27.8 μs <

 Memory estimate: 55.83 KiB, allocs estimate: 15.

julia> @benchmark ReverseDiff.gradient($g, $(randn(1000)))
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (min … max):  151.875 μs …   3.504 ms  ┊ GC (min … max):  0.00% … 94.23%
 Time  (median):     158.858 μs               ┊ GC (median):     0.00%
 Time  (mean ± σ):   183.434 μs ± 253.436 μs  ┊ GC (mean ± σ):  12.42% ±  8.45%

  ▄▆▆█▇▇▅▄▃▂▁▁▁                                                 ▂
  ██████████████▇▆▇▆▆▇▆▆▅▅▄▅▆▃▄▄▃▁▄▄▃▁▃▄▁▁▁▁▁▁▄▃▃▅▄▄▅▄▄▃▅▅▃▆▆▆▆ █
  152 μs        Histogram: log(frequency) by time        259 μs <

 Memory estimate: 374.53 KiB, allocs estimate: 8009.

The overhead can of course be lowered if the tape is compiled:

julia> x = randn(1000);

julia> inputs = (x,); results = (similar(x),); cfg = ReverseDiff.GradientConfig(inputs);

julia> g_tape = ReverseDiff.GradientTape(g, inputs);

julia> compiled_g_tape = ReverseDiff.compile(g_tape)
typename(ReverseDiff.CompiledTape)(g)

julia> @benchmark ReverseDiff.gradient!($results, $compiled_g_tape, $inputs)
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (min … max):  51.798 μs … 117.899 μs  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     54.679 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   55.396 μs ±   4.165 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

    ▃▂▅▁▁▇█▂ ▃▃       ▁              ▂    ▁▁                   ▂
  ▇▇████████▄███▆▇▇▆▆▇█▆▄▄▇▇▄▃▃▆▇▃▁▁██▆▁▁▃██▄▁▄▃▆▅▁▅▃▁▅▅▅▅▃▃▄▅ █
  51.8 μs       Histogram: log(frequency) by time      73.1 μs <

 Memory estimate: 0 bytes, allocs estimate: 0.

but it's still slower than ForwardDiff (this of course varies wrt. input size, etc. but I'm guessing this perf difference is well-established given that basically all reverse-AD frameworks in Julia make use of ForwardDiff for broadcasting).

Would it be possible to record broadcasted instead of broadcast in ReverseDiff.jl (this is the way it's done in Zygote.jl)?

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reproducing the issue's Julia examples with ReverseDiff.GradientTape, ReverseDiff.compile, gradient, and gradient!, then inspect how broadcast and broadcasted are recorded. Done means lazy Broadcast.broadcasted expressions are recorded without falling back to per-element ReverseDiff.TrackedReal tracing, while gradients remain correct and the reported performance gap is reduced.

Written by the indexing model from the issue text.

Assessment

Domain
performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.