force-algorithm-mulithread.jl

using Molly 
using LinearAlgebra
using BenchmarkTools
using Test
using Profile
using ProfileView
using LoopVectorization
using SIMD
using OhMyThreads
using Polyester
using ChunkSplitters

import Base.*

# Teach Julia how to multiply a SIMD vector against a 3D Spatial vector!
@inline Base.:*(c::SIMD.Vec, v::SVector{3}) = SVector(c * v[1], c * v[2], c * v[3])


### ------------ ###
### System Setup ###
### ------------ ###

data_dir = joinpath(dirname(pathof(Molly)), "..", "data")
T = Float64
ff = MolecularForceField(
    T,
    joinpath(data_dir, "force_fields", "ff99SBildn.xml"),
    joinpath(data_dir, "force_fields", "tip3p_standard.xml"),
)

sys = System(
    joinpath(data_dir, "6mrr_equil.pdb"),
    ff;
    nonbonded_method=:pme,
    loggers=(
        energy=TotalEnergyLogger(10),
        writer=TrajectoryWriter(10, "traj_6mrr_5ps.dcd"),
    ),
    array_type=Array,
)

### --------------------------------- ###
### Setup for only LJ interactions ###
### --------------------------------- ###

sys2 = System(
    sys;
    #atoms=[Atom(charge=0.0u"q", mass=a.mass, σ=0.3u"nm", ϵ=1.0u"kJ * mol^-1") for a in sys.atoms],
    specific_inter_lists=(),
    general_inters=(),
    pairwise_inters=(LennardJones(use_neighbors=true, cutoff=DistanceCutoff(1.0u"nm")),),
)

neighbors_sys2 = find_neighbors(sys2, sys2.neighbor_finder; n_threads=1) # get the neighbors list 


### ------------------ ###
### NAIVE Force Algorithm ###
### ------------------ ###

function calculate_forces_naive!(forces, coords, atoms, boundary, neighbors, cutoff)

    fill!(forces, zero(eltype(forces)))

    # Precompute cutoff squared
    cutoff_2 = cutoff^2

    # Loop through neighbors list (ignoring 1-4 interaction)
    # DISABLE BOUNDS CHECKING
    @inbounds for (i, j, _) in neighbors.list

        # get params 
        sigma_i = atoms[i].σ
        eps_i   = atoms[i].ϵ
        
        sigma_j = atoms[j].σ
        eps_j   = atoms[j].ϵ

        # mix
        sigma_ij = (sigma_i + sigma_j) / 2.0
        eps_ij = sqrt(eps_i * eps_j)

        force_prefactor = 24 * eps_ij
        sigma_2 = sigma_ij^2

        r_ij = Molly.vector(coords[i], coords[j], boundary) # vector distance
        dist_2 = r_ij[1]^2 + r_ij[2]^2 + r_ij[3]^2  # calculate squared distance to avoid square roots

        mask = dist_2 < cutoff_2     
        ratio_2 = sigma_2 / dist_2 # Get LJ ratio but squared
        ratio_6 = ratio_2^3 # get ^6 and ^12 terms 
        ratio_12 = ratio_6^2
            
        f_mag_raw = (force_prefactor / dist_2) * (2 * ratio_12 - ratio_6) # Get force magnitude
        f_mag = f_mag_raw * mask
        f_vec = -f_mag * r_ij

        # Newton's Third Law
        forces[i] += f_vec
        forces[j] -= f_vec

    end
    return forces
end


###---------------------------###
### SIMD.JL WORKING ALGORITHM ###
###---------------------------###

function build_padded_adj_list(n_atoms, molly_neighbors, N_SIMD)
    # Create an ampty list of neighbours for each atom
    adj_list = [Int[] for _ in 1:n_atoms]

    # Populate both atom lists while looping through neighbor list
    for (i,j,_) in molly_neighbors.list
        push!(adj_list[i], j)
        push!(adj_list[j], i)
    end
    
    # Chunk size as double N_SIMD for 2x loop unrolling
    chunk_size = 2 * N_SIMD 
    
    for i in 1:n_atoms
        rem = length(adj_list[i]) % chunk_size
        if rem != 0
            pad_count = chunk_size - rem
            for _ in 1:pad_count
                push!(adj_list[i], i)
            end
        end
    end
    return adj_list
end

# function build_packed_adj_list(n_atoms, molly_neighbors, N_SIMD, soa_params, my_inter)
#     # 1. Initialize three parallel arrays
#     adj_list      = [Int[]     for _ in 1:n_atoms]
#     packed_sigmas = [Float64[] for _ in 1:n_atoms]
#     packed_eps    = [Float64[] for _ in 1:n_atoms]
#     #packed_special = [Bool[] for _ in 1:n_atoms]
#     packed_weights = [Float64[] for _ in 1:n_atoms]


#     # 2. Populate all three lists while looping through the neighbor list
#     for (i, j, is_special) in molly_neighbors.list
        
#         w = is_special ? my_inter.weight_special : 1.0

