particle_size_kernels.cuh

Go to the documentation of this file.

/*
 * University of Illinois Open Source License
 * Copyright 2024 Luthey-Schulten Group,
 * All rights reserved.
 *
 * CUDA kernels for particle size-aware RDME operations
 */
 
#ifndef LM_CUDA_PARTICLE_SIZE_KERNELS_CUH
#define LM_CUDA_PARTICLE_SIZE_KERNELS_CUH
 
#include "config.h"
#include "core/Types.h"
#include "cuda/constant.cuh"
 
// Ensure particle_t is defined for CUDA compilation
#ifndef particle_t
typedef uint32_t particle_t;
#endif
 
#ifndef lattice_size_t
typedef uint32_t lattice_size_t;
#endif
 
// Maximum number of particle types supported in constant memory
#define MAX_PARTICLE_TYPES_CONSTANT 256
 
// Constant memory for particle sizes (accessible from all kernels) - must be global
extern __constant__ uint32_t particle_sizes_constant[MAX_PARTICLE_TYPES_CONSTANT];
 
namespace lm {
namespace cuda {
 
/**
 * @brief Device function to get particle size from constant memory
 * @param particleType Particle type ID
 * @return Size of the particle (default: 1)
 */
__device__ __forceinline__ uint32_t getParticleSize(particle_t particleType) {
    return (particleType < MAX_PARTICLE_TYPES_CONSTANT) ?
           particle_sizes_constant[particleType] : 1;
}
 
/**
 * @brief Template device function to calculate total size occupancy at a site
 * @param particles Array of particles at the site
 * @param count Number of particles
 * @return Total size of all particles
 */
template<typename ParticleType>
__device__ uint32_t calculateSiteOccupancy(const ParticleType* particles, uint32_t count) {
    uint32_t totalSize = 0;
    for (uint32_t i = 0; i < count; ++i) {
        totalSize += getParticleSize(particles[i]);
    }
    return totalSize;
}
 
__device__ void validateKernelParams(uint32_t latticeXSize, uint32_t latticeYSize, 
    uint32_t latticeZSize, uint32_t particlesPerSite);
 
 
/**
 * @brief Memory-safe site occupancy calculation with overflow protection
 * @param particles Array of particles at the site
 * @param count Number of particles
 * @param maxCount Maximum safe count to prevent buffer overflow
 * @return Total size of all particles, UINT32_MAX if overflow detected
 */
template<typename ParticleType>
__device__ uint32_t calculateSiteOccupancySafe(const ParticleType* particles, uint32_t count, uint32_t maxCount);
 
/**
 * @brief Template device function to check if adding a particle would exceed capacity
 * @param particles Current particles at site
 * @param count Current particle count
 * @param newParticle Particle to add
 * @param maxCapacity Maximum site capacity
 * @return True if particle can be added
 */
template<typename ParticleType>
__device__ bool canAddParticleDevice(const ParticleType* particles, uint32_t count,
                                    particle_t newParticle, uint32_t maxCapacity) {
    uint32_t currentOccupancy = calculateSiteOccupancy(particles, count);
    uint32_t newParticleSize = getParticleSize(newParticle);
    return (currentOccupancy + newParticleSize) <= maxCapacity;
}
 
// Explicit instantiations for common types
__device__ uint32_t calculateSiteOccupancy(const uint8_t* particles, uint32_t count);
__device__ uint32_t calculateSiteOccupancy(const uint32_t* particles, uint32_t count);
__device__ bool canAddParticleDevice(const uint8_t* particles, uint32_t count,
                                    particle_t newParticle, uint32_t maxCapacity);
__device__ bool canAddParticleDevice(const uint32_t* particles, uint32_t count,
                                    particle_t newParticle, uint32_t maxCapacity);
 
/**
 * @brief Kernel to validate size constraints across entire lattice
 * @param lattice Particle lattice data
 * @param siteLattice Site type lattice
 * @param violationFlags Output array for violation flags
 * @param latticeSize Total lattice size
 * @param particlesPerSite Maximum particles per site
 */
__global__ void validateSizeConstraintsKernel(
    const uint8_t* lattice,
    const uint8_t* siteLattice,
    uint32_t* violationFlags,
    uint32_t latticeSize,
    uint32_t particlesPerSite
);
 
/**
 * @brief Kernel to calculate occupancy statistics
 * @param lattice Particle lattice data
 * @param occupancyCounts Output array for occupancy histogram
 * @param latticeSize Total lattice size
 * @param particlesPerSite Maximum particles per site
 * @param maxOccupancy Maximum occupancy value to track
 */
__global__ void calculateOccupancyStatsKernel(
    const uint8_t* lattice,
    uint32_t* occupancyCounts,
    uint32_t latticeSize,
    uint32_t particlesPerSite,
    uint32_t maxOccupancy
);
 
 
/**
 * @brief Host function to copy particle sizes to constant memory
 * @param particleSizes Host array of particle sizes
 * @param numTypes Number of particle types
 * @return cudaError_t error code
 */
cudaError_t copyParticleSizesToConstantMemory(const uint32_t* particleSizes, size_t numTypes);
 
/**
 * @brief Host function to validate lattice size constraints (ByteLattice)
 * @param lattice Device pointer to lattice data
 * @param siteLattice Device pointer to site data
 * @param latticeSize Total lattice size
 * @param particlesPerSite Maximum particles per site
 * @param stream CUDA stream for async execution
 * @return Number of sites violating constraints
 */
uint32_t validateLatticeConstraints(const uint8_t* lattice, const uint8_t* siteLattice,
                                   uint32_t latticeSize, uint32_t particlesPerSite,
                                   cudaStream_t stream = 0);
 
/**
 * @brief Host function to validate lattice size constraints (IntLattice)
 * @param lattice Device pointer to lattice data
 * @param siteLattice Device pointer to site data
 * @param latticeSize Total lattice size
 * @param particlesPerSite Maximum particles per site
 * @param stream CUDA stream for async execution
 * @return Number of sites violating constraints
 */
uint32_t validateLatticeConstraints(const uint32_t* lattice, const uint8_t* siteLattice,
                                   uint32_t latticeSize, uint32_t particlesPerSite,
                                   cudaStream_t stream = 0);
 
/**
 * @brief Host function to get occupancy statistics from GPU (ByteLattice)
 * @param lattice Device pointer to lattice data
 * @param latticeSize Total lattice size
 * @param particlesPerSite Maximum particles per site
 * @param occupancyCounts Host output array for histogram
 * @param maxOccupancy Maximum occupancy to track
 * @param stream CUDA stream for async execution
 * @return cudaError_t error code
 */
cudaError_t getOccupancyStatistics(const uint8_t* lattice, uint32_t latticeSize,
                                  uint32_t particlesPerSite, uint32_t* occupancyCounts,
                                  uint32_t maxOccupancy, cudaStream_t stream = 0);
 
/**
 * @brief Host function to get occupancy statistics from GPU (IntLattice)
 * @param lattice Device pointer to lattice data
 * @param latticeSize Total lattice size
 * @param particlesPerSite Maximum particles per site
 * @param occupancyCounts Host output array for histogram
 * @param maxOccupancy Maximum occupancy to track
 * @param stream CUDA stream for async execution
 * @return cudaError_t error code
 */
cudaError_t getOccupancyStatistics(const uint32_t* lattice, uint32_t latticeSize,
                                  uint32_t particlesPerSite, uint32_t* occupancyCounts,
                                  uint32_t maxOccupancy, cudaStream_t stream = 0);
 
 
 
// =============== NEW EFFICIENT SIZE TRACKING KERNELS ===============
 
/**
 * @brief Initialize site sizes lattice from current particle lattice
 * @param particles Particle lattice data
 * @param siteSizes Output site sizes lattice (uint32_t per site)
 * @param particleSizes Particle sizes lookup table
 * @param latticeSize Total number of sites
 * @param particlesPerSite Maximum particles per site
 */
__global__ void initializeSiteSizesKernel(
    const uint8_t* particles,
    uint32_t* siteSizes,
    const uint32_t* particleSizes,
    uint32_t latticeSize,
    uint32_t particlesPerSite
);
 
/**
 * @brief Check if particles can be added to sites (batch operation)
 * @param siteSizes Current site sizes lattice
 * @param particleSizes Particle sizes lookup table
 * @param subvolumes Array of subvolume indices to check
 * @param particleTypes Array of particle types to add
 * @param results Output array of boolean results
 * @param maxCapacity Maximum capacity per site
 * @param count Number of operations to perform
 */
__global__ void canAddParticlesBatchKernel(
    const uint32_t* siteSizes,
    const uint32_t* particleSizes,
    const lattice_size_t* subvolumes,
    const particle_t* particleTypes,
    bool* results,
    uint32_t maxCapacity,
    uint32_t count
);
 
/**
 * @brief Add particles to sites and update size lattice (batch operation)
 * @param siteSizes Site sizes lattice to update
 * @param particleSizes Particle sizes lookup table
 * @param subvolumes Array of subvolume indices
 * @param particleTypes Array of particle types to add
 * @param maxCapacity Maximum capacity per site
 * @param count Number of operations to perform
 * @param results Output array indicating success/failure
 */
__global__ void addParticlesBatchUpdateSizesKernel(
    uint32_t* siteSizes,
    const uint32_t* particleSizes,
    const lattice_size_t* subvolumes,
    const particle_t* particleTypes,
    uint32_t maxCapacity,
    uint32_t count,
    bool* results
);
 
/**
 * @brief Remove particles from sites and update size lattice (batch operation)
 * @param siteSizes Site sizes lattice to update
 * @param particleSizes Particle sizes lookup table
 * @param subvolumes Array of subvolume indices
 * @param particleTypes Array of particle types to remove
 * @param count Number of operations to perform
 */
__global__ void removeParticlesBatchUpdateSizesKernel(
    uint32_t* siteSizes,
    const uint32_t* particleSizes,
    const lattice_size_t* subvolumes,
    const particle_t* particleTypes,
    uint32_t count
);
 
/**
 * @brief Fast size-aware diffusion with dedicated site size lattice
 * @param inParticles Input particle lattice
 * @param inSiteSizes Input site sizes lattice
 * @param outParticles Output particle lattice
 * @param outSiteSizes Output site sizes lattice
 * @param inSites Site types
 * @param particleSizes Particle sizes lookup table
 * @param timestepHash Random seed
 * @param latticeXSize X dimension
 * @param latticeYSize Y dimension
 * @param latticeZSize Z dimension
 * @param particlesPerSite Maximum particles per site
 * @param maxCapacity Maximum site capacity
 */
__global__ void fastSizeAwareDiffusionKernel(
    const uint8_t* inParticles,
    const uint32_t* inSiteSizes,
    uint8_t* outParticles,
    uint32_t* outSiteSizes,
    const uint8_t* inSites,
    const uint32_t* particleSizes,
    unsigned long long timestepHash,
    uint32_t latticeXSize,
    uint32_t latticeYSize,
    uint32_t latticeZSize,
    uint32_t particlesPerSite,
    uint32_t maxCapacity
);
/**
 * @brief Perform size-aware propagation
 * @param outLattice Output lattice
 * @param outSizeLattice Output size lattice
 * @param window Window
 * @param sizeWindow Size window
 * @param choices Choices
 * @param latticeIndex Lattice index
 * @param latticeXYSize Lattice XY size
 * @param latticeXYZSize Lattice XYZ size
 * @param windowIndex Window index
 * @param particlesPerSite Maximum particles per site
 * @param siteOverflowList Site overflow list
 */
__device__ void performSizeAwarePropagation(
    unsigned int* __restrict__ outLattice,
    uint32_t* __restrict__ outSizeLattice,
    const unsigned int* __restrict__ window,
    const uint32_t* __restrict__ sizeWindow,
    const uint8_t* __restrict__ choices,
    const uint32_t latticeIndex,
    const uint32_t latticeXYSize,
    const uint32_t latticeXYZSize,
    const uint32_t windowIndex,
    const uint32_t particlesPerSite,
    uint32_t* __restrict__ siteOverflowList);
 
// Host wrapper functions for the new kernels
cudaError_t initializeSiteSizes(const uint8_t* particles, uint32_t* siteSizes,
                               const uint32_t* particleSizes, uint32_t latticeSize,
                               uint32_t particlesPerSite, cudaStream_t stream = 0);
 
cudaError_t canAddParticlesBatch(const uint32_t* siteSizes, const uint32_t* particleSizes,
                                const lattice_size_t* subvolumes, const particle_t* particleTypes,
                                bool* results, uint32_t maxCapacity, uint32_t count,
                                cudaStream_t stream = 0);
 
cudaError_t addParticlesBatchUpdateSizes(uint32_t* siteSizes, const uint32_t* particleSizes,
                                        const lattice_size_t* subvolumes, const particle_t* particleTypes,
                                        uint32_t maxCapacity, uint32_t count, bool* results,
                                        cudaStream_t stream = 0);
 
cudaError_t removeParticlesBatchUpdateSizes(uint32_t* siteSizes, const uint32_t* particleSizes,
                                           const lattice_size_t* subvolumes, const particle_t* particleTypes,
                                           uint32_t count, cudaStream_t stream = 0);
 
} // namespace cuda
} // namespace lm
 
