<commit_before>#include "chrono_parallel/physics/ChMPM.cuh"
#include "chrono_parallel/physics/MPMUtils.h"
#include "chrono_parallel/ChCudaHelper.cuh"
#include "chrono_parallel/ChGPUVector.cuh"
#include "thirdparty/cub/cub.cuh"
namespace chrono {

struct Bounds {
    real minimum[3];
    real maximum[3];
};

real3 min_bounding_point;
real3 max_bounding_point;
vec3 bins_per_axis;

MPM_Settings host_settings;

std::vector<int> particle_node_mapping;
std::vector<int> node_particle_mapping;
std::vector<int> node_start_index;
std::vector<int> particle_number;
uint num_mpm_nodes_active;
std::vector<Mat33> volume_Ap_Fe_transpose;

// GPU Things
real3* lower_bound;
real3* upper_bound;

gpu_vector<real3> pos, vel;
gpu_vector<real> node_mass;
gpu_vector<real> marker_volume;
gpu_vector<real> grid_vel, delta_v;
gpu_vector<real> rhs;
gpu_vector<Mat33> marker_Fe, marker_Fe_hat, marker_Fp, marker_delta_F, PolarR, PolarS;
gpu_vector<real> old_vel_node_mpm;
gpu_vector<real> ml, mg, mg_p, ml_p;
gpu_vector<real> dot_g_proj_norm;

CUDA_CONSTANT MPM_Settings device_settings;
CUDA_CONSTANT Bounds system_bounds;

cudaEvent_t start;
cudaEvent_t stop;
float time_measured = 0;
/////// BB Constants
__device__ real alpha = 0.0001;
__device__ real dot_ms_ms = 0;
__device__ real dot_ms_my = 0;
__device__ real dot_my_my = 0;

#define a_min 1e-13
#define a_max 1e13
#define neg_BB1_fallback 0.11
#define neg_BB2_fallback 0.12

#define LOOP_TWO_RING_GPU(X)                                                                                         \
    const real bin_edge = device_settings.bin_edge;                                                                  \
    const real inv_bin_edge = 1.f / bin_edge;                                                                        \
                                                                                                                     \
    const int cx = GridCoord(xi.x, inv_bin_edge, system_bounds.minimum[0]);                                          \
    const int cy = GridCoord(xi.y, inv_bin_edge, system_bounds.minimum[1]);                                          \
    const int cz = GridCoord(xi.z, inv_bin_edge, system_bounds.minimum[2]);                                          \
    for (int i = cx - 2; i <= cx + 2; ++i) {                                                                         \
        for (int j = cy - 2; j <= cy + 2; ++j) {                                                                     \
            for (int k = cz - 2; k <= cz + 2; ++k) {                                                                 \
                const int current_node = GridHash(i, j, k, device_settings.bins_per_axis_x,                          \
                                                  device_settings.bins_per_axis_y, device_settings.bins_per_axis_z); \
                real3 current_node_location;                                                                         \
                current_node_location.x = i * bin_edge + system_bounds.minimum[0];                                   \
                current_node_location.y = j * bin_edge + system_bounds.minimum[1];                                   \
                current_node_location.z = k * bin_edge + system_bounds.minimum[2];                                   \
                X                                                                                                    \
            }                                                                                                        \
        }                                                                                                            \
    }

#define LOOP_TWO_RING_GPUSP(X)                                                                                         \
    cx = GridCoord(xi.x, inv_bin_edge, system_bounds.minimum[0]);                                                      \
    cy = GridCoord(xi.y, inv_bin_edge, system_bounds.minimum[1]);                                                      \
    cz = GridCoord(xi.z, inv_bin_edge, system_bounds.minimum[2]);                                                      \
    for (int i = cx - 2; i <= cx + 2; ++i) {                                                                           \
        for (int j = cy - 2; j <= cy + 2; ++j) {                                                                       \
            for (int k = cz - 2; k <= cz + 2; ++k) {                                                                   \
                int current_node = GridHash(i, j, k, device_settings.bins_per_axis_x, device_settings.bins_per_axis_y, \
                                            device_settings.bins_per_axis_z);                                          \
                real current_node_locationx = i * bin_edge + system_bounds.minimum[0];                                 \
                real current_node_locationy = j * bin_edge + system_bounds.minimum[1];                                 \
                real current_node_locationz = k * bin_edge + system_bounds.minimum[2];                                 \
                X                                                                                                      \
            }                                                                                                          \
        }                                                                                                              \
    }
//////========================================================================================================================================================================
////
CUDA_GLOBAL void kComputeBounds(const real3* pos,  // input
                                real3* lower,      // output
                                real3* upper       // output
                                ) {
    typedef cub::BlockReduce<real3, num_threads_per_block> BlockReduce;

    __shared__ typename BlockReduce::TempStorage temp_storage;

    const int block_start = blockDim.x * blockIdx.x;
    const int num_valid = min(device_settings.num_mpm_markers - block_start, blockDim.x);

    const int index = block_start + threadIdx.x;
    if (index < device_settings.num_mpm_markers) {
        real3 data = pos[index];
        real3 blockUpper = BlockReduce(temp_storage).Reduce(data, real3Max(), num_valid);
        __syncthreads();
        real3 blockLower = BlockReduce(temp_storage).Reduce(data, real3Min(), num_valid);
        if (threadIdx.x == 0) {
            AtomicMax(upper, blockUpper);
            AtomicMin(lower, blockLower);
        }
    }
}
////========================================================================================================================================================================
CUDA_GLOBAL void kRasterize(const real3* sorted_pos,  // input
                            const real3* sorted_vel,  // input
                            real* grid_mass,          // output
                            real* grid_vel) {         // output
    const int p = blockIdx.x * blockDim.x + threadIdx.x;
    if (p < device_settings.num_mpm_markers) {
        const real3 xi = sorted_pos[p];
        const real3 vi = sorted_vel[p];
        LOOP_TWO_RING_GPU(                                                                     //
            real weight = N(xi - current_node_location, inv_bin_edge) * device_settings.mass;  //
            ATOMIC_ADD(&grid_mass[current_node], weight);                                      //
            ATOMIC_ADD(&grid_vel[current_node * 3 + 0], weight * vi.x);
            ATOMIC_ADD(&grid_vel[current_node * 3 + 1], weight * vi.y);
            ATOMIC_ADD(&grid_vel[current_node * 3 + 2], weight * vi.z);)
    }
}
CUDA_GLOBAL void kRasterize(const real3* sorted_pos,  // input
                            real* grid_mass) {        // output
    const int p = blockIdx.x * blockDim.x + threadIdx.x;
    if (p < device_settings.num_mpm_markers) {
        const real3 xi = sorted_pos[p];
        LOOP_TWO_RING_GPU(                                                                     //
            real weight = N(xi - current_node_location, inv_bin_edge) * device_settings.mass;  //
            ATOMIC_ADD(&grid_mass[current_node], weight);                                      //
            )
    }
}
//
////========================================================================================================================================================================
//
CUDA_GLOBAL void kNormalizeWeights(real* grid_mass,   // input
                                   real* grid_vel) {  // output
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < device_settings.num_mpm_nodes) {
        real n_mass = grid_mass[i];
        if (n_mass > C_EPSILON) {
            grid_vel[i * 3 + 0] /= n_mass;
            grid_vel[i * 3 + 1] /= n_mass;
            grid_vel[i * 3 + 2] /= n_mass;
        }
    }
}
//////========================================================================================================================================================================
////
CUDA_GLOBAL void kComputeParticleVolumes(const real3* sorted_pos,  // input
                                         real* grid_mass,          // output
                                         real* volume) {
    const int p = blockIdx.x * blockDim.x + threadIdx.x;
    if (p < device_settings.num_mpm_markers) {
        const real3 xi = sorted_pos[p];
        real particle_density = 0;
        LOOP_TWO_RING_GPU(                                              //
            real weight = N(xi - current_node_location, inv_bin_edge);  //
            particle_density += grid_mass[current_node] * weight;       //
            )
        // Inverse density to remove division
        particle_density = (bin_edge * bin_edge * bin_edge) / particle_density;
        volume[p] = device_settings.mass * particle_density;
    }
}
CUDA_GLOBAL void kFeHat(const real3* sorted_pos,  // input
                        const Mat33* marker_Fe,   // input
                        const real* grid_vel,     // input
                        Mat33* marker_Fe_hat) {   // output

    const int p = blockIdx.x * blockDim.x + threadIdx.x;
    if (p < device_settings.num_mpm_markers) {
        const real3 xi = sorted_pos[p];
        marker_Fe_hat[p] = Mat33(1.0);
        Mat33 Fe_hat_t(1.0);
        LOOP_TWO_RING_GPU(                                                             //
            real3 vel(grid_vel[i * 3 + 0], grid_vel[i * 3 + 1], grid_vel[i * 3 + 2]);  //
            real3 kern = dN(xi - current_node_location, inv_bin_edge);                 //
            Fe_hat_t += OuterProduct(device_settings.dt * vel, kern);)
        marker_Fe_hat[p] = Fe_hat_t * marker_Fe[p];
    }
}