#         # A -> B interaction
#         push!(adj_list[i], j)
#         push!(packed_sigmas[i], soa_params.σ[j])
#         push!(packed_eps[i],    soa_params.ϵ[j])
        
#         push!(packed_weights[i], w)
                
#         # B -> A interaction (Newton's Third Law)
#         push!(adj_list[j], i)
#         push!(packed_sigmas[j], soa_params.σ[i])
#         push!(packed_eps[j],    soa_params.ϵ[i])
#         push!(packed_weights[j], w)
#     end
    
#     # 3. Chunk size as 1 * N_SIMD (Testing the Register Spilling hypothesis!)
#     chunk_size = 2 * N_SIMD 
    
#     # 4. Pad all three arrays
#     for i in 1:n_atoms
#         rem = length(adj_list[i]) % chunk_size
#         if rem != 0
#             pad_count = chunk_size - rem
#             for _ in 1:pad_count
#                 # Pad index with itself (dist_2 = 0.0 will trigger the mask!)
#                 push!(adj_list[i], i)
                
#                 # Pad parameters with safe dummy values (1.0 prevents any NaN zero-division bugs)
#                 push!(packed_sigmas[i], 1.0)
#                 push!(packed_eps[i],    1.0)
#                 push!(packed_weights[i], 1.0)
#             end
#         end
#     end
    
#     # Return all three perfectly aligned, contiguous arrays
#     return adj_list, packed_sigmas, packed_eps, packed_weights
# end

struct PackedFlatSoA{T}
    offsets::Vector{Int}
    adj_list::Vector{Int}
    sigmas::Vector{T}
    eps::Vector{T}
    weights::Vector{T}
end

function build_packed_adj_list(atoms, molly_neighbors, N_SIMD, soa_params, my_inter)
    
    n_atoms = length(atoms)

    # Temporary ragged arrays to catch the unsorted data
    temp_adj = [Int[] for _ in 1:n_atoms]
    temp_sig = [Float64[] for _ in 1:n_atoms]
    temp_eps = [Float64[] for _ in 1:n_atoms]
    temp_w = [Float64[] for _ in 1:n_atoms]

    for (i, j, is_special) in molly_neighbors.list

        is_skipped = Molly.shortcut_pair(my_inter.shortcut, atoms[i], atoms[j], is_special)

        if is_skipped
            w = 0.0 
        else
            w = is_special ? my_inter.weight_special : 1.0
        end

        # A --> B interaction
        push!(temp_adj[i], j)
        push!(temp_sig[i], soa_params.σ[j])
        push!(temp_eps[i], soa_params.ϵ[j])
        push!(temp_w[i], w)

        # B --> A interaction
        push!(temp_adj[j], i)
        push!(temp_sig[j], soa_params.σ[i])
        push!(temp_eps[j], soa_params.ϵ[i])
        push!(temp_w[j], w)
    end
    
    # The final FLAT arrays
    offsets = zeros(Int, n_atoms + 1)
    flat_adj = Int[]
    flat_sigmas = Float64[]
    flat_eps = Float64[]
    flat_weights = Float64[]
    
    chunk_size = 2 * N_SIMD 
    current_offset = 1
    
    # Squash and pad each flat array with dummy atoms to make the width right
    for i in 1:n_atoms
        offsets[i] = current_offset # Record where Atom i starts
        
        len = length(temp_adj[i])
        append!(flat_adj, temp_adj[i])
        append!(flat_sigmas, temp_sig[i])
        append!(flat_eps, temp_eps[i])
        append!(flat_weights, temp_w[i])
        
        # pad each array with dummy atoms
        rem = len % chunk_size
        if rem != 0
            pad_count = chunk_size - rem
            for _ in 1:pad_count
                push!(flat_adj, i) # triggers dist_2 = 0.0 mask
                push!(flat_sigmas, 1.0)
                push!(flat_eps, 1.0)
                push!(flat_weights, 1.0)
            end
        end
        
        current_offset += len + (rem != 0 ? chunk_size - rem : 0) # move offset forward by the true padded length
    end
    offsets[n_atoms + 1] = current_offset # cap the end
    
    return PackedFlatSoA{Float64}(offsets, flat_adj, flat_sigmas, flat_eps, flat_weights)
end


const N_SIMD = 8
const VFloat = Vec{N_SIMD, T}
const VInt   = Vec{N_SIMD, Int}

full_adj_list_2simd = build_padded_adj_list(length(sys2), neighbors_sys2, N_SIMD)

