Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Experiment] Add deferred_with #599

Closed
wants to merge 15 commits into from

Conversation

vchuravy
Copy link
Member

Currently the pattern to embedd Enzyme into a GPU kernel looks something like this:

function kernel(#=....=#)
    # ...
end

function dkernel(#=...=#)
    autodiff_deferred(Reverse, kernel, #=...=#)
end

Internally this would create a job for Enzyme to later process using the deferred compilation
mechanism originally built for CUDA.jl

In #582 I am reworking this to embedd the necessary information into the IR,
instead of relying on a runtime dictionary inside GPUCompiler.

For deferred compilation ala CUDA.jl we have it relativly easy. We just need to process the discovered method-instance
and link it into the output module.

For the Enzyme use-case it becomes more tricky. We can't embed a CompilerJob since that captures a lot of superflous information,
yet we do need to capture enough information to construct an approprate CompilerJob.

Taking a note from JuliaLang/julia#52964 I propose that we add something like

struct EnzymeContext
   mode
   f_activity
   rt_activity
   args_activity
   width
   #...#
end

@wsmoses essentially everything that the NonGenABI requires.

And essentially rename Enzyme.thunk to GPUCompiler.var"gpuc.deferred.with"(EnzymeContext(#=...=#), f, args...)
(actually lookup.with) since that will have the mi. An issue might be ReverseModePrimal/ReverseModeGradient since that
returns more than a single return value.

The question arises, why not just have gpuc.deferred.with and remove gpuc.deferred?
I am not sure, how the worklist integration will look like, but I for me there is a symmetry to invoke_within and invoke
in the Compiler plugins proposal. (I should probably call this gpuc.deferred.within xD )

Opening this up so that I can unload my brain for now and have it percolate a bit.

Copy link

codecov bot commented Jul 16, 2024

Codecov Report

Attention: Patch coverage is 0% with 10 lines in your changes missing coverage. Please review.

Project coverage is 0.00%. Comparing base (3c80a5d) to head (0d5bcaf).
Report is 1 commits behind head on vc/rework_deferred_codegen.

Files Patch % Lines
src/jlgen.jl 0.00% 10 Missing ⚠️
Additional details and impacted files
@@                    Coverage Diff                     @@
##           vc/rework_deferred_codegen    #599   +/-   ##
==========================================================
  Coverage                        0.00%   0.00%           
==========================================================
  Files                              24      24           
  Lines                            3415    3423    +8     
==========================================================
- Misses                           3415    3423    +8     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@vchuravy vchuravy marked this pull request as draft July 16, 2024 19:59
@maleadt maleadt force-pushed the vc/rework_deferred_codegen branch from d2d7a73 to e532812 Compare July 24, 2024 08:23
@vchuravy vchuravy force-pushed the vc/rework_deferred_codegen branch from 927fc30 to b54b5e4 Compare August 9, 2024 13:59
@vchuravy vchuravy force-pushed the vc/rework_deferred_codegen branch 2 times, most recently from 4f14a5e to a018c17 Compare September 26, 2024 07:16
@vchuravy vchuravy closed this Sep 26, 2024
@vchuravy vchuravy deleted the vc/deferred_with branch September 26, 2024 14:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant