type-instability with repeat and cat?

[kellertuer]

I want to use repeat (or cat) for my own type or to be precise for a field of my own type, in this application to take the array that is stored in my own type, create as many copies as there are dimensions and create a new A1 with this array internally. However, whenever trying to use that, I get type instabilities. So I also tried a second version, since I first only need a specific repeat, that employs cat, but I run into the same problem, that used within a function, it seems not to be stable.

Note that copying both following runExamples() to REPL or running the body in REPL does not yield a type-instability (a prior version of this was posted here)

import Base: size, repeat, ndims
abstract type A end

struct A1{T,N}<:A where {T<:A,N}
    value::Array{T,N}
end
A1{T}(a::Array{T,N}) where {T<:A,N} = A1{T,N}(a)
A1(a::Array{T,N}) where {T<:A,N} = A1{T,N}(a)

struct A2{T} <: A
    value::T
end

size(a::A1{T,N}) where {T<:A,N} = size(a.value)
ndims(a::A1{T,N}) where {T<:A,N} = ndims(a.value)
function repeat(a::A1{T,N}; inner=nothing, outer=nothing) where {T<:A,N}
    b = repeat(a.value; inner=inner, outer=outer)
    return A1(b)
end
repeat(a::A1, counts::Integer...) = repeat(a; outer=counts)
# --- Repeat ---
function runExampleRepeat()
    a = A1( [ A2(0.2) A2(0.1); A2(0.1) A2(0.3) ])
    s = ndims(a)
    d = fill(1,s+1)
    d[s+1] = s
    return repeat(a, d...)
end

# unstable
@code_warntype runExampleRepeat()
# stable
a = A1{A2,2}([ A2(0.2) A2(0.1); A2(0.1) A2(0.3) ])
d = fill(1,ndims(a)+1)
d[end] = ndims(a) 
@code_warntype repeat(a,d...)

# --- Cat ---

function runExampleCat()
    a = A1( [ A2(0.2) A2(0.1); A2(0.1) A2(0.3) ])
    n = ndims(a)
    aT = fill(a.value,n)
    bT = cat(aT..., dims=n+1)
    return A1(bT)
end
# unstable
@code_warntype runExampleCat()
# stable
@code_warntype A1( cat(fill(a.value,ndims(a))...,dims=ndims(a)+1) )

where the @code_warntypes yield for the first example

Body::A1{_1,_2} where _2 where _1
3 1 ─ %1  = %new(A2{Float64}, 0.2)::A2{Float64}                                                                                                                             │╻╷ Type
  │   %2  = %new(A2{Float64}, 0.1)::A2{Float64}                                                                                                                             ││╻  Type
  │   %3  = %new(A2{Float64}, 0.1)::A2{Float64}                                                                                                                             ││╻  Type
  │   %4  = %new(A2{Float64}, 0.3)::A2{Float64}                                                                                                                             ││╻  Type
  │   %5  = invoke Base.typed_hvcat(A2{Float64}::Type{A2{Float64}}, (2, 2)::Tuple{Int64,Int64}, %1::A2{Float64}, %2::Vararg{A2{Float64},N} where N, %3, %4)::Any            │╻  hvcat
  │   %6  = (Main.A1)(%5)::A1{_1,_2} where _2 where _1                                                                                                                      │