# 'where {V <: Vec} means this function takes SIMD vectors and not scalars 
@inline function lj_kernel_simd_ultimate(x_i, y_i, z_i, neigh_x::V, neigh_y::V, neigh_z::V, 
    box_x, box_y, box_z, cutoff_2, 
    sigma_sum, eps_ij, inv_box_x, inv_box_y, inv_box_z) where {V <: Vec}

    # GET THE SPATIAL DIFFERENCES 
    # x_i is a scalar and neigh_x is a vector of 8/16 neighbours 
    # The CPU broadcasts x_i across the vector and subtracts all 8/16 neighbors at once.
    dx = x_i - neigh_x
    dy = y_i - neigh_y
    dz = z_i - neigh_z

    # MINIMUM IMAGE CONVENTION
    # round(dx * inv_box_x) finds how many box-lengths apart the atoms are.
    # muladd(-box_x, ..., dx) pulls them back into the nearest periodic box.
    dx = muladd(-box_x, round(dx * inv_box_x), dx)
    dy = muladd(-box_y, round(dy * inv_box_y), dy)
    dz = muladd(-box_z, round(dz * inv_box_z), dz)

    # Calculates dx^2 + dy^2 + dz^2 using fused multiply-add hardware instructions.
    dist_2 = muladd(dx, dx, muladd(dy, dy, dz * dz))

    # SIMD MASK
    # SIMD cannot skip atoms. It MUST calculate math for all 8/16 slots.
    # We create a Mask (a vector of True/False) to flag which atoms are actually within the cutoff.
    mask = (dist_2 < cutoff_2) & (dist_2 > 0.0)

    # SAFE MATH
    # In the case of padded atoms
    # vifelse says: "If masked True, keep dist_2. If False, temporarily set it to 1.0 to prevent crash."
    safe_dist_2 = vifelse(mask, dist_2, one(V))
    inv_dist_2 = one(V) / safe_dist_2

    # LENNARD JONES MATH 
    # 4 * eps * ((sigma/r)^12 - (sigma/r)^6) using the minimum hardware multiplications

    # R^2 = (sigma_sum / r)^2
    base_ratio_2 = (sigma_sum * sigma_sum) * inv_dist_2
    # R^4 = R^2 * R^2
    br2_sq = base_ratio_2 * base_ratio_2
    # R^6 = R^4 * R^2
    base_ratio_6 = br2_sq * base_ratio_2

    # Factored polynomial. The constants (0.01171875 and -0.375) have absorbed the 
    # '4 * eps' and mixing rule coefficients to save extra multiplication steps.
    term = base_ratio_6 * muladd(0.01171875, base_ratio_6, -0.375)

    # Final raw force magnitude (still includes garbage data for atoms outside cutoff)
    f_mag_raw = eps_ij * inv_dist_2 * term

    # We use vifelse one last time. If the mask is True, keep the force. 
    # If False (it was outside the cutoff), force it to exactly 0.0.
    f_mag = vifelse(mask, f_mag_raw, zero(V))

    return f_mag * dx, f_mag * dy, f_mag * dz
end



