improve inference for getproperty #909

simeonschaub · 2021-02-23T23:22:43Z

This has regressed quite a bit due to #848. With this PR, we should be able to get back the same performance as before, assuming there is no custom implementation or pullback for getproperty. Still need to add tests.

willtebbutt · 2021-02-24T17:19:42Z

This is cool.

assuming there is no custom implementation or pullback for getproperty

Could you elaborate a bit on this? Why is this necessary? I'm sure it could be ascertained from the PR itself, but I could use a helping hand 😂

DhairyaLGandhi · 2021-02-24T17:23:38Z

I am still going through the pr, it'd be good to get an idea if the generated function is too low level for Julia ir so as to break with newer releases.

simeonschaub · 2021-02-24T17:29:56Z

assuming there is no custom implementation or pullback for getproperty

Could you elaborate a bit on this? Why is this necessary? I'm sure it could be ascertained from the PR itself, but I could use a helping hand 😂

Zygote's previous assumption was that getproperty could always be treated just like getfield, unless a custom adjoint was defined. This is obviously wrong if getproperty is overloaded, but it of course opens up a lot more opportunities for optimization. 😄 In the case where a user has overloaded getproperty, there is not really a way around recursing into the actual definition, but this PR special-cases the case where getproperty just calls the fallback definition, so we can basically retain the previous behavior.

willtebbutt · 2021-02-24T17:32:48Z

Ohhh I see. This makes sense. Will comment on the PR itself.

willtebbutt · 2021-02-24T17:35:42Z

src/lib/lib.jl

+    is_getfield_fallback = which(_pullback, pb_sig(x)) == which(_pullback, pb_sig(Any)) && 
+        which(getproperty, sig(x)) == which(getproperty, sig(Any))
+
+    if is_getfield_fallback


Another high-level question, what's the reason that this entire case can't be written as

# we've ensured that we're fine to short-circuit and call `getfield` directly, return :(Zygote._pullback(cx, literal_getfield, x, f))

rather than pulling up the IR for that method and pasting it here as is currently proposed?

The main reason is that we can't (at least that I know of) manually attach backedges to regular expressions, so we need to return a CodeInfo object. We could return a CodeInfo that just contains that call, but if we already go to lengths like this, we might as well just directly inline the whole thing.

This also makes sense to me.

simeonschaub · 2021-02-24T17:37:05Z

I am still going through the pr, it'd be good to get an idea if the generated function is too low level for Julia ir so as to break with newer releases.

It is definitely fairly low-level, so it is not impossible this might cause problems with future Julia versions. It is basically the same mechanism IRTools and Cassette already rely upon, so while I'd definitely prefer a higher-level solution, I think it is a necessary evil if we want the inference improvements. At least in my testing, more aggressive inline annotations didn't really help.

src/lib/lib.jl

pevnak · 2021-03-07T11:37:40Z

This is awesome. I have been tracking the regression. I was originally suspicious to Julia 1.6. and reported here
JuliaLang/julia#39586
which many people confirmed. Bud digging deeper, it seems to me that the problem is the getproperty

This was my MWE

using Flux
using BenchmarkTools
x = randn(Float32, 10, 10);
m = Dense(10,10);
ps = Flux.params(m);
f(w,b,x) = w * x .+ b
w, b = m.W, m.b

@btime gradient(() -> sum(m(x)), ps)
@btime gradient(() -> sum(f(w,b,x)), ps)

which on Zygote 0.5.17 with Flux 11.3 runs as

julia> @btime gradient(() -> sum(m(x)), ps)
  5.785 μs (47 allocations: 4.81 KiB)
Grads(...)

julia> @btime gradient(() -> sum(f(w,b,x)), ps)
  6.145 μs (52 allocations: 4.72 KiB)
Grads(...)

But on [email protected], Flux v0.11.6 it is fairly slow

julia> @btime gradient(() -> sum(m(x)), ps)
  25.046 μs (102 allocations: 7.39 KiB)
Grads(...)

julia> @btime gradient(() -> sum(f(w,b,x)), ps)
  6.247 μs (58 allocations: 4.91 KiB)
Grads(...)

and testing this PR with Flux v0.11.5 gives

julia> @btime gradient(() -> sum(m(x)), ps)
  5.895 μs (54 allocations: 5.02 KiB)
Grads(...)

julia> @btime gradient(() -> sum(f(w,b,x)), ps)
  6.361 μs (58 allocations: 4.91 KiB)
Grads(...)

So this is awesome. Can it be merged and tagged?

simeonschaub · 2021-03-07T12:10:23Z

