Merge pull request #541 from shkodm/sampler_arb

Implemented sampler on Arbitrary grid and fixes for cartesian one

Author: llaniewski

src/ArbLattice.cpp

@@ -22,7 +22,6 @@ ArbLattice::ArbLattice(size_t num_snaps_, const UnitEnv& units_, const std::map<
 }
 
 void ArbLattice::initialize(size_t num_snaps_, const std::map<std::string, int>& setting_zones, pugi::xml_node arb_node) {
-    const int rank = mpitools::MPI_Rank(comm);
     sizes.snaps = num_snaps_;
 #ifdef ADJOINT
     sizes.snaps += 2;  // Adjoint snaps are appended to the total snap allocation
@@ -927,6 +926,36 @@ void ArbLattice::resetAverage(){
     }
 }
 
+void ArbLattice::getSample(int quant, unsigned int lid, real_t scale, real_t *buf) {
+    setSnapIn(Snap);
+#ifdef ADJOINT
+    setAdjSnapIn(aSnap);
+#endif
+    launcher.sampleQuantity(quant, lid, buf, scale, data);
+}
+
+
+void ArbLattice::updateAllSamples(){
+    const int rank = mpitools::MPI_Rank(comm);
+    if (sample->size != 0) {
+	    for (size_t j = 0; j < sample->spoints.size(); j++) {
+		    if (rank == sample->spoints[j].rank) {
+		        for(const Model::Quantity& q : model->quantities) {
+                    if (sample->quant->in(q.name.c_str())){
+                        double v = sample->units->alt(q.unit.c_str());
+                        getSample(q.id, sample->spoints[j].lid, 1/v,
+                                  &sample->gpu_buffer[sample->location[q.name.c_str()]+(data.iter - sample->startIter)*sample->size + sample->totalIter*j*sample->size]);
+                    }
+                }
+	        }
+        }
+    }
+}
+
+unsigned int ArbLattice::getCartesianCoordinateLid(vector_t point) const {
+    return launcher.getCartesianCoordinateLid(point, data);
+}
+
 ArbLattice::~ArbLattice()
 {
 	RFI.Close();

src/ArbLattice.hpp

@@ -96,14 +96,18 @@ class ArbLattice : public LatticeBase {
     virtual void setFlags(const std::vector<big_flag_t>& x) override;
     virtual void setField(const Model::Field& f, const std::vector<real_t>& x) override;
     virtual void setFieldAdjZero(const Model::Field& f) override;
+    virtual void updateAllSamples() override;
 
     const ArbVTUGeom& getVTUGeom() const { return vtu_geom; }
     Span<const flag_t> getNodeTypes() const { return {node_types_host.data(), node_types_host.size()}; }  /// Get host view of node types (permuted)
     const ArbLatticeConnectivity& getConnectivity() const { return connect; }
     const std::vector<unsigned>& getLocalPermutation() const { return local_permutation; }
 
+    unsigned int getCartesianCoordinateLid(vector_t point) const; 
+    void getSample(int quant, unsigned int lid, real_t scale, real_t *buf);
+
     void resetAverage();
-    lbRegion getLocalBoundingBox() const override;                                                                                          /// Compute local bounding box, assuming the arbitrary lattice is a subset of a Cartesian lattice 
+    lbRegion getLocalBoundingBox() const override;                                                                                  /// Compute local bounding box, assuming the arbitrary lattice is a subset of a Cartesian lattice
 
    protected:
     ArbLatticeLauncher launcher;  /// Launcher responsible for running CUDA kernels on the lattice

src/ArbLatticeLauncher.h.Rt

@@ -22,6 +22,10 @@ struct ArbLatticeLauncher {
 
     void getQuantity(int quant, real_t* host_tab, real_t scale, const LatticeData& data) const;
 
+    unsigned int getCartesianCoordinateLid(const vector_t point, const LatticeData& data) const;
+
+    void sampleQuantity(int quant, unsigned int lid, real_t* host_tab, real_t scale, const LatticeData &data) const;
+
 private:
 <?R for (q in rows(Quantities)) { ifdef(q$adjoint);
 if (q$adjoint) {
@@ -31,6 +35,8 @@ if (q$adjoint) {
 }
 ?>
     void getQuantity<?%s q$name ?>(<?%s q$type ?>* tab, real_t scale, const LatticeData& data) const;
+
+    void getSample<?%s q$name ?>(<?%s q$type ?>* tab, real_t scale, const LatticeData& data, unsigned int lid) const;
 <?R }
 ifdef() ?>
 };

src/ArbLatticeLauncher.hpp.Rt

@@ -34,7 +34,8 @@ struct ArbLatticeExecutor : public LinearExecutor {
 
 template <eOperationType I, eCalculateGlobals G, eStage S>
 void ArbLatticeLauncher::RunBorder(CudaStream_t stream, const LatticeData& data) const {
-    const ArbLatticeExecutor<I, G, S> executor{{container.num_border_nodes}, container, data, 0};
+    const ArbLatticeExecutor<I, G, S> executor{{container.num_border_nodes},
+            container, data, 0};
     LaunchExecutorAsync(executor, stream);
 }
 
@@ -99,6 +100,7 @@ struct GetQuantityArbExecutor<?%s q$name?> : public LinearExecutor {
     LatticeData data;
     <?%s q$type ?>* buf;
     real_t scale;
+    unsigned int offset;  // Starting offset for iteration space
 
     CudaDeviceFunction void Execute() const {
         using LA = ArbLatticeAccess; <?R
@@ -109,7 +111,8 @@ if (q$adjoint) { ?>
 }?>
         const int i = threadID(CudaThread, CudaBlock, CudaNumberOfThreads);
         if (inRange(i)) {
-            LA acc(i, container);
+            const int index = i + offset;
+            LA acc(index, container);
             N now(acc, data);
             acc.pop(now); <?R
 if (q$adjoint) { ?>
@@ -129,10 +132,68 @@ if (q$type == "vector_t") {
 };
 
 void ArbLatticeLauncher::getQuantity<?%s q$name ?>(<?%s q$type ?>* tab, real_t scale, const LatticeData& data) const {
-    const GetQuantityArbExecutor<?%s q$name?> executor{{container.num_border_nodes + container.num_interior_nodes}, container, data, tab, scale};
+    const GetQuantityArbExecutor<?%s q$name?> executor{{container.num_border_nodes + container.num_interior_nodes}, container, data, tab, scale, 0};
+    LaunchExecutor(executor);
+}
+
+void ArbLatticeLauncher::getSample<?%s q$name ?>(<?%s q$type ?>* tab, real_t scale, const LatticeData& data, unsigned int lid) const {
+    const GetQuantityArbExecutor<?%s q$name?> executor{{1}, container, data, tab, scale, lid};
     LaunchExecutor(executor);
 }
+
 <?R }
 ifdef() ?>
 
+struct CartesianCoordinateMapperArbExecutor: public LinearExecutor {
+    ArbLatticeContainer container; 
+    LatticeData data;
+    vector_t point; // cartesian point coordinates
+    real_t epsilon2;  // should be a small value (e.g. < 1e-10) to match with only one point
+    unsigned int *matched_lid; 
+
+    CudaDeviceFunction void Execute() const {
+        using LA = ArbLatticeAccess; 
+        const int i = threadID(CudaThread, CudaBlock, CudaNumberOfThreads);
+        if (inRange(i)) {
+            LA acc(i, container);
+            real_t X  = acc.getX();
+            real_t Y = acc.getY();
+            real_t Z = acc.getZ(); 
+            real_t distance = (X - point.x)*(X - point.x) + (Y - point.y)*(Y - point.y) + (Z - point.z)*(Z - point.z); 
+            if (distance < epsilon2) {
+                *matched_lid= i;
+            }
+    }
+  }
+};
+
+
+unsigned int ArbLatticeLauncher::getCartesianCoordinateLid(const vector_t point, const LatticeData& data) const {
+    const auto gpu_lid = cudaMakeUnique<unsigned int>(1);
+    const real_t tol2 = 1e-10;
+    const CartesianCoordinateMapperArbExecutor executor{{container.num_border_nodes + container.num_interior_nodes}, container, data, point, tol2,
+                                                        gpu_lid.get()};
+    LaunchExecutor(executor);
+
+    unsigned int matched_lid;
+    CudaMemcpy(&matched_lid, gpu_lid.get(), sizeof(unsigned int), CudaMemcpyDeviceToHost);
+    return matched_lid; 
+}
+
+void ArbLatticeLauncher::sampleQuantity(int quant, unsigned int lid, real_t* host_tab, real_t scale, const LatticeData &data) const {
+    switch(quant) {	<?R
+for (q in rows(Quantities)) { ifdef(q$adjoint);
+?>
+        case <?%s q$Index ?>: {
+            const auto gpu_tab = cudaMakeUnique<<?%s q$type ?>>(1);
+            getSample<?%s q$name ?>(gpu_tab.get(), scale, data, lid);
+            CudaMemcpy(host_tab, gpu_tab.get(), sizeof(<?%s q$type ?>), CudaMemcpyDeviceToHost);
+            break;
+        } <?R
+}
+ifdef();
+?>
+	}
+}
+
 #endif  // ARBLATTICELAUNCHER_HPP

src/CartLattice.cpp.Rt

@@ -63,7 +63,6 @@ CartLattice::CartLattice(CartConnectivity connect, int ns, const UnitEnv& units_
 {
 	DEBUG_M;
     AllocContainer(launcher.container, getLocalRegion().nx, getLocalRegion().ny, getLocalRegion().nz);
-	sample = std::make_unique<Sampler>(model.get(), units, connectivity.mpi_rank);
 	Snaps = std::make_unique<FTabs[]>(num_snaps);
 	setPosition(0.0,0.0,0.0);
 	DEBUG_M;
@@ -721,8 +720,9 @@ void CartLattice::GetSample<?%s q$name ?>(const lbRegion& over, real_t scale,rea
     launcher.container.in = Snaps[Snap];
 <?R if (q$adjoint) { ?>
     launcher.container.adjin = aSnaps[aSnap]; <?R } ?>
-    lbRegion small = getLocalRegion().intersect(over);
-    launcher.SampleQuantity<?%s q$name ?>(small, (<?%s q$type ?>*)buf, scale, data);
+    lbRegion local_reg = getLocalRegion();
+    lbRegion small_local = over.shift(local_reg.dx, local_reg.dy, local_reg.dz);
+    launcher.SampleQuantity<?%s q$name ?>(small_local, (<?%s q$type ?>*)buf, scale, data);
 }
 <?R } ;ifdef() ?>
 void CartLattice::updateAllSamples(){

src/CartLattice.h.Rt

@@ -8,7 +8,6 @@
 #include "LatticeBase.hpp"
 #include "CartLatticeLauncher.h"
 #include "CartConnectivity.hpp"
-#include "Sampler.h"
 #include "Geometry.h"
 
 #include <memory>
@@ -73,7 +72,6 @@ protected:
 #endif
 
 public:
-  std::unique_ptr<Sampler> sample; //initializing sampler with zero size
   real_t px, py, pz;
 
   CartLattice (CartConnectivity connect, int ns, const UnitEnv& units_);
@@ -148,7 +146,7 @@ public:
   int getPar(const ParStruct& par_struct, double * wb);
   int setPar(const ParStruct& par_struct, double * w);
 
-  void updateAllSamples();
+  virtual void updateAllSamples() override;
   void resetAverage();
 };
 

src/CartLatticeLauncher.hpp.Rt

@@ -307,7 +307,7 @@ public:
   GetQuantitySampleExecutor<?%s q$name ?>(Args&&... args) : GetQuantityExecutor<?%s q$name ?>Base(std::forward<Args>(args)...) {}
 
   CudaHostFunction LaunchParams ComputeLaunchParams(dim3) const {
-    return LaunchParams{dim3(small.nx, small.ny), dim3(1)};
+    return LaunchParams{dim3(small.nx, small.ny, small.nz), dim3(1)};
   }
 };
 

src/Handlers/cbSample.cpp

@@ -16,7 +16,6 @@ int cbSample::Init () {
 		else {
 			s.add_from_string("all",',');
 		}
-                const auto lattice = solver->getCartLattice();
 		for (pugi::xml_node par = node.first_child(); par; par = par.next_sibling()) {
 			if (strcmp(par.name(),"Point") == 0) {
 				lbRegion loc;
@@ -32,33 +31,44 @@ int cbSample::Init () {
 				if (attr) {
 					loc.dz = solver->units.alt(attr.value());
 				}
-				loc = lattice->getLocalRegion().intersect(loc);
-				if (loc.nx == 1)  lattice->sample->addPoint(loc, solver->mpi_rank);
+				loc = solver->lattice->getLocalBoundingBox().intersect(loc);
+
+				if (loc.nx == 1) {
+					unsigned int lid = 0;
+					auto variant = solver->getLatticeVariant();
+					if (auto* lattice = std::get_if<Lattice<ArbLattice>*>(&variant)) {
+						// cache lid for arbitrary lattice
+						const real_t offset = 0.5;
+						vector_t point{real_t(loc.dx) + offset, real_t(loc.dy) + offset, real_t(loc.dz) + offset};
+						lid = (*lattice)->getCartesianCoordinateLid(point);
+					}
+					solver->lattice->sample->addPoint(loc, solver->mpi_rank, lid);
+				}
+
 			} else {
 				error("Uknown element in Sampler\n");
 				return -1;
 			}
 		}
 		filename = solver->outIterFile(nm, ".csv");
-		lattice->sample->units = &solver->units;
-		lattice->sample->mpi_rank = solver->mpi_rank;
-		lattice->sample->Allocate(&s,startIter,everyIter);
-		lattice->sample->initCSV(filename.c_str());
+		solver->lattice->sample->units = &solver->units;
+		solver->lattice->sample->mpi_rank = solver->mpi_rank;
+		solver->lattice->sample->Allocate(&s,startIter,everyIter);
+		solver->lattice->sample->initCSV(filename.c_str());
 		return 0;
 		}
 
 
 int cbSample::DoIt () {
 		Callback::DoIt();
-                const auto lattice = solver->getCartLattice();
-		lattice->sample->writeHistory(solver->iter);
-		lattice->sample->startIter = solver->iter;
+		solver->lattice->sample->writeHistory(solver->iter);
+		solver->lattice->sample->startIter = solver->iter;
 		return 0;
 		}
 
 
 int cbSample::Finish () {
-	   solver->getCartLattice()->sample->Finish();
+	   solver->lattice->sample->Finish();
 	   return Callback::Finish();
 	 }	 
 

src/Lattice.hpp.Rt

@@ -109,8 +109,7 @@ struct Lattice : public LatticeType {
         CudaDeviceSynchronize();
         LatticeType::Snap = tab_out;
         LatticeType::MarkIteration();
-        if constexpr(std::is_same_v<LatticeType, CartLattice>) /// TODO
-            LatticeType::updateAllSamples();
+        LatticeType::updateAllSamples();
         DEBUG_PROF_POP();
     }
 

src/LatticeBase.cpp.Rt

@@ -5,6 +5,7 @@ c_header();
 
 #include "LatticeBase.hpp"
 #include "utils.h"
+#include "Global.h"
 
 LatticeBase::LatticeBase(int zonesettings, int zones, int num_snaps_, const UnitEnv& units_)
   : model(std::make_unique<Model_m>()), zSet(zonesettings, zones), num_snaps(num_snaps_), units(&units_) {
@@ -24,6 +25,8 @@ LatticeBase::LatticeBase(int zonesettings, int zones, int num_snaps_, const Unit
     data.ConstZoneSettings = zSet.gpuConst;
     std::fill_n(iSnaps.get(), maxSnaps, -1);
 
+	sample = std::make_unique<Sampler>(model.get(), units, D_MPI_RANK);
+
 	// Setting settings to default
 <?R for (v in rows(Settings)) {
         if (is.na(v$derived)) { ?>