lkct
commented
1 year ago Book-keeping -- concatenation and indexing
1-D indexing for simple folding, 2-D indexing for mixing (include a dimension for input)
code
import itertools
import subprocess
import sys
import time
from typing import List
import torch
from torch import Tensor
from benchmark.utils.gpu_benchmark import timer
device = "cuda"
B = 128
K = 512
F = 784
f = 288
sizes = {
"B": B,
"K": K,
"F": F,
"f": f,
"2": 2,
}
N = 200
def bench(inputs: List[str], idx_shape: str, dim: int) -> float:
input_data = [torch.rand([sizes[x] for x in y], device=device) for y in inputs]
idx_data = torch.randint(f, [sizes[x] for x in idx_shape], device=device)
idx = [idx_data if i == dim else slice(None) for i in range(len(inputs[0]))]
def proc() -> Tensor:
result = torch.cat(input_data, dim=dim)
result = result.__getitem__(idx)
return result
_, t = timer(proc)
return t
def run(inputs: str, idx_shape: str) -> None:
if "F" in inputs:
dim = inputs.index("F")
args = ([inputs], idx_shape, dim)
else:
dim = inputs.index("f")
args = ([inputs, inputs.replace("f", "F")], idx_shape, dim)
for _ in range(100): # long warm-up is essential?
bench(*args)
s = 0.0
for _ in range(N):
s += bench(*args)
print(f"{s / N * 1000:9.3f}", end="")
def main() -> None:
ans = []
for inputs in itertools.permutations("BKF"):
for idx_shape in ("f",):
inputs = "".join(inputs)
result = subprocess.run(
["python", "bench.py", inputs, idx_shape],
capture_output=True,
check=False,
text=True,
)
assert not result.returncode, result.stderr
ans.append(["i", inputs, "idx", idx_shape, result.stdout])
time.sleep(1)
print(ans[-1], file=sys.stderr)
ans.sort(key=lambda x: x[-1])
for item in ans:
print(*item)
if __name__ == "__main__":
if len(sys.argv) == 1:
main()
else:
run(*sys.argv[1:])1-concat, 1-D index
| input | time |
|---|---|
| BFK | 793.532 |
| KFB | 794.392 |
| FKB | 831.200 |
| FBK | 831.399 |
| KBF | 844.904 |
| BKF | 845.616 |
2-concat, 1-D index
code differ by
for inputs in itertools.permutations("BKf"):| input | time |
|---|---|
| KFB | 1041.293 |
| BFK | 1043.517 |
| KBF | 1090.914 |
| BKF | 1091.461 |
| FKB | 1104.790 |
| FBK | 1109.793 |
0-concat, 2-D index
code differ by
def proc() -> Tensor:
result = input_data[0]
result = result.__getitem__(idx)
return result
...
for idx_shape in ("2f", "f2"):| input | idx | time |
|---|---|---|
| KFB | 2f | 322.121 |
| KFB | f2 | 322.174 |
| BFK | 2f | 323.775 |
| BFK | f2 | 324.232 |
| KBF | f2 | 341.624 |
| BKF | 2f | 341.690 |
| KBF | 2f | 341.697 |
| BKF | f2 | 341.883 |
| FBK | 2f | 440.158 |
| FKB | 2f | 440.267 |
| FBK | f2 | 440.390 |
| FKB | f2 | 440.422 |
This issue is the tracker for benchmarks on pytorch ops with different axes order.
All time counted in $\mu s$.
Notations:
B: batchK(alsoI,J): unitF: foldingR: rankC: componentConclusions:
torch.bmm(possibly with aTensor.expandto manually broadcast) andtorch.matmulis the same astorch.einsumin underlying calculations (all reduced tobmm->baddbmm-> bgemm in BLAS).