Repository: pytorch
Issue: Unexpected poor performance of C++ extension / wish for a fast `operator[]`
Labels: module: performance,module: cpp,triaged

### 🚀 The feature, motivation and pitch

_I had asked this as a [question on SO](https://stackoverflow.com/questions/75580832/why-is-pytorch-c-extension-much-slower-than-its-equivalent-numba-version) some time ago. The lack of answers suggest that I was hopefully not doing something obviously wrong._

The premise of converting a piece of Python code into an C++ extension is that it gives a performance benefit. What I have observed so far in my use cases is that the performance of the C++ extension was surprisingly bad. I had minified one of these use cases into this example of a "variable length delay line" (mixing comb filter). The Python code implementing such a filter is:

```py
def delay_line_impl(samples, delays):
    for i in range(len(samples)):
        delay = int(delays[i].item())
        index_delayed = i - delay
        if index_delayed < 0:
            index_delayed = 0

        samples[i] = 0.5 * (samples[i] + samples[index_delayed])
```

Turning the filter into a C++ extension looks like this:

```c++
void delay_line_forward(torch::Tensor samples, torch::Tensor delays) {
  int64_t input_size = samples.size(-1);

  for (int64_t i = 0; i < input_size; ++i) {
    int64_t delay = delays[i].item<int64_t>();
    int64_t index_delayed = i - delay;
    if (index_delayed < 0) {
      index_delayed = 0;
    }

    samples[i] = 0.5 * (samples[i] + samples[index_delayed]);
  }
}
```

I have benchmarked this filter in all variants:
- The Python version as plain Python
- The Python version converted into TorchScript
- The Python version converted into Numba
- The C++ extension
and also running it both against CPU and GPU.

Here is the full reproduction code:

<details>
<summary>Full benchmark reproduction code</summary>

**`delay_line.cpp`**

```c++
#include <torch/extension.h>

void delay_line_forward(torch::Tensor samples, torch::Tensor delays) {

  int64_t input_size = samples.size(-1);

  for (int64_t i = 0; i < input_size; ++i) {
    int64_t delay = delays[i].item<int64_t>();
    int64_t index_delayed = i - delay;
    if (index_delayed < 0) {
      index_delayed = 0;
    }

    samples[i] = 0.5 * (samples[i] + samples[index_delayed]);
  }
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  py::module::import("torch");
  m.def("delay_line_forward", &delay_line_forward);
}
```

**`benchmark_reproduction.py`**

```py
import json
import time
from contextlib import contextmanager
from typing import Literal

import numba
import numba.cuda
import numpy as np
import torch
from torch.jit._script import script as torchscript
from torch.utils.cpp_extension import load

delay_line_cpp = load(name="delay_line", sources=["delay_line.cpp"])


def delay_line_impl(samples, delays):
    for i in range(len(samples)):
        delay = int(delays[i].item())
        index_delayed = i - delay
        if index_delayed < 0:
            index_delayed = 0

        samples[i] = 0.5 * (samples[i] + samples[index_delayed])


delay_line_numba_cpu = numba.njit()(delay_line_impl)
delay_line_numba_cuda = numba.cuda.jit(delay_line_impl)
delay_line_torchscript = torchscript(delay_line_impl)

Mode = Literal["plain_python", "torchscript", "numba", "cpp"]


def delay_line(samples: torch.Tensor, delays: torch.Tensor, mode: Mode):
    if mode == "plain_python":
        return delay_line_impl(samples, delays)
    elif mode == "torchscript":
        return delay_line_torchscript(samples, delays)  # pyright: ignore
    elif mode == "numba":
        if samples.is_cuda:
            samples_view = numba.cuda.as_cuda_array(samples.detach())
            delays_view = numba.cuda.as_cuda_array(delays.detach())
            return delay_line_numba_cuda[1, 1](samples_view, delays_view)  # pyright: ignore
        else:
            return delay_line_numba_cpu(
                samples.detach().numpy(), delays.detach().numpy()  # pyright: ignore
            )
    elif mode == "cpp":
        delay_line_cpp.delay_line_forward(samples, delays)  # pyright: ignore
    else:
        raise AssertionError("unreachable")


def main():
    size = 1024
    modes: list[Mode] = ["plain_python", "torchscript", "numba", "cpp"]

    results = []

    for use_gpu in [False, True]:
        device = torch.device("cuda:0" if use_gpu else "cpu")
        device_name = "GPU" if use_gpu else "CPU"
        print(f" *** Testing on {device_name}")

        for mode in modes:

            timer = MeasureTime(f"{device_name} / {mode}", report_every=10)

            for _ in range(100):
                samples = torch.tensor(
                    np.random.uniform(-1, +1, size=(size,)), device=device, dtype=torch.float32
                )
                delays = torch.tensor(
                    np.random.randint(1, 100, size=(size,)), device=device, dtype=torch.float32
                )

                with timer.timed():
                    delay_line(samples, delays, mode)

            results.append(
                {"Method": mode, "Device": device_name, "Median time [ms]": timer.median * 1000}
            )
            print()

    try:
        import tabulate

        print(
            tabulate.tabulate(results, headers="keys", tablefmt="rounded_outline", floatfmt=".3f")
        )
    except ImportError:
        print(json.dumps(results, indent=2))


class MeasureTime:
    def __init__(self, name: str, *, report_every: int):
        self.name = name
        self.every = report_every
        self.times: list[float] = []

    @contextmanager
    def timed(self):
        t1 = time.monotonic()
        yield
        t2 = time.monotonic()
        self.times.append(t2 - t1)
        if len(self.times) % self.every == 0:
            self.show_stats()

    def show_stats(self):
        mean = np.mean(self.times_array)
        std = np.std(self.times_array)
        median = np.median(self.times_array)
        unit = "sec"
        if mean < 1:
            mean *= 1000
            std *= 1000
            median *= 1000
            unit = "ms"
        if mean < 1:
            mean *= 1000
            std *= 1000
            median *= 1000
            unit = "us"
        msg = f"{self.name:<30s}    mean: {mean:.1f} ± {std:.1f} {unit}    median: {median:.3f} {unit}"
        print(msg)

    @property
    def times_array(self) -> np.ndarray:
        return np.array(self.times)

    @property
    def median(self) -> float:
        return float(np.median(self.times_array))


if __name__ == "__main__":
    main()
```

</details>

Re-running the benchmark today yields:

```
╭──────────────┬──────────┬────────────────────╮
│ Method       │ Device   │   Median time [ms] │
├──────────────┼──────────┼────────────────────┤
│ plain_python │ CPU      │             18.003 │
│ torchscript  │ CPU      │             10.872 │
│ numba        │ CPU      │              0.013 │
│ cpp          │ CPU      │             11.344 │
│ plain_python │ GPU      │             57.349 │
│ torchscript  │ GPU      │             40.636 │
│ numba        │ GPU      │              0.209 │
│ cpp          │ GPU      │             39.016 │
╰──────────────┴──────────┴────────────────────╯
```

As can be seen the C++ extension is still several orders of magnitudes slower than the Numba versions. Considering that the Python and TorchScript versions are capable of auto-grad, there seems little justification to give up auto-grad with the C++ extension, because the performance gains are tiny (or even slower than TorchScript).

Back in February I had spend some time narrowing down where the bad performance is coming from, and it looked it mostly comes down to the `operator[]` accessed themselves.

Of course, "scalar-like" accesses to tensors are a kind of an anti-pattern, because they don't vectorize well. However, the use case of writing a C++ extension often arises exactly in these situations where one is facing a problem that cannot easily be written in vectorized form. If the problem could easily be written in a more vectorized / non-scalar fashion, the problem could be solved in plain `pytorch` already. For instance, I'd assume this problem arises for any kind of IIR filter that needs to build its output "one-by-one" given the nature of how the underlying non-reducible recurrence relation.

Thus, it would be great if there was a way to implement such filters as C++ extensions with performance comparable to Numba. I'm wondering if this can be accomplished by introducing a kind of low-level `operator[]` that provides better performance.

### Alternatives

Use plain Python, TorchScript, or Numba.

However, there may be reason why one would actually like to implement these things in C++.

### Additional context

_No response_

cc @jbschlosser