[inductor] Lower aten.index_add_/aten.index_add

[desertfire]

Should be fairly similar to how aten.index_put_/aten.index_put is currently handled, see

• torchdynamo/torchinductor/lowering.py

  Line 1445 in 7cc850c

  def index_put_(self, indices, values, accumulate=False):

• torchdynamo/torchinductor/decomposition.py

  Line 357 in 7cc850c

  def index_put(self, indices, values, accumulate=False):

[eellison]

I looked into this a bit, I'm not sure it has that much in common with index_put, since index_add can accumulate arbitrarily many indices of the same value. Also, the usage seems to be pretty different -
for fastNLP_Bert you have inputs: Operator: aten.index_put_.default cnt: 1, ((T([6, 476], i64), [T([6], i64), T([6], i64)], T([], i64)), {}) (aka indices of length 6).

For e.g. hf_Longformer you have the following: ((T([2359296], f16), 0, T([4718592], i64), T([4718592], f16)), {}) (aka, indices of length 4718592).

I can't think of a lowering of index_add within inductor that makes a ton of sense at least within the current IR. Maybe a custom template op (similar to matmul/conv) might make sense ? Or maybe there's a way to generalize this pattern with embedding_bag. Both are pretty common in huggingface.

cc @ngimel @jansel

[jansel]

You should be able to map this to ir.Scatter(..., scatter_mode="atomic_add").

ir.Scatter is very similar to ir.Pointwise, except:

ir.Pointwise computes

out[x] = inner_fn(x)

ir.Scatter computes

out[output_indexer(x)] = inner_fn(x)

If you set scatter_mode="atomic_add" it instead does:

out[output_indexer(x)] += inner_fn(x)

In a thread-safe way.

The iteration ranges of scatter then become the iteration space of the scatter loops (which will be different than the size of the output).

The lowering should be something like this (warning, this is untested):

def index_add_(self, dim, index, source, *, alpha=1):
    index_loader = index.make_loader()
    source_loader = source.make_loader()

    def output_indexer(index):
        index = list(index)
        index[dim] = ops.indirect_indexing(index_loader([index[dim]]))
        return index

    def fn(index):
        return ops.mul(
            ops.constant(alpha, source.get_dtype()), 
            source_loader(index)
        )
    
    scatter = ir.Scatter(
        device=self.get_device(),
        dtype=self.get_dtype(),
        inner_fn=fn,
        ranges=list(source.get_size()),
        output_indexer=output_indexer,
        scatter_mode="atomic_add"
    )
    buffer = ir.ComputedBuffer(
        None,
        ir.MutationLayout(self),
        scatter,
    )
    buffer.name = V.graph.register_buffer(buffer)
    return self

[ngimel]

For completeness, we should also add index_select (which can be done as Pointwise, similar toindex)

[ngimel]

@eellison index_put also accumulates arbitrarily many indices into the same location

[lezcano]

This and index_copy can be done as a decomposition in terms of index_put doing something like x[:, :, ..., idx] = tensor. Writing a lowering should also be easy, as it's pretty much the same as index_select which also has a lowering. I was going to implement the two and see how they compare, but a priori I think they should be comparable, so the decomposition should be a better option.

[lezcano]

The decompositions are here: pytorch/pytorch#85002.

@eellison if you implement the lowering, you can benchmark it against the decomposition see how they fare.

[lezcano]

Of course, if you guys think that the decomposition is more convenient than the lowering, I could implement the decomposition for index_select and remove the lowering.