@inline function task_kernel(forces, i, adj_list, coords, atoms, boundary, cutoff_2, inv_box_x, inv_box_y, inv_box_z, box_x, box_y, box_z, flat_coords, flat_atoms, ::Val{N_SIMD}, ::Val{VFloat}, ::Val{VInt}) where {N_SIMD, VFloat, VInt}

    # Get the list of all neighbor IDs for Atom I
    neigh_list = adj_list[i]
    n_neighbors = length(neigh_list)

    # Hoist Atom I's data out of the loop (it doesn't change)
    xi = ustrip(coords[i][1])
    yi = ustrip(coords[i][2])
    zi = ustrip(coords[i][3])

    # Hoist target parameters from the un-flattened atoms array
    sigma_i = ustrip(atoms[i].σ)
    eps_i = ustrip(atoms[i].ϵ)

    # --- 2. VERTICAL ACCUMULATORS ---
    # These are empty SIMD registers holding 0.0s. 
    # Instead of adding forces to a scalar total, we add them vertically in the vector registers.
    # We use two sets (chunk 1 and chunk 2) because we are unrolling the loop by 2x.
    f_ix_vec_1 = zero(VFloat)
    f_iy_vec_1 = zero(VFloat)
    f_iz_vec_1 = zero(VFloat)

    f_ix_vec_2 = zero(VFloat)
    f_iy_vec_2 = zero(VFloat)
    f_iz_vec_2 = zero(VFloat)
    
    # --- 3. THE SIMD LOOP ---
    # Step forward by 2 * N_SIMD (e.g., if N_SIMD=8, we process 16 neighbors per loop).
    @inbounds for j in 1:(2 * N_SIMD):n_neighbors

        # vload reads N_SIMD integers directly from the neighbor list at exactly the same time.
        # neigh_idxs_1 contains the first 8 neighbors. neigh_idxs_2 contains the next 8.
        neigh_idxs_1 = vload(VInt, neigh_list, j)
        neigh_idxs_2 = vload(VInt, neigh_list, j + N_SIMD)

        # --- 4. THE MEMORY STRIDE TRAP (AoS -> SoA) ---
        # flat_atoms is an Array of Structs (AoS). To find specific variables, you must do pointer math.
        # * 6 assumes the Atom struct is EXACTLY 6 Float64s wide in memory.
        # - 1 finds the 5th element (sigma). No offset finds the 6th element (epsilon).
        idx_sigma_1 = neigh_idxs_1 * 6 - 1
        idx_eps_1 = neigh_idxs_1 * 6 
        idx_sigma_2 = neigh_idxs_2 * 6 - 1
        idx_eps_2 = neigh_idxs_2 * 6

        # vgather tells the CPU: "Go to flat_atoms and grab the numbers at these specific 8 scattered memory indices."
        neigh_sigmas_1 = vgather(flat_atoms, idx_sigma_1)
        neigh_eps_1 = vgather(flat_atoms, idx_eps_1)
        neigh_sigmas_2 = vgather(flat_atoms, idx_sigma_2)
        neigh_eps_2 = vgather(flat_atoms, idx_eps_2)

        # --- 5. MIXING RULES ---
        # This calculates the combined sigma and epsilon for 8/16 pairs simultaneously.
        sigma_sum_1 = sigma_i + neigh_sigmas_1
        sigma_sum_2 = sigma_i + neigh_sigmas_2
        eps_ij_1 = sqrt(eps_i * neigh_eps_1)
        eps_ij_2 = sqrt(eps_i * neigh_eps_2)

        # Find the X, Y, Z coordinates for all 16 neighbors using the same stride logic (* 3)
        idx_x_1 = neigh_idxs_1 * 3 - 2
        idx_y_1 = neigh_idxs_1 * 3 - 1
        idx_z_1 = neigh_idxs_1 * 3
        idx_x_2 = neigh_idxs_2 * 3 - 2
        idx_y_2 = neigh_idxs_2 * 3 - 1
        idx_z_2 = neigh_idxs_2 * 3
        
        neigh_x_1 = vgather(flat_coords, idx_x_1)
        neigh_y_1 = vgather(flat_coords, idx_y_1)
        neigh_z_1 = vgather(flat_coords, idx_z_1)
        neigh_x_2 = vgather(flat_coords, idx_x_2)
        neigh_y_2 = vgather(flat_coords, idx_y_2)
        neigh_z_2 = vgather(flat_coords, idx_z_2)
        
        # --- 6. FIRE THE KERNEL ---
        # We pass our fully loaded SIMD vectors into the pure math function.
        f_x_chunk_1, f_y_chunk_1, f_z_chunk_1 = lj_kernel_simd_ultimate(
            xi, yi, zi, 
            neigh_x_1, neigh_y_1, neigh_z_1, 
            box_x, box_y, box_z, cutoff_2, sigma_sum_1, eps_ij_1, inv_box_x, inv_box_y, inv_box_z
        )
        
        f_x_chunk_2, f_y_chunk_2, f_z_chunk_2 = lj_kernel_simd_ultimate(
            xi, yi, zi, 
            neigh_x_2, neigh_y_2, neigh_z_2, 
            box_x, box_y, box_z, cutoff_2, sigma_sum_2, eps_ij_2, inv_box_x, inv_box_y, inv_box_z
        )
        
        # Add the 8 forces from chunk 1 to our running total registers.
        f_ix_vec_1 += f_x_chunk_1
        f_iy_vec_1 += f_y_chunk_1
        f_iz_vec_1 += f_z_chunk_1
        
        # Add the 8 forces from chunk 2 to our running total registers.
        f_ix_vec_2 += f_x_chunk_2
        f_iy_vec_2 += f_y_chunk_2
        f_iz_vec_2 += f_z_chunk_2
    end
    
    # --- 8. HORIZONTAL SUM (REDUCTION) ---
    # At the end of the loop, f_ix_vec_1 looks like: [1.2, -0.4, 3.1, 0.0, ...]
    # The sum() function squashes those 8 vertical slots into one single scalar number.
    f_ix = sum(f_ix_vec_1 + f_ix_vec_2)
    f_iy = sum(f_iy_vec_1 + f_iy_vec_2)
    f_iz = sum(f_iz_vec_1 + f_iz_vec_2)

    @inbounds forces[i] = SVector(f_ix, f_iy, f_iz) * u"kJ * mol^-1 * nm^-1"

end

function calculate_forces_simd_tasks!(forces, sys::System{D, AT, T}, adj_list, cutoff, ::Val{N_SIMD}, ::Val{VFloat}, ::Val{VInt}) where {D, AT, T, N_SIMD, VFloat, VInt}

    coords = sys.coords
    atoms = sys.atoms
    boundary = sys.boundary

    # Fill array with zeroes, but with no memory alllocation overhead
    #fill!(forces, zero(eltype(forces)))

    # Do everything in squares 
    cutoff_2 = ustrip(cutoff)^2
    
    box_x = ustrip(boundary.side_lengths[1])
    box_y = ustrip(boundary.side_lengths[2])
    box_z = ustrip(boundary.side_lengths[3])

    inv_box_x = 1.0 / box_x
    inv_box_y = 1.0 / box_y
    inv_box_z = 1.0 / box_z

    # zero allocation flattening of the atoms and coords as flat float64 array
    flat_coords = reinterpret(T, coords)
    flat_atoms = reinterpret(T, atoms)

    n_t = Threads.nthreads()
    num_chunks = 8 * n_t

    # Allocate chunks
    chunk_list = collect(chunks(1:length(coords); n=num_chunks))

    # Hardware counter 
    counter = Threads.Atomic{Int}(1)

    # Spawn the right number of tasks per cores
    @sync for _ in 1:n_t
        Threads.@spawn begin
            
            while true
                # Get chunk then add 1 
                chunk_id = Threads.atomic_add!(counter, 1)
                
                # Break if it exceeds
                if chunk_id > num_chunks
                    break
                end
                
                # Extract the specific chunk assigned to loop iteration
                i_range = chunk_list[chunk_id]
                
                # Run the math for the chunk
                for i in i_range
                    task_kernel(forces, i, adj_list, coords, atoms, boundary, cutoff_2, inv_box_x, inv_box_y, inv_box_z, 
                        box_x, box_y, box_z, flat_coords, flat_atoms, Val(N_SIMD), Val(VFloat), Val(VInt))
                end
            end
        end
    end
    
    return forces
