CudaByteLatticeExtended.h

Go to the documentation of this file.

/*
 * University of Illinois Open Source License
 * Copyright 2024 Luthey-Schulten Group,
 * All rights reserved.
 *
 * Extended CUDA ByteLattice with particle size support
 */
 
#ifndef LM_RDME_CUDABYTELATTICEEXTENDED_H_
#define LM_RDME_CUDABYTELATTICEEXTENDED_H_
 
#include <vector>
#include <map>
#include <memory>
#include "core/Exceptions.h"
#include "cuda/lm_cuda.h"
#include "rdme/ByteLatticeExtended.h"
#include "rdme/CudaByteLattice.h"
#include "rdme/ParticleSizeManager.h"
 
namespace lm {
// Forward declaration
class MpdRdmeSolverExtended;
 
namespace rdme {

class CudaByteLatticeExtended : public CudaByteLattice
{
    friend class lm::MpdRdmeSolverExtended;
 
public:
    CudaByteLatticeExtended(lattice_coord_t size, si_dist_t spacing, uint particlesPerSite,
                           bool enableSizeChecking = true);

    CudaByteLatticeExtended(lattice_size_t xSize, lattice_size_t ySize, lattice_size_t zSize,
                           si_dist_t spacing, uint particlesPerSite, bool enableSizeChecking = true);

    virtual ~CudaByteLatticeExtended();
 
    // GPU memory management overrides
    virtual void copyToGPU() override;
    virtual void copyFromGPU() override;
    virtual void * getGPUMemorySiteSizes();
    virtual void * getGPUMemoryParticleSizes();
    virtual void * getGPUMemorySizeLattice();
    virtual void * getGPUMemorySizeDest();
    virtual void swapSizeSrcDest();
 
    // Particle size management methods
    void setParticleSize(particle_t particleType, uint32_t size);
    uint32_t getParticleSize(particle_t particleType) const;
 
    // Size-aware lattice operations
    uint32_t getSizeBasedOccupancy(lattice_size_t x, lattice_size_t y, lattice_size_t z) const;
    uint32_t getSizeBasedOccupancy(lattice_size_t subvolume) const;
    bool canAddParticle(lattice_size_t x, lattice_size_t y, lattice_size_t z, particle_t particle) const;
    bool canAddParticle(lattice_size_t subvolume, particle_t particle) const;
 
    // GPU-accelerated size-aware operations
    virtual void addParticleWithSizeCheck(lattice_size_t x, lattice_size_t y, lattice_size_t z,
                                         particle_t particle);
    virtual void addParticleWithSizeCheck(lattice_size_t subvolume, particle_t particle);
 
    // Override methods that can cause GPU memory to become stale
    virtual void addParticle(lattice_size_t x, lattice_size_t y, lattice_size_t z, particle_t particle) override;
    virtual void addParticle(lattice_size_t index, particle_t particle) override;
    virtual void removeParticles(lattice_size_t x, lattice_size_t y, lattice_size_t z) override;
    virtual void removeParticles(lattice_size_t index) override;
    virtual void removeAllParticles() override;
 
    // Size checking control
    void setSizeCheckingEnabled(bool enable);
    bool isSizeCheckingEnabled() const;
 
    // Access to size manager
    const ParticleSizeManager& getParticleSizeManager() const;
 
    // Validation and statistics
    void validateSizeConstraints() const;
    std::map<uint32_t, uint32_t> getSizeBasedOccupancyStats() const;
 
    // GPU-specific utilities
    void updateGPUSiteSizes();
    void syncSiteSizesFromGPU();
 
protected:
    virtual void allocateCudaMemory() override;
    virtual void deallocateCudaMemory() override;
 
    // Update site size when particles are added/removed
    void updateSiteSize(lattice_size_t subvolume, particle_t particle, bool adding);
 
protected:
    // CPU-side data structures
    std::unique_ptr<ParticleSizeManager> sizeManager_;
    bool sizeCheckingEnabled_;
    uint32_t maxSiteCapacity_;
    std::vector<uint32_t> siteSizes_;  // Current size occupancy per site
 
    // GPU memory for size tracking
    size_t cudaSiteSizesSize_;
    void * cudaSiteSizes_[2];        // Double-buffered size occupancy lattices on GPU (uint32_t*)
    uint cudaSiteSizesCurrent_;      // Current buffer index (0 or 1)
    size_t cudaParticleSizesSize_;
    void * cudaParticleSizes_;       // Particle sizes lookup table on GPU (uint32_t*)
 
    bool isGPUSiteSizesSynched_;
    bool isGPUParticleSizesSynched_;
 
    // Fast GPU operations for size checking
    void initializeGPUSiteSizes();
    void recalculateAllSiteSizes();
    bool canAddParticleGPU(lattice_size_t subvolume, particle_t particle) const;
    void addParticleGPU(lattice_size_t subvolume, particle_t particle);
    void removeParticleGPU(lattice_size_t subvolume, particle_t particle);
 
    // Helper methods
    site_size_t getParticlesAtSite(lattice_size_t subvolume, particle_t* particles,
                                  site_size_t maxCount) const;
};
 
}
}
 
#endif // LM_RDME_CUDABYTELATTICEEXTENDED_H_