Summary:
This together with #66580 and #68376 will remove all syncs in embedding.

This PR includes #68376, please review after merging #68376

This PR introduces perf regressions and increases memory usage:
- `exclusive_sum` is now computing the entire `numel` elements instead of `num_of_segments` elements, and the trailing `numel - num_of_segments` results will be discarded.
- Some memory allocation now needs `numel` spaces instead of `num_of_segments` or `num_of_partial_segments`.

These are the prices we must pay in order to get a sync-free implementation.

I haven't done any benchmark yet. I will do it later.

Pull Request resolved: #70943

Reviewed By: H-Huang

Differential Revision: D34881660

Pulled By: ngimel

fbshipit-source-id: b0760fa33608c46cd4145ceb09878bf94a9f959d

`numel` is a too loose upper bound for `num_of_segments` and `num_of_partial_segments`. It causes perf regressions.

This PR moves to a tighter upper bound.

Benchmark with jupyter notebook:
```python
import torch
num_embeddings = 1024
embedding_dim = 512
e = torch.nn.Embedding(num_embeddings, embedding_dim).cuda()

size = 1*1024*1024
i = torch.arange(size, device='cuda') % num_embeddings
o = e(i)
g = torch.randn_like(o)
torch.cuda.synchronize()
```
```python
%%timeit
o.backward(g, retain_graph=True)
torch.cuda.synchronize()
```

Before #70943: 3.6 ms
After #70943: 6.9 ms
With this PR: 3.55 ms
Pull Request resolved: #78588
Approved by: https://github.com/ngimel

Summary:
`numel` is a too loose upper bound for `num_of_segments` and `num_of_partial_segments`. It causes perf regressions.

This PR moves to a tighter upper bound.

Benchmark with jupyter notebook:
```python
import torch
num_embeddings = 1024
embedding_dim = 512
e = torch.nn.Embedding(num_embeddings, embedding_dim).cuda()

size = 1*1024*1024
i = torch.arange(size, device='cuda') % num_embeddings
o = e(i)
g = torch.randn_like(o)
torch.cuda.synchronize()
```
```python
%%timeit
o.backward(g, retain_graph=True)
torch.cuda.synchronize()
```

Before #70943: 3.6 ms
After #70943: 6.9 ms
With this PR: 3.55 ms

Pull Request resolved: #78588
Approved by: https://github.com/ngimel

Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/5bcbad76148a2808ae8331b8cfbd42885269a8d9

Reviewed By: seemethere

Differential Revision: D36815606

Pulled By: seemethere

fbshipit-source-id: 0d4c349c7bd8aa8e2c9209d6ec15415867f722a0