Improve the performance of btriunpack

[bamos]

The part of btriunpack that extract pivots has been causing some unexpected performance bottlenecks in qpth. Here's a newer version I've tried that uses gather/scatter operations across a batched vector instead of row interchanges on a batched matrix. I think it's a step towards a better method but the current form is just as slow. I want to do what the LAPACK LASWP function provides but with a batch so maybe we could use some knowledge from those implementations, like this one in OpenBLAS.

Slightly improved pivot matrix extraction but still slow version using gather/scatter

Pidx = type(LU_data)(range(sz)).repeat(nBatch, 1).long()

for i in range(sz):
    k = LU_pivots[:, i] - 1
    t = Pidx[:, i].clone()
    Pidx[:, i] = torch.gather(Pidx, 1, k.unsqueeze(1).long())
    Pidx.scatter_(1, k.unsqueeze(1).long(), t.unsqueeze(1))

P = type(LU_data)(nBatch, sz, sz).zero_()
for i in range(nBatch):
    P[i].scatter_(0, Pidx[i].unsqueeze(0), 1.0)

[apaszke]

So it seems to me that a simple way to improve the perf is to:

1. get rid of all masks
   1a. do batched triu, tril (should be fairly simple to add)
   1b. do L[::L.size(1) + 1] = 1 to fill the diagonal (zero memory copy)
2. use scatter_ (perhaps it's possible to batch the ops using some tricks)

[IssamLaradji]

Was this fixed ? is it possible to do batched triu, tril now ? Cheers!

[apaszke]

I'm not 100% sure, but I think we haven't changed the code since then. Is this function a bottleneck for you?

[IssamLaradji]

@apaszke I was hoping to use batch .triu for a different algorithm. I am implementing Variational Gaussian Processes and one of the requirements is obtaining triu for n matrices. I am currently using the following code where I am applying .triu to qVar across its first axis (qVar is n x m x m):

def apply(func, M):
    tList = [func(m) for m in torch.unbind(M, dim=0)]
    res = torch.stack(tList, dim=0)
    return res 

apply(torch.triu,  qVar)

but as n gets larger, this operation becomes a large bottleneck. Are there plans to have a generic batched function mode ? In bayesian methods it is often necessary to apply batch functions such as Inverse, cholesky, triu, etc. Sorry for posting this here which is not completely relevant to btriunpack, I might have to open a new issue for this. Thanks!

[apaszke]

This function seems fairly inefficient. How about this:

B, N = 4, 10
qVar = torch.triu(torch.Tensor(N, N)).expand(B, N, N)

[IssamLaradji]

Good idea! So I could define indices as ind = torch.triu(torch.ones(N, N)).expand(B, N, N), then to extract the upper triangles for each of the B qVar matrices I could do:
qVar = qVar * ind

[vishwakftw]

Closing in favour of #15286; this should have been addressed in the same.