mitsuba-renderer/drjit

[Feature Request] Add a scatter atomic operation

DoeringChristian opened this issue · 0 comments

For constructing hash grids or similar data structures it is useful to have access to the previous value of a Read Modify Write (RMW) operation (such as scatter_reduce).
Such data structures could in theory be constructed using external modules but I thought this might be helpful for prototyping renderers using purely Mitsuba3/Dr.Jit.
I propose to add a scatter_atomic function.
This function would work similar to the scatter_reduce function but return the previous value of the value at the address, that was modified.
Here is an example of how a potential scatter_atomic function could be used (in the context constructing a hashgrid):

cell = cell_index(pos)

cell_size = dr.zeros(mi.UInt, n_cells)

index_in_cell = dr.scatter_atomic(dr.ReduceOp.Add, cell_size, 1, cell)

The backends used by Dr.Jit (both LLVM and CUDA) already support this operation.
The PTX standard for the CUDA backend specifies two different RMW operations red and atom where the scatter_reduce function uses the red operation and the scatter_atomic function would simply use the atom instruction.
For the LLVM backend Dr.Jit already utilizes the atomicrmw instruction, which returns the previous value (I don't know if LLVM optimizes this in case the output is not used).

I have implemented a prototype implementation for the CUDA backend as I know more about it than LLVM.

Thanks for considering the issue.