Skip to content

Commit

Permalink
optimizer: enhance SROA, enable PhiNodefy load forwarding with unin…
Browse files Browse the repository at this point in the history
…itialized allocations

This commit addresses this [TODO comment](https://github.com/JuliaLang/julia/blob/3a5146f6b94e32bef5a812619068f82218b5482c/base/compiler/ssair/passes.jl#L831-L832).
Now SROA can eliminate this sort of allocation:
```julia
let src = code_typed((Bool,)) do cond
        r = Ref{Any}()
        if cond
            r[] = 42
        else
            r[] = 32
        end
        return r[]
    end |> only |> first
    @test !any(isnew, src.code)
end
```
  • Loading branch information
aviatesk committed Nov 25, 2021
1 parent 2c38b36 commit 22bbb15
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 17 deletions.
49 changes: 34 additions & 15 deletions base/compiler/ssair/passes.jl
Original file line number Diff line number Diff line change
Expand Up @@ -93,23 +93,50 @@ function compute_value_for_use(ir::IRCode, domtree::DomTree, allblocks::Vector{I
end
end

# even when the allocation contains an uninitialized field, we try an extra effort to check
# if this load at `idx` have any "safe" `setfield!` calls that define the field
function has_safe_def(
ir::IRCode, domtree::DomTree, allblocks::Vector{Int}, du::SSADefUse,
newidx::Int, idx::Int, inclusive::Bool = false)
def = first(find_def_for_use(ir, domtree, allblocks, du, idx, inclusive))

# this field is supposed to be defined at the `:new` site (but it's not and thus this load will throw)
def == newidx && return false

def 0 && return true # found a "safe" definition

# we may be able to replace this load with `PhiNode` if all the predecessors have "safe" definitions
idxblock = block_for_inst(ir, idx)
for pred in ir.cfg.blocks[idxblock].preds
lastidx = last(ir.cfg.blocks[pred].stmts)
# NOTE `lastidx` isn't a load, thus we can use inclusive coondition within the `find_def_for_use`
has_safe_def(ir, domtree, allblocks, du, newidx, lastidx, true) || return false
end
return true
end

# find the first dominating def for the given use
function find_def_for_use(ir::IRCode, domtree::DomTree, allblocks::Vector{Int}, du::SSADefUse, use_idx::Int)
stmtblock = block_for_inst(ir.cfg, use_idx)
curblock = find_curblock(domtree, allblocks, stmtblock)
function find_def_for_use(
ir::IRCode, domtree::DomTree, allblocks::Vector{Int}, du::SSADefUse, use::Int, inclusive::Bool=false)
useblock = block_for_inst(ir.cfg, use)
curblock = find_curblock(domtree, allblocks, useblock)
local def = 0
for idx in du.defs
if block_for_inst(ir.cfg, idx) == curblock
if curblock != stmtblock
if curblock != useblock
# Find the last def in this block
def = max(def, idx)
else
# Find the last def before our use
def = max(def, idx >= use_idx ? 0 : idx)
if inclusive
def = max(def, idx use ? idx : 0)
else
def = max(def, idx < use ? idx : 0)
end
end
end
end
return def, stmtblock, curblock
return def, useblock, curblock
end

function simple_walk(compact::IncrementalCompact, @nospecialize(defssa#=::AnySSAValue=#),
Expand Down Expand Up @@ -833,16 +860,8 @@ function sroa_pass!(ir::IRCode)
allblocks = sort(vcat(phiblocks, ldu.def_bbs))
blocks[fidx] = phiblocks, allblocks
if fidx + 1 > length(defexpr.args)
# even if the allocation contains an uninitialized field, we try an extra effort
# to check if all uses have any "solid" `setfield!` calls that define the field
# although we give up the cases below:
# `def == idx`: this field can only defined at the allocation site (thus this case will throw)
# `def == 0`: this field comes from `PhiNode`
# we may be able to traverse control flows of PhiNode values, but it sounds
# more complicated than beneficial under the current implementation
for use in du.uses
def = find_def_for_use(ir, domtree, allblocks, du, use)[1]
(def == 0 || def == newidx) && @goto skip
has_safe_def(ir, domtree, allblocks, du, newidx, use) || @goto skip
end
end
end
Expand Down
32 changes: 30 additions & 2 deletions test/compiler/irpasses.jl
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,6 @@ let src = code_typed1((Bool,)) do cond
end
@test !any(isnew, src.code)
end
# FIXME to handle this case, we need a more strong alias analysis
let src = code_typed1((Bool,)) do cond
r = Ref{Any}()
if cond
Expand All @@ -157,7 +156,22 @@ let src = code_typed1((Bool,)) do cond
end
return r[]
end
@test_broken !any(isnew, src.code)
@test !any(isnew, src.code)
end
let src = code_typed1((Bool,Bool,Any,Any,Any)) do c1, c2, x, y, z
r = Ref{Any}()
if c1
if c2
r[] = x
else
r[] = y
end
else
r[] = z
end
return r[]
end
@test !any(isnew, src.code)
end
let src = code_typed1((Bool,)) do cond
r = Ref{Any}()
Expand All @@ -169,6 +183,20 @@ let src = code_typed1((Bool,)) do cond
# N.B. `r` should be allocated since `cond` might be `false` and then it will be thrown
@test any(isnew, src.code)
end
let src = code_typed1((Bool,Bool,Any,Any)) do c1, c2, x, y
r = Ref{Any}()
if c1
if c2
r[] = x
end
else
r[] = y
end
return r[]
end
# N.B. `r` should be allocated since `c2` might be `false` and then it will be thrown
@test any(isnew, src.code)
end

# should include a simple alias analysis
struct ImmutableOuter{T}; x::T; y::T; z::T; end
Expand Down

0 comments on commit 22bbb15

Please sign in to comment.