Allow for ndims query (#551)

Author: vchuravy

src/compiler.jl

@@ -30,3 +30,4 @@ end
 
 @inline groupsize(ctx) = __groupsize(ctx)
 @inline ndrange(ctx) = __ndrange(ctx)
+@inline Base.ndims(ctx) = ndims(__iterspace(ctx))

src/nditeration.jl

@@ -63,6 +63,7 @@ end
 @inline workitems(range::NDRange{N, B, W}) where {N, B, W <: StaticSize} = CartesianIndices(get(W))::CartesianIndices{N}
 @inline blocks(range::NDRange{N, B}) where {N, B <: DynamicSize} = range.blocks::CartesianIndices{N}
 @inline blocks(range::NDRange{N, B}) where {N, B <: StaticSize} = CartesianIndices(get(B))::CartesianIndices{N}
+@inline Base.ndims(::NDRange{N}) where {N} = N
 
 import Base.iterate
 @inline iterate(range::NDRange) = iterate(blocks(range))

test/compiler.jl

@@ -41,6 +41,7 @@ function compiler_testsuite(backend, ArrayT)
     kernel = index(CPU(), DynamicSize(), DynamicSize())
     iterspace = NDRange{1, StaticSize{(128,)}, StaticSize{(8,)}}();
     ctx = KernelAbstractions.mkcontext(kernel, 1, nothing, iterspace, Val(KernelAbstractions.NoDynamicCheck()))
+    @test ndims(ctx) == 1
     @test KernelAbstractions.__index_Global_NTuple(ctx, CartesianIndex(1)) == (1,)
 
     A = ArrayT{Int}(undef, 1)

test/nditeration.jl

@@ -7,10 +7,12 @@ function nditeration_testsuite()
         let ndrange = NDRange{2, DynamicSize, DynamicSize}(CartesianIndices((256, 256)), CartesianIndices((32, 32)))
             @test length(ndrange) == 256 * 256
             @test all(p -> p[1] == p[2], zip(ndrange, CartesianIndices((256, 256))))
+            @test ndims(ndrange) == 2
         end
         let ndrange = NDRange{2, StaticSize{(256, 256)}, DynamicSize}(nothing, CartesianIndices((32, 32)))
             @test length(ndrange) == 256 * 256
             @test all(p -> p[1] == p[2], zip(ndrange, CartesianIndices((256, 256))))
+            @test ndims(ndrange) == 2
         end
     end
 
@@ -40,13 +42,15 @@ function nditeration_testsuite()
                 I = Tuple(I)
                 @test check(idx, i - 1, ntuple(i -> I[i] - 1, length(I))..., Dim_x, Dim_y)
             end
+            @test ndims(ndrange) == 2
         end
         let ndrange = NDRange{2, DynamicSize, DynamicSize}(CartesianIndices((4, 4)), CartesianIndices((Dim_x, Dim_y)))
             idx = linear_iteration(ndrange)
             for (i, I) in zip(1:length(blocks(ndrange)), blocks(ndrange))
                 I = Tuple(I)
                 @test check(idx, i - 1, ntuple(i -> I[i] - 1, length(I))..., Dim_x, Dim_y)
             end
+            @test ndims(ndrange) == 2
         end
 
         Dim_x = 32
@@ -58,13 +62,15 @@ function nditeration_testsuite()
                 I = Tuple(I)
                 @test check(idx, i - 1, ntuple(i -> I[i] - 1, length(I))..., Dim_x, Dim_y)
             end
+            @test ndims(ndrange) == 2
         end
         let ndrange = NDRange{2, DynamicSize, DynamicSize}(CartesianIndices((4, 4 * 32)), CartesianIndices((Dim_x, Dim_y)))
             idx = linear_iteration(ndrange)
             for (i, I) in zip(1:length(blocks(ndrange)), blocks(ndrange))
                 I = Tuple(I)
                 @test check(idx, i - 1, ntuple(i -> I[i] - 1, length(I))..., Dim_x, Dim_y)
             end
+            @test ndims(ndrange) == 2
         end
 
         Dim_x = 1
@@ -76,13 +82,15 @@ function nditeration_testsuite()
                 I = Tuple(I)
                 @test check(idx, i - 1, ntuple(i -> I[i] - 1, length(I))..., Dim_x, Dim_y)
             end
+            @test ndims(ndrange) == 2
         end
         let ndrange = NDRange{2, DynamicSize, DynamicSize}(CartesianIndices((4 * 32, 4)), CartesianIndices((Dim_x, Dim_y)))
             idx = linear_iteration(ndrange)
             for (i, I) in zip(1:length(blocks(ndrange)), blocks(ndrange))
                 I = Tuple(I)
                 @test check(idx, i - 1, ntuple(i -> I[i] - 1, length(I))..., Dim_x, Dim_y)
             end
+            @test ndims(ndrange) == 2
         end
     end
 end