Optional CPIC and per-node particle lists, plus opt-in kernel compilation

Cargo.toml

@@ -26,7 +26,7 @@ stensor = "0.4.2"
 
 [workspace.lints]
 rust.unexpected_cfgs = { level = "warn", check-cfg = [
-    'cfg(feature, values("dim2", "dim3"))',
+    'cfg(feature, values("dim2", "dim3", "cpic", "node_particle_lists"))',
 ] }
 
 [patch.crates-io]

crates/slosh2d/Cargo.toml

@@ -16,8 +16,10 @@ path = "../../src/lib.rs"
 required-features = ["dim2"]
 
 [features]
-default = ["dim2"]
+default = ["dim2", "cpic", "node_particle_lists"]
 dim2 = []
+cpic = []
+node_particle_lists = []
 
 comptime = ["slosh_testbed2d/comptime", "stensor/comptime"]
 runtime = ["slosh_testbed2d/runtime", "stensor/runtime"]

crates/slosh3d/Cargo.toml

@@ -16,8 +16,10 @@ path = "../../src/lib.rs"
 required-features = ["dim3"]
 
 [features]
-default = ["dim3"]
+default = ["dim3", "cpic", "node_particle_lists"]
 dim3 = []
+cpic = []
+node_particle_lists = []
 
 comptime = ["stensor/comptime", "slosh_testbed3d/comptime"]
 runtime = ["stensor/runtime", "slosh_testbed3d/runtime"]

shaders/slosh/grid/grid.slang

@@ -2,6 +2,20 @@ module grid;
 
 import slosh.aliases;
 
+// SLOSH_CPIC gates the CPIC parts of Node: the incompatible-momentum lane and the cdf collision
+// field. register_shaders sets it from the cpic cargo feature so this layout matches the Rust
+// GpuGridNode. Default on so a compile that never sets the macro still builds the full layout.
+#ifndef SLOSH_CPIC
+#define SLOSH_CPIC 1
+#endif
+
+// SLOSH_NODE_PARTICLE_LISTS gates the per-node particle linked lists (built in
+// finalize_particles_sort, cleared in reset). Set from the node_particle_lists cargo feature; the
+// bound buffers and the Rust GridArgs fields share it so everything stays in sync.
+#ifndef SLOSH_NODE_PARTICLE_LISTS
+#define SLOSH_NODE_PARTICLE_LISTS 1
+#endif
+
 // TODO: a lot if what exposed from this module should be methods of Grid (or other structs)
 //       rather than free function. (For now they have just been converted from WGSL to Slang.)
 
