Interpreter fixes (prerequisite to CUDA support)

[wsmoses]

and also should fix the nested gradient overlay, and ideally improve compile times (but not necessarily time to first @ compile ]

[github-actions]

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

JuliaFormatter

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/utils.jl

Lines 286 to 288 in ff729ce

	mi = ccall(:jl_specializations_get_linfo, Ref{Core.MethodInstance},
	(Any, Any, Any), match.method, match.spec_types, match.sparams)

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/utils.jl

Line 290 in ff729ce

frame = Core.Compiler.InferenceState(result, #=cache_mode=#:local, interp)

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/utils.jl

Lines 307 to 346 in ff729ce

	opt = Core.Compiler.OptimizationState(frame, interp)
	caller = frame.result
	@static if VERSION < v"1.11-"
	ir = Core.Compiler.run_passes(opt.src, opt, caller)
	else
	ir = Core.Compiler.run_passes_ipo_safe(opt.src, opt, caller)
	Core.Compiler.ipo_dataflow_analysis!(interp, ir, caller)
	end
	# Rewrite type unstable calls to recurse into call_with_reactant to ensure
	# they continue to use our interpreter. Reset the derived return type
	# to Any if our interpreter would change the return type of any result.
	# Also rewrite invoke (type stable call) to be :call, since otherwise apparently
	# screws up type inference after this (TODO this should be fixed).
	any_changed = false
	for (i, inst) in enumerate(ir.stmts)
	@static if VERSION < v"1.11"
	changed, next = rewrite_inst(inst[:inst], ir)
	Core.Compiler.setindex!(ir.stmts[i], next, :inst)
	else
	changed, next = rewrite_inst(inst[:stmt], ir)
	Core.Compiler.setindex!(ir.stmts[i], next, :stmt)
	end
	if changed
	any_changed = true
	Core.Compiler.setindex!(ir.stmts[i], Any, :type)
	end
	end
	Core.Compiler.finish(interp, opt, ir, caller)
	src = Core.Compiler.ir_to_codeinf!(opt)
	# Julia hits various internal errors trying to re-perform type inference
	# on type infered code (that we undo inference of), if there is no type unstable
	# code to be rewritten, just use the default methodinstance (still using our methodtable),
	# to improve compatibility as these bugs are fixed upstream.
	if !any_changed
	src = Core.Compiler.retrieve_code_info(mi, world)
	@show "post non change", src
	end

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/utils.jl

Line 362 in ff729ce

code_info.slotnames = Any[:call_with_reactant, REDUB_ARGUMENTS_NAME, code_info.slotnames...]

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/utils.jl

Line 372 in ff729ce

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/utils.jl

Line 389 in ff729ce

actual_argument = Expr(:call, Core.GlobalRef(Core, :getfield), overdub_args_slot, offset)

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/utils.jl

Lines 394 to 396 in ff729ce

	#push!(overdubbed_code, actual_argument)
	push!(fn_args, Core.SSAValue(length(overdubbed_code)))

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/utils.jl

Lines 405 to 406 in ff729ce

	pop!(overdubbed_codelocs)
	pop!(fn_args)

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/utils.jl

Line 410 in ff729ce

push!(overdubbed_code, Expr(:call, Core.GlobalRef(Core, :getfield), overdub_args_slot, offset - 1))

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/utils.jl

Lines 415 to 417 in ff729ce

	push!(overdubbed_code, Expr(:(=), Core.SlotNumber(n_method_args + n_prepended_slots), trailing_arguments))
	push!(overdubbed_codelocs, code_info.codelocs[1])
	push!(fn_args, Core.SSAValue(length(overdubbed_code)))

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/utils.jl

Lines 423 to 424 in ff729ce

	arg_partially_inline!(code_info.code, fn_args, method.sig, Any[static_params...],
	n_prepended_slots, n_prepended_slots, length(overdubbed_code), :propagate)

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/src/utils.jl

Line 445 in ff729ce

$(Expr(:meta, :generated, call_with_reactant_generator))

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/cuda.jl

Line 21 in ff729ce

@show @code_hlo optimize=false square!(A)

[wsmoses]

Currently I want to ensure this doesn't break anything.

Then I will potentially revamp the Enzyme overload to use more of this (and already this should hopefully fix the use of the Enzyme overload in type unstable code), as well as CUDA and perhaps pythoncall stuff

[wsmoses]

....and naturally I segfault julia. @aviatesk @vtjnash @gbaraldi @vchuravy @topolarity do you have any thoughts here?

[mofeing]

now that we're adding it, should we rename it to @reactant_overlay? because it's just a partial macro to @overlay

[mofeing]

I think we're putting too much into "utils.jl"? Maybe this code makes more sense in "Compiler.jl"?

Also, with these fixes we'll be able to write diff rules from the Julia side?

[wsmoses]

Closing in favor of #365