Allow for deferred codegen in !toplevel

[vchuravy]

Fixes EnzymeAD/Enzyme.jl#1173

Need to come up with a standalone test-case.

[jgreener64]

Could this get merged?

[wsmoses]

@vchuravy gentle bump

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 73.45%. Comparing base (288b613) to head (0e00885).

❗ Current head 0e00885 differs from pull request most recent head da0e767. Consider uploading reports for the commit da0e767 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #556      +/-   ##
==========================================
- Coverage   82.86%   73.45%   -9.41%     
==========================================
  Files          24       24              
  Lines        3361     3330      -31     
==========================================
- Hits         2785     2446     -339     
- Misses        576      884     +308     

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

[vchuravy]

I tried reproducing the Enzyme MWE with examples/jit.jl but the stars didn't align.

@maleadt do you recall why you limited this to toplevel originally? I assume because it wouldn't make sense for device side launched kernels.

[maleadt]

The reasoning is that deferred functions called from the entrypoint (which is what the compiler driver is architected around) can be discovered recursively without having to expand them during deferred compilation. I guess that this doesn't hold with Enzyme because the deferred compilation jobs aren't called by the main entrypoint? To support this, the compiler driver will need to be changed. For example, if you simply let them expand from within a deferred context as proposed here, the jobs variable at top level won't list all compilation jobs, and the post-processing that happens at top level won't be executing on all functions.

[vchuravy]

Hm... Doesn't the inner call resolve its own deferred codegen first and then both of them get linked into the outer one?

[maleadt]

Yes, but that doesn't correctly populate the jobs variable that's used at top level:

GPUCompiler.jl/src/driver.jl

Lines 423 to 435 in 288b613

	# finish the module
	#
	# we want to finish the module after optimization, so we cannot do so
	# during deferred code generation. instead, process the deferred jobs
	# here.
	if toplevel
	entry = finish_ir!(job, ir, entry)
	for (job′, fn′) in jobs
	job′ == job && continue
	finish_ir!(job′, ir, functions(ir)[fn′])
	end
	end

We could work around this by passing jobs to the inner deferred codegen generators, but that's a hack. Essentially the current system relies on a compilation job having a single point of entry, even into deferred jobs.

[vchuravy]

Hm I am a bit confused how

https://github.com/JuliaGPU/CUDA.jl/blob/c2d444b0f5a76f92c5ba6bc1534a53319218b563/test/core/execution.jl#L974-L1002

works. I just pushed a commit that returns the job variable from the !toplevel compilation.

[maleadt]

Hm I am a bit confused how

https://github.com/JuliaGPU/CUDA.jl/blob/c2d444b0f5a76f92c5ba6bc1534a53319218b563/test/core/execution.jl#L974-L1002

We probably don't really rely on the finish_ir phase right now.

I'm not particularly happy with threading yet more state through, but if you need this... The compiler driver should be reworked to support compiling multiple translation units without relying on a single entry-point.

[vchuravy]

Yeah I can try to improve this while working on #582

[maleadt]

CUDA.jl CI failure looks related.

[vchuravy]

Ah I see for https://github.com/JuliaGPU/CUDA.jl/blob/c2d444b0f5a76f92c5ba6bc1534a53319218b563/test/core/execution.jl#L974-L975

We have an infinite loop since we don't pass through the list of already codegen'd functions. So we recurse into the original top-level and carry on from there.