@@ -192,11 +206,11 @@ void reset_hmap(
 ) {
     let id = invocation_id.x;
     if (id < grid[0].hmap_capacity) {
+        // Only `state` needs clearing: insertion CAS-probes on `state` and the winner writes
+        // `key`, and `find_block_header_id` reads `value` only after matching `state` (which is
+        // always written before any lookup). Skipping key/value saves 44 of the 48 bytes/entry
+        // on this clear, which sweeps the whole fixed-size hashmap every substep.
         hmap_entries[id].state = NONE;
-        // Resetting the following two isn’t necessary for correctness,
-        // but it makes debugging easier.
-        hmap_entries[id].key = BlockVirtualId(vector<uint, DIM>(0));
-        hmap_entries[id].value = BlockHeaderId(0);
     }
     if (id == 0) {
         grid[0].num_active_blocks = 0u;
@@ -309,12 +323,14 @@ public struct Node {
     /// (depending on the context). The fourth component contains the cell’s mass.
     // TODO: maybe we don’t really need to pack ourself.
     public vector<float, DIM + 1> momentum_velocity_mass;
+#if SLOSH_CPIC
     // If a particle is incompatible with this node (as pepr CPIC’s concept of compatibility
     // based on affinities), it will contribute to this field instead of `momentum_velocity_mass`.
     // That way the P2G/G2P transfers on incompatible nodes still work properly and we don’t
     // loose the contributions of other compatible particles influencing this incompatible node too.
     public vector<float, DIM + 1> momentum_velocity_mass_incompatible;
     public NodeCdf cdf;
+#endif
 }
 
 #if DIM == 2
@@ -442,8 +458,10 @@ func reset(
     uint3 invocation_id: SV_DispatchThreadID,
     StructuredBuffer<Grid> grid,
     RWStructuredBuffer<Node> nodes,
+#if SLOSH_NODE_PARTICLE_LISTS
     RWStructuredBuffer<NodeLinkedList> nodes_linked_lists,
     RWStructuredBuffer<NodeLinkedList> rigid_nodes_linked_lists,
+#endif
 ) {
 // TODO: slang doesn’t have a way to retrieve the number of workgroups?
 //   let num_threads = num_workgroups.x * GRID_WORKGROUP_SIZE * num_workgroups.y * num_workgroups.z;
@@ -452,12 +470,16 @@ func reset(
 //   for (var i = invocation_id.x; i < num_nodes; i += num_threads) {
     if (i < num_nodes) {
        nodes[i].momentum_velocity_mass = vector<float, DIM + 1>(0.0);
+#if SLOSH_CPIC
        nodes[i].momentum_velocity_mass_incompatible = vector<float, DIM + 1>(0.0);
        nodes[i].cdf = NodeCdf(0.0, 0, NONE);
+#endif
+#if SLOSH_NODE_PARTICLE_LISTS
        nodes_linked_lists[i].head = NONE;
        nodes_linked_lists[i].len = 0u;
        rigid_nodes_linked_lists[i].head = NONE;
        rigid_nodes_linked_lists[i].len = 0u;
+#endif
    }
 //   }
 }

shaders/slosh/grid/sort.slang

@@ -3,6 +3,12 @@ module sort;
 import slosh.grid.grid;
 import slosh.solver.particle;
 
+// SLOSH_NODE_PARTICLE_LISTS gates the per-node particle linked-list build (see grid.slang).
+// Default on so a standalone compile keeps the full behavior.
+#ifndef SLOSH_NODE_PARTICLE_LISTS
+#define SLOSH_NODE_PARTICLE_LISTS 1
+#endif
+
 // Returns the within-block sort bucket for a particle counted/inserted into its primary
 // block: one bucket per associated-cell slab along the slowest-varying node axis (y in 2D,
 // z in 3D).
@@ -303,8 +309,10 @@ func finalize_particles_sort(
     StructuredBuffer<GridHashMapEntry> hmap_entries,
     StructuredBuffer<Position> particles_pos,
     ConstantBuffer<uint> particles_len,
+#if SLOSH_NODE_PARTICLE_LISTS
     RWStructuredBuffer<AtomicNodeLinkedList> nodes_linked_lists,
     RWStructuredBuffer<uint> particle_node_linked_lists,
+#endif
     RWStructuredBuffer<uint> sorted_particle_ids,
     RWStructuredBuffer<AtomicActiveBlockHeader> active_blocks,
 
@@ -337,12 +345,14 @@ func finalize_particles_sort(
             }
         }
 
+#if SLOSH_NODE_PARTICLE_LISTS
         // Setup the per-node particle linked-list (still consumed by the gather P2G and the
         // rigid/CDF paths).
         let node_global_id = node_id(block_header_id_to_physical_id(active_block_id_0), assoc);
         let prev_head = nodes_linked_lists[node_global_id.id].head.exchange(id);
         nodes_linked_lists[node_global_id.id].len.add(1u);
         particle_node_linked_lists[id] = prev_head;
+#endif
     }
 }
 
@@ -353,8 +363,10 @@ func sort_rigid_particles(
     StructuredBuffer<Grid> grid,
     StructuredBuffer<GridHashMapEntry> hmap_entries,
     StructuredBuffer<Position> rigid_particles_pos,
+#if SLOSH_NODE_PARTICLE_LISTS
     RWStructuredBuffer<AtomicNodeLinkedList> rigid_nodes_linked_lists,
     RWStructuredBuffer<uint> rigid_particle_node_linked_lists,
+#endif
 ) {
     let id = invocation_id.x;
     if (id < rigid_particles_pos.getCount()) {
@@ -367,6 +379,7 @@ func sort_rigid_particles(
 
         // NOTE: if the rigid particle doesn’t map to any block, we can just ignore it
         //       is it won’t affect the simulation.
+#if SLOSH_NODE_PARTICLE_LISTS
         if (active_block_id.id != NONE) {
             // Setup the per-node rigid particle linked-list.
             let node_local_id = associated_cell_index_in_block_off_by_one(particle, cell_width);
@@ -375,5 +388,6 @@ func sort_rigid_particles(
             rigid_nodes_linked_lists[node_global_id.id].len.add(1u);
             rigid_particle_node_linked_lists[id] = prev_head;
         }
+#endif
     }
 }

shaders/slosh/solver/grid_update.slang

@@ -7,6 +7,12 @@ import slosh.solver.params;
 import slosh.collision.collide;
 import slosh.solver.boundary_condition;
 
+// SLOSH_CPIC gates the incompatible-momentum lane on Node (see grid.slang). Default on so a
+// standalone compile matches the crate's default layout.
+#ifndef SLOSH_CPIC
+#define SLOSH_CPIC 1
+#endif
+
 #if DIM == 2
 static const uint WORKGROUP_SIZE_X = 8;
 static const uint WORKGROUP_SIZE_Y = 8;
@@ -73,10 +79,12 @@ func grid_update(
      * Step 1: momentum -> velocity update (gravity + clamping).
      */
     let momentum_velocity_mass = nodes[global_id].momentum_velocity_mass;
-    let momentum_velocity_mass_incompatible = nodes[global_id].momentum_velocity_mass_incompatible;
     var velocity_mass = update_single_cell(sim_params, cell_width, cell_pos, momentum_velocity_mass);
+#if SLOSH_CPIC
+    let momentum_velocity_mass_incompatible = nodes[global_id].momentum_velocity_mass_incompatible;
     let velocity_mass_incompatible = update_single_cell(sim_params, cell_width, cell_pos, momentum_velocity_mass_incompatible);
     nodes[global_id].momentum_velocity_mass_incompatible = velocity_mass_incompatible;
+#endif
 
     /*
      * Step 2: rigid-body collision boundary condition.

src/grid/grid.rs

@@ -34,7 +34,11 @@ struct GridArgs<'a, B: Backend> {
     n_block_groups: &'a GpuVector<[u32; 3], B>,
     n_g2p_p2g_groups: &'a GpuVector<[u32; 3], B>,
     nodes: &'a GpuVector<GpuGridNode, B>,
+    // Per-node particle linked lists, bound only under the node_particle_lists feature (the
+    // shader params are #if SLOSH_NODE_PARTICLE_LISTS).
+    #[cfg(feature = "node_particle_lists")]
     nodes_linked_lists: &'a GpuVector<[u32; 2], B>,
+    #[cfg(feature = "node_particle_lists")]
     rigid_nodes_linked_lists: &'a GpuVector<[u32; 2], B>,
     scan_values: &'a GpuVector<u32, B>,
     // Snapshot of `num_active_blocks` after the primary-block touch pass (the base-block count),
@@ -47,6 +51,7 @@ struct GridArgs<'a, B: Backend> {
     rigid_particles_pos: &'a GpuVector<Vector, B>,
     rigid_particle_needs_block: &'a GpuVector<u32, B>,
     sorted_particle_ids: &'a GpuVector<u32, B>,
+    #[cfg(feature = "node_particle_lists")]
     particle_node_linked_lists: &'a GpuVector<u32, B>,
 }
 
@@ -80,14 +85,17 @@ impl<B: Backend> WgGrid<B> {
         prefix_sum: &mut PrefixSumWorkspace<B>,
         sort_module: &'a WgSort<B>,
         prefix_sum_module: &'a WgPrefixSum<B>,
+        reset_nodes: bool,
     ) -> Result<(), B::Error> {
         let args = GridArgs {
             grid: &grid.meta,
             hmap_entries: &grid.hmap_entries,
             n_block_groups: &grid.indirect_n_blocks_groups,
             n_g2p_p2g_groups: &grid.indirect_n_g2p_p2g_groups,
             nodes: &grid.nodes,
+            #[cfg(feature = "node_particle_lists")]
             nodes_linked_lists: &grid.nodes_linked_lists,
+            #[cfg(feature = "node_particle_lists")]
             rigid_nodes_linked_lists: &grid.rigid_nodes_linked_lists,
             scan_values: &grid.scan_values,
             num_base_blocks: &grid.active_blocks_snapshot,
@@ -97,6 +105,7 @@ impl<B: Backend> WgGrid<B> {
             rigid_particles_pos: &rigid_particles.sample_points,
             rigid_particle_needs_block: &rigid_particles.rigid_particle_needs_block,
             sorted_particle_ids: particles.sorted_ids(),
+            #[cfg(feature = "node_particle_lists")]
             particle_node_linked_lists: particles.node_linked_lists(),
         };
 
@@ -202,14 +211,22 @@ impl<B: Backend> WgGrid<B> {
             .copy_scan_values_to_first_particles_and_prepare_for_finalize
             .launch_indirect(backend, pass, &args, grid.indirect_n_blocks_groups.buffer())?;
 
-        // Reset here so the linked list heads get reset before `finalize_particles_sort` which
-        // also setups the per-node linked list.
-        self.reset.launch_indirect(
-            backend,
-            pass,
-            &args,
-            grid.indirect_n_g2p_p2g_groups.buffer(),
-        )?;
+        // `reset` zeroes the node momentum lane, but it's also the only thing that initializes
+        // the per-node data the features depend on: the cpic incompatible/cdf lanes that
+        // grid_update reads, and the linked-list heads/len that finalize_particles_sort builds
+        // onto. Those buffers start uninitialized, so we can only honor `reset_nodes: false` when
+        // both features are off and the momentum lane (overwritten by every P2G) is all that's
+        // left. It has to run before finalize_particles_sort, which populates the linked list.
+        let must_reset =
+            reset_nodes || cfg!(feature = "cpic") || cfg!(feature = "node_particle_lists");
+        if must_reset {
+            self.reset.launch_indirect(
+                backend,
+                pass,
+                &args,
+                grid.indirect_n_g2p_p2g_groups.buffer(),
+            )?;
+        }
         sort_module.finalize_particles_sort.launch_capped(
             backend,
             pass,
@@ -241,7 +258,12 @@ pub struct GpuGridMetadata {
 #[repr(C)]
 pub struct GpuGridNode {
     momentum_velocity_mass: glam::Vec4,
+    // Gated by the cpic feature (default on). Turning cpic off drops the node from 48 to 16
+    // bytes, cutting grid traffic ~3x, and gives up collision-detection-field support. Keep in
+    // lockstep with Node in grid.slang, which SLOSH_CPIC gates the same way.
+    #[cfg(feature = "cpic")]
     momentum_velocity_mass_incompatible: glam::Vec4,
+    #[cfg(feature = "cpic")]
     cdf: GpuGridNodeCdf,
 }
 
@@ -343,6 +365,7 @@ pub struct GpuNodeCollision {
 /// Collision detection field data for a grid node.
 ///
 /// Stores signed distance and affinity information for MPM-rigid body coupling.
+#[cfg(feature = "cpic")]
 #[derive(Copy, Clone, PartialEq, Default, Debug, ShaderType)]
 #[repr(C)]
 pub struct GpuGridNodeCdf {
@@ -386,9 +409,12 @@ pub struct GpuGrid<B: Backend> {
     pub active_blocks_snapshot: GpuVector<u32, B>,
     /// Workspace for prefix sum operations.
     pub scan_values: GpuVector<u32, B>,
-    /// Per-node linked lists for MPM particles.
+    /// Per-node linked lists for MPM particles. Only built/reset under the node_particle_lists
+    /// feature, but always allocated: the P2G launchers still bind this buffer (the gather P2G
+    /// reads it; the scatter-style P2G binds it without reading). Gating the allocation means
+    /// dropping that unused scatter-style binding first.
     pub nodes_linked_lists: GpuVector<[u32; 2], B>,
-    /// Per-node linked lists for rigid body particles.
+    /// Per-node linked lists for rigid body particles. See `nodes_linked_lists`.
     pub rigid_nodes_linked_lists: GpuVector<[u32; 2], B>,
     /// Indirect dispatch arguments for block-parallel kernels.
     pub indirect_n_blocks_groups: Arc<GpuScalar<[u32; 3], B>>,

src/grid/sort.rs

@@ -38,8 +38,10 @@ pub struct WgSort<B: Backend> {
 struct SortArgs<'a, B: Backend> {
     grid: &'a GpuScalar<GpuGridMetadata, B>,
     hmap_entries: &'a GpuScalar<GpuGridHashMapEntry, B>,
+    #[cfg(feature = "node_particle_lists")]
     rigid_nodes_linked_lists: &'a GpuScalar<[u32; 2], B>,
     rigid_particles_pos: &'a GpuScalar<Vector, B>,
+    #[cfg(feature = "node_particle_lists")]
     rigid_particle_node_linked_lists: &'a GpuScalar<u32, B>,
 }
 
@@ -64,8 +66,10 @@ impl<B: Backend> WgSort<B> {
         let args = SortArgs {
             grid: &grid.meta,
             hmap_entries: &grid.hmap_entries,
+            #[cfg(feature = "node_particle_lists")]
             rigid_nodes_linked_lists: &grid.rigid_nodes_linked_lists,
             rigid_particles_pos: &rigid_particles.sample_points,
+            #[cfg(feature = "node_particle_lists")]
             rigid_particle_node_linked_lists: &rigid_particles.node_linked_lists,
         };
 

src/lib.rs

@@ -89,6 +89,18 @@ pub const SLANG_SRC_DIR: include_dir::Dir<'_> =
 pub fn register_shaders(compiler: &mut SlangCompiler) {
     stensor::register_shaders(compiler);
     compiler.add_dir(SLANG_SRC_DIR.clone());
+    // Mirror the cpic and node_particle_lists cargo features into shader macros so the shader
+    // Node layout and sort kernels match GpuGridNode and GridArgs on the Rust side. slosh sets
+    // these itself (unlike DIM, which the consumer picks) since they're crate-layout invariants.
+    compiler.set_global_macro("SLOSH_CPIC", if cfg!(feature = "cpic") { 1 } else { 0 });
+    compiler.set_global_macro(
+        "SLOSH_NODE_PARTICLE_LISTS",
+        if cfg!(feature = "node_particle_lists") {
+            1
+        } else {
+            0
+        },
+    );
 }
 
 /// Mathematical types and utilities for physics simulation.