Type container equivalence

I have struct which encapsulates a value with timestamp, like this.

struct Timestampped{T}
     time:: UInt64
     data:: T
end

Julia seem to have trouble with type stability when it has Timestampped{Union{A, Nothing}}. I tried a bit with covariance, @profview keeps showing them as dynamic dispatch. I had to manually unpack them everywhere and repack them like Union{Timestampped{A}, Timestampped{Nothing}}. Is there a better way to do this?

Can you provide code with a specific example of A? If it is concrete, small union-splitting can remove dynamic dispatch.

I got a function current_value(context) which returns Union{Float32, Nothing} I cannot do Timestampped(current_timestamp, current_value(context)) I had to split it manually with

@inline process(:: Nothing) = Timestampped(current_timestamp, nothing)
@inline process(data:: T) = Timestampped(current_timestamp, data)

otherwise the consumer of the function is flagged with dynamic_call on the line which uses the return value. In essense, its not able to split union on type variable.

Does it only do small union-splitting if the outer-most type has few unions?

What is the function call code for splitting it manually? If it is simply process(current_value(context)) as implied, then union-splitting of the output of current_value(context) is already removing a dynamic dispatch of process.

Then it would be strange to me that you have to go as far as making process. I would have expected Timestampped(current_timestamp, current_value(context)) to first infer that current_value(context) is either an instance of Float32 or Nothing, hence Union{Float32, Nothing} in @code_warntype. Then infer the output as Timestampped{Float32} or Timestampped{Nothing}, reported as Union{Timestampped{Float32}, Timestampped{Nothing}} in @code_warntype. The inferred union of types should not be different from using process.

I’m also a little confused by your mention of Timestampped{Union{A, Nothing}} at the beginning. That is a different concrete type, not a union of inferred types, and I haven’t seen any code so far using it. Union{Float32, Nothing} is an abstract type and thus cannot be a return type, just a union of possible return types according to the compiler’s inference.

If you worry about type stability, then never use Union type!
There are two workaround:

  1. Wrap Union{...} in a sum type and access the value through type-stable function. For example:
struct Maybe{T}
    val::Union{T, Nothing}
end
isNothing(m::Maybe{T})::Bool = m.val isa Nothing
# When you want to use the value
function cast(m::Maybe{T})::T
    mval = m.val
    if mval isa Nothing
        error("Not a valid maybe value")
    end
    return mval
end

This is a bit of inconvenience, though…

  1. Use negative value to represent a invalid timestamp, so there is no type instablity. However this might be dangerous, because it’s easy to forget validating the negative value (which is also a common mistake in C/C++ language).
  1. Like I said, it doesn’t seem like the code he gave could result in Timestampped{Union{Float32, Nothing}}. The Union seems to only be in type inference reports.

  2. It’s hyperbolic to state things like that. A variable inferred to be one of a small (3 or 4) union of types is optimized via union-splitting to conditional branches of compile-time dispatches. It’s the reason why so many functions return nothing as a sentinel value, including iterate. The thing to avoid is combinatorial explosion of such inferred unions in methods like foo(::Union{A,B}, ::Union{C,D,E}) returning any of Union{AC, AD, AE, BC, BD, BE}, which overwhelms the compiler and prevents union-splitting optimization.

iterate is different, because in most case iterate is used in for-loop and followed by a if ... isa Nothing splitting inserted by compiler, so type stability is trivial there. Other base functions like findfirst do return Union{...}, but I never use these functions if type stability is my concern. They are more or less designed for interactive usage (“convenience”), so users don’t need to validate and unpack the resultsin REPL. Even so, I think iterate should return sum type instead of Union{...}.

A variable inferred to be one of a small (3 or 4) union of types is optimized via union-splitting to conditional branches of compile-time dispatches.

Union splitting by compiler can still cause path explosions. Every time you split a union, you need to specialize the following code path (this is basically how abstract interpretation works : specialize basic block until you reach a fixed point). This can slow down compilation (cost of method dispatch is also non-trival). Also, people may forget to handle Union somewhere, which might accumulate graduately along code paths and function calls (if you have recursive function calls or union in a for loop, then you might be unlucky). Since many Julia codes have no type annotations, even a minor type unstability can impact a lot of codes unintentionally.

