Possible type instability in Base.reducedim_init

[pevnak]

Hi All,

I have been hunting a weird type instability in Flux's softmax function and I think the culprit is in Base.reducedim_init. the following snippet

using JETTest
x = randn(Float32, 2, 2)
@report_dispatch Base.reducedim_init(identity, max, x, 1)

outputs

No errors !
(Union{Matrix{Float32}, Matrix{Float64}}, 0)

but I would expect this to return Matrix{Float32} as there should not be a possibility for this to automatically promote to Matrix{Float64}. The @code_typed Base.reducedim_init(identity, max, x, 1) has the same output as @report_dispatch returning => Union{Matrix{Float32}, Matrix{Float64}}.

I am running MacOSX

julia> versioninfo()
Julia Version 1.7.0-beta4
Commit d0c90f37ba (2021-08-24 12:35 UTC)
Platform Info:
  OS: macOS (x86_64-apple-darwin19.6.0)
  CPU: Intel(R) Core(TM) i5-8279U CPU @ 2.40GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-12.0.1 (ORCJIT, skylake)

with JETTest v0.1.4.

Tomas

[timholy]

Probably from 76952a8, CC @cmcaine. Inf is a Float64,

julia/base/reducedim.jl

Line 128 in 3794f9a

for (f1, f2, initval, typeextreme) in ((:min, :max, :Inf, :typemax), (:max, :min, :(-Inf), :typemin))

[pevnak]

I have tried to change it to

for (f1, f2, initval, typeextreme) in ((:min, :max, :(typemax(T)), :typemax), (:max, :min, :(typemin(T)), :typemin))

but it did not help.

[KristofferC]

Are you sure? It seems like it fixed it for me.

[cmcaine]

Some types (e.g. BigInt) don't support typemin/typemax, which is why Inf and -Inf are used.

This line should maybe be changed:

julia/base/reducedim.jl

Line 153 in 3794f9a

v0 = $initval

Tr($initval) is probably what it should be.

[pevnak]

It has fixed @code_typed but not the JETTest. But I might be doing something wrong.

[cmcaine]

But also I can't think of a type that would have NaN and not have Inf, so if we assume that we can simplify this code at the cost of running nonmissingtype, but probably that gets done at compile time anyway.

[cmcaine]

There's another type instability here that depends on whether the first slice of the array A is empty:

julia/base/reducedim.jl

Lines 138 to 166 in 3794f9a

	A1 = view(A, ri...)
	if isempty(A1)
	# If the slice is empty just return non-view version as the initial array
	return copy(A1)
	else
	# otherwise use the min/max of the first slice as initial value
	v0 = mapreduce(f, $f2, A1)
	T = _realtype(f, promote_union(eltype(A)))
	Tr = v0 isa T ? T : typeof(v0)
	# but NaNs and missing need to be avoided as initial values
	if (v0 == v0) === false
	# v0 is NaN
	v0 = $initval
	elseif isunordered(v0)
	# v0 is missing or a third-party unordered value
	Tnm = nonmissingtype(Tr)
	# TODO: Some types, like BigInt, don't support typemin/typemax.
	# So a Matrix{Union{BigInt, Missing}} can still error here.
	v0 = $typeextreme(Tnm)
	end
	# v0 may have changed type.
	Tr = v0 isa T ? T : typeof(v0)
	return reducedim_initarray(A, region, v0, Tr)
	end
	end

In the first part of the if/else we copy a view of the array (so we use return an array matching the type of A).

In the else we return an array of type Tr, which depends on the heuristic _realtype and if that heuristic fails it just uses the type of v0.

We can't fix this perfectly because f could do something wacky, but we should probably use the heuristic in both or neither branches.

[cmcaine]

I never really understood what this function was doing, so I tried not to change it much, but I think we could maybe replace else with return map(f, A1)?

Maybe we could replace the whole if/else with that and remove the type guessing altogether?

[timholy]

I also wonder if we could just eliminate the init and perform the first reduction operation. Kind of like this:

result = first(X)
for x in Iterators.drop(X, 1)
    result += x
end

In principle this could also be a performance enhancement: I've seen cases where that fill! operation is noticeable in profiles.

Need to think about empty X and case where the operator changes the type though.

[pevnak]

I like Tim's approach, as the current implementation overengineered, but likely I do not see all details. For example what happens when first(X) returns missing. It is not clear to me, how missing should be handled. Therefore something like

result, skip_items = first_non_missing(x)
for x in Iterators.drop(X, skip_items)
    result += x
end

If operator changes the type, then the operation would by unstable by definition, I would expect.