Extra allocations associated with unused ArgumentError

[BioTurboNick]

I've encountered a weird bug causing additional memory to be allocated when running a method that never executes the changed branch.

In 1.5.2-pre.29 and earlier, both functions have 1 allocation. In the current master, the top function has 1 additional allocation.

function hvncat1(dims::Tuple{Vararg{Int, N}}, xs::T...) where T <:Number where N
    a = Array{Int,N}(undef, dims...)
    if prod(dims) != length(xs)
        throw(ArgumentError("argument count does not match specified shape (expected $(prod(dims)), got $(length(xs)))"))
    end

    nr = dims[1]
    nc = dims[2]
    k = 1
    na = prod(dims[3:N])
    @inbounds for d=1:na
        for i=1:nr
            for j=1:nc
                a[i+nr*(j-1)+nr*nc*(d-1)] = xs[k]
                k += 1
            end
        end
    end
    a
end

function hvncat2(dims::Tuple{Vararg{Int, N}}, xs::T...) where T <:Number where N
    a = Array{Int,N}(undef, dims...)
    size1 = prod(dims)
    size2 = length(xs)
    if size1 != size2
        throw(ArgumentError("argument count does not match specified shape (expected $(size1), got $(size2)))"))
    end

    nr = dims[1]
    nc = dims[2]
    k = 1
    na = prod(dims[3:N])
    @inbounds for d=1:na
        for i=1:nr
            for j=1:nc
                a[i+nr*(j-1)+nr*nc*(d-1)] = xs[k]
                k += 1
            end
        end
    end
    a
end

using BenchmarkTools
@btime hvncat1((2, 2, 2), 1, 3, 2, 3, 3, 5, 4, 6) # 2 allocations: 224 bytes in master; 1 allocation: 144 bytes in 1.5.2
@btime hvncat2((2, 2, 2), 1, 3, 2, 3, 3, 5, 4, 6) # 1 allocation: 144 bytes

[KristofferC]

Maybe fixed by #37572. Try on very latest master.

[BioTurboNick]

I rebased onto that commit just before, so it doesn't appear to have fixed this.

[BioTurboNick]

Just noting that this issue is still present in 1.6.0-beta1

[BioTurboNick]

This appears to be a type inference issue.

In the latter case with 1 allocation, the lowering looks like this:

│    %8  = Core.tuple(Main.undef)::Core.Const((UndefInitializer(),))
│          (a = Core._apply_iterate(Base.iterate, %7, %8, dims))
│          (size1 = Main.prod(dims))
│          (size2 = Main.length(xs))
│    %12 = (size1 != size2::Core.Const(8))::Bool
└───       goto #3 if not %12
2 ── %14 = Base.string("argument count does not match specified shape (expected ", size1, ", got ", size2::Core.Const(8), "))")::Any
│    %15 = Main.ArgumentError(%14)::Any
└───       Main.throw(%15)

In the former case with 2 allocations:

│    %8  = Core.tuple(Main.undef)::Core.Const((UndefInitializer(),))
│          (a = Core._apply_iterate(Base.iterate, %7, %8, dims))
│    %10 = Main.prod(dims)::Int64
│    %11 = Main.length(xs)::Core.Const(8)
│    %12 = (%10 != %11)::Bool
└───       goto #3 if not %12
2 ── %14 = Main.prod(dims)::Any
│    %15 = Main.length(xs)::Any
│    %16 = Base.string("argument count does not match specified shape (expected ", %14, ", got ", %15, ")")::Any
│    %17 = Main.ArgumentError(%16)::Any
└───       Main.throw(%17)

What's extra weird is that it seems to properly find the type on one line (%10, %11), and then it forgets it a few lines later (%14, %15).

@KristofferC - if you have an idea of what code would be responsible for this, I could poke around.

[BioTurboNick]

After looking into it more, I think this is actually a difference between the released Julia packages and the development copy I compile myself. Should they differ in this way, though? Are there extra optimizations turned on there? EDIT: But 1.6.0-beta1 release does have this behavior, and 1.5.3 does not.

[vchuravy]

The inference result is due to #35982, we are deoptimizing branches that yield to throw. I haven't investigated, but I would hypothesize that escape analysis no longer sees either dims or xs as non-escaping and therefore we don't ellide those allocations

[BioTurboNick]

Thanks! Is that something that can be fixed readily, or will it just be advised for people to take care not to do things like call functions within a throw expression?

[BioTurboNick]

