#1544: collection: Implement default_map for indices with more than 3…

… dimensions
DARMA-tasking · Apr 19, 2022 · bde766a · bde766a
1 parent 7a5bbfb
commit bde766a
Show file tree

Hide file tree

Showing 5 changed files with 42 additions and 6 deletions.
diff --git a/src/vt/topos/index/dense/dense_array.h b/src/vt/topos/index/dense/dense_array.h
@@ -102,7 +102,7 @@ struct DenseIndexArray : BaseIndex, serialization::ByteCopyTrait {
   DenseIndexArray(std::array<IndexType, ndim> in_array);
   DenseIndexArray(DenseIndexArraySingleInitTag, IndexType const& init_value);
 
-  NumDimensionsType ndims() const { return ndim; }
+  static constexpr NumDimensionsType ndims() { return ndim; }
 
   IndexType& operator[](IndexType const& index);
   IndexType const& operator[](IndexType const& index) const;

diff --git a/src/vt/topos/index/index.h b/src/vt/topos/index/index.h
@@ -60,6 +60,7 @@ template <typename T = IdxBase> using Index1D  = DenseIndexArray<T, 1>;
 template <typename T = IdxBase> using Index2D  = DenseIndexArray<T, 2>;
 template <typename T = IdxBase> using Index3D  = DenseIndexArray<T, 3>;
 template <typename T, int8_t N> using IdxType = DenseIndexArray<T, N>;
+template <typename T, int8_t N> using IndexND  = DenseIndexArray<T, N>;
 
 static_assert(IndexTraits<Index1D<IdxBase>>::is_index, "Does not conform");
 static_assert(IndexTraits<Index2D<IdxBase>>::is_index, "Does not conform");
@@ -76,11 +77,14 @@ using IdxBase    = index::IdxBase;
 using Index1D  = index::Index1D<index::IdxBase>;
 using Index2D  = index::Index2D<index::IdxBase>;
 using Index3D  = index::Index3D<index::IdxBase>;
+template <int8_t N>
+using IndexND  = index::IndexND<index::IdxBase, N>;
 
 template <typename T, int8_t N> using IdxType   = index::IdxType<T, N>;
 template <typename T>           using IdxType1D = index::Index1D<T>;
 template <typename T>           using IdxType2D = index::Index2D<T>;
 template <typename T>           using IdxType3D = index::Index3D<T>;
+template <typename T, int8_t N> using IdxTypeND = index::IndexND<T, N>;
 
 }  // end namespace vt
 

diff --git a/src/vt/topos/mapping/dense/dense.h b/src/vt/topos/mapping/dense/dense.h
@@ -74,7 +74,7 @@ using IdxPtr = Index*;
 template <typename T = IdxBase> using Idx1DPtr = IdxType1D<T>*;
 template <typename T = IdxBase> using Idx2DPtr = IdxType2D<T>*;
 template <typename T = IdxBase> using Idx3DPtr = IdxType3D<T>*;
-template <typename T, typename N> using IdxNDPtr = vt::index::IdxType<T, N>*;
+template <typename T, int8_t N> using IdxNDPtr = vt::index::IdxType<T, N>*;
 
 template <typename Idx, index::NumDimensionsType ndim>
 NodeType denseBlockMap(IdxPtr<Idx> idx, IdxPtr<Idx> max_idx, NodeType nnodes);
@@ -85,26 +85,31 @@ template <typename T = IdxBase>
 NodeType defaultDenseIndex2DMap(Idx2DPtr<T> idx, Idx2DPtr<T> max, NodeType n);
 template <typename T = IdxBase>
 NodeType defaultDenseIndex3DMap(Idx3DPtr<T> idx, Idx3DPtr<T> max, NodeType n);
-template <typename T, typename N>
-NodeType defaultDenseIndex3DMap(IdxNDPtr<T,N> idx, IdxNDPtr<T,N> max, NodeType n);
+template <typename T, int8_t N>
+NodeType defaultDenseIndexNDMap(IdxNDPtr<T,N> idx, IdxNDPtr<T,N> max, NodeType n);
 
 template <typename T = IdxBase>
 NodeType dense1DRoundRobinMap(  Idx1DPtr<T> idx, Idx1DPtr<T> max, NodeType n);
 template <typename T = IdxBase>
 NodeType dense2DRoundRobinMap(  Idx2DPtr<T> idx, Idx2DPtr<T> max, NodeType n);
 template <typename T = IdxBase>
 NodeType dense3DRoundRobinMap(  Idx3DPtr<T> idx, Idx3DPtr<T> max, NodeType n);
+template <typename T, int8_t N>
+NodeType denseNDRoundRobinMap(  IdxNDPtr<T,N> idx, IdxNDPtr<T,N> max, NodeType n);
 
 template <typename T = IdxBase>
 NodeType dense1DBlockMap(       Idx1DPtr<T> idx, Idx1DPtr<T> max, NodeType n);
 template <typename T = IdxBase>
 NodeType dense2DBlockMap(       Idx2DPtr<T> idx, Idx2DPtr<T> max, NodeType n);
 template <typename T = IdxBase>
 NodeType dense3DBlockMap(       Idx3DPtr<T> idx, Idx3DPtr<T> max, NodeType n);
