codegen: add a pass for late conversion of known modify ops to call atomicrmw #57010

vtjnash · 2025-01-10T01:02:22Z

The ExpandAtomicModify can recognize our pseudo-intrinsic julia.atomicmodify and convert it into some of known atomicrmw expressions, or simplify it with more inlining, as applicable. Ideally we could get this pass upstreamed, since there's nothing specific to julia about this pass, and LLVM's IR cannot express this quite correctly without making it a new intrinsic.

This ensures that now our @atomic modify is as fast as Threads.Atomic!

julia> @code_llvm Threads.atomic_add!(r, 10)
; Function Signature: atomic_add!(Base.Threads.Atomic{Int64}, Int64)
;  @ atomics.jl:307 within `atomic_add!`
define i64 @"julia_atomic_add!_2680"(ptr noundef nonnull align 8 dereferenceable(8) %"x::Atomic", i64 signext %"v::Int64") #0 {
top:
; ┌ @ Base_compiler.jl:94 within `modifyproperty!`
   %0 = atomicrmw add ptr %"x::Atomic", i64 %"v::Int64" acq_rel, align 8
; └
; ┌ @ Base_compiler.jl:54 within `getproperty`
   ret i64 %0
; └
}

base/atomics.jl

gbaraldi · 2025-01-10T02:28:50Z

Why is there such a large aotcompile.cpp diff? Looks unrelated

src/aotcompile.cpp

vtjnash · 2025-01-13T21:48:47Z

Why is there such a large aotcompile.cpp diff? Looks unrelated

It is mostly support code for this, since if you don't have the IR emitted for the target, then it has to use a loop&call, which isn't want people want to see, so we need to make sure to emit the target code too into all of the correct compile units. We could land some of it separately, but it wouldn't do anything (be mostly untested) until this landed

oscardssmith · 2025-01-27T20:20:08Z

Is this ready to merge? Edit: no. CI seems pretty unhappy with this currently.

DilumAluthge · 2025-01-27T21:38:15Z

Also there's a merge conflict.

vtjnash · 2025-01-28T16:05:13Z

Please review, but hold off on merging. I think there might be one other PR I try to merge first

oscardssmith · 2025-01-28T16:07:03Z

CI failures look real: Specifically the following seems to be a multithreading test that breaks with this PR:

threads_exec.jl with JULIA_NUM_THREADS == 1,0: Test Failed at /cache/build/tester-demeter6-15/julialang/julia-master/julia-67ad4b675e/share/julia/test/threads_exec.jl:444
  Expression: accum[] + Int(var[]) === sum(0:nloops)
   Evaluated: 129821317250305000 === 500500

vchuravy · 2025-02-10T11:28:27Z

src/codegen.cpp

@@ -9476,7 +9552,7 @@ static jl_llvm_functions_t
    Instruction &prologue_end = ctx.builder.GetInsertBlock()->back();

    // step 11a. Emit the entry safepoint
-    if (JL_FEAT_TEST(ctx, safepoint_on_entry))
+    if (params.safepoint_on_entry && JL_FEAT_TEST(ctx, safepoint_on_entry))


This looks weird, shouldn't ctx be derived from params?

ctx derives from params.params not params (I agree the naming choice here was a choice)

This broke GPUCompiler.jl's tests, together with

julia/src/aotcompile.cpp

Lines 797 to 798 in 9b63fd9

if (jl_ir_inlining_cost((jl_value_t*)src) < UINT16_MAX)

params.safepoint_on_entry = false; // ensure we don't block ExpandAtomicModifyPass from inlining this code if applicable

... which seems like a very ad hoc way of avoiding safepoints? It also introduces an inconsistency wrt. code_llvm, which doesn't go through jl_emit_native.

Yeah, we might just want to emit two separate copies of the function (for inlining and not for inlining) which is what code_llvm essentially does. I went back and forth when I was making the PR with which strategy made more sense.

src/llvm-expand-atomic-modify.cpp

…tomicrmw The ExpandAtomicModify can recognize our pseudo-intrinsic julia.atomicmodify and convert it into some of known atomicrmw expressions, or simplify it with more inlining, as applicable. This ensures that now our `@atomic` modify is as fast as `Threads.Atomic` for the cases we implement now.

vtjnash · 2025-03-28T14:37:34Z

