[Operator] Add diagonal backward

[awayzjj]

PR Category

Type of Change

Description

Issue

Close #314

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

[awayzjj]

@StrongSpoon Hi, I have 2 questions.

1. How to registered backward independently into aten library, I could not find a reference in the repo, so I still implemented a class with forward and backward function as a draft.
2. I implemented the backward function, and the UT passed.

    def backward(ctx, out_grad):
        logging.debug("GEMS DIAGONAL BACKWARD")
        (inp,) = ctx.saved_tensors
        grad_input = torch.zeros_like(inp)
        diag = torch.diagonal(grad_input, ctx.offset, ctx.dim1, ctx.dim2)
        diag.copy_(out_grad)
        return grad_input, None, None, None

it is simple, use torch.diagonal to get a view of the grad_input, and copy the value from the out_grad, I wonder if we need a triton kernel here.

Thank you very much!

[awayzjj]

@StrongSpoon A gentle reminder.

@StrongSpoon Hi, I have 2 questions.

1. How to registered backward independently into aten library, I could not find a reference in the repo, so I still implemented a class with forward and backward function as a draft.
2. I implemented the backward function, and the UT passed.

    def backward(ctx, out_grad):
        logging.debug("GEMS DIAGONAL BACKWARD")
        (inp,) = ctx.saved_tensors
        grad_input = torch.zeros_like(inp)
        diag = torch.diagonal(grad_input, ctx.offset, ctx.dim1, ctx.dim2)
        diag.copy_(out_grad)
        return grad_input, None, None, None

it is simple, use torch.diagonal to get a view of the grad_input, and copy the value from the out_grad, I wonder if we need a triton kernel here.

Thank you very much!

[StrongSpoon]

Hi awayzjj,

We used to implement operator with both forward and backward as a subclass of torch.autograd.Function, and register it to the forward interface using AutogradCUDA key. But recently we found that AutogradCUDA could not work with torch.compile perfectly. As a solution, we recommend to implement forward function and backward function, and register them into aten library respectively.
Take tanh as an example, its forward interface is defined in https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml. In the previous practice, we implemented class Tanh and registered it into tanh interface. But it's better to register forward function to tanh and backward function to tanh_backward, using CUDA as dispatch key.

- func: tanh(Tensor self) -> Tensor
  device_check: NoCheck   # TensorIterator
  structured_delegate: tanh.out
  variants: function, method
  dispatch:
    QuantizedCPU: tanh_quantized_cpu
    MkldnnCPU: mkldnn_tanh
    SparseCPU, SparseCUDA: tanh_sparse
    SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: tanh_sparse_csr
    NestedTensorCPU, NestedTensorCUDA: NestedTensor_tanh
  tags: [core, pointwise]
- func: tanh_backward(Tensor grad_output, Tensor output) -> Tensor
  python_module: nn
  structured_delegate: tanh_backward.grad_input

[StrongSpoon]

As is the same, there is a definition of diagonal_backward interface, which is expected to be reimplemented by developers.

[StrongSpoon]

Besides, we require developer to implement the function by writing a Triton kernel function instead of torch apis. If you are confused about the format, please refer to our source code.

[awayzjj]

Besides, we require developer to implement the function by writing a Triton kernel function instead of torch apis. If you are confused about the format, please refer to our source code.

Please review my PR, thanks!

[StrongSpoon]

please provide the benchmark results.

[StrongSpoon]

it's okay to fuse backward and diagonal_backward into one.

[StrongSpoon]

since torch.zeros also calls for a kernel, there exist two kernels indeed. I wonder if it's feasible to initialize grad_input as an empty tensor, and assign it in one kernel function.

[awayzjj]

Got it, I'll give it a try.

[StrongSpoon]

why not require them equal? I thought backward function doesn't change the value.

[awayzjj]

Yes, I fixed it.

[awayzjj]

Hi, @StrongSpoon
CI failed, but the failed UT is not relavent with my PR(I can reproduce it after runing the UT several times on the latest master branch)

[StrongSpoon]

lgtm