4 │   %7  = (Main.ndims)(%6)::Any                                                                                                                                           │
5 │   %8  = (%7 + 1)::Any                                                                                                                                                   │
  │   %9  = (isa)(%8, Tuple{})::Bool                                                                                                                                        │
  └──       goto #3 if not %9                                                                                                                                               │
  2 ─ %11 = π (%8, Tuple{})                                                                                                                                                 │
  │   %12 = $(Expr(:foreigncall, :(:jl_new_array), Array{Int64,0}, svec(Any, Any), :(:ccall), 2, Array{Int64,0}, :(%11)))::Array{Int64,0}                                   ││╻  Type
  │   %13 = invoke Base.fill!(%12::Array{Int64,0}, 1::Int64)::Array{Int64,0}                                                                                                ││
  └──       goto #4                                                                                                                                                         │
  3 ─ %15 = (Main.fill)(1, %8)::Any                                                                                                                                         │
  └──       goto #4                                                                                                                                                         │
  4 ┄ %17 = φ (#2 => %13, #3 => %15)::Any                                                                                                                                   │
6 │   %18 = (%7 + 1)::Any                                                                                                                                                   │
  │         (Base.setindex!)(%17, %7, %18)                                                                                                                                  │
7 │   %20 = (Core.tuple)(%6)::Tuple{A1{_1,_2} where _2 where _1}                                                                                                            │
  │   %21 = (Core._apply)(Main.repeat, %20, %17)::A1{_1,_2} where _2 where _1                                                                                               │
  └──       return %21  

while for the second (the repeat) one correctly gets Body::A1{A2,3}and similarly for the approach with cat. Bot do yield a 3D array within A1 where the third dimension has 2 (ndims(a)) entries, but even with a lot of reading and research the function does not get type stable.

So why can't repeat or cat infer the proper type of its result?

[deahhh]

It's weird!

# code1
function add_gpu_test()
    xs = Vector{AbstractArray{Float32, 4}}([rand(Float32, 2,3,3,1) for _ in 1:3])
    @code_warntype cat(xs...; dims=Val(3))
end
add_gpu_test()

#code2
function add_gpu_test()
    xs = Vector{AbstractArray{Float32, 4}}([rand(Float32, 2,3,3,1) for _ in 1:3])
    cat(xs...; dims=Val(3))
end
@code_warntype add_gpu_test()

code1 is type-stable while code2 is not. Why?

[deahhh]

@kellertuer , @KristofferC
Why nobody repair the bug?

[KristofferC]

Same reason you don't do it I guess? People have other things to do that take priority.

[deahhh]

Same reason you don't do it I guess? People have other things to do that take priority.

I don't have the skills to do that yet.
Both the runtime efficiency issue and the memory efficiency issue will be solved sooner or later, right? I'm losing faith. I've tried everything I can to optimise my program, and it's still 30 times slower. The most efficiency is lost between statements and even the big guys who developed Flux haven't solved it. Oh my god.

[deahhh]

It's weird!

# code1
function add_gpu_test()
    xs = Vector{AbstractArray{Float32, 4}}([rand(Float32, 2,3,3,1) for _ in 1:3])
    @code_warntype cat(xs...; dims=Val(3))
end
add_gpu_test()

#code2
function add_gpu_test()
    xs = Vector{AbstractArray{Float32, 4}}([rand(Float32, 2,3,3,1) for _ in 1:3])
    cat(xs...; dims=Val(3))
end
@code_warntype add_gpu_test()

code1 is type-stable while code2 is not. Why?

Today, I solve it partly by using Base.cat_t. But, it is still type-unstable, outputs Union{CUDA.CuArray{Float32, 3, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}}

[kellertuer]

I think I rephrased my problem back then, though I did not write that in the thread linked in the original post here.

Why nobody repair the bug?

I think Kristopher is right: Probably there is not much demand for this to be addressed and I for myself do this only as a small part of my job, so mainly in my free time and I am not clever enough to fix this myself, I can not contribute to a solution here but

I'm losing faith.

faith in what? That this specific small thing will be fixed by someone? If someone comes along needing this and being able to fix this case, sure it will be fixed, but again, open source means (a) a lot pof people putting in their free time and (b) anyone who sees how to fix it can fix it (and thats for example not me).

[deahhh]

faith in what?

it's "Both the runtime efficiency issue and the memory efficiency issue will be solved sooner or later, right?"
For I've tried everything I can to optimise my program, and it's still 30 times slower (than pytorch version). The most efficiency is lost between statements and even the big guys who developed Flux haven't solved it. He guess it's caused by GC.