Add getSize static method to calculate the size of sources in SliceUnion

This commit is contained in:
Alejandro Gallo 2023-01-25 16:25:09 +01:00
parent 933d556c84
commit 4e2d1143e5
2 changed files with 65 additions and 0 deletions

View File

@ -387,6 +387,22 @@ template <typename F=double>
} }
} }
static size_t
getSize(const std::vector<size_t> sliceLength,
const std::vector<size_t> paramLength,
const size_t np,
const MPI_Comm global_world) {
const RankMap<F> rankMap(paramLength, np, global_world);
const size_t
nSources = rankMap.nSources(),
sliceSize = std::accumulate(sliceLength.begin(),
sliceLength.end(),
1UL,
std::multiplies<size_t>());
return nSources * sliceSize;
}
// CONSTRUCTOR // CONSTRUCTOR
SliceUnion( std::vector<typename Slice<F>::Type> sliceTypes_ SliceUnion( std::vector<typename Slice<F>::Type> sliceTypes_
, std::vector<size_t> sliceLength_ , std::vector<size_t> sliceLength_

View File

@ -235,11 +235,54 @@ Atrip::Output Atrip::run(Atrip::Input<F> const& in) {
MPI_Comm_size(child_comm, &child_size); MPI_Comm_size(child_comm, &child_size);
} }
// a, b, c, d, e, f and P => Nv
// H => No
// total_source_sizes contains a list of the number of elements
// in all sources of every tensor union, therefore nSlices * sliceSize
const std::vector<size_t> total_source_sizes = {
// ABPH
SliceUnion<F>::getSize({Nv, No}, {Nv, Nv}, (size_t)np, universe),
// ABHH
SliceUnion<F>::getSize({No, No}, {Nv, Nv}, (size_t)np, universe),
// TABHH
SliceUnion<F>::getSize({No, No}, {Nv, Nv}, (size_t)np, universe),
// TAPHH
SliceUnion<F>::getSize({Nv, No, No}, {Nv}, (size_t)np, universe),
// HHHA
SliceUnion<F>::getSize({No, No, No}, {Nv}, (size_t)np, universe),
};
const size_t
total_source_size = sizeof(DataFieldType<F>)
* std::accumulate(total_source_sizes.begin(),
total_source_sizes.end(),
0UL);
#if defined(HAVE_CUDA)
DataPtr<F> all_sources_pointer;
cuMemAlloc(&all_sources_pointer, total_source_size);
#else
DataPtr<F>
all_sources_pointer = (DataPtr<F>)malloc(total_source_size);
#endif
size_t _source_pointer_idx = 0;
// BUILD SLICES PARAMETRIZED BY NV x NV =============================={{{1 // BUILD SLICES PARAMETRIZED BY NV x NV =============================={{{1
WITH_CHRONO("nv-nv-slices", WITH_CHRONO("nv-nv-slices",
LOG(0,"Atrip") << "building NV x NV slices\n"; LOG(0,"Atrip") << "building NV x NV slices\n";
// TODO
// DataPtr<F> offseted_pointer = all_sources_pointer
// * total_source_sizes[_source_pointer_idx++];
ABPH<F> abph(*in.Vppph, (size_t)No, (size_t)Nv, (size_t)np, child_comm, universe); ABPH<F> abph(*in.Vppph, (size_t)No, (size_t)Nv, (size_t)np, child_comm, universe);
// TODO
// DataPtr<F> offseted_pointer = all_sources_pointer
// * total_source_sizes[_source_pointer_idx++];
ABHH<F> abhh(*in.Vpphh, (size_t)No, (size_t)Nv, (size_t)np, child_comm, universe); ABHH<F> abhh(*in.Vpphh, (size_t)No, (size_t)Nv, (size_t)np, child_comm, universe);
// TODO
// DataPtr<F> offseted_pointer = all_sources_pointer
// * total_source_sizes[_source_pointer_idx++];
TABHH<F> tabhh(*in.Tpphh, (size_t)No, (size_t)Nv, (size_t)np, child_comm, universe); TABHH<F> tabhh(*in.Tpphh, (size_t)No, (size_t)Nv, (size_t)np, child_comm, universe);
) )
@ -251,7 +294,13 @@ Atrip::Output Atrip::run(Atrip::Input<F> const& in) {
// BUILD SLICES PARAMETRIZED BY NV ==================================={{{1 // BUILD SLICES PARAMETRIZED BY NV ==================================={{{1
WITH_CHRONO("nv-slices", WITH_CHRONO("nv-slices",
LOG(0,"Atrip") << "building NV slices\n"; LOG(0,"Atrip") << "building NV slices\n";
// TODO
// DataPtr<F> offseted_pointer = all_sources_pointer
// * total_source_sizes[_source_pointer_idx++];
TAPHH<F> taphh(*in.Tpphh, (size_t)No, (size_t)Nv, (size_t)np, child_comm, universe); TAPHH<F> taphh(*in.Tpphh, (size_t)No, (size_t)Nv, (size_t)np, child_comm, universe);
// TODO
// DataPtr<F> offseted_pointer = all_sources_pointer
// * total_source_sizes[_source_pointer_idx++];
HHHA<F> hhha(*in.Vhhhp, (size_t)No, (size_t)Nv, (size_t)np, child_comm, universe); HHHA<F> hhha(*in.Vhhhp, (size_t)No, (size_t)Nv, (size_t)np, child_comm, universe);
) )