Regression in memory allocation optimization of a mutable StaticArray

[cdsousa]

Julia version 1.6.2 was able to optimize out the memory allocation of the following code, but version1.7.0-beta3.0 is not:

using StaticArrays
function cpu_kernel(a, i)
    SIZE = 15
    v = ones(MVector{SIZE, UInt32})
    for x in v
        a[i] += x
    end
end
a = zeros(UInt32, 1)
cpu_kernel(a, 1)
@time cpu_kernel(a, 1)

This has been noticed originally here: JuliaGPU/CUDA.jl#38 (comment)

(Maybe related to #41512 ?)

[cdsousa]

This also allocates in Version 1.8.0-DEV.310 (2021-08-06).

[KristofferC]

A bisection might be beneficial here to find the offending commit.

[maleadt]

I'm bisecting this.

EDIT:

7ba81e852695bb42a1c6d5f543bd55f02d799a2c is the first bad commit
commit 7ba81e852695bb42a1c6d5f543bd55f02d799a2c
Author: Valentin Churavy <v.churavy@gmail.com>
Date:   Mon May 10 12:53:54 2021 -0400

    [LLVM] Upgrade to LLVM 12

[maleadt]

Reduced Julia example:

import Base: getindex, ones
struct SOneTo{a} <: AbstractUnitRange{Int}
end
SOneTo(a) = SOneTo{a}()
Base.first(::SOneTo) = 1
Base.last(::SOneTo{a}) where a = a
function Base.iterate(::SOneTo{a}, b) where a
    if b < a
        c = b + 1
        c, c
    end
end
abstract type StaticArray{a, c, b} <: AbstractArray{c, b} end
struct Size{a}
    end
Size(b) = Size{tuple(b.parameters...)}()
Size(::c) where c<: StaticArray{a} where a= @isdefined(a) ? Size(a) : d
mutable struct MArray{a , b, c, d} <: StaticArray{a, b, c}
    function MArray{a,b,c,d}(e) where {a,b,c,d}
        new()
    end
end
function getindex(f::MArray, h)
    checkbounds(f,h)
      g
end
MVector{a, b} = MArray{Tuple{a}, b, 1, a}
Base.axes(a::StaticArray) = b(Size(a))
function b(::Size{c}) where c
    map(SOneTo, c) end
ones(::Type{a}) where a = b(Size, a)
@generated function b(::c, ::Type{a}) where {c, a }
    quote
        a(0)
    end
end
function b(a, c)
    d = 15
    e = ones(MVector{d, UInt32})
    for f in e
    end
end
a = using InteractiveUtils
@code_llvm b(a, 1)

And the IR that fails to optimize:

@a = external global {}*

define void @b() {
d:
  %c = call {} addrspace(10)* @julia.gc_alloc_obj(i8* undef, i64 0, {} addrspace(10)* null)
  br label %e

e:
  %ac3 = phi i64 [ 2, %d ], [ %ac, %f ]
  %g = phi i64 [ 2, %d ], [ %ab, %f ]
  %ac.off = add i64 %ac3, -1
  %0 = icmp ugt i64 %ac.off, 14
  br i1 %0, label %h, label %i

h:
  %1 = call {} addrspace(10)* @aj({} addrspace(10)* %c)
  unreachable

i:
  %j = load atomic {}*, {}** @a unordered, align 8
  %k.l = icmp eq {}* %j, null
  br i1 %k.l, label %al, label %f

f:
  %aa = icmp ugt i64 %g, 4
  %ab = add i64 %g, 1
  %ac = select i1 %aa, i64 undef, i64 %ab
  %exitcond = icmp eq i64 %ab, 16
  br i1 %exitcond, label %ap.critedge, label %e

ap.critedge:
  ret void

al:
  %2 = call {}* @aq()
  br label %f
}

declare {} addrspace(10)* @julia.gc_alloc_obj(i8*, i64, {} addrspace(10)*)

declare {}* @aq()

declare {} addrspace(10)* @aj({} addrspace(10)*)

Contains an allocation after optimization on LLVM 12, not on LLVM 11. Can be checked with LLVM.jl on 1.6 and 1.7 using:

using LLVM

Context() do ctx
    ir = read(first(ARGS), String)
    mod = parse(LLVM.Module, ir; ctx)

    ModulePassManager() do pm
        ccall(:jl_add_optimization_passes, Cvoid,
                (LLVM.API.LLVMPassManagerRef, Cint, Cint),
                pm, Base.JLOptions().opt_level, #=lower_intrinsics=# 1)
        run!(pm, mod)
    end

    println(mod)
end

[vtjnash]

It is impressive that LLVM used to be able to remove the bounds check. We may want to start emitting explicit semantic copies on error paths though, to make that less challenging for it.

[JeffBezanson]

Seems to be due to the value escaping via a BoundsError, and LLVM no longer removing a bounds check. So we could work around it by adding a copy when throwing a BoundsError for this type. Obviously it would be great to fix this in LLVM, but that seems like it could be quite difficult.

[maleadt]

We may want to start emitting explicit semantic copies on error paths though, to make that less challenging for it.

How would that look here with the MWE (which doesn't rely on StaticArrays.jl)? Base.throw_boundserror(A::MArray, I) = throw(BoundsError(deepcopy(A), I)) doesn't seem to cut it.

[vtjnash]

deepcopy needs to escape the object (to detect reference cycles). Maybe a regular array copy?

[maleadt]

Maybe a regular array copy?

Still allocates. throw_boundserror doesn't even need to capture the array, @inline Base.throw_boundserror(f::MArray, I) = throw(I) also causes an allocation.

[vtjnash]

true, I don't think we would optimize (or inline) that, because it is on the error path

[JeffBezanson]

So far this seems too difficult to fix to be release-blocking. We absolutely should fix it though, and can backport the fix to a future 1.7.x patch release.