Fuse `concat` with `reduce_by_index`

[FluxusMagna]

This optimisation is conceptually very similar to how concat is fused into scatter. This is important for algorithms that produce many points to be reduced per item. A fairly minimal example of something that should fuse but currently does not is

entry test n =
    let (is, vs) = tabulate n (\i -> ((i, (i+1) % n), (1f32, 1f32))) |> unzip
    let unzip_concat as = unzip as |> uncurry concat
    in hist (+) 0f32 n (unzip_concat is) (unzip_concat vs)

here the lack off fusion is not too problematic, as the number of items to reduce scales with the array, but once the number of items to reduce increases beyond that there can be some really bad scaling issues, as (is, vs) needs to be manifested due to the lack of fusion. Unlike scatter where the number of points to be scattered are usually less than or equal to the size of the destination array, in reduce_by_index this is often not the case.

I have some code used to transfer point-value pairs a onto coarser grid using polynomial interpolation, and this produces many items to reduce per point to be transferred, especially in higher dimensions. This is necessary for implementing Particle Mesh Ewald (PME) summation which is a common method in Molecular Dynamics (MD) and more generally for solving linear partial differential equations with complex geometry.

As an example of the scaling issues, this is the results for 3d cubic interpolation where each point is reduced to 4x4x4 bins with different weights. The grid size is kept constant, only the number of points to be reduced is changed.

grid_test.fut:cubic_bench (no tuning file):
[3][100000]i32:         4475μs (95% CI: [    4464.6,     4510.0])
[3][1000000]i32:       40286μs (95% CI: [   39761.9,    42289.3])
[3][10000000]i32:   69413528μs |███       |  3 runs ^C

The smaller systems are also slower than expected, but as the array that needs to be manifested grows, the function becomes unusably slow.

In the worst case here, the intermediate data, that would be fused away, is far larger than the input data or result grid.

[athas]

This is the implementation of concat-scatter fusion, and it should be straightforward to do the same thing for Hists as for Scremas.

[FluxusMagna]

Do you want me to give it a go? It doesn't seem too daunting given the similarity.

[athas]

If you find it fun, but otherwise I'll do it tomorrow. It won't be exactly the same, as we do not represent scatters and histograms the same way.

[FluxusMagna]

In that case I'll let you do it. Thanks a lot!

[athas]

The following rewrite (by hand) would suffice for the example:

entry test n =
  let (is, vs) = tabulate n (\i -> ((i, (i + 1) % n), (1f32, 1f32))) |> unzip
  in reduce_by_index (hist (+) 0f32 n (map (.0) is) (map (.0) vs))
                     (+)
                     0
                     (map (.1) is)
                     (map (.1) vs)

Would that not work for the real application?

[FluxusMagna]

I had not considered that composition since I assume it would lead to two kernels. The example is really not very similar to the real code, and the points share a lot of computation (it's different components of a matrix vector product, where the vector is dependent on the input), so unless they fuse there would be redundant computation or other issues. That aside it would make modular composition significantly more difficult if not impossible.

[athas]

It does lead to two kernels, but kernel launch overhead generally isn't the problem - loading the indexes/values from memory is (or the values used to construct them). The trick above only works if the concatenated arrays are completely independent. I'll continue trying to figure out how to fuse more generally, although I realise that we can't do it just like we do scatter, as the IR does not allow multiple histogram updates from a single input element.

[FluxusMagna]

Kernel launches might not be an issue here, but in the worst case it would have 6x6x6 kernels (although this higher order interpolant is probably overkill for this specific application). 4x4x4 is what I intend to use, and 64 kernels doesn't seem very appealing either.

[FluxusMagna]

https://gitlab.com/Gusten_Isfeldt/solver/-/blob/main/lib/gitlab.com/Gusten_Isfeldt/solver/interpolation/grid.fut?ref_type=heads

For reference this is what the code looks like. The module system allows some pretty neat composition so that different order interpolation and dimensionality reuses code.

[FluxusMagna]

By the way, why wouldn't accumulators work for this case?

[athas]

They would - if the histogram was expressed with accumulators, the concat-scatter rule would fire. It is an open and fair question why we maintain the Hist construct in the IR. The reason is that code generation for histograms is quite elaborate, and depends on being able to detect histograms as a structured construct, rather than as ad-hoc uses of accumulators. In order to do this high quality code generation, we would need to be able to detect that some instances of accumulator usage in fact correspond to structured histograms, and I am uncertain that we can do so reliably - although we do have a pass that tries, for the benefit of AD.

[FluxusMagna]

I wanted to see what the potential speedup would be so I tried using accumulators directly. The same benchmark as above performed like this. So a decent speedup for the cases that were already 'okay' and as expected much better scaling for the larger case(it's about 500 times faster than before).

[3][100000]i32:         2572μs (95% CI: [    2556.5,     2641.8])
[3][1000000]i32:       14371μs (95% CI: [   14361.2,    14378.4])
[3][10000000]i32:     125431μs (95% CI: [  125069.8,   125802.6])

[athas]

Well damn. I also put in some more thought, and I think the best solution is to switch histograms entirely to accumulators. I think they are not so difficult to detect as I first imagined, but it's not a one-afternoon project for sure.

[FluxusMagna]

Seems like a reasonable way to do it. I don't know how big your repository of histogram tests are, but I suppose it should be fairly easy to see potential performance regressions and (potentially less easy) address those cases with better detection.

[FluxusMagna]

It should be noted though that the benchmark here uses a dest that is too large for the subhistogram optimisation to be viable, so it's not necessarily representative of smaller histograms.