Type instability with Type{T} selectors in keyword arguments

General Usage
1

I’m trying to write a function that outputs arrays that are filled according to a certain computation. I want the user to be able to select the type of the array according to an argument.

Consider these two example functions:

function my_fill1(val, array_constructor::Type{AC}=Array) where {AC<:AbstractArray}
    T = typeof(val)
    result = AC{T, 1}(undef, 10)
    result .= val
    result
end

function my_fill2(val; array_constructor::Type{AC}=Array) where {AC<:AbstractArray}
    T = typeof(val)
    result = AC{T, 1}(undef, 10)
    result .= val
    result
end

In both versions, I’m using the Type{T} selector to avoid type instabilities. The only difference between them is that, in the first version, array_constructor is a positional argument, while, in the second, it is a keyword argument. For my use case, the function has many more arguments, so I would really prefer use the version with keyword arguments.

When I test both versions for type stability with the default values for array_constructor, I get positive results:

@code_warntype my_fill1(1) # type stable
@code_warntype my_fill2(1) # type stable

Nonetheless, when I overwrite the default values, my_fill1 remains stable while my_fill2 becomes unstable:

@code_warntype my_fill1(1, Array) # type stable
@code_warntype my_fill2(1, array_constructor=Array) # type unstable (output type is Any)

What is going on here? Am I doing something wrong? Thanks in advance for any help!

2

Whenever you’re testing something sensitive with types, it’s easy for @code_warntype to get the specialization and constant propagation heuristics wrong. If Array is often written as a literal in the cases you’re interested in, the more “typical” result you’d see would be reflected with:

julia> f() = my_fill2(1, array_constructor=Array)
f (generic function with 1 method)

julia> @code_warntype f()
MethodInstance for f()
MethodInstance for f()
  from f() @ Main REPL[4]:1
Arguments
  #self#::Core.Const(Main.f)
Body::Vector{Int64}

cf Can we support a way to make `@code_typed` stop lying to us all? :) · Issue #32834 · JuliaLang/julia · GitHub

3

Just to be sure I understood what you said: you mean that there is no problem with the way the functions are implemented, just with @code_warntype?

4

Yeah, mostly. But it’s also tricky because a function’s performance depends not just upon the function definition itself but also how it’s called. @code_warntype is trying to guess how it might get called.

Putting it in a simple function that represents how you know it’ll get called is always a good answer.

5

I see. In the setting that I’m imagining, the user will just set this keyword value as CuArray or some other GPU array type according to their hardware. Then this should be safe?

6

Separate issues.

  1. The code_ reflection methods redo type inference and compilation given the input types, which isn’t always what Julia does. It’s possible for Julia to compile a method once for a wider argument type annotation instead of for each specific input type (@nospecialize, some cases of ::Function, ::Type, and ::Vararg), and it’s possible for Julia to infer the call’s return type given the type annotation instead of each input type (Base.@nospecializeinfer). The former can occur without the latter, in other words the compiler can figure out a call’s return type but choose to use compiled code that didn’t use the same degree of type inference. @code_warntype won’t show those nuances, and while it’s arguable that we shouldn’t assume that the inferred types are always relevant to execution, @code_llvm only generates compiled code that specializes completely, hence the “lying.”

  2. That’s not what’s making the difference for keyword arguments here, in fact @code_warntype is doing the opposite of over-specialization. Keyword input types are specialized for compilation (this is different from the “method specialization” where you write multiple methods with varied positional arguments), but ::Type{AC} annotations get lost in the process. The implementation puts the keyword arguments in a NamedTuple like (array_constructor = AC,), which only retains AC’s type (DataType or UnionAll) for type inference instead of the specific AC. Making a positional argument or passing in a dummy instance for an inner typeof call are workarounds, though the latter isn’t always feasible like in the Array case.

1 Like
7

That’s technically true (I think), but it’s deceiving because — as demonstrated above in the @code_warntype f() — this can be recovered through constant propagation/inlining.

It’s probably also worth noting that AC{T,1} isn’t the generic operation that you’d hope, even though AC <: AbstractArray. While AbstractArray itself has the two type parameters like this — and while all subtypes need to report what those parameters are — the exact parameters on the concrete type can and do vary wildly… and the constructors themselves aren’t a part of the array interface. You might as well pass a function or flag instead of the type in any case.

1 Like
8

Constant propagation is an extra layer of nuance that I think was getting missed next to the code_ limitations (“there is no problem with the way the functions are implemented”), which is why I pointed out the fundamental loss of Type{AC} annotations of keyword arguments via NamedTuples. I don’t actually know how fixable that is, but it’s hypothetically possible to pass AC.

@code_warntype f(), or as I prefer to avoid using up names with @code_warntype (()->my_fill2(1, array_constructor=Array))(), is the appropriate usage of @code_warntype when you expect the inputs to be constants or static parameters known by the compiler. The calls in the original post assume the opposite. In either case, I don’t think @code_warntype is lying about type inference this time.

9

About the lack of generality, I know that the restriction AC <: AbstractArray will not solve it. It will accept array_constructor=Matrix but then the constructor AC{T, 1} will error. I’m ok with that. The actual calculation that I’m interested only makes sense if can call AC{T, N}.

If I used functions or flags, then the result would be type unstable, no? Isn’t that the whole point of the Type selector?

10

Yes, it would be if the flags were not distinguished by their type like array_type = :gpu. But if you pass in a function (different functions are different types) or a dedicated struct GPU; end, with array_type = GPU() then you’d be in the clear.

11

I guess in this case I’d need to include the corresponding GPU packages as dependencies, or at least as extensions, of my package. I’d prefer to avoid that. The actual calculations I’m doing involve KernelAbstractions, which allow me to write the kernels without depending on the GPU packages.

12

Ideally you’re not writing your own fill-like function at all, just calling genericfill(sufficient_array_type::AbstractArray, value, dims::Integer...), but that doesn’t exist. As pointed out, abstract types don’t impose any structure on the type parameters of their subtypes, so it’s not feasible for you to write 1 method or constructor call for arbitrary dense array types, even if just limited to 1D vectors. That’s what multiple dispatch is for.

fill!(A, x) however does exist, I just don’t know if all the array types you’re interested in support that, and the users would have to instantiate (a likely uninitialized) A themselves.

13

If you’re already using KernelAbstractions.jl, you can use what they provide, no? See, e.g., KernelAbstractions.allocate and the Backend flags.

1 Like
14

Yes, this is actually exactly what I need, thank you for pointing that out!

1 Like