Yes, sorry, I had other stuff to do and didn't come around to finishing this. I will try to finish this up some time later today.

willtebbutt · 2021-04-22T05:35:02Z

@simeonschaub how's this going? :)

simeonschaub · 2021-04-26T20:25:24Z

Please don't merge this yet, it currently breaks custom adjoints for getproperty.

simonmandlik · 2021-05-27T10:50:21Z

Is there a chance of this being merged in the near future? :)

simeonschaub · 2021-05-27T21:36:00Z

Yes, feel free to bug me about it from time to time, which probably increases the chances of me eventually finishing it. 😄 I will try to get to it soon again but every time I work on it, I discover some new issue...

DhairyaLGandhi · 2021-06-25T07:19:29Z

bump :)

1681: Support NamedTuples for Chain + Parallel r=mcabbott a=mcabbott Closes #1680, WIP. Todo list includes: - [x] add Parallel too - [ ] ~~worry about whether any of this will upset Zygote, like FluxML/Zygote.jl#909 or, kick that can down the road. - [x] add tests Co-authored-by: Michael Abbott <[email protected]>

1681: Support NamedTuples for Chain + Parallel r=DhairyaLGandhi a=mcabbott Closes #1680, WIP. Todo list includes: - [x] add Parallel too - [ ] ~~worry about whether any of this will upset Zygote, like FluxML/Zygote.jl#909 or, kick that can down the road. - [x] add tests Co-authored-by: Michael Abbott <[email protected]>

1681: Support NamedTuples for Chain + Parallel r=darsnack a=mcabbott Closes #1680, WIP. Todo list includes: - [x] add Parallel too - [ ] ~~worry about whether any of this will upset Zygote, like FluxML/Zygote.jl#909 or, kick that can down the road. - [x] add tests Co-authored-by: Michael Abbott <[email protected]>

willtebbutt · 2021-09-24T10:48:09Z

@simeonschaub bump :)

This has regressed quite a bit due to FluxML#848. With this PR, we should be able to get back the same performance as before, assuming there is no custom implementation or pullback for `getproperty`. Still need to add tests.

simeonschaub · 2021-10-05T15:25:30Z

Ok, sorry for taking so long to getting around for this. I have pushed some more changes and now the only broken case left is the backedges for overloading Zygote._pullback(::AContext, ::typeof(getproperty), x, ::Symbol). I am beginning to suspect this might actually be a bug in Base, but it may be enough of an edge case that we could go forward with this nevertheless.

simeonschaub · 2021-10-05T15:32:09Z

Any ideas why the GH actions jobs are not running?

DhairyaLGandhi · 2021-10-05T15:39:09Z

I see they're going

willtebbutt

Thanks for pushing this forward. A couple of test-related comments.

test/compiler.jl

willtebbutt

Thanks for the fantastic work @simeonschaub . I think it just needs a patch bump, then it'll be good to go.

willtebbutt · 2021-10-14T17:06:51Z

@simeonschaub can we bors this?

simeonschaub · 2021-10-14T17:10:23Z

Sure, please go ahead!

willtebbutt · 2021-10-14T17:16:24Z

bors r+

bors · 2021-10-14T17:41:29Z

Build succeeded:

buildkite/zygote-dot-jl

willtebbutt reviewed Feb 24, 2021

View reviewed changes

src/lib/lib.jl Outdated Show resolved Hide resolved

willtebbutt reviewed Feb 24, 2021

View reviewed changes

src/lib/lib.jl Outdated Show resolved Hide resolved

DhairyaLGandhi mentioned this pull request Mar 22, 2021

accumulation of gradients #922

Closed

willtebbutt mentioned this pull request Apr 23, 2021

Optimise posteiror marginals and Zygote JuliaGaussianProcesses/TemporalGPs.jl#67

Merged

agerlach mentioned this pull request May 7, 2021

[WIP] ExpectationProblem interface SciML/SciMLExpectations.jl#55

Merged

9 tasks

pevnak mentioned this pull request Jul 2, 2021

The gradients utilizing getproperty takes ages #1014

Open

mcabbott mentioned this pull request Jul 29, 2021

Support NamedTuples for Chain + Parallel FluxML/Flux.jl#1681

Merged

3 tasks

This was referenced Oct 2, 2021

getproperty performance workaround #1091

Closed

Possible to write rules for methods not collections of methods? JuliaDiff/ChainRulesCore.jl#471

Open

Type instability with structs #1094

Closed

WIP: improve inference for getproperty

a3f8dc4