Just confirmed, if I comment out these lines from abstract_call_gf_by_type in abstractinterpretation.jl, the issue is gone:

    if sv.params.unoptimize_throw_blocks && sv.currpc in sv.throw_blocks
        return CallMeta(Any, false)
    end

[BioTurboNick]

I looked at the LLVM generated for the optimized and unoptimized version and I found a few differences prior to the branch, at the very top of the method:

Unoptimized:

%gcframe84 = alloca [4 x {}*], align 16
%gcframe84.sub = getelementptr inbounds [4 x {}*], [4 x {}*]* %gcframe84, i64 0, i64 0
%10 = bitcast [4 x {}*]* %gcframe84 to i8*
call void @llvm.memset.p0i8.i32(i8* nonnull align 16 dereferenceable(32) %10, i8 0, i32 32, i1 false)
%11 = bitcast [4 x {}*]* %gcframe84 to i64*
store i64 8, i64* %11, align 16
%12 = getelementptr inbounds [4 x {}*], [4 x {}*]* %gcframe84, i64 0, i64 1

Optimized:

%gcframe84 = alloca [3 x {}*], align 16
%gcframe84.sub = getelementptr inbounds [3 x {}*], [3 x {}*]* %gcframe84, i64 0, i64 0
%10 = bitcast [3 x {}*]* %gcframe84 to i8*
call void @llvm.memset.p0i8.i32(i8* nonnull align 16 dereferenceable(24) %10, i8 0, i32 24, i1 false)
%11 = bitcast [3 x {}*]* %gcframe84 to i64*
store i64 4, i64* %11, align 16
%12 = getelementptr inbounds [3 x {}*], [3 x {}*]* %gcframe84, i64 0, i64 1

I'm not entirely sure how to read LLVM instructions with respect to how allocations are counted. Are the differences here relevant?

Obviously more changes in the throw section itself, and I see some alloc calls there, but it should be completely skipped when the throw condition isn't met.

[BioTurboNick]

Hmm... it still happens even if I try to extract the throw function?

t(dims, xs) = throw(ArgumentError("argument count does not match specified shape (expected $(prod(dims)), got $(length(xs)))"))
function hvncat3(dims::Tuple{Vararg{Int, N}}, xs::T...) where T <:Number where N
           a = Array{Int,N}(undef, dims...)
           prod(dims) != length(xs) && t(dims, xs)
           nr = dims[1]
           nc = dims[2]
           k = 1
           na = prod(dims[3:N])
               @inbounds for d=1:na
               for i=1:nr
                   for j=1:nc
                       a[i+nr*(j-1)+nr*nc*(d-1)] = xs[k]
                       k += 1
                   end
               end
           end
           a
end

Adding @noinline to t() doesn't seem to affect it either. @code_warntype looks good. All the LLVM lines I identified above are missing from the LLVM code in this version, so they must not be related?

I am quite confused.

[BioTurboNick]

I used --track-allocations=all and the only allocations I see in the function are the expected ones...

        - function hvncat3(dims::Tuple{Vararg{Int, N}}, xs::T...) where T <:Number where N
      144     a = Array{Int,N}(undef, dims...)
        0     prod(dims) != length(xs) && t(dims, xs)
        0     nr = dims[1]
        0     nc = dims[2]
        -     k = 1
        0     na = prod(dims[3:N])
        0         @inbounds for d=1:na
        0         for i=1:nr
        0             for j=1:nc
        0                 a[i+nr*(j-1)+nr*nc*(d-1)] = xs[k]
        0                 k += 1
        -             end
        -         end
        -     end
        0     a
        - end
        - t(dims, xs) = throw(ArgumentError("argument count does not match specified shape (expected $(prod(dims)), got $(length(xs)))"))

[BioTurboNick]

I'm assuming track_allocations is accurate... Is @btime making a mistake?

EDIT: Ah, found the 80 bytes allocated associated with the function call itself(?)

 80 hhh() = hvncat3((2,2,2), 4, 5, 6, 4, 2, 1, 3, 5)
  - function hvncat3(dims::Tuple{Vararg{Int, N}}, xs::T...) where T <:Number where N
144     a = Array{Int,N}(undef, dims...)
...

[BioTurboNick]

This is starting to get over my head, but one more observation: The issue seems only related to referencing the xs... vararg inside the throw block. dims does not trigger the issue.

[timholy]

Dumb question, but if you're happy sidestepping this rather than fixing it, can you put the string interpolation in a separate function? We do this frequently, try grep -R "noinline throw" *. (The@noinline is no longer necessary but a legacy from a time when it was.)