Accept index ranges in ndrange

docs/src/api.md

@@ -91,4 +91,10 @@ KernelAbstractions.@context
 KernelAbstractions.argconvert
 KernelAbstractions.NDIteration.DynamicSize
 KernelAbstractions.NDIteration.StaticSize
+KernelAbstractions.NDIteration.NDRange
+KernelAbstractions.NDIteration.StaticOffset
+KernelAbstractions.NDIteration.DynamicOffset
+KernelAbstractions.NDIteration.extents
+KernelAbstractions.NDIteration.offsets
+KernelAbstractions.NDIteration.linear_index
 ```

docs/src/index.md

@@ -132,6 +132,11 @@ but users must avoid the use of `@index(Global)` and instead use their own deriv
 end
 ```
 
+### 0.10
+- `ndrange` entries may be index ranges, given statically (`kernel(backend, workgroupsize, (-2:N+3, 0:M+1))`)
+  or at launch (`ndrange=(-2:N+3, 0:M+1)`, a single range, or a `CartesianIndices`).
+  `@index(Global, Cartesian)` and `@index(Global, NTuple)` return the shifted indices.
+
 ## Semantic differences
 
 ### To CUDA.jl/AMDGPU.jl

docs/src/kernels.md

@@ -232,5 +232,26 @@ kernel = my_kernel(backend, 32, size(A))
 kernel(A)
 ```
 
+### Index ranges
+
+Each entry of `ndrange` is either an extent (indices `1:n`) or a range of indices, so a kernel
+can iterate over a region whose indices do not start at 1. `ndrange` can also be given as a
+single range or as a `CartesianIndices`:
+
+```julia
+# static ndrange over the indices -2:N+3 along x and 0:M+1 along y
+kernel = my_kernel(backend, (16, 16), (-2:N+3, 0:M+1))
+kernel(A)
+
+# dynamic ndrange
+kernel = my_kernel(backend, (16, 16))
+kernel(A, ndrange=(-2:N+3, 0:M+1))
+kernel(A, ndrange=CartesianIndices(A))   # e.g. for an OffsetArray
+```
+
+Inside the kernel `@index(Global, Cartesian)` and `@index(Global, NTuple)` return the shifted
+indices, `@index(Global, Linear)` counts the indices of the region from 1 in column-major order,
+and `@ndrange()` returns the extents.
+
 Obtain the backend from an array with [`get_backend`](@ref) and always call [`synchronize`](@ref) before reading results on the host.
 See the [Quickstart](@ref) for a full walkthrough and the Examples section of the manual for larger patterns.

src/KernelAbstractions.jl

@@ -414,7 +414,7 @@ end
 @inline function __index_Global_Linear(ctx)
     I = @inbounds expand(__iterspace(ctx), KI.get_group_id().x, KI.get_local_id().x)
     # TODO: This is unfortunate, can we get the linear index cheaper
-    return @inbounds LinearIndices(__ndrange(ctx))[I]
+    return linear_index(__ndrange(ctx), I)
 end
 
 @inline function __index_Local_Cartesian(ctx)
@@ -544,6 +544,8 @@ last (possibly partial) workgroup. Primarily used by backend implementations and
 @inline function partition(kernel, ndrange, workgroupsize)
     static_ndrange = KernelAbstractions.ndrange(kernel)
     static_workgroupsize = KernelAbstractions.workgroupsize(kernel)
+    ndrange = NDIteration.normalize_ndrange(ndrange)
+    workgroupsize = NDIteration.normalize_workgroupsize(workgroupsize)
 
     if ndrange === nothing && static_ndrange <: DynamicSize ||
             workgroupsize === nothing && static_workgroupsize <: DynamicSize
@@ -562,7 +564,7 @@ last (possibly partial) workgroup. Primarily used by backend implementations and
     end
 
     if static_ndrange <: StaticSize
-        if ndrange !== nothing && ndrange != get(static_ndrange)
+        if ndrange !== nothing && !NDIteration.same_axes(ndrange, get(static_ndrange))
             error("Static NDRange ($static_ndrange) and launch NDRange ($ndrange) differ")
         end
         ndrange = get(static_ndrange)
@@ -577,14 +579,16 @@ last (possibly partial) workgroup. Primarily used by backend implementations and
 
     @assert workgroupsize !== nothing
     @assert ndrange !== nothing
-    blocks, workgroupsize, dynamic = NDIteration.partition(ndrange, workgroupsize)
+    blocks, workgroupsize, dynamic = NDIteration.partition(extents(ndrange), workgroupsize)
 
     if static_ndrange <: StaticSize
         static_blocks = StaticSize{blocks}
         blocks = nothing
+        mapping = NDIteration.static_mapping(ndrange)
     else
         static_blocks = DynamicSize
         blocks = CartesianIndices(blocks)
+        mapping = NDIteration.dynamic_mapping(ndrange)
     end
 
     if static_workgroupsize <: StaticSize
@@ -594,7 +598,7 @@ last (possibly partial) workgroup. Primarily used by backend implementations and
         workgroupsize = CartesianIndices(workgroupsize)
     end
 
-    iterspace = NDRange{length(ndrange), static_blocks, static_workgroupsize}(blocks, workgroupsize)
+    iterspace = NDRange{length(ndrange), static_blocks, static_workgroupsize}(blocks, workgroupsize, mapping)
     return iterspace, dynamic
 end
 

src/compiler.jl

@@ -5,21 +5,24 @@ struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace}
 
     # CPU variant
     function CompilerMetadata{NDRange, CB}(idx, ndrange, iterspace) where {NDRange, CB}
-        if ndrange !== nothing
-            ndrange = CartesianIndices(ndrange)
-        end
+        ndrange = cartesian(ndrange)
         return new{NDRange, CB, typeof(idx), typeof(ndrange), typeof(iterspace)}(idx, ndrange, iterspace)
     end
 
     # GPU variante: index is given implicit
     function CompilerMetadata{NDRange, CB}(ndrange, iterspace) where {NDRange, CB}
-        if ndrange !== nothing
-            ndrange = CartesianIndices(ndrange)
-        end
+        ndrange = cartesian(ndrange)
         return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace)}(nothing, ndrange, iterspace)
     end
 end
 
+# `CartesianIndices` covering a launch `ndrange` (any form accepted by `partition`).
+cartesian(::Nothing) = nothing
+cartesian(ci::CartesianIndices) = ci
+cartesian(n::Integer) = CartesianIndices((Int(n),))
+cartesian(r::AbstractUnitRange) = CartesianIndices((r,))
+cartesian(t::Tuple) = CartesianIndices(t)
+
 @inline __iterspace(cm::CompilerMetadata) = cm.iterspace
 @inline __groupindex(cm::CompilerMetadata) = cm.groupindex
 @inline __groupsize(cm::CompilerMetadata) = size(workitems(__iterspace(cm)))

src/nditeration.jl

@@ -2,13 +2,82 @@ module NDIteration
 
 export _Size, StaticSize, DynamicSize, get
 export NDRange, blocks, workitems, expand
+export StaticOffset, DynamicOffset, offsets, extents, linear_index
 export DynamicCheck, NoDynamicCheck
 
 import Base.@pure
 
 struct DynamicCheck end
 struct NoDynamicCheck end
 
+# An axis of an `ndrange` is either an extent (`Int`) or a range of indices (`UnitRange{Int}`).
+axis(n::Integer) = Int(n)
+axis(r::Base.OneTo) = Int(length(r))
+axis(r::AbstractUnitRange) = UnitRange{Int}(r)
+
+extent(n::Integer) = Int(n)
+extent(r::AbstractUnitRange) = length(r)
+
+axis_offset(::Integer) = 0
+axis_offset(r::AbstractUnitRange) = first(r) - 1
+
+"""
+    extents(ndrange)
+
+Number of indices along each axis of `ndrange`, given as a tuple of extents and/or ranges,
+a `CartesianIndices`, a single range, or an integer.
+"""
+extents(t::Tuple) = map(extent, t)
+extents(ci::CartesianIndices) = size(ci)
+extents(r::AbstractUnitRange) = (length(r),)
+extents(n::Integer) = (Int(n),)
+
+"""
+    offsets(ndrange)
+
+Offset of the first index along each axis of `ndrange` relative to 1.
+"""
+offsets(t::Tuple) = map(axis_offset, t)
+
+"""
+    normalize_ndrange(ndrange)
+
+Canonical form of a launch `ndrange`: `nothing`, or a tuple of `Int` extents and
+`UnitRange{Int}` axes.
+"""
+normalize_ndrange(::Nothing) = nothing
+normalize_ndrange(n::Integer) = (Int(n),)
+normalize_ndrange(r::AbstractUnitRange) = (axis(r),)
+normalize_ndrange(ci::CartesianIndices) = map(axis, ci.indices)
+normalize_ndrange(t::Tuple) = map(axis, t)
+
+"""
+    normalize_workgroupsize(workgroupsize)
+
+Canonical form of a launch `workgroupsize`: `nothing`, or a tuple of `Int` extents.
+"""
+normalize_workgroupsize(::Nothing) = nothing
+normalize_workgroupsize(n::Integer) = (Int(n),)
+normalize_workgroupsize(t::Tuple) = extents(t)
+
+# Two ndranges denote the same indices.
+same_axes(a::Tuple, b::Tuple) = extents(a) == extents(b) && offsets(a) == offsets(b)
+
+"""
+    linear_index(ndrange::CartesianIndices, I::CartesianIndex)
+
+Column-major position of `I` within `ndrange`, counted from 1.
+"""
+@inline function linear_index(ndrange::CartesianIndices{N}, I::CartesianIndex{N}) where {N}
+    lo = map(first, ndrange.indices)
+    sz = size(ndrange)
+    idx = I.I[N] - lo[N]
+    for d in (N - 1):-1:1
+        idx = idx * sz[d] + (I.I[d] - lo[d])
+    end
+    return idx + 1
+end
+
 abstract type _Size end
 
 """
@@ -22,29 +91,54 @@ struct DynamicSize <: _Size end
     StaticSize{S}
 
 Marker type encoding a compile-time workgroup size or `ndrange` as a tuple `S`.
+Each entry of `S` is an `Int` extent or, for an `ndrange` axis whose indices do not start
+at 1, a `UnitRange{Int}`.
 """
 struct StaticSize{S} <: _Size
     function StaticSize{S}() where {S}
-        return new{S::Tuple{Vararg{Int}}}()
+        return new{S::Tuple{Vararg{Union{Int, UnitRange{Int}}}}}()
     end
 end
 
 @pure StaticSize(s::Tuple{Vararg{Int}}) = StaticSize{s}()
 @pure StaticSize(s::Int...) = StaticSize{s}()
 @pure StaticSize(s::Type{<:Tuple}) = StaticSize{tuple(s.parameters...)}()
+StaticSize(s::Tuple{Vararg{Union{Integer, AbstractUnitRange{<:Integer}}}}) = StaticSize{map(axis, s)}()
+StaticSize(ci::CartesianIndices) = StaticSize(ci.indices)
 
 # Some @pure convenience functions for `StaticSize`
 @pure get(::Type{StaticSize{S}}) where {S} = S
 @pure get(::StaticSize{S}) where {S} = S
 @pure Base.getindex(::StaticSize{S}, i::Int) where {S} = i <= length(S) ? S[i] : 1
 @pure Base.ndims(::StaticSize{S}) where {S} = length(S)
-@pure Base.length(::StaticSize{S}) where {S} = prod(S)
+@pure Base.length(::StaticSize{S}) where {S} = prod(extents(S))
 
+"""
+    StaticOffset{O}
+
+Compile-time offset `O::NTuple{N, Int}` added to the indices produced by an [`NDRange`](@ref).
+"""
+struct StaticOffset{O}
+    function StaticOffset{O}() where {O}
+        return new{O::Tuple{Vararg{Int}}}()
+    end
+end
+
+"""
+    DynamicOffset{N}
+
+Runtime offset added to the indices produced by an [`NDRange`](@ref).
+"""
+struct DynamicOffset{N}
+    offset::NTuple{N, Int}
+end
 
 """
     NDRange
 
-Encodes a blocked iteration space. 
+Encodes a blocked iteration space. The `mapping` field relates blocked indices to
+`ndrange` indices: `nothing` for the identity, or a [`StaticOffset`](@ref)/[`DynamicOffset`](@ref)
+for an `ndrange` whose indices do not start at 1.
 
 # Example
 ```