+template <typename T, int8_t N>
+NodeType denseNDBlockMap(       IdxNDPtr<T,N> idx, IdxNDPtr<T,N> max, NodeType n);
 
 template <typename T = IdxBase>   using i1D   = IdxType1D<T>;
 template <typename T = IdxBase>   using i2D   = IdxType2D<T>;
 template <typename T = IdxBase>   using i3D   = IdxType3D<T>;
+template <typename T, int8_t N>   using iND   = IdxTypeND<T,N>;
 
 template <typename F, F* f, typename IndexT>
 using Adapt = MapFunctorAdapt<F, f, IndexT>;
@@ -115,18 +120,26 @@ template <typename T = IdxBase>
 using dense2DMapFn    = Adapt<MapAdapter<i2D<T>>, defaultDenseIndex2DMap<T>, i2D<T> >;
 template <typename T = IdxBase>
 using dense3DMapFn    = Adapt<MapAdapter<i3D<T>>, defaultDenseIndex3DMap<T>, i3D<T> >;
+template <typename T, int8_t N>
+using denseNDMapFn    = Adapt<MapAdapter<iND<T, N>>, defaultDenseIndexNDMap<T>, iND<T, N> >;
+
 template <typename T = IdxBase>
 using dense1DRRMapFn  = Adapt<MapAdapter<i1D<T>>, dense1DRoundRobinMap<T>, i1D<T>>;
 template <typename T = IdxBase>
 using dense2DRRMapFn  = Adapt<MapAdapter<i2D<T>>, dense2DRoundRobinMap<T>, i2D<T>>;
 template <typename T = IdxBase>
 using dense3DRRMapFn  = Adapt<MapAdapter<i3D<T>>, dense3DRoundRobinMap<T>, i3D<T>>;
+template <typename T, int8_t N>
+using denseNDRRMapFn  = Adapt<MapAdapter<iND<T, N>>, denseNDRoundRobinMap<T>, iND<T, N> >;
+
 template <typename T = IdxBase>
 using dense1DBlkMapFn = Adapt<MapAdapter<i1D<T>>, dense1DBlockMap<T>, i1D<T>>;
 template <typename T = IdxBase>
 using dense2DBlkMapFn = Adapt<MapAdapter<i2D<T>>, dense2DBlockMap<T>, i2D<T>>;
 template <typename T = IdxBase>
 using dense3DBlkMapFn = Adapt<MapAdapter<i3D<T>>, dense3DBlockMap<T>, i3D<T>>;
+template <typename T, int8_t N>
+using denseNDBlkMapFn = Adapt<MapAdapter<iND<T, N>>, denseNDBlockMap<T>, iND<T, N> >;
 
 }}  // end namespace vt::mapping
 

diff --git a/src/vt/topos/mapping/dense/dense.impl.h b/src/vt/topos/mapping/dense/dense.impl.h
@@ -68,6 +68,11 @@ NodeType defaultDenseIndex3DMap(Idx3DPtr<T> idx, Idx3DPtr<T> max, NodeType nx) {
   return dense3DBlockMap<T>(idx, max, nx);
 }
 
+template <typename T, int8_t N>
+NodeType defaultDenseIndexNDMap(IdxNDPtr<T,N> idx, IdxNDPtr<T,N> max, NodeType nx) {
+  return denseNDBlockMap<T>(idx, max, nx);
+}
+
 // Default round robin mappings
 template <typename T>
 NodeType dense1DRoundRobinMap(Idx1DPtr<T> idx, Idx1DPtr<T> max, NodeType nx) {
@@ -104,6 +109,11 @@ NodeType dense3DBlockMap(Idx3DPtr<T> idx, Idx3DPtr<T> max, NodeType nx) {
   return denseBlockMap<IdxType3D<T>, 3>(idx, max, nx);
 }
 
+template <typename T, int8_t N>
+NodeType denseNDBlockMap(IdxNDPtr<T,N> idx, IdxNDPtr<T,N> max, NodeType nx) {
+  return denseBlockMap<IdxTypeND<T,N>, N>(idx, max, nx);
+}
+
 template <typename IndexElmType, typename PhysicalType>
 inline NodeType blockMapDenseFlatIndex(
   IndexElmType* flat_idx_ptr, IndexElmType* num_elems_ptr,

diff --git a/src/vt/vrt/collection/defaults/default_map.h b/src/vt/vrt/collection/defaults/default_map.h
@@ -61,8 +61,17 @@ struct DefaultMapBase {
   using MapParamPackType = std::tuple<IndexPtrType,IndexPtrType,NodeType>;
 };
 
-template <typename CollectionT, typename Enable=void>
-struct DefaultMap;
+template <typename CollectionT, typename Enable = void>
+struct DefaultMap : DefaultMapBase<CollectionT> {
+  using BaseType = typename CollectionT::IndexType::DenseIndexType;
+  using BlockMapType =
+    ::vt::mapping::denseNDMapFn<BaseType, CollectionT::IndexType::ndims()>;
+  using RRMapType =
+    ::vt::mapping::denseNDRRMapFn<BaseType, CollectionT::IndexType::ndims()>;
+  using DefaultMapType =
+    ::vt::mapping::denseNDMapFn<BaseType, CollectionT::IndexType::ndims()>;
+  using MapType = DefaultMapType;
+};
 
 /*
  * Default mappings for Index1D: RR, Block, etc.