Add softmax_csr implementation

[DamianSzwichtenberg]

This PR adds forward and backward implementation of sparse softmax operation as defined here.

In the pytorch_geometric implementation we cannot take advantage of model compilation when groups are defined via ptr. softmax_csr introduced here provides a well-performing kernel for such a scenario.

Performance boost (achieved on 28C, single socket machine):
~7x for forward pass
~8x for backward pass
Additionally, GAT training time was reduced by ~5%.

[codecov]

Codecov Report

Attention: 4 lines in your changes are missing coverage. Please review.

Comparison is base (2b9af1c) 85.65% compared to head (40c8f52) 86.19%.

Files	Patch %	Lines
pyg_lib/csrc/ops/cpu/softmax_kernel.cpp	91.11%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #264      +/-   ##
==========================================
+ Coverage   85.65%   86.19%   +0.54%     
==========================================
  Files          32       34       +2     
  Lines        1115     1188      +73     
==========================================
+ Hits          955     1024      +69     
- Misses        160      164       +4     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[DamianSzwichtenberg]

@pyg-team/intel-team Please take a look.

[kgajdamo]

I wonder why ptr is optional because if you don't provide it, you get an error.

[DamianSzwichtenberg]

I wanted to make API in its final form, otherwise, each change here would require a change in pytorch_geometric. I'll add support for index in the near future.

[DamianSzwichtenberg]

@kgajdamo, after rethinking your suggestion, I decided to change the API and create a specialized softmax_csr operation that accepts ptr only. Rationale:
torch.compile gives nice results for softmax with groups defined via index, hence I don't see a reason to have a specialized kernel for that option.

[DamianSzwichtenberg]

@kgajdamo @rusty1s Please take a look. I made the softmax implementation a bit more general, so now it covers any src dimensionality as well as any given dim. I also restricted groups to be defined via ptr (rationale here).

[yanbing-j]

Hi @DamianSzwichtenberg , this PR looks good to me. The overall structure of softmax kernel with sparse input is similar with that in softmax kernel of dense input in PyTorch. With the sparsity, the performance boost is from parallelism, right?

And will this PR upstream to PyTorch later? Since there is no SparseCsr support for softmax yet in PyTorch.

[DamianSzwichtenberg]

Hi @DamianSzwichtenberg , this PR looks good to me. The overall structure of softmax kernel with sparse input is similar with that in softmax kernel of dense input in PyTorch. With the sparsity, the performance boost is from parallelism, right?

These kernels differ quite a bit. In softmax_csr groups are created from ptr across given dim and then, for each dimension different from dim we also create seperate group. If you take a look at tests, to achieve the same result with torch.nn.Softmax you need to slice the tensor. As for performance gains, parallelization is made through the groups defined by user (which also differs from Softmax). The other thing is that, for most common scenario (dim=0), we access data in contiguous manner, despite the fact that contiguous elements do not belong to the same group. So performance comes from parallelization and good memory access pattern.

And will this PR upstream to PyTorch later? Since there is no SparseCsr support for softmax yet in PyTorch.

There are no plans to upstream this operation to PyTorch. As above, softmax_csr differs from Softmax operation defined in torch.

[kgajdamo]

Looks good to me.

[rusty1s]

Can we define the autograd function directly in C++?

[DamianSzwichtenberg]

Should be possible, will check.

[DamianSzwichtenberg]

Change available at #282