/**
 * @brief Kernel for windowed size-aware diffusion
 * @param inLattice Input lattice state
 * @param inSites Site types
 * @param inSizeLattice Input size lattice
 * @param outLattice Output lattice state
 * @param outSizeLattice Output size lattice
 * @param timestepHash Random seed
 * @param siteOverflowList List for tracking overflow sites
 * @param latticeXSize X dimension
 * @param latticeYSize Y dimension
 * @param latticeZSize Z dimension
 * @param particlesPerSite Maximum particles per site
 */
__global__ void sizeAwareWindowedDiffusionKernel(
    const unsigned int* inLattice,
    const uint8_t* inSites,
    const uint32_t* inSizeLattice,
    unsigned int* outLattice,
    uint32_t* outSizeLattice,
    const unsigned long long timestepHash,
    uint32_t* siteOverflowList,
    uint32_t latticeXSize,
    uint32_t latticeYSize,
    uint32_t latticeZSize,
    uint32_t particlesPerSite);
 
/**
 * @brief Kernel for size-aware reaction processing
 * @param inLattice Input lattice state
 * @param inSites Site types
 * @param outLattice Output lattice state
 * @param timestepHash Random seed for this timestep
 * @param siteOverflowList List for tracking overflow sites
 * @param reactionMatrix Reaction stoichiometry matrix
 * @param reactionLocationMatrix Reaction location matrix
 * @param latticeXSize X dimension
 * @param latticeYSize Y dimension
 * @param latticeZSize Z dimension
 * @param particlesPerSite Maximum particles per site
 * @param numReactions Number of reactions
 * @param numSpecies Number of species
 */
