How to implement grid-based particle collision detection as described in this Nvidia whitepaper?

[bcolloran]

Hi, apologies if this is a foolish question, I'm still new to GPU programming and Tai Chi (which is really amazing btw!).

I was looking at this old whitepaper from Nvidia that is about using a grid to detect particle collisions, and it seems to have a pattern that I didn't think was possible on GPU. In the section "Building the Grid using Atomic Operations" (page 7-8), they describe using a "scattered global write" to write the indices of all the particles in a grid cell to a fixed-length array (one array per grid cell) without race conditions in writing to the array; the races are avoided by using atomicAdds to track the number of particles in the cell so far. If I understand correctly, in Tai Chi's notation, it would work something like this:

grid = ti.Struct.field(
    {
        "num_in_cell": ti.i32,  # number of particles in grid cell
        "indices": ti.types.vector(4, ti.i32),  # indices of particles in this grid cell
    }
)

...

@ti.kernel
def bin_particles():
    for p in particles:
        ij = particles[p].x.cast(int)
        num_in_cell = grid[ij].num_in_cell
        grid[ij].indices[num_in_cell] = p
        grid[ij].num_in_cell += 1

I thought I would go ahead and just try this to see if it would work, and I'm not surprised to find that it doesn't, because it seems like a weird pattern. I get the error:

TaichiSyntaxError: The 0-th index of a Matrix/Vector must be a compile-time constant integer, got <class 'taichi.lang.expr.Expr'>.
This is because matrix operations will be **unrolled** at compile-time for performance reason.
If you want to *iterate through matrix elements*, use a static range:

So I guess I have a couple questions:

1. am I even understanding the description of the Nvidia algorithm correctly?
2. is there a way to implement this algorithm with Tai Chi that would be better than the naive approach i tried?
3. if there is not a way to implement this particular grid-based collision detection algorithm, does anyone have suggestions about other fast particle collision detection algorithms that can be implemented with Tai Chi? (Or perhaps even example code?)

Thanks for you help! :-)

[k-ye]

Oops. The immediate issue I can tell is that you cannot index into a vector (matrix) using a runtime integer variable yet. (The .indices[num_in_cell] part, as num_in_cell is a runtime variable). The current situation is a bit unfortunate, but we're really close on supporting it (@strongoier is driving this feature). For now a workaround is to do a static loop:

for vi in static(range(4)):
  if num_in_cell == vi:
    grid[ij].indices[vi] = p

[bcolloran]

Wow, I was afraid I was going to need a totally different approach, but this is such an easy workaround! Thanks @k-ye, this is great :-)

Also exciting to hear that improved support for this kind of operation is in progress -- I've been very impressed by the super rapid pace of this project. @strongoier, is there an issue I could watch to follow your work on this? (I did a search, but I did not manage to find anything that looked related, sorry if I missed it). Thanks!

[strongoier]

@bcolloran You can watch the progress in #2590, which tracks all related PRs. More details about the design roadmap will also come soon :-)