I’ve noticed that certain PyTorch functions run slower when I make calls to scipy.linalg.solve_triangular.
I have managed to reproduce this issue with the following script:
import time

import numpy as np
import torch
import scipy.linalg

def torch_fn(input_size):
    with torch.no_grad():
        a = torch.randn(input_size, 4, 4, dtype=torch.float64)
        b = torch.randn(input_size, 4, 5, dtype=torch.float64)
        begin = time.perf_counter()
        c = torch.bmm(a, b)
        return time.perf_counter() - begin


def scipy_fn(input_size):
    a = np.asarray([[1, 0], [1, 1]], dtype=np.float64)
    b = np.random.randn(input_size, 2).astype(np.float64)
    begin = time.perf_counter()
    c = scipy.linalg.solve_triangular(a, b.T, lower=True).T
    return time.perf_counter() - begin


def avg_time(torch_input_size, scipy_input_size):
    tot = 0
    for _ in range(100):
        tot += torch_fn(torch_input_size)
        if scipy_input_size > 0:
            scipy_fn(scipy_input_size)
    return tot / 100


if __name__ == '__main__':
    print(avg_time(5, 0))
    print(avg_time(5, 3))
    print(avg_time(5, 75))
    print(avg_time(30, 0))
    print(avg_time(30, 3))
    print(avg_time(30, 75))

The output I get is
0.0002978778630495071
8.901087567210197e-06
0.006054648053832352
0.0008917823247611523
8.631939999759198e-06
0.004081738111563027

So it seems like making a scipy.linalg.solve_triangular call with sufficiently large inputs results in torch.bmm running much slower. The runtime difference may not seem too significant, but in our codebase, we call other torch functions and we observe much larger differences (e.g. 30ms vs. < 0.5ms if we call/do not call scipy).
I suspect the issue may have to do with the underlying linear algebra library that scipy calls (in the above case, the trtrs function in lapack/blas). Beyond this, I can’t explain why this runtime discrepancy occurs.
I’m wondering if someone might have an explanation for this?