[RFC] Lowering aten::_index_put_impl_ into tm_tensor.scatter op

[erman-gurses]

(#614)

[cathyzhyi]

This op is quite tricky. We need to transform the operands meet the requirement for for tm_tensor.scatter

torch-mlir/external/llvm-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Dialect/TMTensor/IR/TMTensorOps.td

Line 99 in ba29d4f

def TMTensor_ScatterOp : TMTensor_Op<"scatter",

(assuming readers are familiar with both aten.index_put and tm_tensor.scatter)

For all indices tensors:

1. broadcast to the same shape (call it indices_shape for now)
2. flatten to 1d tensor
3. unsqueeze to make the last dimension size one
4. concatenate all the tensors together

For values operand:

1. broadcast to the values_shape. values_shape is consisted of 2 parts. First part is the same as the indices_shape. The second part is input tensor shape with dimensions matching the given indices tensors removed.
2. collapse the indices_shape part to be 1 dimension.

Some steps above are more convenient to be decomposed to other torch ops (like flatten to 1d tensor, unsqueeze, concatenate), some are more convenient to be converted directly to tm_tensor or linalg ops. The conversion should go into ConvertTorchToTMTensor pass to be added by #574.

We might not be able to handle all the cases aten.index_put support. For example, it allows a mixed of None and tensors for the indices list in any order. We might want to start with no None indices first.

[silvasean]

We will need a new pass called something like DecomposeOrConvertTorchOp to convert Torch ops to a mixed of torch and aten ops.

Do you mean a mix of torch and TMTensor ops? I think this pass would naturally be the scope of a ConvertTorchToTMTensor pass.

[cathyzhyi]

sry, I meant a mixed of torch and linalg/tmtensor ops.

ConvertTorchToTMTensor

We have another pass called lowerAlgorithmicOps #574 which lowers algorithmic torch ops to linalg and tm_tensor ops. I feel calling ConvertTorchToTMTensor for this new pass would cause confusion as it is not explicit the torch can also be decomposed to Torch ops.

[silvasean]

I left a comment in #574. I think that there should be LowerAlgorithmicOps (similar to DecomposeComplexOps) as a Torch->Torch lowering. Then ConvertTorchToTMTensor which happens in backends (i.e. after the Torch backend contract, just before ConvertTorchToLinalg).

[cathyzhyi]

I left a comment in #574. I think that there should be LowerAlgorithmicOps (similar to DecomposeComplexOps) as a Torch->Torch lowering. Then ConvertTorchToTMTensor which happens in backends (i.e. after the Torch backend contract, just before ConvertTorchToLinalg).

Ha, I see. There seems to be confusion about passes naming. I updated my comments. Thanks for clarifying this!

[vivekkhandelwal1]

Hi @cathyzhyi @silvasean
I am working on the implementation of this op. I need this op for the functorch decomposition of aten::embedding_dense_backward.
https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.py#L318-L333

@register_decomposition(aten.embedding_dense_backward)
def embedding_dense_backward(grad_output: Tensor, indices: Tensor, num_weights: int, padding_idx: int, scale_grad_by_freq: bool):
    numel = indices.numel()
    grad = grad_output.view(numel, grad_output.size(-1))
    grad_weight = aten.new_zeros(grad_output, (num_weights, grad_output.shape[-1]))
    indices_rank1 = indices.view(numel)
    if scale_grad_by_freq:
        counts = aten.new_zeros(indices, (num_weights,))
        ones = aten.new_ones(indices, (numel,))
        counts = aten.index_put(counts, [indices_rank1], ones, accumulate=True)
        grad_weights_scale = aten.index(counts, [indices_rank1])
        grad = grad / grad_weights_scale.unsqueeze(1)
    skip_padding = (indices_rank1 != padding_idx).unsqueeze(1)
    skip_padding = skip_padding.expand_as(grad)
    zero_grad = aten.full_like(grad, 0)
    return aten.index_put(grad_weight, [indices_rank1], aten.where(skip_padding, grad, zero_grad), accumulate=True)

I have gone through the comments mentioned above. Since this op seems a bit complex, should we implement this op as of now for the cases required, which is a 1-d tensor of indices and a 1-d tensor of values with
indices.size() == values.size()
We can add the remaining cases as a TO-DO and add them as and when required. The current implementation of tm_tensor.scatter op easily supports the 1-d case.

[silvasean]

Starting with the 1D version of this and emitting a clear diagnostic for the unsupported cases seems like a good incremental first step.

[cathyzhyi]

closed by #650