diff --git a/src/looping.jl b/src/looping.jl index 08018cdd8..d3f2a0451 100644 --- a/src/looping.jl +++ b/src/looping.jl @@ -366,9 +366,17 @@ eval(quote loop_ranges_store[dims] = LoopRanges(; parallel_dims=dims, rank0 = rank0, ranges...) else - # Use the same ranges as serial loops + # Loop over all indices for non-parallelised dimensions (dimensions not in + # `dims`), but only loop over parallel dimensions (dimensions in `dims`) on + # rank0. + this_ranges = Dict(d=>1:n for (d,n) in dim_sizes) + if !rank0 + for d ∈ dims + this_ranges[d] = 1:0 + end + end loop_ranges_store[dims] = LoopRanges(; - parallel_dims=dims, rank0 = rank0, serial_ranges...) + parallel_dims=dims, rank0 = rank0, this_ranges...) end end diff --git a/test/loop_setup_tests.jl b/test/loop_setup_tests.jl index c34ea677b..725714d96 100644 --- a/test/loop_setup_tests.jl +++ b/test/loop_setup_tests.jl @@ -197,26 +197,26 @@ function runtests() @test loop_ranges_store[()].z == 1:0 @test loop_ranges_store[(:s,)].s == 1:0 - @test loop_ranges_store[(:s,)].r == 1:0 - @test loop_ranges_store[(:s,)].z == 1:0 + @test loop_ranges_store[(:s,)].r == 1:3 + @test loop_ranges_store[(:s,)].z == 1:4 - @test loop_ranges_store[(:r,)].s == 1:0 + @test loop_ranges_store[(:r,)].s == 1:2 @test loop_ranges_store[(:r,)].r == 1:0 - @test loop_ranges_store[(:r,)].z == 1:0 + @test loop_ranges_store[(:r,)].z == 1:4 - @test loop_ranges_store[(:z,)].s == 1:0 - @test loop_ranges_store[(:z,)].r == 1:0 + @test loop_ranges_store[(:z,)].s == 1:2 + @test loop_ranges_store[(:z,)].r == 1:3 @test loop_ranges_store[(:z,)].z == 1:0 @test loop_ranges_store[(:s,:r)].s == 1:0 @test loop_ranges_store[(:s,:r)].r == 1:0 - @test loop_ranges_store[(:s,:r)].z == 1:0 + @test loop_ranges_store[(:s,:r)].z == 1:4 @test loop_ranges_store[(:s,:z)].s == 2:2 @test loop_ranges_store[(:s,:z)].r == 1:3 @test loop_ranges_store[(:s,:z)].z == 1:4 - @test loop_ranges_store[(:r,:z)].s == 1:0 + @test loop_ranges_store[(:r,:z)].s == 1:2 @test loop_ranges_store[(:r,:z)].r == 1:0 @test loop_ranges_store[(:r,:z)].z == 1:0 @@ -269,23 +269,23 @@ function runtests() @test loop_ranges_store[()].z == 1:0 @test loop_ranges_store[(:s,)].s == 1:0 - @test loop_ranges_store[(:s,)].r == 1:0 - @test loop_ranges_store[(:s,)].z == 1:0 + @test loop_ranges_store[(:s,)].r == 1:3 + @test loop_ranges_store[(:s,)].z == 1:4 - @test loop_ranges_store[(:r,)].s == 1:0 + @test loop_ranges_store[(:r,)].s == 1:2 @test loop_ranges_store[(:r,)].r == 1:0 - @test loop_ranges_store[(:r,)].z == 1:0 + @test loop_ranges_store[(:r,)].z == 1:4 - @test loop_ranges_store[(:z,)].s == 1:0 - @test loop_ranges_store[(:z,)].r == 1:0 + @test loop_ranges_store[(:z,)].s == 1:2 + @test loop_ranges_store[(:z,)].r == 1:3 @test loop_ranges_store[(:z,)].z == 1:0 @test loop_ranges_store[(:s,:r)].s == 1:0 @test loop_ranges_store[(:s,:r)].r == 1:0 - @test loop_ranges_store[(:s,:r)].z == 1:0 + @test loop_ranges_store[(:s,:r)].z == 1:4 @test loop_ranges_store[(:s,:z)].s == 1:0 - @test loop_ranges_store[(:s,:z)].r == 1:0 + @test loop_ranges_store[(:s,:z)].r == 1:3 @test loop_ranges_store[(:s,:z)].z == 1:0 @test loop_ranges_store[(:r,:z)].s == 1:2 @@ -305,23 +305,23 @@ function runtests() @test loop_ranges_store[()].z == 1:0 @test loop_ranges_store[(:s,)].s == 1:0 - @test loop_ranges_store[(:s,)].r == 1:0 - @test loop_ranges_store[(:s,)].z == 1:0 + @test loop_ranges_store[(:s,)].r == 1:3 + @test loop_ranges_store[(:s,)].z == 1:4 - @test loop_ranges_store[(:r,)].s == 1:0 + @test loop_ranges_store[(:r,)].s == 1:2 @test loop_ranges_store[(:r,)].r == 1:0 - @test loop_ranges_store[(:r,)].z == 1:0 + @test loop_ranges_store[(:r,)].z == 1:4 - @test loop_ranges_store[(:z,)].s == 1:0 - @test loop_ranges_store[(:z,)].r == 1:0 + @test loop_ranges_store[(:z,)].s == 1:2 + @test loop_ranges_store[(:z,)].r == 1:3 @test loop_ranges_store[(:z,)].z == 1:0 @test loop_ranges_store[(:s,:r)].s == 1:0 @test loop_ranges_store[(:s,:r)].r == 1:0 - @test loop_ranges_store[(:s,:r)].z == 1:0 + @test loop_ranges_store[(:s,:r)].z == 1:4 @test loop_ranges_store[(:s,:z)].s == 1:0 - @test loop_ranges_store[(:s,:z)].r == 1:0 + @test loop_ranges_store[(:s,:z)].r == 1:3 @test loop_ranges_store[(:s,:z)].z == 1:0 @test loop_ranges_store[(:r,:z)].s == 1:2 @@ -341,23 +341,23 @@ function runtests() @test loop_ranges_store[()].z == 1:0 @test loop_ranges_store[(:s,)].s == 1:0 - @test loop_ranges_store[(:s,)].r == 1:0 - @test loop_ranges_store[(:s,)].z == 1:0 + @test loop_ranges_store[(:s,)].r == 1:3 + @test loop_ranges_store[(:s,)].z == 1:4 - @test loop_ranges_store[(:r,)].s == 1:0 + @test loop_ranges_store[(:r,)].s == 1:2 @test loop_ranges_store[(:r,)].r == 1:0 - @test loop_ranges_store[(:r,)].z == 1:0 + @test loop_ranges_store[(:r,)].z == 1:4 - @test loop_ranges_store[(:z,)].s == 1:0 - @test loop_ranges_store[(:z,)].r == 1:0 + @test loop_ranges_store[(:z,)].s == 1:2 + @test loop_ranges_store[(:z,)].r == 1:3 @test loop_ranges_store[(:z,)].z == 1:0 @test loop_ranges_store[(:s,:r)].s == 1:0 @test loop_ranges_store[(:s,:r)].r == 1:0 - @test loop_ranges_store[(:s,:r)].z == 1:0 + @test loop_ranges_store[(:s,:r)].z == 1:4 @test loop_ranges_store[(:s,:z)].s == 1:0 - @test loop_ranges_store[(:s,:z)].r == 1:0 + @test loop_ranges_store[(:s,:z)].r == 1:3 @test loop_ranges_store[(:s,:z)].z == 1:0 @test loop_ranges_store[(:r,:z)].s == 1:2 @@ -377,26 +377,26 @@ function runtests() @test loop_ranges_store[()].z == 1:0 @test loop_ranges_store[(:s,)].s == 1:0 - @test loop_ranges_store[(:s,)].r == 1:0 - @test loop_ranges_store[(:s,)].z == 1:0 + @test loop_ranges_store[(:s,)].r == 1:3 + @test loop_ranges_store[(:s,)].z == 1:4 - @test loop_ranges_store[(:r,)].s == 1:0 + @test loop_ranges_store[(:r,)].s == 1:2 @test loop_ranges_store[(:r,)].r == 1:0 - @test loop_ranges_store[(:r,)].z == 1:0 + @test loop_ranges_store[(:r,)].z == 1:4 - @test loop_ranges_store[(:z,)].s == 1:0 - @test loop_ranges_store[(:z,)].r == 1:0 + @test loop_ranges_store[(:z,)].s == 1:2 + @test loop_ranges_store[(:z,)].r == 1:3 @test loop_ranges_store[(:z,)].z == 1:0 @test loop_ranges_store[(:s,:r)].s == 1:0 @test loop_ranges_store[(:s,:r)].r == 1:0 - @test loop_ranges_store[(:s,:r)].z == 1:0 + @test loop_ranges_store[(:s,:r)].z == 1:4 @test loop_ranges_store[(:s,:z)].s == 1:0 - @test loop_ranges_store[(:s,:z)].r == 1:0 + @test loop_ranges_store[(:s,:z)].r == 1:3 @test loop_ranges_store[(:s,:z)].z == 1:0 - @test loop_ranges_store[(:r,:z)].s == 1:0 + @test loop_ranges_store[(:r,:z)].s == 1:2 @test loop_ranges_store[(:r,:z)].r == 1:0 @test loop_ranges_store[(:r,:z)].z == 1:0