This has regressed quite a bit due to FluxML#848. With this PR, we should be able to get back the same performance as before, assuming there is no custom implementation or pullback for `getproperty`. Still need to add tests.

simeonschaub force-pushed the sds/getprop branch from 24bcd83 to a3f8dc4 Compare October 5, 2021 15:19

simeonschaub closed this Oct 5, 2021

simeonschaub reopened this Oct 5, 2021

simeonschaub marked this pull request as ready for review October 5, 2021 15:31

willtebbutt reviewed Oct 6, 2021

View reviewed changes

test/compiler.jl Show resolved Hide resolved

test/compiler.jl Outdated Show resolved Hide resolved

address review comments

bdbb369

simeonschaub force-pushed the sds/getprop branch from 3d56f79 to bdbb369 Compare October 12, 2021 16:21

simeonschaub changed the title ~~WIP: improve inference for getproperty~~ improve inference for getproperty Oct 12, 2021

willtebbutt approved these changes Oct 12, 2021

View reviewed changes

simeonschaub added 2 commits October 12, 2021 12:26

Merge remote-tracking branch 'upstream/master' into sds/getprop

812b8c8

bump patch version

1d189fb

bors bot merged commit 5887e46 into FluxML:master Oct 14, 2021

willtebbutt mentioned this pull request Oct 14, 2021

getproperty type-stability work-around FluxML/ZygoteRules.jl#21

Closed

simeonschaub deleted the sds/getprop branch October 14, 2021 17:53

mcabbott mentioned this pull request Oct 16, 2021

Insert _project into getproperty's gradient, and then improve z2d etc. to restore stability #1104

Merged

ToucheSir mentioned this pull request Jan 10, 2022

Zygote's compilation scales badly with the number of ~ statements TuringLang/Turing.jl#1754

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

improve inference for getproperty #909

improve inference for getproperty #909

simeonschaub commented Feb 23, 2021

willtebbutt commented Feb 24, 2021

DhairyaLGandhi commented Feb 24, 2021

simeonschaub commented Feb 24, 2021

willtebbutt commented Feb 24, 2021

willtebbutt Feb 24, 2021 •

edited

Loading

simeonschaub Feb 24, 2021

willtebbutt Feb 24, 2021

simeonschaub commented Feb 24, 2021

pevnak commented Mar 7, 2021

simeonschaub commented Mar 7, 2021

willtebbutt commented Apr 22, 2021

simeonschaub commented Apr 26, 2021

simonmandlik commented May 27, 2021 •

edited

Loading

simeonschaub commented May 27, 2021

DhairyaLGandhi commented Jun 25, 2021

willtebbutt commented Sep 24, 2021

simeonschaub commented Oct 5, 2021

simeonschaub commented Oct 5, 2021

DhairyaLGandhi commented Oct 5, 2021

willtebbutt left a comment

willtebbutt left a comment

willtebbutt commented Oct 14, 2021

simeonschaub commented Oct 14, 2021

willtebbutt commented Oct 14, 2021

bors bot commented Oct 14, 2021

improve inference for getproperty #909

improve inference for getproperty #909

Conversation

simeonschaub commented Feb 23, 2021

willtebbutt commented Feb 24, 2021

DhairyaLGandhi commented Feb 24, 2021

simeonschaub commented Feb 24, 2021

willtebbutt commented Feb 24, 2021

willtebbutt Feb 24, 2021 • edited Loading

Choose a reason for hiding this comment

simeonschaub Feb 24, 2021

Choose a reason for hiding this comment

willtebbutt Feb 24, 2021

Choose a reason for hiding this comment

simeonschaub commented Feb 24, 2021

pevnak commented Mar 7, 2021

simeonschaub commented Mar 7, 2021

willtebbutt commented Apr 22, 2021

simeonschaub commented Apr 26, 2021

simonmandlik commented May 27, 2021 • edited Loading

simeonschaub commented May 27, 2021

DhairyaLGandhi commented Jun 25, 2021

willtebbutt commented Sep 24, 2021

simeonschaub commented Oct 5, 2021

simeonschaub commented Oct 5, 2021

DhairyaLGandhi commented Oct 5, 2021

willtebbutt left a comment

Choose a reason for hiding this comment

willtebbutt left a comment

Choose a reason for hiding this comment

willtebbutt commented Oct 14, 2021

simeonschaub commented Oct 14, 2021

willtebbutt commented Oct 14, 2021

bors bot commented Oct 14, 2021

willtebbutt Feb 24, 2021 •

edited

Loading

simonmandlik commented May 27, 2021 •

edited

Loading