JuliaGPU / Adapt.jl

Other
86 stars 24 forks source link

Specialize adapt_structure for small tuples #81

Closed charleskawczynski closed 3 months ago

charleskawczynski commented 3 months ago

Closes #79.

codecov[bot] commented 3 months ago

Codecov Report

Attention: Patch coverage is 87.50000% with 1 lines in your changes are missing coverage. Please review.

Project coverage is 92.85%. Comparing base (9c1b146) to head (5e85c7c).

Files Patch % Lines
src/base.jl 87.50% 1 Missing :warning:
Additional details and impacted files ```diff @@ Coverage Diff @@ ## master #81 +/- ## ========================================== - Coverage 93.65% 92.85% -0.80% ========================================== Files 6 6 Lines 63 70 +7 ========================================== + Hits 59 65 +6 - Misses 4 5 +1 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

maleadt commented 3 months ago

Can't you do this with a simple if, assuming it would get compiled away, instead of inserting many methods?

This also reminds me of https://github.com/JuliaLang/julia/pull/53665.

charleskawczynski commented 3 months ago

I was going to suggest that we could use an if-statement, is that your preference?

maleadt commented 3 months ago

Yes, you're now putting additional pressure on the dispatch system where the compiler should be able to handle this just fine.

charleskawczynski commented 3 months ago

Alright, I've updated the PR. This passes inference for me in the #79 example. Does 20 qualify as "small", @maleadt?

charleskawczynski commented 3 months ago

bump 🙂

maleadt commented 3 months ago

I don't like that we don't have a test for this, but OK I guess.

Seeing how you went with the if approach, I take it using afold (which handles this in Base, https://github.com/JuliaLang/julia/blob/bc7ba3d5c8b2dab1c0e19537739b67c2da902d11/base/operators.jl#L579-L625) didn't solve the allocations?

charleskawczynski commented 3 months ago

I don't like that we don't have a test for this, but OK I guess.

Seeing how you went with the if approach, I take it using afold (which handles this in Base, https://github.com/JuliaLang/julia/blob/bc7ba3d5c8b2dab1c0e19537739b67c2da902d11/base/operators.jl#L579-L625) didn't solve the allocations?

Thank you for merging!

Yeah, I agree about the tests. I did try to write a simple reproducer, but I wasn't able to recreate the failure. I'm not sure what ingredients are needed. I'll take another look to see if I can cook something up.

I tend to try to avoid using Julia internals, but I can check to see if afold works if you prefer that over the existing solution.

maleadt commented 3 months ago

I tend to try to avoid using Julia internals

Yeah, that's fair.

charleskawczynski commented 3 months ago

I've updated, and I'll try to continue updating/simplifying, the reproducer in #79. It's drastically simpler now, although there's still a lot more to go.

charleskawczynski commented 3 months ago

Thanks again for fixing this! This fixed 955 inference failures per timestep 😬 found by JET! 🚀