`DiskIndex` is not inferred

[tiemvanderdeure]

For views of more than 2 dimensions, DiskIndex is not inferred, which then spreads - so even getindex becomes Any.

I think this is because of the recursive call to _resolve_indices (there are multiple dispatches, in the case of the MWE below it is this one):

DiskArrays.jl/src/diskindex.jl

Lines 52 to 56 in 716f509

	function _resolve_indices(chunks, i, indices_pre::DiskIndex, strategy::BatchStrategy)
	inow = first(i)
	indices_new, chunksrem = process_index(inow, chunks, strategy)
	_resolve_indices(chunksrem, tail(i), merge_index(indices_pre, indices_new), strategy)
	end

MWE:

import DiskArrays
using DiskArrays.TestTypes

a = AccessCountDiskArray(reshape(1:24, 2, 3, 4), chunksize=(2, 2, 2))
a_view = @view a[:,1:2,2]
a_view2= @view a[:,1:2,2:4]

@code_warntype DiskArrays.DiskIndex(a_view, (1:1,1:1), DiskArrays.NoBatch()) # all good
@code_warntype DiskArrays.DiskIndex(a_view2, (1:1,1:1,1:1), DiskArrays.NoBatch()) #DiskArrays.DiskIndex

@code_warntype a_view[1:1,1:1] # All good
@code_warntype a_view2[1:1,1:1] # Any

I've tried just adding an @inline and a Base.@constprop :aggressive but those didn't work.

@lupemba ran into this one in JuliaGeo/CommonDataModel.jl#35

[tiemvanderdeure]

Tested on some old versions and I'm getting the same results back to at least 0.4.0

[rafaqz]

Did you have a look with Cthulhu.jl @descend ?

[tiemvanderdeure]

I did! But then when I look into the second call to resolve_indices, everything is blue and I am getting a warning