Like I have said before, if you really concern type stability and you don’t know exactly how Julia compiler works or you have only vague impression, then just avoid anything that could cause problems (Union{…}, Any, Vector{Any}). So you don’t need tools like JET, profview or Cthulu to trace type instability bugs, which hardly produce readable reports because they rely on Julia’s lowered type IR to detect errors, where after lowering a bunch of meaningful debug info has lost. Don’t return Union{...} because you might forget to split the Union. Use sum type instead, which also forces you checking the validity of results). If you have to use them, then always manually split the Union using if x isa ... elseif x isa ...end.

I know it’s boring and tedious to wrap things in sum type and manually split union/sum type everywhere. That’s why other languages like Haskell or Rust have syntax sugar like optional chaining or do-notation to simplify optional value handling. In Julia we don’t have such things (and we don’t even have pattern match or true sum type), so we have to do things the old way.

Maybe this is a bit too extreme, but only a bit.
Whenever you put these unions into containers, type stability can easily be lost. Here’s a small example using only Base types — floats, missings, tuples, and arrays:

Despite the simplicity of this example, the type-stable (float-only) code is 10x faster than float+missing.

Mixing Missing with (eg) Float64 is not a case of type instability. The compiler knows that all values are exactly one of those two types and no dynamic dispatch is necessary. However, the branching necessary to handle the inhomogenous types can have substantial performance implications. Branches (especially those that are not predicted well) have considerable cost on their own, and also usually preclude SIMD optimizations.

You only get massive speedups when you can exploit native NaN behavior directly while using SIMD. As soon as you want to apply special-casing to the NaN/missing case (so that branching is unavoidable) or otherwise cannot use SIMD, the performance gap shrinks considerably. See the following example:

xnan = randn(1000)
@. xnan = ifelse(xnan < -2, NaN, xnan)
xmiss = @. ifelse(isnan(xnan), missing, xnan)
julia> using BenchmarkTools

julia> @btime sum($xnan)
  61.978 ns (0 allocations: 0 bytes)
NaN

julia> @btime sum($xmiss)
  950.000 ns (0 allocations: 0 bytes)
missing

julia> @btime foldl(+,$xnan)
  860.274 ns (0 allocations: 0 bytes)
NaN

julia> @btime foldl(+,$xmiss)
  970.370 ns (0 allocations: 0 bytes)
missing

sum is much faster for NaN than missing. But note that if we instead use foldl(+,x), we see a 14x slowdown for the NaN array due to SIMD and other optimizations being forbidden. Meanwhile the missing array sees almost no degradation for switching from sum to foldl as those optimizations were already impossible. The performance gap becomes relatively small when using foldl.

1 Like

Usually, anyway. In aplavin’s linked example, the array containing heterogenous tuples containing arrays has type Vector{Tuple{Vector}}, Tuple{Vector} being abstract and not a small union. It’s not exactly the same as inference of a variable’s possible types, it’s how Julia determines an array’s element type, which has an extra layer of promotions and avoidance of excessive specificity, in this case of tuples. For example, the type of [(3,), (4.5,)] is Vector{Tuple{Real}}. So, keeping stability to 1 type is still safer, but union-splitting provides an idiomatic flexibility that should be leveraged.

In any case, this has gotten somewhat removed from OP’s problem. Nobody really needs to sum NaNs and missings, rather skip over them with skipmissing, NaNStatistics.jl, or Skipper.jl (authored by aplavin in fact). Even the more general discussion of how to handle small inferred type unions is premature because OP has not said how they wanted to do that, let alone skipping or erroring on sentinel values.

The stated problem is that a processing method wrapping a constructor call seemed to be necessary to restore a Union-splitting optimization and to avoid a dynamic dispatch. Runnable MWEs that replicate the issue and the process fix could reveal some other source of type instability or something that needs a compiler improvement. The available code is not runnable and provides little clarity, in fact @inline process(data:: T) = Timestampped(current_timestamp, data) would just throw UndefVarError: T not defined.

2 Likes