add LoopInfo + LICM

[vchuravy]

In recent discussion as well as in the past the topic came up that LLVM LICM pass is currently not effectual for calls to Julia code and Julia specific constructs. In #36786 (comment) I explained my understanding that we don't need to change the semantics of math functions to enable LICM. In general it seems that informing LLVM about the pureness of Julia functions might not be valid (#29566 (comment), #32368) since Julia and LLVM have differing definitions, this will fixe #29285

We could also do #36809 (Custom-LICM pass for GC.@preserve on the LLVM level) as part of this pass.

In order to make this more useful we might want to track more properties of Julia methods. For simplicity for now I have limited myself to @pure functions.

TODO:

• Create a pre-header properly, this is a CFG manipulation and might need updating domtrees dynamically, removing all unreachable blocks #33730 to update the domtree dynamically
• Iterate over nested loops correctly, innermost to outermost, while maintaining LoopInfo
• Various todos in the code

cc: @yhls, @jpsamaroo

@timholy Together with a Loop Induction Variables deduction pass the loopinfo pass could be used to implement #35074 without special casing CartesianIndices

I am putting this up mostly to get help with figuring out the CFG manipulation part. I looked through inlining, but it seems we could make the insertion of BBs an independent util.

[vchuravy]

With this PR:

Base.@pure @noinline mysin(x) = sin(x)
function example(A, x)
           for i in 1:length(A)
               A[i] = A[i] + (mysin(x) + mysin(x))
           end
       end

turns into:

julia> @code_typed example(zeros(3), 3.0)
CodeInfo(
1 ── %1  = Base.arraylen(A)::Int64
│    %2  = Base.sle_int(1, %1)::Bool
│    %3  = Base.ifelse(%2, %1, 0)::Int64
│    %4  = Base.slt_int(%3, 1)::Bool
└───       goto #3 if not %4
2 ──       goto #4
3 ──       goto #4
4 ┄─ %8  = φ (#2 => true, #3 => false)::Bool
│    %9  = φ (#3 => 1)::Int64
│    %10 = φ (#3 => 1)::Int64
│    %11 = Base.not_int(%8)::Bool
│    %12 = invoke Main.mysin(_3::Float64)::Float64
│    %13 = invoke Main.mysin(_3::Float64)::Float64
│    %14 = Base.add_float(%12, %13)::Float64
└───       goto #10 if not %11
5 ┄─ %16 = φ (#4 => %9, #9 => %26)::Int64
│    %17 = φ (#4 => %10, #9 => %27)::Int64
│    %18 = Base.arrayref(true, A, %16)::Float64
│    %19 = Base.add_float(%18, %14)::Float64
│          Base.arrayset(true, A, %19, %16)::Vector{Float64}
│    %21 = (%17 === %3)::Bool
└───       goto #7 if not %21
6 ──       goto #8
7 ── %24 = Base.add_int(%17, 1)::Int64
└───       goto #8
8 ┄─ %26 = φ (#7 => %24)::Int64
│    %27 = φ (#7 => %24)::Int64
│    %28 = φ (#6 => true, #7 => false)::Bool
│    %29 = Base.not_int(%28)::Bool
└───       goto #10 if not %29
9 ──       goto #5
10 ┄       return nothing
) => Nothing

Instead of:

CodeInfo(
1 ── %1  = Base.arraylen(A)::Int64
│    %2  = Base.sle_int(1, %1)::Bool
│    %3  = Base.ifelse(%2, %1, 0)::Int64
│    %4  = Base.slt_int(%3, 1)::Bool
└───       goto #3 if not %4
2 ──       goto #4
3 ──       goto #4
4 ┄─ %8  = φ (#2 => true, #3 => false)::Bool
│    %9  = φ (#3 => 1)::Int64
│    %10 = φ (#3 => 1)::Int64
│    %11 = Base.not_int(%8)::Bool
└───       goto #10 if not %11
5 ┄─ %13 = φ (#4 => %9, #9 => %26)::Int64
│    %14 = φ (#4 => %10, #9 => %27)::Int64
│    %15 = Base.arrayref(true, A, %13)::Float64
│    %16 = invoke Main.mysin(_3::Float64)::Float64
│    %17 = invoke Main.mysin(_3::Float64)::Float64
│    %18 = Base.add_float(%16, %17)::Float64
│    %19 = Base.add_float(%15, %18)::Float64
│          Base.arrayset(true, A, %19, %13)::Vector{Float64}
│    %21 = (%14 === %3)::Bool
└───       goto #7 if not %21
6 ──       goto #8
7 ── %24 = Base.add_int(%14, 1)::Int64
└───       goto #8
8 ┄─ %26 = φ (#7 => %24)::Int64
│    %27 = φ (#7 => %24)::Int64
│    %28 = φ (#6 => true, #7 => false)::Bool
│    %29 = Base.not_int(%28)::Bool
└───       goto #10 if not %29
9 ──       goto #5
10 ┄       return nothing
) => Nothing

[yuyichao]

We could also do #36809 (Custom-LICM pass for GC.@preserve on the LLVM level) as part of this pass.

For sure. Though I wouldn't use one to replace the other as long as we still uses any LLVM loop passes. In the llvm pass, I'm counting on LLVM LICM to hoist some generation of the pointer value out of the loop (e.g. loading from a array/object that is loop invariant).

Also, there are some julia specific LICM that can only be done at LLVM level, including everything GC/allocation related. It can certainly be done in julia code in the future but it'll require a much lower level IR.

[yuyichao]

Also, I feel like it wouldn't hurt to mark some functions that are believed to be "pure"/LICM safe with some custom metadata so that our own LLVM pass can pick it up.

[vtjnash]

In general it seems that informing LLVM about the pureness of Julia functions might not be valid (#29566 (comment), #32368) since Julia and LLVM have differing definitions, this will fixe #29285

FWIW, the issue was that many of our uses of @pure aren't valid to LICM, not that it was bad to set the attribute in LLVM.

[yuyichao]

I’m not really commenting on what is safe to LICM, I just mean that whatever identified here as safe should also be safe to do so in llvm. I assume there isn't a case where the code is safe to licm in Julia but not in llvm?

[vchuravy]

Now that we have effect analysis this is now feasible and thanks #45305 added the gnarly details of how to do CFG manipulations.

Still some details to figure out, in particular how to deal with loop-nests and multiple loops in a function since the act of mutating the CFG invalidates the collect loop information.

Also need to check multiple invariant statements doing the right thing.

Developed this out-of-tree in https://github.com/vchuravy/Loops.jl and https://vchuravy.dev/talks/licm/

Will fix the long-standing #29285