__global__ void sizeAwareReactionKernel(
    const uint8_t* inLattice,
    const uint32_t* inSizeLattice,
    const uint8_t* inSites,
    uint32_t* outSizeLattice,
    uint8_t* outLattice,
    unsigned long long timestepHash,
    uint32_t* siteOverflowList,
    const int8_t* reactionMatrix,
    const uint8_t* reactionLocationMatrix,
    uint32_t latticeXSize,
    uint32_t latticeYSize,
    uint32_t latticeZSize,
    uint32_t particlesPerSite,
    uint32_t numReactions,
    uint32_t numSpecies,
    uint32_t numSiteTypes
);
 
/**
 * Optimized size-aware reaction kernel with precomputed propensities
 * Multiple signatures based on compile-time macros (MPD_GLOBAL_S_MATRIX, MPD_GLOBAL_R_MATRIX)
 */
#ifdef MPD_GLOBAL_S_MATRIX
#ifdef MPD_GLOBAL_R_MATRIX
__global__ void sizeAwarePrecompReactionKernel(
    const unsigned int* inLattice,
    const uint32_t* inSizeLattice,
    const uint8_t* inSites,
    uint32_t* outSizeLattice,
    unsigned int* outLattice,
    const unsigned long long timestepHash,
    uint32_t* siteOverflowList,
    const int8_t* __restrict__ SG,
    const uint8_t* __restrict__ RLG,
    const unsigned int* __restrict__ reactionOrdersG,
    const unsigned int* __restrict__ reactionSitesG,
    const unsigned int* __restrict__ D1G,
    const unsigned int* __restrict__ D2G,
    const float* __restrict__ reactionRatesG,
    const float* __restrict__ qp0,
    const float* __restrict__ qp1,
    const float* __restrict__ qp2,
    uint32_t latticeXSize,
    uint32_t latticeYSize,
    uint32_t latticeZSize,
    uint32_t particlesPerSite,
    uint32_t numSpecies);
