[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.