end

##################################################################################################################################
##################################################################################################################################
##################################################################################################################################
##################################################################################################################################
##################################################################################################################################

struct FastLennardJones{C, SC, S, E, W} <: PairwiseInteraction
    cutoff::C
    shortcut::SC
    σ_mixing::S
    ϵ_mixing::E
    weight_special::W
end

# Very bottom level maths function - mirroring pairwise_force(::LennardJones) or pairwise_force(::Coulomb) except since it takes r^2 it outputs F/r 
@inline function pairwise_force_div_r(inter::FastLennardJones, dist_2, sigma_ij, eps_ij)

    inv_dist_2 = 1.0 / dist_2
    base_ratio_2 = (sigma_ij * sigma_ij) * inv_dist_2
    br2_sq = base_ratio_2 * base_ratio_2
    base_ratio_6 = br2_sq * base_ratio_2

    term = base_ratio_6 * muladd(48.0, base_ratio_6, -24.0)
    
    return eps_ij * inv_dist_2 * term
end

# Cutoff function that will call the bottom level maths kernel, but also handles cutoff logic
@inline function force_apply_cutoff(cutoff::Molly.DistanceCutoff, inter, dist_2, sigma_ij, eps_ij, cutoff_2)
    return pairwise_force_div_r(inter, dist_2, sigma_ij, eps_ij)
end

@inline function force_apply_cutoff(cutoff::Molly.ShiftedForceCutoff, inter, dist_2, sigma_ij, eps_ij, cutoff_2)
    fdr_real = pairwise_force_div_r(inter, dist_2, sigma_ij, eps_ij)
    fdr_cut = pairwise_force_div_r(inter, cutoff_2, sigma_ij, eps_ij)
    return fdr_real - fdr_cut
end

@inline function force_apply_cutoff(cutoff::Molly.NoCutoff, inter, dist_2, sigma_ij, eps_ij, cutoff_2)
    # No cutoff means no boundaries
    return pairwise_force_div_r(inter, dist_2, sigma_ij, eps_ij)
end

@inline function custom_force(inter, dr, safe_dist_2, atom_i, atom_j, neigh_weights, cutoff_2)
    
    # Mixing rules
    sigma_ij = Molly.σ_mixing(inter.σ_mixing, atom_i, atom_j)
    eps_ij = Molly.ϵ_mixing(inter.ϵ_mixing, atom_i, atom_j)

    # Cutoff
    f_div_r = force_apply_cutoff(inter.cutoff, inter, safe_dist_2, sigma_ij, eps_ij, cutoff_2)

    # Shortcut AND Special all in one
    f_div_r_weighted = f_div_r * neigh_weights

    return f_div_r_weighted * dr
end


# 'where {V <: Vec} means this function takes SIMD vectors and not scalars 
# MIDDLE MAN OF THE SIMULATION #
@inline function lj_kernel_simd_molly(x_i, y_i, z_i, neigh_x::V, neigh_y::V, neigh_z::V, 
    sim_params, atom_i, atom_j, inter, neigh_weights) where {V <: Vec}

    # NOTE: every variable in this function except the central atom i and the box dimensions, is a vector of 8 numbers

    # Boundary math - neigh_x is a vector of 8 numbers 
    dx = x_i - neigh_x
    dy = y_i - neigh_y
    dz = z_i - neigh_z

    # Calculate how many box lengths apart and wrap the neighbour to the closest image 
    dx = muladd(-sim_params.box_x, round(dx * sim_params.inv_box_x), dx)
    dy = muladd(-sim_params.box_y, round(dy * sim_params.inv_box_y), dy)
    dz = muladd(-sim_params.box_z, round(dz * sim_params.inv_box_z), dz)

    # Calculate the square of the distance 
    dist_2 = muladd(dx, dx, muladd(dy, dy, dz * dz))

    # Create 8 true or false values 
    mask = (dist_2 < sim_params.cutoff_2) & (dist_2 > 0.0)
    
   # Swap any NaN distances to be one(V) so there is no divide-by-zero 
   # After the function, use the mask values to check if the force should be zero
   safe_dist_2 = vifelse(mask, dist_2, one(V)) 

   # Package as an svector to allow dr * force syntax 
   dr = SVector(dx, dy, dz)

   # Pass to the user defined force function which will calculate garbage forces for any dummy atoms 
   fdr_raw = custom_force(inter, dr, safe_dist_2, atom_i, atom_j, neigh_weights, sim_params.cutoff_2)
   
   # Zero out forces for atoms outside the cutoff (or padded dummy atoms)
   f_x = vifelse(mask, fdr_raw[1], zero(V))
   f_y = vifelse(mask, fdr_raw[2], zero(V))
   f_z = vifelse(mask, fdr_raw[3], zero(V))

    return f_x, f_y, f_z