#else
__global__ void sizeAwarePrecompReactionKernel(
    const uint8_t* inLattice,
    const uint32_t* inSizeLattice,
    const uint8_t* inSites,
    uint32_t* outSizeLattice,
    uint8_t* outLattice,
    unsigned long long timestepHash,
    uint32_t* siteOverflowList,
    const int8_t* __restrict__ SG,
    const uint8_t* __restrict__ RLG,
    const unsigned int* __restrict__ reactionOrdersG,
    const unsigned int* __restrict__ reactionSitesG,
    const unsigned int* __restrict__ D1G,
    const unsigned int* __restrict__ D2G,
    const float* __restrict__ reactionRatesG,
    const float* __restrict__ qp0,
    const float* __restrict__ qp1,
    const float* __restrict__ qp2,
    uint32_t latticeXSize,
    uint32_t latticeYSize,
    uint32_t latticeZSize,
    uint32_t particlesPerSite,
    uint32_t numSpecies);
#endif
#else
__global__ void sizeAwarePrecompReactionKernel(
    const uint8_t* inLattice,
    const uint32_t* inSizeLattice,
    const uint8_t* inSites,
    uint32_t* outSizeLattice,
    uint8_t* outLattice,
    unsigned long long timestepHash,
    uint32_t* siteOverflowList,
    const float* __restrict__ qp0,
    const float* __restrict__ qp1,
    const float* __restrict__ qp2,
    uint32_t latticeXSize,
    uint32_t latticeYSize,
    uint32_t latticeZSize,
    uint32_t particlesPerSite,
    uint32_t numSpecies);
