NVIDIA / NVIDIA/MatX

[FEA] Detect transposed tensor on reductions and switch to fast path

Open
#482 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
C++
Stars
1.4k
Forks
127
Avg merge
4d 1h
Merged PRs (30d)
11

Description

Currently a reduction across columns vs rows is slow since it naively takes in a transpose operator and indexes it using a random access iterator. This causes adjacent threads to have a strided access equal to the second dimension.

We experimented with an einsum transpose, followed by a reduction, and the results were as follows on an A30:

permute/sum

|  T  | Tensor Size | NumElements |  DataSize   | Samples | CPU Time  | Noise | GPU Time  | Noise | Elem/s  | GlobalMem BW | BWUtil | Samples | Batch GPU |
|-----|-------------|-------------|-------------|---------|-----------|-------|-----------|-------|---------|--------------|--------|---------|-----------|
| F32 |    2^6 = 64 |          64 |  64.000 MiB |    608x |  1.534 ms | 6.09% |  1.446 ms | 0.63% | 44.248K |  46.397 GB/s |  4.97% |    609x |  1.457 ms |
| F32 |   2^7 = 128 |         128 |   1.000 GiB |     17x | 31.070 ms | 0.31% | 30.981 ms | 0.09% |  4.132K |  34.658 GB/s |  3.71% |     18x | 30.965 ms |
| F64 |    2^6 = 64 |          64 | 128.000 MiB |    254x |  2.058 ms | 4.40% |  1.971 ms | 0.18% | 32.467K |  68.088 GB/s |  7.30% |    267x |  1.944 ms |
| F64 |   2^7 = 128 |         128 |   2.000 GiB |     16x | 32.480 ms | 0.29% | 32.391 ms | 0.07% |  3.952K |  66.298 GB/s |  7.11% |     17x | 32.367 ms |

einsum/sum

|  T  | Tensor Size | NumElements |  DataSize   | Samples |  CPU Time  | Noise  |  GPU Time  | Noise  |  Elem/s  | GlobalMem BW | BWUtil | Samples | Batch GPU  |
|-----|-------------|-------------|-------------|---------|------------|--------|------------|--------|----------|--------------|--------|---------|------------|
| F32 |    2^6 = 64 |          64 |  64.000 MiB |    944x | 585.716 us | 14.30% | 577.762 us | 14.06% | 110.772K | 116.153 GB/s | 12.45% |    991x | 504.905 us |
| F32 |   2^7 = 128 |         128 |   1.000 GiB |   1846x |   8.025 ms |  0.82% |   8.017 ms |  0.81% |  15.966K | 133.933 GB/s | 14.35% |   1847x |   8.007 ms |
| F64 |    2^6 = 64 |          64 | 128.000 MiB |   3136x | 668.751 us |  1.46% | 661.259 us |  0.91% |  96.785K | 202.973 GB/s | 21.75% |   3137x | 649.249 us |
| F64 |   2^7 = 128 |         128 |   2.000 GiB |   1380x |  10.776 ms |  1.44% |  10.768 ms |  1.44% |  11.887K | 199.424 GB/s | 21.37% |   1381x |  10.760 ms |

The einsum version is quite a bit faster since it hits SoL on the transpose. However, there can be room for improvement where a kernel aware of both transpose and reductions can be faster by tiling. This issue is to implement that kernel and compare performance.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by locating the reduction and transpose paths described in the issue, including the random-access iterator and einsum-based transpose route. Compare the existing behavior with the proposed tiled kernel approach on the A30 benchmarks; done means the transpose-aware reduction is implemented and its performance is compared against the reported baselines.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
hpc, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.