Avoid cloning gradient tensor in embedding backward pass

[jhadidjojo]

Summary: I found memory spike during embedding kernel backward split_embedding_backward_codegen_rowwise_adagrad_unweghted_exact_cuda, which was traced into the below code making a clone of the gradient tensor. This logic didn't seem to be there in the original code: https://github.com/pytorch/FBGEMM/pull/2347/files#diff-944ab49dcbcf54826cc3e1eab5e3c0c787b5a195f602c2d3052adae14c506d78.

Differential Revision: D56420646

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D56420646

[netlify]

✅ Deploy Preview for pytorch-fbgemm-docs ready!

Name	Link
🔨 Latest commit	3f16122
🔍 Latest deploy log	https://app.netlify.com/sites/pytorch-fbgemm-docs/deploys/6629059bd7f6f200086dbaa1
😎 Deploy Preview	https://deploy-preview-2526--pytorch-fbgemm-docs.netlify.app
📱 Preview on mobile	Toggle QR Code... Use your smartphone camera to open QR code link.

---

To edit notification comments on pull requests, go to your Netlify site configuration.

[flaviotruzzi]

This will break the signature for const, will it not? The reason for that if I recall it was that we moved this memory alignment into the operator, inside the cuda kernel, the signatures ask that to be const. If we don’t do it inside the operator we need to figure out a way to adjust the meta kernel which don’t have access to the data pointers to check the alignment.

[jhadidjojo]

This will break the signature for const, will it not?

No, I don't think so. You're right that grad_output is a const reference, and if you try to assign that to a new value it will error (that's probably the error you saw?). Here I'm creating a new (non-const) reference called aligned_grad_output. You can assign it to the grad_output by reference without copy.

[flaviotruzzi]

This will break the signature for const, will it not?

No, I don't think so. You're right that grad_output is a const reference, and if you try to assign that to a new value it will error (that's probably the error you saw?). Here I'm creating a new (non-const) reference called aligned_grad_output. You can assign it to the grad_output by reference without copy.

Cool, I am only on my phone, and still on leave, can you check with someone else to approve it?

[jhadidjojo]

@pytorchbot merge

[pytorch-bot]

This PR needs to be approved by an authorized maintainer before merge.

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D56420646

[facebook-github-bot]

This pull request has been merged in a75037b.