Support for `allocate(backend, desired_axes)` in addition to `allocate(backend, desired_size)`

[dzhang314]

Hey KA.jl team, thanks for this amazing package! I'm very impressed at how uniquely easy it is to write portable GPU code using KA.jl, and I wish other programming ecosystems had this awesome power.

Recently, I've been rewriting a lot of my Julia code to use axes instead of size, and I noticed that allocate(backend, T, desired_axes) doesn't work, while allocate(backend, T, desired_size) does. This would be a nice feature to bring the allocate interface in line with Base.similar, which accepts size or axes. Are there any plans to support this? I would be happy to help implement this myself if someone can point me to the right code paths to looks at.

Note that this would be helpful for API compatibility even if KA.jl doesn't plan to support offset axes or symbolic axes -- we could just throw on any axis that isn't a 1:n or Base.OneTo(n).

[vchuravy]

KA is split across multiple packages, there are the backends in CUDA/AMDGPU/oneAPI/OpenCL... and then of course the core here.

One question for now would be what does similar do with offset axes on Array?

In any case maybe your second suggestion you could look at:

KernelAbstractions.jl/src/KernelAbstractions.jl

Lines 526 to 535 in f600e50

	"""
	allocate(::Backend, Type, dims...)::AbstractArray
	Allocate a storage array appropriate for the computational backend.
	!!! note
	Backend implementations **must** implement `allocate(::NewBackend, T, dims::Tuple)`
	"""
	allocate(backend::Backend, T::Type, dims...) = allocate(backend, T, dims)
	allocate(backend::Backend, T::Type, dims::Tuple) = throw(MethodError(allocate, (backend, T, dims)))

[dzhang314]

Base.similar only accepts axes of type Base.OneTo. You get a MethodError if you try to pass in axes of a different type.
OffsetArrays.jl modifies this behavior through type-piracy on Base.similar, adding methods that accept UnitRange axes.

               _
   _       _ _(_)_     |  Documentation: https://docs.julialang.org
  (_)     | (_) (_)    |
   _ _   _| |_  __ _   |  Type "?" for help, "]?" for Pkg help.
  | | | | | | |/ _` |  |
  | | |_| | | | (_| |  |  Version 1.10.9 (2025-03-10)
 _/ |\__'_|_|_|\__'_|  |  Official https://julialang.org/ release
|__/                   |

julia> similar(Float64[], Base.OneTo(4), Base.OneTo(4))
4×4 Matrix{Float64}:
 0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0

julia> similar(Float64[], Base.OneTo(4), 1:4)
ERROR: MethodError: no method matching similar(::Vector{Float64}, ::Type{Float64}, ::Tuple{Int64, UnitRange{Int64}})

Closest candidates are:
  similar(::Vector{T}, ::Type) where T
   @ Base array.jl:417
  similar(::Array, ::Type, ::Tuple{Vararg{Int64, N}}) where N
   @ Base array.jl:420
  similar(::Base.Broadcast.Broadcasted{Base.Broadcast.ArrayConflict}, ::Type{ElType}, ::Any) where ElType
   @ Base broadcast.jl:229
  ...

Stacktrace:
 [1] similar(::Vector{Float64}, ::Base.OneTo{Int64}, ::UnitRange{Int64})
   @ Base ./abstractarray.jl:831
 [2] top-level scope
   @ REPL[2]:1

julia> using OffsetArrays

julia> similar(Float64[], Base.OneTo(4), 1:4)
4×4 OffsetArray(::Matrix{Float64}, 1:4, 1:4) with eltype Float64 with indices 1:4×1:4:
 0.0           0.0           0.0           0.0
 0.0           0.0           6.75045e-310  0.0
 0.0           6.75045e-310  6.75045e-310  0.0
 6.75045e-310  0.0           0.0           0.0

I think this would be a reasonable default behavior for KernelAbstractions.jl. We can manually handle Base.OneTo axes with the standard size-based allocate, then leave it up to CUDA.jl, AMDGPU.jl, etc. if they want to support other axis types by adding methods to allocate.

[vchuravy]

Yeah essentially matching:

julia> zeros(Base.OneTo(10))
10-element Vector{Float64}:
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0