Port cholesky_inverse to ATen

[IvanYashchuk]

Now we can remove _th_potri!

Compared to the original TH-based cholesky_inverse, complex (#33152) and batched inputs (#7500) are now supported both on CPU and CUDA.

Closes #24685.
Closes #24543.

Ref. #49421, #42666

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit f117e84 (more details on the Dr. CI page):

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions to the (internal) Dr. CI Users group.

[IvanYashchuk]

Benchmarks for torch.float64:
Before this PR:

[-------------------- cholesky_inverse (TH) torch.float64 --------------------]
                              |  cholesky_inverse CUDA  |  cholesky_inverse CPU
 --------------------------------------------------------------------
      torch.Size([32, 32])    |          5198.2         |            7.4       
      torch.Size([512, 512])  |          8669.2         |         5555.0   

After:

[---------------------- cholesky_inverse (ATen) torch.float64 -----------------------]
                                   |  cholesky_inverse CUDA  |  cholesky_inverse CPU
 -------------------------------------------------------------------------
      torch.Size([32, 32])         |           5402.7        |            16.2      
      torch.Size([512, 512])       |           6949.7        |          7204.2   

      torch.Size([1, 32, 32])      |             84.1        |            16.4      
      torch.Size([32, 32, 32])     |             87.7        |           241.9      
      torch.Size([256, 32, 32])    |            191.9        |          1627.3         
      torch.Size([1, 512, 512])    |           2491.6        |          7307.1      
      torch.Size([32, 512, 512])   |          48022.1        |        180274.9      
      torch.Size([256, 512, 512])  |         375196.1        |       1432507.8      

Times are in microseconds (us).

[mruberry]

So this PR pretty significantly regresses performance on both CUDA and CPU? That's worrying. What do you think is happening and can we address this?

cc @ngimel

[mruberry]

Wait -- doesn't this PR implemented batched CUDA, though?

[IvanYashchuk]

Oh, this comment is a bit outdated. I need to change it. At that time I thought that batched MAGMA routine would require infos to live on GPU.
Now CUDA path is implemented in terms of MAGMA's apply_cholesky_solve. It is different from other batched functions in that it doesn't take an array of ints infos argument, it doesn't raise any errors related to the algorithm. It returns only a single integer to tell whether all passed arguments were ok or not.
If cuSOLVER would be used instead in the future, then infos would need to be created on GPU.

I will change it to:

• if input on CPU then we need infos of size batchsize(input) to fill with error codes from LAPACK for each matrix in the batched tensor.
• if input on GPU then we need only one integer living on CPU for storing error code from MAGMA

It also means that for some inputs CPU could raise an error while GPU would output something because the operations used are different.

[mruberry]

It also means that for some inputs CPU could raise an error while GPU would output something because the operations used are different.

Are there some specific cases you're thinking of?

[IvanYashchuk]

Checked the docs. It will happen for inputs with zero on the diagonal. CPU version raises an error with the index of the zero diagonal element, while CUDA gives inf. I need to add this to tests.

In [1]: import torch

In [2]: a = torch.randn(2, 2)

In [3]: a
Out[3]: 
tensor([[-0.6556, -0.4479],
        [-0.9347,  0.8169]])

In [4]: a[1, 1] = 0

In [5]: a
Out[5]: 
tensor([[-0.6556, -0.4479],
        [-0.9347,  0.0000]])

In [6]: torch.cholesky_inverse(a)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-6-680561e9e67c> in <module>
----> 1 torch.cholesky_inverse(a)

RuntimeError: Lapack Error potri : A(2,2) is 0, A cannot be factorized at ../aten/src/TH/generic/THTensorLapack.cpp:245

In [7]: torch.cholesky_solve(torch.eye(2), a)
Out[7]: 
tensor([[inf, -inf],
        [-inf, inf]])

In [8]: torch.cholesky_solve(torch.eye(2, device='cuda'), a.cuda())
Out[8]: 
tensor([[inf, -inf],
        [-inf, inf]], device='cuda:0')

[mruberry]

Testing (and, in the near future, documenting) this behavior sounds great.

[mruberry]

Great work, @IvanYashchuk. There are a couple follow-ups (adjustment based on future OpInfo fixes, docs) as @anjali411 points out, but I think they're separable. I've added this to the list of operators to review in our 1.8 scrub.

When you're happy with this PR ping me and let's merge it.

[facebook-github-bot]

@anjali411 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[anjali411]

@IvanYashchuk please rebase the PR and let me know once it's ready for merge.

[IvanYashchuk]

@anjali411 I resolved the conflict.

[facebook-github-bot]

@anjali411 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[facebook-github-bot]

@anjali411 merged this pull request in 6e4746c.

[dncliss]

After this issue was merged, the builds on ppc64le platform started failing again, and for the same reason they had
been failing just recently before #51217 was merged shortly before this. That merge fixed the issue, and then this merge
introduced it again.
I believe it probably simply means that this feature was missing an equivalent line in aten/src/ATen/native/BatchLinearAlgebraKernel.cpp at about line 177 .... that most likely the following line should be present (compare to the REGISTER_VSX_DISPATCH lines shortly below this point):

REGISTER_VSX_DISPATCH(cholesky_inverse_stub, &cholesky_inverse_kernel_impl);

[anjali411]

@dncliss thanks for reporting the issue, and yeah the fix should be similar. @IvanYashchuk could you create a fix for this?

[dncliss]

@anjali411 Thanks for confirming. I didn't open a new issue (I figure maybe @IvanYashchuk can just push the 1-line fix) but if you want it reported that way I could do so.

For reference, the regular ppc64le builds seen here: https://powerci.osuosl.org/job/pytorch-master-nightly-py3-linux-ppc64le-gpu/ older than build 1047 had the error fixed by PR 51217, then 1047 itself was a successful build, and 1048 onward introduced the failure after the merging of this cholesky_inverse feature PR. The error, as before, is an "undefined reference" during linking.

[IvanYashchuk]

@anjali411, @dncliss I submitted the fix. Sorry for the trouble!