end

@inline function indiv_task_kernel(forces, i, packed_data, soa_params, sim_params, coords, flat_coords, inter, ::Val{N_SIMD}, ::Val{VFloat}, ::Val{VInt}) where {N_SIMD, VFloat, VInt}

    # Use the offsets to find where Atom i's data lives in the flat 1d arrays
    start_idx = packed_data.offsets[i]
    end_idx   = packed_data.offsets[i+1] - 1

    # For some reason quicker than indexing from flat_coords
    xi = ustrip(coords[i][1])
    yi = ustrip(coords[i][2])
    zi = ustrip(coords[i][3])

    # Central atom proxy
    atom_i_proxy = (σ = soa_params.σ[i], ϵ = soa_params.ϵ[i])

    # Two sets of accumulators for 2x ILP
    f_ix_vec_1 = zero(VFloat); f_iy_vec_1 = zero(VFloat); f_iz_vec_1 = zero(VFloat)
    f_ix_vec_2 = zero(VFloat); f_iy_vec_2 = zero(VFloat); f_iz_vec_2 = zero(VFloat)
    
    # Loop through from 1 until length of offsets which is length of atoms minus the 1 difference added at the start - no pointer chasing
    @inbounds for j in start_idx:(2 * N_SIMD):end_idx
        
        # Load the 8 atom IDs of the different neighbors for both chunks 
        neigh_idxs_1 = vload(VInt, packed_data.adj_list, j)
        neigh_idxs_2 = vload(VInt, packed_data.adj_list, j + N_SIMD)

        # Load the sigmas and epsilons of the 8 neighbours into the vector registers for both chunks 
        neigh_sigmas_1 = vload(VFloat, packed_data.sigmas, j)
        neigh_eps_1    = vload(VFloat, packed_data.eps, j)
        neigh_sigmas_2 = vload(VFloat, packed_data.sigmas, j + N_SIMD)
        neigh_eps_2    = vload(VFloat, packed_data.eps, j + N_SIMD)

        neigh_weights_1 = vload(VFloat, packed_data.weights, j)
        neigh_weights_2 = vload(VFloat, packed_data.weights, j + N_SIMD)

        # Build Proxies
        atom_j_proxy_1 = (σ = neigh_sigmas_1, ϵ = neigh_eps_1)
        atom_j_proxy_2 = (σ = neigh_sigmas_2, ϵ = neigh_eps_2)

        # Calculate the x y and z indices using each of the atoms indexes (neigh_idx) from the SIMD chunk of the list
        idx_x_1 = neigh_idxs_1 * 3 - 2
        idx_y_1 = neigh_idxs_1 * 3 - 1
        idx_z_1 = neigh_idxs_1 * 3
        idx_x_2 = neigh_idxs_2 * 3 - 2
        idx_y_2 = neigh_idxs_2 * 3 - 1
        idx_z_2 = neigh_idxs_2 * 3
        
        # Gather the x y and z coordinates of the 8 neighboours using the generated x/y/z idxs - this gathers 8 of each coordinate
        neigh_x_1 = vgather(flat_coords, idx_x_1)
        neigh_y_1 = vgather(flat_coords, idx_y_1)
        neigh_z_1 = vgather(flat_coords, idx_z_1)
        neigh_x_2 = vgather(flat_coords, idx_x_2)
        neigh_y_2 = vgather(flat_coords, idx_y_2)
        neigh_z_2 = vgather(flat_coords, idx_z_2)
        
        # Pass the scalar values for atom i and then the vectors of 8 values for the neighbours, alongside the box params
        f_x_chunk_1, f_y_chunk_1, f_z_chunk_1 = lj_kernel_simd_molly(xi, yi, zi, neigh_x_1, neigh_y_1, neigh_z_1, 
            sim_params, atom_i_proxy, atom_j_proxy_1, inter, neigh_weights_1)
        
        # Do the same for the second chunk
        f_x_chunk_2, f_y_chunk_2, f_z_chunk_2 = lj_kernel_simd_molly(xi, yi, zi, neigh_x_2, neigh_y_2, neigh_z_2, 
            sim_params, atom_i_proxy, atom_j_proxy_2, inter, neigh_weights_2)
        
        # Adds the 8 force values into running totals 
        f_ix_vec_1 += f_x_chunk_1; f_iy_vec_1 += f_y_chunk_1; f_iz_vec_1 += f_z_chunk_1
        f_ix_vec_2 += f_x_chunk_2; f_iy_vec_2 += f_y_chunk_2; f_iz_vec_2 += f_z_chunk_2
    end
    
    # Horizontal sum across both chunks to get a single scalar force for the central atom
    f_ix = sum(f_ix_vec_1 + f_ix_vec_2)
    f_iy = sum(f_iy_vec_1 + f_iy_vec_2)
    f_iz = sum(f_iz_vec_1 + f_iz_vec_2)

    @inbounds forces[i] = SVector(f_ix, f_iy, f_iz) * u"kJ * mol^-1 * nm^-1"
