dsharlet / slinky

Optimize pipelines for locality
MIT License
8 stars 2 forks source link

Use `slice_dim` instead of `crop_dim` if possible #275

Open dsharlet opened 6 months ago

dsharlet commented 6 months ago

In large real pipelines, there are many crops for pointwise dimensions, corresponding to "batch dimensions". For example:

loop(d4, serial, [0, buffer_max(v2, 4)], 1) {
 crop_dim(v2, 4, [d4, d4]) {
  crop_dim(v1, 4, [d4, d4]) {
   loop(d3, serial, [0, buffer_max(v2, 3)], 1) {
    crop_dim(v2, 3, [d3, d3]) {
     crop_dim(v1, 3, [d3, d3]) {
      loop(d2, serial, [0, buffer_max(v2, 2)], 1) {
       crop_dim(v2, 2, [d2, d2]) {
        crop_dim(v1, 2, [d2, d2]) {
         // Other non-pointwise stuff here

This leaves the pointwise dimensions in place, so the callbacks get the full rank of the inputs, and then have to do work to skip over the batch dimensions.

If we could generate slices instead of crops, we would avoid this inefficiency. To do this, we'd need some kind of indication that the callback is going to treat some dimensions as batch dimensions, and we could only do this if all uses of the dimension are treated as batch dimensions by the callback.

In the meantime, I think we can at least optimize our buffer traversal helpers for the common case of many trailing dimensions.

dsharlet commented 6 months ago

I had an interesting thought about this: if we say that buffers are always infinite rank, and the "missing" dimensions are always broadcast dimensions, then cropping the last dimension to extent 1 is equivalent to a slice and we could just dynamically implement the crop that way. for_each_element and related helpers already implement this implicit broadcasting logic...