CuIterator doesn't specify IteratorSize but has no length()

[mentics]

Describe the bug

I’m trying to use CuIterator and Flux.train!. Attempting to do so gives an error because train! tries to get the length of the CuIterator. So this:

iter = CuIterator(batch(i) for i in 1:10)
Flux.train!(loss, ps, iter, opt)

results in an error:

MethodError: no method matching length(::CuIterator{Base.Generator{UnitRange{Int64}, var"#17#18"}})

It appears that HasLength() is the default for Base.IteratorSize and CuIterator doesn’t specify it, so it uses the default.

To reproduce

The Minimal Working Example (MWE) for this bug:

data = CuIterator(x for x in .1:.1:1.)
itrsz = Base.IteratorSize(typeof(data))
n = (itrsz == Base.HasLength()) || (itrsz == Base.HasShape{1}()) ? length(data) : 0

Manifest.toml

CUDA v3.12.0

Expected behavior

As suggested here, CuIterator should probably delegate HasLength and length to the wrapped iterator.

Version info

Details on Julia:

Julia Version 1.8.0
Commit 5544a0fab7 (2022-08-17 13:38 UTC)
Platform Info:
  OS: Windows (x86_64-w64-mingw32)
  CPU: 24 × AMD Ryzen 9 5900X 12-Core Processor
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-13.0.1 (ORCJIT, znver3)
  Threads: 24 on 24 virtual cores
Environment:
  JULIA_EDITOR = code

Details on CUDA:

CUDA toolkit 11.7, artifact installation
NVIDIA driver 516.94.0, for CUDA 11.7
CUDA driver 11.7

Libraries:
- CUBLAS: 11.10.1
- CURAND: 10.2.10
- CUFFT: 10.7.1
- CUSOLVER: 11.3.5
- CUSPARSE: 11.7.3
- CUPTI: 17.0.0
- NVML: 11.0.0+516.94
- CUDNN: 8.30.2 (for CUDA 11.5.0)
- CUTENSOR: 1.4.0 (for CUDA 11.5.0)

Toolchain:
- Julia: 1.8.0
- LLVM: 13.0.1
- PTX ISA support: 3.2, 4.0, 4.1, 4.2, 4.3, 5.0, 6.0, 6.1, 6.3, 6.4, 6.5, 7.0, 7.1, 7.2
- Device capability support: sm_35, sm_37, sm_50, sm_52, sm_53, sm_60, sm_61, sm_62, sm_70, sm_72, sm_75, sm_80, sm_86

1 device:
  0: NVIDIA GeForce RTX 2080 SUPER (sm_75, 6.783 GiB / 8.000 GiB available)

[maleadt]

Sounds like a good idea. We should probably cache the length upon construction of the iterator. Care to create a PR? 🙂

[mentics]

Trying this out and looking at https://docs.julialang.org/en/v1/manual/interfaces/, it seems we probably need at least the following:

Base.IteratorSize(::Type{<:CuIterator{B}}) where B = Base.IteratorSize(B)
Base.HasLength(::Type{<:CuIterator{B}}) where B = Base.HasLength(B)
Base.HasShape{N}(::Type{<:CuIterator{B}}) where {B,N} = Base.HasShape{N}(B)
Base.IteratorEltype(::Type{<:CuIterator{B}}) where B = Base.IteratorEltype(B)

Base.size(itr::T) where T<:CuIterator = Base.size(itr.batches)
Base.length(itr::T) where T<:CuIterator = Base.length(itr.batches)
Base.eltype(itr::T) where T<:CuIterator = Base.eltype(itr.batches)

This seemed to work for me, but I'm a little new to Julia, so not sure.

I wouldn't expect caching the length would be needed. Couldn't the compiler at least inline and perhaps compile away the above anyway so there wouldn't be any performance penalty?

[mcabbott]

I think you should only define Base.IteratorSize and Base.IteratorEltype; the other types exist to be returned by that, but aren't directly called on anything.

julia> Base.IteratorSize(UnitRange{Int})
Base.HasShape{1}()

julia> methods(Base.HasLength)
# 1 method for type constructor:
 [1] Base.HasLength()

I would also have guessed that just passing these through without caching would be fine. The traits certainly depend only on the types, so ought to compile away. I suppose there are iterators whose length takes a while to compute but would guess they are rare, and if desired should probably take care of caching this themselves?

julia> @btime length($("abc"));
  min 4.500 ns, mean 4.628 ns (0 allocations)

julia> @btime length($(String(rand('a':'z', 1000))));
  min 713.686 ns, mean 719.636 ns (0 allocations)

[mentics]

I think you should only define Base.IteratorSize and Base.IteratorEltype; the other types exist to be returned by that, but aren't directly called on anything.

When only defining Base.IteratorSize, I got further errors. In your example, if someone has a CuIterator{UnitRange{Int}} then Base.IteratorSize(CuIterator{UnitRange{Int}}) would then delegate and return Base.HasShape{1}(). But the contract says: if HasShape is returned for a type, length and size must be defined for that type. I got an error because code was calling axes(itr::CuIterator{T}) and failed because size(itr::CuIterator{T}) was not defined. By delegating, you're claiming that CuIterator can support all these cases, so I think you have to define the other methods to satisfy that contract.

[mcabbott]

Yes, sorry I meant should only define of the traits those two; length, size, axes are still needed. The trait Base.IteratorSize tells you whether it's safe to assume those exist.

For eltype I think you want to define it on the type not the instance, something like Base.eltype(::Type{ CuIterator{T}}) where T = eltype(T).