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.

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