This is now ready for someone to review, particularly any correctness implications of optimizing atomics that I didn't already address

maleadt · 2025-05-14T07:39:16Z

src/aotcompile.cpp

@@ -823,7 +849,6 @@ void *jl_emit_native_impl(jl_array_t *codeinfos, LLVMOrcThreadSafeModuleRef llvm
    size_t idx = 0;
    for (auto &global : params.global_targets) {
        gvars[idx] = global.second->getName().str();
-        global.second->setInitializer(literal_static_pointer_val(global.first, global.second->getValueType()));


This broke another part of GPUCompiler.jl, relying on the presence of inttoptrs in the IR to reconstruct object pointers. Instead, the IR now only contains nulls (set further down):

function kernel() undefined return end

@"*Main.##280.undefined#41043" = internal global ptr null, !julia.constgv !0 @"jl_sym#undefined#41044" = internal global ptr null, !julia.constgv !0 @"jl_global#41045" = internal global ptr null, !julia.constgv !0 define void @julia_kernel_41041() #0 !dbg !5 { top: %pgcstack = call ptr @julia.get_pgcstack() %"*Main.##280.undefined#41043" = load ptr, ptr @"*Main.##280.undefined#41043" %0 = call ptr addrspace(10) @jl_get_binding_value_seqcst(ptr %"*Main.##280.undefined#41043") %1 = icmp ne ptr addrspace(10) %0, null br i1 %1, label %ok, label %err err: %"jl_sym#undefined#41044" = load ptr, ptr @"jl_sym#undefined#41044" %2 = addrspacecast ptr %"jl_sym#undefined#41044" to ptr addrspace(12) %"jl_global#41045" = load ptr, ptr @"jl_global#41045" %3 = addrspacecast ptr %"jl_global#41045" to ptr addrspace(12) call void @ijl_undefined_var_error(ptr addrspace(12) %2, ptr addrspace(12) %3) unreachable ok: ret void }

Can you explain how we should now use jl_emit_native to get proper IR?

@vtjnash

#57010 overhauled how global variables are initialized in code, breaking GPUCompiler.jl. After talking to @vtjnash, we are apparently supposed to use `jl_get_llvm_gvs` now to get the Julia memory location of a global variable (for a binding or constant value), however, the elements in there don't exactly correspond to the global variables in the module (not all module globals are "managed" by Julia, and the order is different). To make GPUCompiler.jl work again, here I propose renaming `jl_get_llvm_gvs` to `jl_get_llvm_gv_inits`, and making `jl_get_llvm_gvs` instead return the list of the managed global variables, making it possible to perform the global variable initialization that was done by Julia before using something like: ```julia if VERSION >= v"1.13.0-DEV.533" num_gvars = Ref{Csize_t}(0) @CCall jl_get_llvm_gvs(native_code::Ptr{Cvoid}, num_gvars::Ptr{Csize_t}, C_NULL::Ptr{Cvoid})::Nothing gvs = Vector{Ptr{LLVM.API.LLVMOpaqueValue}}(undef, num_gvars[]) @CCall jl_get_llvm_gvs(native_code::Ptr{Cvoid}, num_gvars::Ptr{Csize_t}, gvs::Ptr{LLVM.API.LLVMOpaqueValue})::Nothing inits = Vector{Ptr{Cvoid}}(undef, num_gvars[]) @CCall jl_get_llvm_gv_inits(native_code::Ptr{Cvoid}, num_gvars::Ptr{Csize_t}, inits::Ptr{Cvoid})::Nothing for (gv_ref, init) in zip(gvs, inits) gv = GlobalVariable(gv_ref) ptr = const_inttoptr(ConstantInt(Int64(init)), value_type(gv)) initializer!(gv, ptr) obj = Base.unsafe_pointer_to_objref(init) @safe_info "Resolved $(name(gv)) to $obj" end end ```

oscardssmith reviewed Jan 10, 2025

View reviewed changes

base/atomics.jl Show resolved Hide resolved

oscardssmith added performance Must go faster multithreading Base.Threads and related functionality atomics labels Jan 10, 2025

gbaraldi reviewed Jan 10, 2025

View reviewed changes

src/aotcompile.cpp Outdated Show resolved Hide resolved

vtjnash force-pushed the jn/atomic-modify-opt branch 2 times, most recently from 666985c to c98b43f Compare January 17, 2025 14:32

vtjnash force-pushed the jn/atomic-modify-opt branch from 67ad4b6 to dbecbe8 Compare February 8, 2025 20:38

vchuravy reviewed Feb 10, 2025

View reviewed changes

vchuravy approved these changes Feb 10, 2025

View reviewed changes

src/llvm-expand-atomic-modify.cpp Outdated Show resolved Hide resolved

oscardssmith mentioned this pull request Feb 19, 2025

Teach alloc-opt to handle atomics a bit better #57208

Merged

vtjnash force-pushed the jn/atomic-modify-opt branch from dbecbe8 to 7eb8691 Compare February 21, 2025 19:44

vtjnash added 2 commits March 28, 2025 14:35

remove deprecated Threads.Atomics

5cfdf66

vtjnash force-pushed the jn/atomic-modify-opt branch from 7eb8691 to 5cfdf66 Compare March 28, 2025 14:35

oscardssmith requested a review from gbaraldi March 28, 2025 19:49

vtjnash merged commit 5a62fe0 into master May 7, 2025
4 of 7 checks passed

vtjnash deleted the jn/atomic-modify-opt branch May 7, 2025 15:03

maleadt mentioned this pull request May 8, 2025

Fixes for 1.13 JuliaGPU/GPUCompiler.jl#693

Merged

7 tasks

maleadt reviewed May 14, 2025

View reviewed changes

maleadt mentioned this pull request May 15, 2025

Expose native_code's jl_sysimg_gvars. #58423

Merged

maleadt added a commit to JuliaGPU/GPUCompiler.jl that referenced this pull request May 22, 2025

Adapt to JuliaLang/julia#57010.

8a36324

maleadt added a commit to JuliaGPU/GPUCompiler.jl that referenced this pull request May 22, 2025

Adapt to JuliaLang/julia#57010.

f8d91d7

maleadt mentioned this pull request May 23, 2025

Don't emit module metadata for non-toplevel modules. #58508

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

codegen: add a pass for late conversion of known modify ops to call atomicrmw #57010

codegen: add a pass for late conversion of known modify ops to call atomicrmw #57010

Uh oh!

vtjnash commented Jan 10, 2025

Uh oh!

Uh oh!

gbaraldi commented Jan 10, 2025

Uh oh!

Uh oh!

vtjnash commented Jan 13, 2025

Uh oh!

oscardssmith commented Jan 27, 2025 •

edited

Loading

Uh oh!

DilumAluthge commented Jan 27, 2025

Uh oh!

vtjnash commented Jan 28, 2025

Uh oh!

oscardssmith commented Jan 28, 2025

Uh oh!

vchuravy Feb 10, 2025

Uh oh!

vtjnash Feb 11, 2025

Uh oh!

maleadt May 8, 2025

Uh oh!

vtjnash May 8, 2025

Uh oh!

Uh oh!

vtjnash commented Mar 28, 2025

Uh oh!

Uh oh!

maleadt May 14, 2025

Uh oh!

Uh oh!

	if (jl_ir_inlining_cost((jl_value_t*)src) < UINT16_MAX)
	params.safepoint_on_entry = false; // ensure we don't block ExpandAtomicModifyPass from inlining this code if applicable

Uh oh!

codegen: add a pass for late conversion of known modify ops to call atomicrmw #57010

codegen: add a pass for late conversion of known modify ops to call atomicrmw #57010

Uh oh!

Conversation

vtjnash commented Jan 10, 2025

Uh oh!

Uh oh!

gbaraldi commented Jan 10, 2025

Uh oh!

Uh oh!

vtjnash commented Jan 13, 2025

Uh oh!

oscardssmith commented Jan 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

DilumAluthge commented Jan 27, 2025

Uh oh!

vtjnash commented Jan 28, 2025

Uh oh!

oscardssmith commented Jan 28, 2025

Uh oh!

vchuravy Feb 10, 2025

Choose a reason for hiding this comment

Uh oh!

vtjnash Feb 11, 2025

Choose a reason for hiding this comment

Uh oh!

maleadt May 8, 2025

Choose a reason for hiding this comment

Uh oh!

vtjnash May 8, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

vtjnash commented Mar 28, 2025

Uh oh!

Uh oh!

maleadt May 14, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

oscardssmith commented Jan 27, 2025 •

edited

Loading