// CUDA_GLOBAL void kSVD(Mat33* marker_Fe_hat, Mat33* PolarR, Mat33* PolarS) {
//    const int p = blockIdx.x * blockDim.x + threadIdx.x;
//    if (p < device_settings.num_mpm_markers) {
//        Mat33 U, V, R, S, W;
//        real3 E;
//        SVD(marker_Fe_hat[p], U, E, V);
//        // Perform polar decomposition F = R*S
//        R = MultTranspose(U, V);
//        S = V * MultTranspose(Mat33(E), V);
//
//        PolarR[p] = R;
//        PolarS[p] = S;
//    }
//}

CUDA_GLOBAL void kApplyForces(const real3* sorted_pos,     // input
                              const Mat33* marker_Fe_hat,  // input
                              const Mat33* marker_Fe,      // input
                              const Mat33* marker_Fp,      // input
                              const real* marker_volume,   // input
                              const real* node_mass,       // input
                              Mat33* PolarR,               // input<commit_msg>Fe_hat computation was completely wrong, fixed and improved performance<commit_after>#include "chrono_parallel/physics/ChMPM.cuh"
#include "chrono_parallel/physics/MPMUtils.h"
#include "chrono_parallel/ChCudaHelper.cuh"
#include "chrono_parallel/ChGPUVector.cuh"
#include "thirdparty/cub/cub.cuh"
namespace chrono {

struct Bounds {
    real minimum[3];
    real maximum[3];
};

real3 min_bounding_point;
real3 max_bounding_point;
vec3 bins_per_axis;

MPM_Settings host_settings;

std::vector<int> particle_node_mapping;
std::vector<int> node_particle_mapping;
std::vector<int> node_start_index;
std::vector<int> particle_number;
uint num_mpm_nodes_active;
std::vector<Mat33> volume_Ap_Fe_transpose;

// GPU Things
real3* lower_bound;
real3* upper_bound;

gpu_vector<real3> pos, vel;
gpu_vector<real> node_mass;
gpu_vector<real> marker_volume;
gpu_vector<real> grid_vel, delta_v;
gpu_vector<real> rhs;
gpu_vector<Mat33> marker_Fe, marker_Fe_hat, marker_Fp, marker_delta_F, PolarR, PolarS;
gpu_vector<real> old_vel_node_mpm;
gpu_vector<real> ml, mg, mg_p, ml_p;
gpu_vector<real> dot_g_proj_norm;

CUDA_CONSTANT MPM_Settings device_settings;
CUDA_CONSTANT Bounds system_bounds;

cudaEvent_t start;
cudaEvent_t stop;
float time_measured = 0;
/////// BB Constants
__device__ real alpha = 0.0001;
__device__ real dot_ms_ms = 0;
__device__ real dot_ms_my = 0;
__device__ real dot_my_my = 0;

#define a_min 1e-13
#define a_max 1e13
#define neg_BB1_fallback 0.11
#define neg_BB2_fallback 0.12

#define LOOP_TWO_RING_GPU(X)                                                                                         \
    const real bin_edge = device_settings.bin_edge;                                                                  \
    const real inv_bin_edge = 1.f / bin_edge;                                                                        \
                                                                                                                     \
    const int cx = GridCoord(xi.x, inv_bin_edge, system_bounds.minimum[0]);                                          \
    const int cy = GridCoord(xi.y, inv_bin_edge, system_bounds.minimum[1]);                                          \
    const int cz = GridCoord(xi.z, inv_bin_edge, system_bounds.minimum[2]);                                          \
    for (int i = cx - 2; i <= cx + 2; ++i) {                                                                         \
        for (int j = cy - 2; j <= cy + 2; ++j) {                                                                     \
            for (int k = cz - 2; k <= cz + 2; ++k) {                                                                 \
                const int current_node = GridHash(i, j, k, device_settings.bins_per_axis_x,                          \
                                                  device_settings.bins_per_axis_y, device_settings.bins_per_axis_z); \
                real3 current_node_location;                                                                         \
                current_node_location.x = i * bin_edge + system_bounds.minimum[0];                                   \
                current_node_location.y = j * bin_edge + system_bounds.minimum[1];                                   \
                current_node_location.z = k * bin_edge + system_bounds.minimum[2];                                   \
                X                                                                                                    \
            }                                                                                                        \
        }                                                                                                            \
    }

#define LOOP_TWO_RING_GPUSP(X)                                                                                         \
    cx = GridCoord(xi.x, inv_bin_edge, system_bounds.minimum[0]);                                                      \
    cy = GridCoord(xi.y, inv_bin_edge, system_bounds.minimum[1]);                                                      \
    cz = GridCoord(xi.z, inv_bin_edge, system_bounds.minimum[2]);                                                      \
    for (int i = cx - 2; i <= cx + 2; ++i) {                                                                           \
        for (int j = cy - 2; j <= cy + 2; ++j) {                                                                       \
            for (int k = cz - 2; k <= cz + 2; ++k) {                                                                   \
                int current_node = GridHash(i, j, k, device_settings.bins_per_axis_x, device_settings.bins_per_axis_y, \
                                            device_settings.bins_per_axis_z);                                          \
                real current_node_locationx = i * bin_edge + system_bounds.minimum[0];                                 \
                real current_node_locationy = j * bin_edge + system_bounds.minimum[1];                                 \
                real current_node_locationz = k * bin_edge + system_bounds.minimum[2];                                 \
                X                                                                                                      \
            }                                                                                                          \
        }                                                                                                              \
    }
//////========================================================================================================================================================================
////
CUDA_GLOBAL void kComputeBounds(const real3* pos,  // input
                                real3* lower,      // output
                                real3* upper       // output
                                ) {
    typedef cub::BlockReduce<real3, num_threads_per_block> BlockReduce;

    __shared__ typename BlockReduce::TempStorage temp_storage;

    const int block_start = blockDim.x * blockIdx.x;
    const int num_valid = min(device_settings.num_mpm_markers - block_start, blockDim.x);

    const int index = block_start + threadIdx.x;
    if (index < device_settings.num_mpm_markers) {
        real3 data = pos[index];
        real3 blockUpper = BlockReduce(temp_storage).Reduce(data, real3Max(), num_valid);
        __syncthreads();
        real3 blockLower = BlockReduce(temp_storage).Reduce(data, real3Min(), num_valid);
        if (threadIdx.x == 0) {
            AtomicMax(upper, blockUpper);
            AtomicMin(lower, blockLower);
        }
    }
}
////========================================================================================================================================================================
CUDA_GLOBAL void kRasterize(const real3* sorted_pos,  // input
                            const real3* sorted_vel,  // input
                            real* grid_mass,          // output
                            real* grid_vel) {         // output
    const int p = blockIdx.x * blockDim.x + threadIdx.x;
    if (p < device_settings.num_mpm_markers) {
        const real3 xi = sorted_pos[p];
        const real3 vi = sorted_vel[p];
        LOOP_TWO_RING_GPU(                                                                     //
            real weight = N(xi - current_node_location, inv_bin_edge) * device_settings.mass;  //
            ATOMIC_ADD(&grid_mass[current_node], weight);                                      //
            ATOMIC_ADD(&grid_vel[current_node * 3 + 0], weight * vi.x);
            ATOMIC_ADD(&grid_vel[current_node * 3 + 1], weight * vi.y);
            ATOMIC_ADD(&grid_vel[current_node * 3 + 2], weight * vi.z);)
    }
}
CUDA_GLOBAL void kRasterize(const real3* sorted_pos,  // input
                            real* grid_mass) {        // output
    const int p = blockIdx.x * blockDim.x + threadIdx.x;
    if (p < device_settings.num_mpm_markers) {
        const real3 xi = sorted_pos[p];
        LOOP_TWO_RING_GPU(                                                                     //
            real weight = N(xi - current_node_location, inv_bin_edge) * device_settings.mass;  //
            ATOMIC_ADD(&grid_mass[current_node], weight);                                      //
            )
    }
}
//
////========================================================================================================================================================================
//
CUDA_GLOBAL void kNormalizeWeights(real* grid_mass,   // input
                                   real* grid_vel) {  // output
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < device_settings.num_mpm_nodes) {
        real n_mass = grid_mass[i];
        if (n_mass > C_EPSILON) {
            grid_vel[i * 3 + 0] /= n_mass;
            grid_vel[i * 3 + 1] /= n_mass;
            grid_vel[i * 3 + 2] /= n_mass;
        }
    }
}
//////========================================================================================================================================================================
////
CUDA_GLOBAL void kComputeParticleVolumes(const real3* sorted_pos,  // input
                                         real* grid_mass,          // output
                                         real* volume) {
    const int p = blockIdx.x * blockDim.x + threadIdx.x;
    if (p < device_settings.num_mpm_markers) {
        const real3 xi = sorted_pos[p];
        real particle_density = 0;
        LOOP_TWO_RING_GPU(                                              //
            real weight = N(xi - current_node_location, inv_bin_edge);  //
            particle_density += grid_mass[current_node] * weight;       //
            )
        // Inverse density to remove division
        particle_density = (bin_edge * bin_edge * bin_edge) / particle_density;
        volume[p] = device_settings.mass * particle_density;
    }
}
CUDA_GLOBAL void kFeHat(const real3* sorted_pos,  // input
                        const Mat33* marker_Fe,   // input
                        const real* grid_vel,     // input
                        Mat33* marker_Fe_hat) {   // output

    const int p = blockIdx.x * blockDim.x + threadIdx.x;
    if (p < device_settings.num_mpm_markers) {
        const real3 xi = sorted_pos[p];
        Mat33 Fe_hat_t(1.0f);

        int cx, cy, cz;
        const real bin_edge = device_settings.bin_edge;
        const real inv_bin_edge = 1.f / bin_edge;

        LOOP_TWO_RING_GPUSP(

            real vnx = device_settings.dt * grid_vel[current_node * 3 + 0];  //
            real vny = device_settings.dt * grid_vel[current_node * 3 + 1];  //
            real vnz = device_settings.dt * grid_vel[current_node * 3 + 2];

            real Tx = (xi.x - current_node_locationx) * inv_bin_edge;  //
            real Ty = (xi.y - current_node_locationy) * inv_bin_edge;  //
            real Tz = (xi.z - current_node_locationz) * inv_bin_edge;  //

            real valx = dN(Tx) * inv_bin_edge * N(Ty) * N(Tz);  //
            real valy = N(Tx) * dN(Ty) * inv_bin_edge * N(Tz);  //
            real valz = N(Tx) * N(Ty) * dN(Tz) * inv_bin_edge;  //

            Fe_hat_t[0] += vnx * valx; Fe_hat_t[1] += vny * valx; Fe_hat_t[2] += vnz * valx;  //
            Fe_hat_t[4] += vnx * valy; Fe_hat_t[5] += vny * valy; Fe_hat_t[6] += vnz * valy;  //
            Fe_hat_t[8] += vnx * valz; Fe_hat_t[9] += vny * valz; Fe_hat_t[10] += vnz * valz;

            //            real3 vel(grid_vel[current_node * 3 + 0], grid_vel[current_node * 3 + 1],
            //                      grid_vel[current_node * 3 + 2]);                  //
            //            real3 kern = dN(xi - current_node_location, inv_bin_edge);  //
            //            Fe_hat_t += OuterProduct(device_settings.dt * vel, kern);

            )

        marker_Fe_hat[p] = Fe_hat_t * marker_Fe[p];
    }
}

// CUDA_GLOBAL void kSVD(Mat33* marker_Fe_hat, Mat33* PolarR, Mat33* PolarS) {
//    const int p = blockIdx.x * blockDim.x + threadIdx.x;
//    if (p < device_settings.num_mpm_markers) {
//        Mat33 U, V, R, S, W;
//        real3 E;
//        SVD(marker_Fe_hat[p], U, E, V);
//        // Perform polar decomposition F = R*S
//        R = MultTranspose(U, V);
//        S = V * MultTranspose(Mat33(E), V);
//
//        PolarR[p] = R;
//        PolarS[p] = S;
//    }
//}

CUDA_GLOBAL void kApplyForces(const real3* sorted_pos,     // input
                              const Mat33* marker_Fe_hat,  // input
                              const Mat33* marker_Fe,      // input
                              const Mat33* marker_Fp,      // input
                              const real* marker_volume,   // input
                              const real* node_mass,       // input
                              Mat33* PolarR,               // input
{'fraction_non_alphanumeric': 0.14137096774193547, 'fraction_numerical': 0.011330645161290323, 'mean_word_length': 1.2509529860228716, 'pattern_counts': {'":': 0, '<': 59, '<?xml version=': 0, '>': 35, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 197, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4875257', 'n_tokens_mistral': 6693, 'n_tokens_neox': 6383, 'n_words': 1342}