futhark
futhark copied to clipboard
Proposal: new streaming operators for scatter and reduce_by_index
Inspired by the need for efficient rasterisation (and some long-ago complaints by Cosmin), I would like to add the following operators for efficiently updating multiple elements of some targets for each element in an input array. The interface is heavily inspired by @melsman's expand utility function:
val stream_scatter 'a 'b [k] [n] :
*[k]a
-> (b -> i32)
-> (b -> i32 -> (i32, a))
-> [n]b
-> *[k]a
The imperative semantics of an expression stream_scatter dest f g arr is given by the following pseudocode:
for i < n:
m = f(arr[i])
for j < m:
(l, x) = g(arr[i], j)
dest[l] = x
I expect that a similar core language representation with reasonable fusion rules will be easy to design, and that a stream_reduce_by_index will be a straightforward extension of the same idea.
The core IR implementation of this has been added in #1197, but I still don't see a nice way of safely exposing this to the source language.