Converted the block and loop limit storage from vector to ntuple

smillerc · Dec 9, 2022 · 4a5cf94 · 4a5cf94 · smillerc · Dec 9, 2022
1 parent 08337a9
commit 4a5cf94
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 11 deletions.
diff --git a/src/BlockHaloArrays.jl b/src/BlockHaloArrays.jl
@@ -33,13 +33,13 @@ scaling than multi-threaded loops
  - `globaldims::NTuple{D,Int}`: Dimensions of the array if it were a simple `Array{T,D}`, e.g. `(20,20)`
 
 """
-struct BlockHaloArray{T,N,NH,NBL,AA<:Array{T,N},SA} <: AbstractBlockHaloArray
-    blocks::Vector{AA}
+struct BlockHaloArray{NB,T,N,NH,NBL,AA<:Array{T,N},SA,NL} <: AbstractBlockHaloArray
+    blocks::NTuple{NB,AA}
     block_layout::NTuple{NBL,Int}
     halodims::NTuple{NH,Int}
     global_blockranges::Vector{NTuple{N,UnitRange{Int}}}
     nhalo::Int
-    loop_limits::Vector{Vector{Int}}
+    loop_limits::NTuple{NB, NTuple{NL, Int}}
     globaldims::NTuple{N,Int}
     neighbor_blocks::Vector{Dict{Symbol,Int}}
 
@@ -72,7 +72,7 @@ struct MPIBlockHaloArray{T,N,NH,NBL,AA<:Array{T,N}} <: AbstractBlockHaloArray
     halodims::NTuple{NH,Int}
     global_blockranges::Vector{NTuple{N,UnitRange{Int}}}
     nhalo::Int
-    loop_limits::Vector{Vector{Int}}
+    loop_limits::NTuple{NBL, NTuple{4, Int}}
     globaldims::NTuple{N,Int}
     neighbor_blocks::Vector{Dict{Symbol,Int}}
     _global_halo_send_buf::Vector{Array{T,NH}}
@@ -186,6 +186,7 @@ function BlockHaloArray(dims::NTuple{N,Int}, halodims::NTuple{N2,Int}, nhalo::In
                           for ax in axes(blocks[i])] |> flatten |> collect
     end
 
+    loop_lims = Tuple(Tuple.(loop_limits))
     neighbors = get_neighbor_blocks(tile_dims)
 
     # Pass undef'ed vector to the constructor, since the views
@@ -198,8 +199,8 @@ function BlockHaloArray(dims::NTuple{N,Int}, halodims::NTuple{N2,Int}, nhalo::In
     DonorViewTypes = typeof(first(_donor_views)[:ilo])
     donor_views = Vector{Dict{Symbol, DonorViewTypes}}(undef, nblocks)
 
-    A = BlockHaloArray(blocks, tile_dims, halodims, vec(block_ranges), nhalo,
-        loop_limits, dims, neighbors,
+    A = BlockHaloArray(tuple(blocks...), tile_dims, halodims, vec(block_ranges), nhalo,
+        loop_lims, dims, neighbors,
         cummulative_blocksize_per_dim, LI,
         donor_views, halo_views)
 

diff --git a/src/halo_exchange.jl b/src/halo_exchange.jl
@@ -324,7 +324,7 @@ Copy data from the neighbor block into the current block's halo region
 """
 updateblockhalo!(A, current_block_id) = updateblockhalo!(A, current_block_id, false)
 
-function updateblockhalo!(A::BlockHaloArray{T,N,NH,NBL,AA,SA}, current_block_id::Integer, include_periodic_bc::Bool) where {T,N,NH,NBL,AA,SA}
+function updateblockhalo!(A::BlockHaloArray{NB,T,N,NH,NBL,AA,SA,Y}, current_block_id::Integer, include_periodic_bc::Bool) where {NB,T,N,NH,NBL,AA,SA,Y}
     exchange_map = hmap[NH]
 
     for (dom_id, halo_id) in exchange_map

diff --git a/src/indexing.jl b/src/indexing.jl
@@ -55,7 +55,7 @@ since the actual data within the BlockHaloArray is a series of separate blocks.
 the array in this manner will be slower due to the additional arithmetic need to find the global
 to local index conversion for each block.
 """
-function Base.getindex(A::BlockHaloArray{T,N,NH,NBL,AA}, global_idx::Vararg{Int, N2}) where {T,N,N2,NH,NBL,AA}
+function Base.getindex(A::BlockHaloArray, global_idx::Vararg{Int, N}) where {N}
     block_idx = ntuple(i ->
             get_block_idx(global_idx[A.halodims[i]],
                 A._cummulative_blocksize_per_dim,
@@ -86,7 +86,7 @@ end
 """
 Get the scalar block index give a global index tuple
 """
-function blockindex(A::BlockHaloArray{T,N,NH,NBL,AA}, global_idx::NTuple{N,Int}) where {T,N,NH,NBL,AA}
+function blockindex(A::BlockHaloArray, global_idx::NTuple{N,Int}) where {N}
     block_idx = ntuple(i ->
             get_block_idx(global_idx[A.halodims[i]],
                 A._cummulative_blocksize_per_dim,
@@ -100,7 +100,7 @@ end
 """
 Get the block-local indices give a global index tuple
 """
-function localindex(A::BlockHaloArray{T,N,NH,NBL,AA}, global_idx::NTuple{N,Int}) where {T,N,NH,NBL,AA}
+function localindex(A::BlockHaloArray, global_idx::NTuple{N,Int}) where {N}
 
     block_idx = blockindex(A, global_idx)
     local_idx = ntuple(i ->

diff --git a/test/unit/test_constructors.jl b/test/unit/test_constructors.jl
@@ -70,6 +70,9 @@ end
     nhalo = 2
     nblocks = 4
     A = BlockHaloArray(dims, nhalo, nblocks)
+    @show A.global_blockranges
+    @show A.block_layout
+    @show typeof(A)
 end
 
 @testitem "3D Array, 1D halo dims" begin

diff --git a/test/unit/test_util_funcs.jl b/test/unit/test_util_funcs.jl
@@ -138,7 +138,7 @@ end
     @test globalindices(A, 1, (3, 3)) == (1, 1)
     @test globalindices(A, 4, (3, 3)) == (21, 21)
 
-    @show globalindices(A, 4, (22, 22))
+    # @show globalindices(A, 4, (22, 22))
 end
 
 @testitem "copy!" begin