pytorch
1dee99c9 - LU Solve using cublas and cusolver (#59148)
Commit
3 years agoLU Solve using cublas and cusolver (#59148)
Summary:
This PR introduces cuSOLVER and cuBLAS for the `lu_solve` routine. Solves a part of https://github.com/pytorch/pytorch/issues/47953.
Since usage of cuSOLVER with MAGMA introduces performance regressions in MAGMA (https://github.com/pytorch/pytorch/issues/56590), we use heuristics for determining when to call cuSOLVER, cuBLAS or MAGMA depending on the batch and matrix sizes. The 64-bit cuSOLVER API is not introduced in this PR since there are several problems with the LU factorization using cusolver (https://github.com/pytorch/pytorch/pull/59148).
The following are performance benchmarks using various configurations:
<details>
```
[--------------------------------------------------------- LU solve CUDA torch.float64 ----------------------------------------------------------]
| lu_solve CUSOLVER | lu_solve MAGMA | lu_solve CUBLAS | lu_solve cuSOLVER/MAGMA | lu_solve TEST ALL
1 threads: ---------------------------------------------------------------------------------------------------------------------------------------
torch.Size([1, 1, 1]) | 703.4 | 489.8 | 511.8 | 710.1 | 487.1
torch.Size([2, 1, 1]) | 738.9 | 504.1 | 513.0 | 958.2 | 494.4
torch.Size([4, 1, 1]) | 790.7 | 514.7 | 506.8 | 983.9 | 540.2
torch.Size([8, 1, 1]) | 865.3 | 496.4 | 514.7 | 975.2 | 520.0
torch.Size([16, 1, 1]) | 955.5 | 483.9 | 508.3 | 937.6 | 526.5
torch.Size([32, 1, 1]) | 1167.7 | 495.2 | 511.2 | 934.0 | 528.7
torch.Size([64, 1, 1]) | 1730.0 | 492.1 | 537.8 | 936.4 | 533.2
torch.Size([128, 1, 1]) | 2748.4 | 499.7 | 526.5 | 982.9 | 540.8
torch.Size([1, 2, 2]) | 724.6 | 498.2 | 541.7 | 715.0 | 504.7
torch.Size([2, 2, 2]) | 737.0 | 514.3 | 527.6 | 934.5 | 524.5
torch.Size([4, 2, 2]) | 750.5 | 524.1 | 537.4 | 935.5 | 543.0
torch.Size([8, 2, 2]) | 844.8 | 513.7 | 538.9 | 953.3 | 534.4
torch.Size([16, 2, 2]) | 1013.1 | 521.9 | 530.0 | 932.2 | 537.9
torch.Size([32, 2, 2]) | 1335.8 | 515.1 | 544.4 | 939.9 | 559.5
torch.Size([64, 2, 2]) | 1819.6 | 511.8 | 534.1 | 973.9 | 540.0
torch.Size([128, 2, 2]) | 3018.7 | 526.3 | 546.1 | 979.3 | 543.5
torch.Size([1, 8, 8]) | 732.5 | 524.9 | 532.9 | 762.4 | 516.8
torch.Size([2, 8, 8]) | 771.2 | 514.9 | 538.7 | 1007.5 | 531.1
torch.Size([4, 8, 8]) | 811.3 | 507.7 | 534.6 | 1002.2 | 548.5
torch.Size([8, 8, 8]) | 866.6 | 530.0 | 532.0 | 1016.1 | 562.9
torch.Size([16, 8, 8]) | 991.8 | 533.6 | 548.0 | 1022.6 | 548.5
torch.Size([32, 8, 8]) | 1271.7 | 541.2 | 534.7 | 1013.8 | 545.6
torch.Size([64, 8, 8]) | 1817.2 | 530.2 | 520.6 | 1008.7 | 566.3
torch.Size([128, 8, 8]) | 2678.7 | 531.6 | 552.2 | 1006.2 | 555.0
torch.Size([1, 16, 16]) | 738.2 | 546.1 | 536.6 | 775.6 | 540.1
torch.Size([2, 16, 16]) | 782.6 | 543.5 | 539.6 | 1010.9 | 541.1
torch.Size([4, 16, 16]) | 815.2 | 546.1 | 560.9 | 1012.5 | 553.1
torch.Size([8, 16, 16]) | 877.7 | 543.0 | 547.9 | 1012.8 | 551.5
torch.Size([16, 16, 16]) | 1008.7 | 549.2 | 562.7 | 1016.6 | 546.8
torch.Size([32, 16, 16]) | 1291.9 | 540.8 | 560.3 | 1055.8 | 539.3
torch.Size([64, 16, 16]) | 1846.3 | 553.5 | 556.0 | 1010.8 | 551.9
torch.Size([128, 16, 16]) | 2953.8 | 562.7 | 547.5 | 1026.2 | 555.8
torch.Size([1, 32, 32]) | 789.1 | 590.6 | 590.9 | 790.5 | 579.0
torch.Size([2, 32, 32]) | 806.9 | 596.6 | 600.2 | 1085.6 | 573.8
torch.Size([4, 32, 32]) | 852.0 | 597.9 | 588.2 | 1098.9 | 574.7
torch.Size([8, 32, 32]) | 914.2 | 597.8 | 591.4 | 1090.3 | 585.7
torch.Size([16, 32, 32]) | 1063.0 | 604.6 | 597.3 | 1094.0 | 580.5
torch.Size([32, 32, 32]) | 1302.0 | 602.0 | 598.9 | 1090.3 | 583.6
torch.Size([64, 32, 32]) | 1861.7 | 601.1 | 599.8 | 1113.4 | 588.6
torch.Size([128, 32, 32]) | 3251.0 | 619.6 | 595.3 | 1106.8 | 608.9
torch.Size([1, 64, 64]) | 978.6 | 842.7 | 778.6 | 1071.4 | 825.8
torch.Size([2, 64, 64]) | 1072.3 | 845.7 | 785.4 | 1400.6 | 829.0
torch.Size([4, 64, 64]) | 1051.9 | 842.9 | 796.1 | 1352.2 | 788.2
torch.Size([8, 64, 64]) | 1090.3 | 834.1 | 805.2 | 1382.6 | 804.7
torch.Size([16, 64, 64]) | 1206.9 | 835.7 | 802.2 | 1365.6 | 801.2
torch.Size([32, 64, 64]) | 1671.2 | 846.5 | 794.5 | 1345.1 | 814.2
torch.Size([64, 64, 64]) | 2759.3 | 848.5 | 795.4 | 1409.7 | 832.9
torch.Size([128, 64, 64]) | 4928.6 | 877.4 | 848.3 | 1439.0 | 883.9
torch.Size([1, 128, 128]) | 1315.6 | 1158.4 | 1130.0 | 1301.3 | 1177.1
torch.Size([2, 128, 128]) | 1334.7 | 1198.2 | 1186.6 | 1703.9 | 1209.5
torch.Size([4, 128, 128]) | 1374.6 | 1200.7 | 1266.2 | 1640.6 | 1272.3
torch.Size([8, 128, 128]) | 1453.6 | 1215.9 | 1287.3 | 1669.1 | 1288.7
torch.Size([16, 128, 128]) | 1882.1 | 1244.9 | 1337.6 | 1698.8 | 1347.1
torch.Size([32, 128, 128]) | 2789.0 | 1284.5 | 1398.6 | 1747.6 | 1396.3
torch.Size([64, 128, 128]) | 4763.0 | 1425.2 | 1581.7 | 1921.0 | 1584.1
torch.Size([128, 128, 128]) | 8835.9 | 1808.9 | 1968.7 | 2197.6 | 1961.8
torch.Size([1, 512, 512]) | 4369.9 | 4577.6 | 4804.0 | 4331.4 | 4599.0
torch.Size([2, 512, 512]) | 4635.9 | 4850.1 | 5159.1 | 5315.4 | 4845.5
torch.Size([4, 512, 512]) | 5367.5 | 5261.6 | 6134.7 | 5807.8 | 5345.2
torch.Size([8, 512, 512]) | 7025.2 | 6184.5 | 7065.6 | 6711.6 | 6303.9
torch.Size([16, 512, 512]) | 10221.3 | 7849.7 | 8820.1 | 8323.6 | 7992.1
torch.Size([32, 512, 512]) | 16574.8 | 11208.4 | 12284.3 | 11704.7 | 11394.4
torch.Size([64, 512, 512]) | 29500.1 | 18043.1 | 19249.3 | 18744.0 | 18242.1
torch.Size([128, 512, 512]) | 56783.3 | 33903.9 | 34713.5 | 33893.8 | 34041.8
torch.Size([1, 1024, 1024]) | 14864.5 | 15714.6 | 16128.1 | 14726.7 | 14992.6
torch.Size([2, 1024, 1024]) | 17891.0 | 18553.3 | 19111.6 | 19271.5 | 19283.0
torch.Size([4, 1024, 1024]) | 22143.4 | 21909.2 | 23667.1 | 22698.9 | 22713.8
torch.Size([8, 1024, 1024]) | 30621.1 | 28669.9 | 30822.9 | 29725.0 | 29760.8
torch.Size([16, 1024, 1024]) | 47045.9 | 41900.0 | 44353.8 | 43215.6 | 43237.5
torch.Size([32, 1024, 1024]) | 79245.5 | 68316.9 | 70959.0 | 69506.4 | 69876.7
torch.Size([64, 1024, 1024]) | 147973.9 | 121120.6 | 124601.1 | 122084.4 | 122578.7
torch.Size([128, 1024, 1024]) | 295586.2 | 232871.8 | 237421.8 | 233765.3 | 234704.6
Times are in microseconds (us).
```
</details>
Here's the details of how the tests were performed:
* CUSOLVER - Only call `cusolver` for all problem sizes.
* MAGMA - Only call `magma` for all problem sizes (this is the current master branch).
* CUBLAS - Only call `cublas` for all problem sizes.
* cuSOLVER / MAGMA - Use cusolver for `batch_size == 1` and magma for all others.
* TEST ALL - Employ heuristics to switch between cublas/cusolver/magma. This yields the best overall results (this PR).
Script for reproducing the results:
<details>
``` python
import torch
import pickle
import itertools
from torch.utils.benchmark import Timer
import sys
shapes = [1, 2, 8, 16, 32, 64, 128, 512, 1024]
batches = [(1,), (2,), (4,), (8,), (16,), (32,), (64,), (128,)]
results = []
num_threads = 1
dtype = torch.float64
repeats = 2
from torch.testing._internal.common_utils import random_hermitian_pd_matrix
def lu_factorize_solve(mat, b):
lu_data = torch.lu(mat)
x = torch.lu_solve(b, *lu_data)
for shape, batch in itertools.product(shapes, batches):
mat = torch.randn(*batch, shape, shape, dtype=dtype, device='cuda')
b = torch.randn(*batch, shape, 1, dtype=dtype, device='cuda')
tasks = [("lu_factorize_solve(mat, b)", "lu_solve CUSOLVER")]
print("shape: ", shape, " batch: ", batch)
timers = [Timer(stmt=stmt, num_threads=num_threads, label=f"LU solve CUDA {dtype}",
sub_label=f"{mat.shape}", description=label, globals=globals()) for stmt, label in tasks]
for i, timer in enumerate(timers * repeats):
results.append(
pickle.dumps(timer.blocked_autorange())
)
print(f"\r{i + 1} / {len(timers) * repeats}", end="")
sys.stdout.flush()
f = open("cusolver_lu_solve.pickle", "wb")
pickle.dump(results, f)
f.close()
```
</details>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/59148
Reviewed By: H-Huang
Differential Revision: D29160609
Pulled By: mruberry
fbshipit-source-id: 7280f25db1e66aa650ea15608a6dc5d688fb4db2
Author
Committer
Parents
Loading