┌ Warning: Inference terminated in an incomplete state due to argument-type changes during recursion
└ @ Cthulhu ~/.julia/packages/Cthulhu/kMTSr/src/codeview.jl:164
52 function _resolve_indices(chunks, i, indices_pre::DiskIndex, strategy::BatchStrategy)
53     inow::UnitRange{Int64} = first(i)::UnitRange{Int64}
54     indices_new::DiskArrays.DiskIndex{1, 1, Tuple{Colon}, Tuple{Colon}, Tuple{UnitRange{Int64}}}, chunksrem::Tuple{DiskArrays.RegularChunks} = process_index(inow::UnitRange{Int64}, chunks, strategy)::Tuple{DiskArrays.RegularChunks}
55     _resolve_indices(chunksrem::Tuple{DiskArrays.RegularChunks}, Base.tail(i), merge_index(indices_pre, indices_new::DiskArrays.DiskIndex{1, 1, Tuple{Colon}, Tuple{Colon}, Tuple{UnitRange{Int64}}})::DiskArrays.DiskIndex{2, 2, Tuple{Colon, Colon}, Tuple{Colon, Colon}, Tuple{UnitRange{Int64}, UnitRange{Int64}}}, strategy)
56 end
Select a call to descend into or ↩ to ascend. [q]uit. [b]ookmark.
Toggles: [w]arn, [h]ide type-stable statements, [t]ype annotations, [s]yntax highlight for Source/LLVM/Native, [j]ump to source always, [v]scode: inlay types, [V]scode: diagnostics.
Show: [S]ource code, [A]ST, [T]yped code, [L]LVM IR, [N]ative code
Actions: [E]dit source code, [R]evise and redisplay
 • first(i)
   process_index(inow, chunks, strategy)
   process_index(inow, chunks, strategy)
   process_index(inow::UnitRange{Int64}, chunks, strategy)
   %14 = tail(::Tuple{UnitRange{Int64}, UnitRange{Int64}})::Tuple{UnitRange{Int64}}
   merge_index(indices_pre, indices_new::DiskArrays.DiskIndex{1, 1, Tuple{Colon}, Tuple{Colon}, Tuple{UnitRange{Int64}}})
   _resolve_indices(chunksrem::Tuple{DiskArrays.RegularChunks}, Base.tail(i), merge_index(indices_pre, indices_new::DiskArrays.DiskIndex{1, 1, Tuple{Colon}, Tuple{Colon}, Tuple…

Seems like the compiler is giving up because of the recursive calls, and Cthulhu is having a hard time showing that?

[rafaqz]

Yeah its actually missing the problem somehow, which is process_index being type unstable - a few lines above the first time Cthulhu shows any type instability. But if you go in and look at the process_index call its unstable.

But its actually really simple - broadcasts over Tuples is just always bad. Fixed immediately with map.

OMG revise keeps somehow fixing the type stability but there is something very weird about how the return values of process_index are being split into indices_new, chunksrem that just isn't type stable as its iterating and leaving a Union type for some reason.

[tiemvanderdeure]

Yeah I was just looking at it and couldn't find the broadcast over Tuples. Very weird

[rafaqz]

There are quite a few broadcasts over tuples in there, I'm cleaning all that up, but it seems the compiler has some directive to stop caring after the third recursion (or something like that) and even chtulhu doesn't know whats happening after that

[tiemvanderdeure]

Locally I just tried to tweak the code such that process_index just returns indices_new and get chunksrem in a separate line, but it's still type unstable. Plus @code_warntype a_view[1:1,1:1] also hits that path and it actually is type stable. So maybe that's a red herring and it actually is the recursion that's causing it?

[tiemvanderdeure]

it seems the compiler has some directive to stop caring after the third recursion (or something like that) and even chtulhu doesn't know whats happening after that

I think that's probably what's happening. Seems there have been some attempts to address this JuliaLang/julia#48059

[rafaqz]

Yeah - we can actually kick the can down the road by doing this:

function _resolve_indices(chunks::Tuple{<:Any}, i::Tuple{<:Any}, indices_pre::DiskIndex, strategy::BatchStrategy)
    inow = first(i)
    indices_new, chunksrem = process_index(inow, chunks, strategy)
    return merge_index(indices_pre, indices_new)
end

Which makes the ndims 3 view type stable. But then 4 isn't type stable. But if we can kick it out to 5 then it hardly ever matters

[tiemvanderdeure]

Or we can use some macro trick similar to in base apparently JuliaLang/julia#48059 (comment)

[rafaqz]

Ugh thats kind of brutal tho

As dumb as it is, adding these methods works for both 3 and 4 dimensional arrays.

# This makes 3 arg type stable
function _resolve_indices(chunks::Tuple{<:Any}, i::Tuple{<:Any}, indices_pre::DiskIndex, strategy::BatchStrategy)
    inow = first(i)
    indices_new, chunksrem = process_index(inow, chunks, strategy)
    return merge_index(indices_pre, indices_new)
end
# This makes 4 arg type stable
function _resolve_indices(chunks::Tuple{<:Any,<:Any}, i::Tuple{<:Any,<:Any}, indices_pre::DiskIndex, strategy::BatchStrategy)
    inow = first(i)
    indices_new, chunksrem = process_index(inow, chunks, strategy)
    return _resolve_indices(chunksrem, tail(i), merge_index(indices_pre, indices_new), strategy)
end

I guess it tricks the recursion detection by giving it different methods to look at?

[rafaqz]

And, of course this even works for 5!

# This makes 5 arg type stable
function _resolve_indices(chunks::Tuple{<:Any,<:Any,<:Any}, i::Tuple{<:Any,<:Any,<:Any}, indices_pre::DiskIndex, strategy::BatchStrategy)
    inow = first(i)
    indices_new, chunksrem = process_index(inow, chunks, strategy)
    return _resolve_indices(chunksrem, tail(i), merge_index(indices_pre, indices_new), strategy)
end

How far do we need to go?? (turns out 6 works too, I will stop there...)

[tiemvanderdeure]

In practice I'm sure anything more than 5 will be super rare - but it's not a very satisfying solution. Maybe we can get some ideas for a more proper fix from the slack helpdesk.