end


function calculate_forces_simd_molly_tasks!(forces, sys::System{D, AT, T}, packed_data, inter, soa_params, ::Val{N_SIMD}, ::Val{VFloat}, ::Val{VInt}) where {D, AT, T, N_SIMD, VFloat, VInt}

    coords = sys.coords
    flat_coords = reinterpret(T, coords) # zero allocation flattening of the coords as flat float64 array for vgather
    boundary = sys.boundary

    # Calculate box metrics for distances
    box_x = ustrip(boundary.side_lengths[1])
    box_y = ustrip(boundary.side_lengths[2])
    box_z = ustrip(boundary.side_lengths[3])
    inv_box_x = 1.0 / box_x
    inv_box_y = 1.0 / box_y
    inv_box_z = 1.0 / box_z

    if hasfield(typeof(inter.cutoff), :dist_cutoff)
        cutoff_2 = ustrip(inter.cutoff.dist_cutoff)^2
    else
        cutoff_2 = Inf
    end

    # Package box metrics 
    sim_params = (
        box_x = box_x, box_y = box_y, box_z = box_z,
        inv_box_x = inv_box_x, inv_box_y = inv_box_y, inv_box_z = inv_box_z,
        cutoff_2 = cutoff_2
    )

    # Multithreading 
    n_t = Threads.nthreads()
    num_chunks = 8 * n_t

    # Allocate chunks
    chunk_list = collect(chunks(1:length(coords); n=num_chunks))

    # Hardware counter 
    counter = Threads.Atomic{Int}(1)

    # Spawn the right number of tasks per cores
    @sync for _ in 1:n_t
        Threads.@spawn begin
            
            while true
                # Get chunk then add 1 
                chunk_id = Threads.atomic_add!(counter, 1)
                
                # Break if it exceeds
                if chunk_id > num_chunks
                    break
                end
                
                # Extract the specific chunk assigned to loop iteration
                i_range = chunk_list[chunk_id]
                
                # Run the math for the chunk
                for i in i_range
                    indiv_task_kernel(forces, i, packed_data, soa_params, sim_params, coords, flat_coords, inter, Val(N_SIMD), Val(VFloat), Val(VInt))
                end
            end
        end
    end
    
    return forces
end



struct SimpleLJ
    cutoff::typeof(1.0u"nm") # Keeps Molly's unit system happy
end

#import Molly: PairwiseInteraction



# Create your dummy inter to pass around
#my_inter = SimpleLJ(1.0u"nm")
my_cutoff = Molly.DistanceCutoff(1.0u"nm")
my_inter = FastLennardJones(my_cutoff, Molly.LJZeroShortcut(), Molly.LorentzMixing(), Molly.GeometricMixing(), 1.0)

my_soa_params = (
    σ = [ustrip(a.σ) for a in sys.atoms], 
    ϵ = [ustrip(a.ϵ) for a in sys.atoms]
)

# 2. Build the padded neighbor list AND the packed parameter arrays
n_atoms = length(sys.atoms)
# full_adj_list, packed_sigmas, packed_eps, packed_weights = build_packed_adj_list(
#     n_atoms, neighbors_sys2, N_SIMD, my_soa_params, my_inter
# )

my_packed_data = build_packed_adj_list(sys.atoms, neighbors_sys2, N_SIMD, my_soa_params, my_inter)

### ------------- ###
### Testing suite ###
### ------------- ###

# Parameters for benchmarking and Profiling
const RUN_BENCHMARKS = true
const RUN_PROFILING = false
const RUN_LLVM = false


