High-level API for torch.sparse.mm with optimized spmm_reduce kernel using CSR format
#6699
Conversation
JakubPietrakIntel
requested review from
rusty1s,
mingfeima,
andreazanetti,
mszarma,
DamianSzwichtenberg,
kgajdamo and
yanbing-j
for more information, see https://pre-commit.ci
…_geometric into spmm_reduce_api
for more information, see https://pre-commit.ci
…_geometric into spmm_reduce_api
for more information, see https://pre-commit.ci
…_geometric into spmm_reduce_api
for more information, see https://pre-commit.ci
…_geometric into spmm_reduce_api
for more information, see https://pre-commit.ci
Removes duplicate edged to the given homogeneous or heterogeneous graph. It will change the original order of dataset by concatenating one of duplicated edges at the end of the dataset. It can be used to clean-up a known repeated self-connecting edges issue in ogbn-products. Reference to ogbn-products Leaderboard: [here](https://ogb.stanford.edu/docs/nodeprop/#:~:text=Note%3A%20A%20very%20small%20number%20of%20self%2Dconnecting%20edges%20are%20repeated%20(see%20here)%3B%20you%20may%20remove%20them%20if%20necessary) Moved this to separate PR from #6699 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Matthias Fey <matthias.fey@tu-dortmund.de>
…into spmm_reduce_api
| reduce = 'sum' if reduce == 'add' else reduce | ||
|
|
||
| if reduce not in ['sum', 'mean', 'min', 'max']: | ||
| raise ValueError(f"`reduce` argument '{reduce}' not supported") | ||
|
|
||
| if isinstance(src, SparseTensor): | ||
| return torch_sparse.matmul(src, other, reduce) |
There was a problem hiding this comment.
Hi @JakubPietrakIntel @rusty1s , May I know the logic here? If src is SparseTensor which is the default data type, it will never go into the spmm optimized impl in PT.
There was a problem hiding this comment.
Yes, I changed this back for now since torch.sparse.mm is missing CUDA support. I think we can either patch this in SparseTensor or patch this here such that we only call torch.sparse.mm in case PyTorch >= 2.0 and CPU. Otherwise, I think this solution is fine, as we support torch.sparse.Tensor now anyway (and are in the process of removing torch-sparse dependency).
There was a problem hiding this comment.
Understand. So for now, we need add a patch to call torch.sparse.mm when PT >= 2.0 and CPU. Do you know the time of adding CUDA support?
There was a problem hiding this comment.
I don‘t have any insights into that. Let me add the optimized routine in a separate PR.
Related to important optimization in Pytorch:
✅ port sparse_mm.reduce to pytorch and optimize it on CPU #83727
Updating simplified high-level API for
spmm_reduce()kernel and tests.The current kernel implementation has limitation to process
srcof typetorch.Tensorintorch.sparse_csrformat, therefore I've added an option to auto-convertsrcto CSR format usingsrc.to_sparse_csr(), which isFalseby default and will result inValueErrorif the input is not provided in the correct format.The conversion from
SparseTensortotorch.Tensoris enabled by default for Pytorch > 1.13.Added transfrom to remove duplicated in ogbn-products dataset, because the new kernel can't handle duplicate entries (useful for benchmarks).
Re-opened this PR because the draft (#6689) needed to be scrapped after a rebase.