#endif
 
// =============== BYTE LATTICE OVERFLOW CORRECTION KERNEL ===============
 
/**
 * Byte-lattice-compatible overflow correction kernel
 * Handles overflows for uint8_t lattices instead of unsigned int lattices
 */
__global__ void correct_byte_overflows(
    uint8_t* lattice,
    uint32_t* siteOverflowList,
    uint32_t latticeXSize,
    uint32_t latticeYSize,
    uint32_t latticeZSize,
    uint32_t particlesPerSite);
 
// =============== NEW 1D SIZE-AWARE DIFFUSION KERNELS ===============
 
/**
 * @brief Size-aware Z-direction diffusion kernel (1D approach)
 */
__global__ void sizeAwareZDiffusionKernel(
    const uint8_t* inLattice,
    const uint8_t* inSites,
    const uint32_t* inSizeLattice,
    uint8_t* outLattice,
    uint32_t* outSizeLattice,
    const unsigned long long timestepHash,
    uint32_t* siteOverflowList,
    uint32_t latticeXSize,
    uint32_t latticeYSize,
    uint32_t latticeZSize,
    uint32_t particlesPerSite
);
 
/**
 * @brief Size-aware Y-direction diffusion kernel (1D approach)
 */
__global__ void sizeAwareYDiffusionKernel(
    const uint8_t* inLattice,
    const uint8_t* inSites,
    const uint32_t* inSizeLattice,
    uint8_t* outLattice,
    uint32_t* outSizeLattice,
    const unsigned long long timestepHash,
    uint32_t* siteOverflowList,
    uint32_t latticeXSize,
    uint32_t latticeYSize,
    uint32_t latticeZSize,
    uint32_t particlesPerSite
);
 
/**
 * @brief Size-aware X-direction diffusion kernel (1D approach)
 */
__global__ void sizeAwareXDiffusionKernel(
    const uint8_t* inLattice,
    const uint8_t* inSites,
    const uint32_t* inSizeLattice,
    uint8_t* outLattice,
    uint32_t* outSizeLattice,
    const unsigned long long timestepHash,
    uint32_t* siteOverflowList,
    uint32_t latticeXSize,
    uint32_t latticeYSize,
    uint32_t latticeZSize,
    uint32_t particlesPerSite
);
 
 
#endif // LM_CUDA_PARTICLE_SIZE_KERNELS_CUH