@testset "MD Kernel Test Suite" begin

    # Pre-initialise force arrays to compare
    my_forces_arr = [zero(sys2.coords[1]) * u"kJ" / u"nm" / u"nm" / u"mol" for i in 1:length(sys2)]
    ref_forces_arr = [zero(sys.coords[1]) * u"kJ" / u"nm" / u"nm" / u"mol" for i in 1:length(sys)] 

    # Additional set-up for Molly in-place kernel
    needs_vir = false
    buffers = Molly.init_buffers!(sys2, Threads.nthreads())

    my_soa_params = (
    σ = [ustrip(a.σ) for a in sys.atoms], 
    ϵ = [ustrip(a.ϵ) for a in sys.atoms]
)


    @testset "Math Correctness" begin
        
        # Calculate forces of both and compare
        calculate_forces_simd_molly_tasks!(my_forces_arr, sys2, my_packed_data, my_inter, my_soa_params, Val(N_SIMD), Val(VFloat), Val(VInt))


        Molly.forces!(ref_forces_arr, sys2, neighbors_sys2, buffers, Val(needs_vir), 0; n_threads=Threads.nthreads())
        
        max_diff = maximum(norm.(my_forces_arr .- ref_forces_arr))
        println("\nMax difference between custom and Molly in-place kernel: ", max_diff)
        @test max_diff < 1e-9 * u"kJ * mol^-1 * nm^-1"
    end

    @testset "Memory Efficiency" begin
        
        # Use a let block to create local scope and stop dynamic dispatch, equivalent to $ in BenchmarkTools
        allocs = let f = my_forces_arr, s = sys2, adj = full_adj_list, cut = 1.0u"nm"
            calculate_forces_simd_tasks!(f, s, adj, cut)
            @allocated calculate_forces_simd_tasks!(f, s, adj, cut)
        end

        println("Bytes allocated in pure custom kernel: ", allocs) # needed?
        @test allocs == 0
    end

    if RUN_BENCHMARKS
        
        println("\n--- Running Benchmarks ---\n")

        n_t = Threads.nthreads()
        println("--- Running Benchmarks (Active Threads: $n_t) ---\n")

        println("\n--- Custom SIMD Kernel (Tasks) ---\n")
        my_simd_tasks_1_bmark = @benchmark calculate_forces_simd_tasks!(f, $sys2, $full_adj_list_2simd, $(1.0u"nm"), Val(N_SIMD), Val(VFloat), Val(VInt)) setup=(f = fill!($my_forces_arr, zero(eltype($my_forces_arr))))
        display(my_simd_tasks_1_bmark)
    
        println("\n--- Custom SIMD MOLLY Kernel (Tasks) ---\n")

        my_simd_tasks_2_bmark = @benchmark calculate_forces_simd_molly_tasks!(f, $sys2, $my_packed_data, $my_inter,  $my_soa_params, Val(N_SIMD), Val(VFloat), Val(VInt)) setup=(f = fill!($my_forces_arr, zero(eltype($my_forces_arr))))
        display(my_simd_tasks_2_bmark)

        println("\n--- Molly In-Place Kernel ---\n")
        ref_bmark = @benchmark Molly.forces!(f, $sys2, $neighbors_sys2, $buffers, Val($needs_vir), 0; n_threads=$n_t) setup=(f = fill!($ref_forces_arr, zero(eltype($ref_forces_arr))))
        display(ref_bmark)
    else
        println("(Benchmarks skipped, set RUN_BENCHMARKS = true to run)")
    end

    if RUN_PROFILING
        
        println("\n--- Running Profiling ---")

        # Run once to ensure compiled 
        #calculate_forces_simd_tasks!(my_forces_arr, sys2, full_adj_list, 1.0u"nm", Val(N_SIMD), Val(VFloat), Val(VInt))
        #Molly.forces!(ref_forces_arr, sys2, neighbors_sys2, buffers, Val(needs_vir), 0; n_threads=Threads.nthreads())
        calculate_forces_simd_molly_tasks!(my_forces_arr, sys2, full_adj_list, 1.0u"nm", my_soa_params, Val(N_SIMD), Val(VFloat), Val(VInt))



        # Clear old data
        Profile.clear()

        ProfileView.@profview for _ in 1:100
            #calculate_forces_simd_tasks!(my_forces_arr, sys2, full_adj_list, 1.0u"nm", Val(N_SIMD), Val(VFloat), Val(VInt))
            #Molly.forces!(ref_forces_arr, sys2, neighbors_sys2, buffers, Val(needs_vir), 0; n_threads=Threads.nthreads())
            calculate_forces_simd_molly_tasks!(my_forces_arr, sys2, full_adj_list, 1.0u"nm", my_soa_params, Val(N_SIMD), Val(VFloat), Val(VInt))





        end
        
    else
        println("(Profiling skipped, set RUN_PROFILING = true to run)")
    end

    if RUN_LLVM

        println("\n--- Running LLVM IR ---")

        #@code_llvm calculate_forces_simd_molly_tasks!(my_forces_arr, sys2, full_adj_list, 1.0u"nm", Val(N_SIMD), Val(VFloat), Val(VInt))
        #@code_llvm calculate_forces_simd_molly_tasks!(my_forces_arr, sys2, full_adj_list, 1.0u"nm", my_soa_params, Val(N_SIMD), Val(VFloat), Val(VInt))

        # open("forces_llvm_simd_as_fast_manual_loop_unrolling.txt", "w") do file
        #     # Temporarily redirect all terminal output into this file
        #     redirect_stdout(file) do
        #         # Run your exact command
        #         @code_llvm debuginfo=:none calculate_forces_simd_single!(my_forces_arr, sys2.coords, sys2.atoms, sys2.boundary, full_adj_list, 1.0u"nm")
        #     end
        # end

    else
        println("(LLVM IR skipped, set RUN_LLVM = true to run)\n")
    end
end