pytorch

mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00

Files

Pearu Peterson d4708a6da7 Add scatter_mm and bsr_scatter_mm operations. (#110396 )

This PR introduces `scatter_mm` operation (compute `mm` of arbitrary pairs of tensors given in batches of tensors) that is used to implement `bsr_scatter_mm` that is equivalent to `bsr_dense_mm` (the `mm` operation on bsr and strided tensors). The implementation is provided both in Triton (when tensor dimensions are multiples of 16) and in PyTorch (otherwise).

The figures below illustrate the performance differences of `bsr_scatter_mm` and `bsr_dense_mm` (GPU: `NVIDIA GeForce RTX 2060 SUPER`). The first figure represents the performance equilibrium point in BSR tensor sparsity at which value `bsr_scatter_mm` or `bsr_dense_mm` have the same performance characteristics as `torch.matmul`. The second figure represents speedups from using `bsr_scatter_mm` at its performance equilibrium points with respect to `bsr_dense_mm`.

<img src="https://github.com/pytorch/pytorch/assets/402156/526d182e-937f-4812-a6c4-904f52d6d5ab" width="48%"> <img src="https://github.com/pytorch/pytorch/assets/402156/ccb606ab-1f3f-4133-887c-b56285f4f168" width="48%">

The same figures for GPU card `NVIDIA A100-SXM4-80GB`:

<img src="https://github.com/pytorch/pytorch/assets/402156/25466f1d-df34-4d1c-a975-afb478e4d9f0" width="48%"> <img src="https://github.com/pytorch/pytorch/assets/402156/6ada91f0-a20f-4f0d-8a48-1f4ccc60d08e" width="48%">

In sum:
- `bsr_scatter_mm` is about 2x faster than `bsr_dense_mm` for small block sizes of 16 and 32 and large tensors [GPU: `NVIDIA GeForce RTX 2060 SUPER`].
- `bsr_scatter_mm` is up to 2x faster than `bsr_dense_mm` for small block sizes of 16 and large tensors [GPU: `NVIDIA A100-SXM4-80GB`].
- `bsr_dense_mm` is up to 20 % faster than `bsr_scatter_mm` for block sizes of 64 or larger [GPU: `NVIDIA GeForce RTX 2060 SUPER`].
- However, `bsr_dense_mm` fails with `OutOfResources` exception for block sizes of 256 or larger whereas `bsr_scatter_mm` succeeds.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/110396
Approved by: https://github.com/cpuhrsch

2023-10-23 19:45:30 +00:00

dlmc

Apply UFMT to all files in benchmarks/ (#105928 )

2023-07-26 01:18:48 +00:00

__init__.py

Apply UFMT to all files in benchmarks/ (#105928 )

2023-07-26 01:18:48 +00:00

benchmark_semi_structured_sparsity.py

[sparse] Add padding for dense matrices in semi-structured sparse (#110583 )

2023-10-13 20:04:23 +00:00

README.md

Add CSR (compressed sparse row) layout for sparse tensors (#50937 )

2021-04-12 10:09:12 -07:00

spmm.py

Apply UFMT to all files in benchmarks/ (#105928 )

2023-07-26 01:18:48 +00:00

spmv.py

Apply UFMT to all files in benchmarks/ (#105928 )

2023-07-26 01:18:48 +00:00

test_csr.sh

[BE] Prefer dash over underscore in command-line options (#94505 )

2023-02-09 20:16:49 +00:00

triton_ops.py

Add scatter_mm and bsr_scatter_mm operations. (#110396 )

2023-10-23 19:45:30 +00:00

utils.py

Apply UFMT to all files in benchmarks/ (#105928 )

2023-07-26 01:18:48 +00:00

README.md

#Sparse benchmarks

These sets of benchmarks are for the sparse matrix functionality. They exist for comparing the performance of sparse matrix routines such as SpMV between various sparse matrix formats and with other frameworks such as TensorFlow.