From cc9bf887df66f2958ce6b0a6b9c2f0401e5c8aee Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Tue, 7 Jan 2025 16:04:14 +0100 Subject: [PATCH 1/2] Allow for ndims query --- src/compiler.jl | 1 + src/nditeration.jl | 1 + test/compiler.jl | 1 + test/nditeration.jl | 8 ++++++++ 4 files changed, 11 insertions(+) diff --git a/src/compiler.jl b/src/compiler.jl index c60bc83be..45b7bc07b 100644 --- a/src/compiler.jl +++ b/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)) diff --git a/src/nditeration.jl b/src/nditeration.jl index ab3fd4ec3..66b307dcc 100644 --- a/src/nditeration.jl +++ b/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)) diff --git a/test/compiler.jl b/test/compiler.jl index cf86386e6..7b0efd7ed 100644 --- a/test/compiler.jl +++ b/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) diff --git a/test/nditeration.jl b/test/nditeration.jl index e00ac3ce5..2258e8d4b 100644 --- a/test/nditeration.jl +++ b/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,6 +42,7 @@ 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) @@ -47,6 +50,7 @@ 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 Dim_x = 32 @@ -58,6 +62,7 @@ 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) @@ -65,6 +70,7 @@ 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 Dim_x = 1 @@ -76,6 +82,7 @@ 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) @@ -83,6 +90,7 @@ 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 end end From 8671ed9d6b13589dce027985fc2b9c481c461d5b Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Thu, 9 Jan 2025 14:18:29 +0100 Subject: [PATCH 2/2] Update src/nditeration.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- src/nditeration.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/nditeration.jl b/src/nditeration.jl index 66b307dcc..8cfdbce95 100644 --- a/src/nditeration.jl +++ b/src/nditeration.jl @@ -63,7 +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 +@inline Base.ndims(::NDRange{N}) where {N} = N import Base.iterate @inline iterate(range::NDRange) = iterate(blocks(range))