@@ -58,16 +152,17 @@ for block in ndrange
 end
 ```
 """
-struct NDRange{N, StaticBlocks, StaticWorkitems, DynamicBlock, DynamicWorkitems}
+struct NDRange{N, StaticBlocks, StaticWorkitems, DynamicBlock, DynamicWorkitems, Mapping}
     blocks::DynamicBlock
     workitems::DynamicWorkitems
+    mapping::Mapping
 
     function NDRange{N, B, W}() where {N, B, W}
-        return new{N, B, W, Nothing, Nothing}(nothing, nothing)
+        return new{N, B, W, Nothing, Nothing, Nothing}(nothing, nothing, nothing)
     end
 
-    function NDRange{N, B, W}(blocks, workitems) where {N, B, W}
-        return new{N, B, W, typeof(blocks), typeof(workitems)}(blocks, workitems)
+    function NDRange{N, B, W}(blocks, workitems, mapping = nothing) where {N, B, W}
+        return new{N, B, W, typeof(blocks), typeof(workitems), typeof(mapping)}(blocks, workitems, mapping)
     end
 end
 
@@ -77,18 +172,29 @@ end
 @inline blocks(range::NDRange{N, B}) where {N, B <: StaticSize} = CartesianIndices(get(B))::CartesianIndices{N}
 @inline Base.ndims(::NDRange{N}) where {N} = N
 
+@inline offsets(::NDRange{N, B, W, DB, DW, Nothing}) where {N, B, W, DB, DW} = ntuple(_ -> 0, Val(N))
+@inline offsets(::NDRange{N, B, W, DB, DW, StaticOffset{O}}) where {N, B, W, DB, DW, O} = O
+@inline offsets(range::NDRange{N, B, W, DB, DW, DynamicOffset{N}}) where {N, B, W, DB, DW} = range.mapping.offset
+
+# Mapping of a partitioned `ndrange` (in canonical form); a plain size tuple has no mapping.
+static_mapping(::Tuple{Vararg{Int}}) = nothing
+static_mapping(t::Tuple) = StaticOffset{offsets(t)}()
+dynamic_mapping(::Tuple{Vararg{Int}}) = nothing
+dynamic_mapping(t::Tuple) = DynamicOffset(offsets(t))
+
 import Base.iterate
 @inline iterate(range::NDRange) = iterate(blocks(range))
 @inline iterate(range::NDRange, state) = iterate(blocks(range), state)
 
 Base.length(range::NDRange) = length(blocks(range))
 
 @inline function expand(ndrange::NDRange{N}, groupidx::CartesianIndex{N}, idx::CartesianIndex{N}) where {N}
+    offset = offsets(ndrange)
     nI = ntuple(Val(N)) do I
         Base.@_inline_meta
         stride = size(workitems(ndrange), I)
         gidx = groupidx.I[I]
-        (gidx - 1) * stride + idx.I[I]
+        (gidx - 1) * stride + idx.I[I] + offset[I]
     end
     return CartesianIndex(nI)
 end
@@ -153,6 +259,8 @@ Returns the number of workgroups necessary and whether the last workgroup
 needs to perform dynamic bounds-checking.
 """
 @inline function partition(ndrange, __workgroupsize)
+    ndrange = extents(ndrange)
+    __workgroupsize = extents(__workgroupsize)
     @assert length(__workgroupsize) <= length(ndrange)
     # pad workgroupsize with ones
     workgroupsize = ntuple(Val(length(ndrange))) do I

src/pocl/backend.jl

@@ -196,7 +196,7 @@ function (obj::KA.Kernel{POCLBackend})(args::Vararg{Any, N}; ndrange = nothing,
     # figure out the optimal workgroupsize automatically
     if KA.workgroupsize(obj) <: KA.DynamicSize && workgroupsize === nothing
         wg_info = cl.work_group_info(kernel.fun, device())
-        wg_size_nd = threads_to_workgroupsize(wg_info.size, ndrange)
+        wg_size_nd = threads_to_workgroupsize(wg_info.size, KA.NDIteration.extents(ndrange))
         iterspace, dynamic = KA.partition(obj, ndrange, wg_size_nd)
         ctx = KA.mkcontext(obj, ndrange, iterspace)
     end

test/nditeration.jl

@@ -16,6 +16,37 @@ function nditeration_testsuite()
         end
     end
 
+    @testset "offsets" begin
+        @test NDIteration.get(StaticSize((1:4, 0:9))) == (1:4, 0:9)
+        @test NDIteration.get(StaticSize(CartesianIndices((3, 0:9)))) == (3, 0:9)
+        @test length(StaticSize((1:4, 0:9))) == 40
+        @test extents((1:4, 0:9, 7)) == (4, 10, 7)
+        @test extents(CartesianIndices((3, 0:9))) == (3, 10)
+        @test extents(0:9) == (10,)
+        @test offsets((1:4, 0:9, 7)) == (0, -1, 0)
+
+        let ndrange = NDRange{2, StaticSize{(4, 4)}, StaticSize{(8, 8)}}(nothing, nothing, StaticOffset{(-8, 3)}())
+            @test offsets(ndrange) == (-8, 3)
+            @test expand(ndrange, CartesianIndex(1, 1), CartesianIndex(1, 1)) == CartesianIndex(-7, 4)
+            @test expand(ndrange, CartesianIndex(4, 4), CartesianIndex(8, 8)) == CartesianIndex(24, 35)
+        end
+        let ndrange = NDRange{2, DynamicSize, DynamicSize}(CartesianIndices((4, 4)), CartesianIndices((8, 8)), DynamicOffset((-8, 3)))
+            @test offsets(ndrange) == (-8, 3)
+            @test expand(ndrange, 1, 1) == CartesianIndex(-7, 4)
+            @test expand(ndrange, 16, 64) == CartesianIndex(24, 35)
+        end
+        let ndrange = NDRange{2, DynamicSize, DynamicSize}(CartesianIndices((4, 4)), CartesianIndices((8, 8)))
+            @test offsets(ndrange) == (0, 0)
+            @test ndrange.mapping === nothing
+        end
+
+        let ci = CartesianIndices((-3:4, 2:11))
+            @test linear_index(ci, CartesianIndex(-3, 2)) == 1
+            @test linear_index(ci, CartesianIndex(4, 2)) == 8
+            @test linear_index(ci, CartesianIndex(4, 11)) == 80
+        end
+    end
+
     # GPU scenario where we get a linear index into workitems/blocks
     function linear_iteration(ndrange)
         idx = Array{CartesianIndex{2}}(undef, length(blocks(ndrange)) * length(workitems(ndrange)))