Port CPU torch.geqrf to ATen by IvanYashchuk · Pull Request #56249 · pytorch/pytorch

IvanYashchuk · 2021-04-16T09:50:24Z

Stack from ghstack:

Fix MAGMA qr for empty batched inputs #56257 Fix MAGMA qr for empty batched inputs
Add cuSOLVER path for torch.linalg.qr #56256 Add cuSOLVER path for torch.linalg.qr
Remove size arguments for internal orgqr and geqrf calls #56255 Remove size arguments for internal orgqr and geqrf calls
Add non-allocating helper function for torch.linalg.qr #56254 Add non-allocating helper function for torch.linalg.qr
Add cuBLAS path for batched torch.geqrf #56253 Add cuBLAS path for batched torch.geqrf
Add cuSOLVER path for torch.geqrf #56252 Add cuSOLVER path for torch.geqrf
Port CUDA torch.geqrf to ATen #56251 Port CUDA torch.geqrf to ATen
Update internal code for torch.geqrf #56250 Update internal code for torch.geqrf
Port CPU torch.geqrf to ATen #56249 Port CPU torch.geqrf to ATen

This PR ports torch.geqrf from TH to ATen. CUDA path will be
implemented in a follow-up PR.
With ATen port support for complex and batched inputs is added.
There were no correctness tests, they are
added in this PR and I added OpInfo for this operation.

We can implement the QR decomposition as a composition of geqrf and
orgqr (torch.linalg.householder_product).
Also we can implement the least squares solver with geqrf + ormqr +
trtrs. So it's useful to have this function renewed at least for the
internal code.

Resolves #24705

Differential Revision: D27907357

This PR ports `torch.geqrf` from TH to ATen. CUDA path will be implemented in a follow-up PR. With ATen port support for complex and batched inputs is added. There were no correctness tests, they are added in this PR and I added OpInfo for this operation. We can implement the QR decomposition as a composition of geqrf and orgqr (torch.linalg.householder_product). Also we can implement the least squares solver with geqrf + ormqr + trtrs. So it's useful to have this function renewed at least for the internal code. Resolves #24705 [ghstack-poisoned]

facebook-github-bot · 2021-04-16T09:50:36Z

💊 CI failures summary and remediations

As of commit dd13e03 (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.

This PR ports `torch.geqrf` from TH to ATen. CUDA path will be implemented in a follow-up PR. With ATen port support for complex and batched inputs is added. There were no correctness tests, they are added in this PR and I added OpInfo for this operation. We can implement the QR decomposition as a composition of geqrf and orgqr (torch.linalg.householder_product). Also we can implement the least squares solver with geqrf + ormqr + trtrs. So it's useful to have this function renewed at least for the internal code. Resolves pytorch#24705 ghstack-source-id: cd264d7 Pull Request resolved: pytorch#56249

This PR ports `torch.geqrf` from TH to ATen. CUDA path will be implemented in a follow-up PR. With ATen port support for complex and batched inputs is added. There were no correctness tests, they are added in this PR and I added OpInfo for this operation. We can implement the QR decomposition as a composition of geqrf and orgqr (torch.linalg.householder_product). Also we can implement the least squares solver with geqrf + ormqr + trtrs. So it's useful to have this function renewed at least for the internal code. Resolves #24705 [ghstack-poisoned]

This PR ports `torch.geqrf` from TH to ATen. CUDA path will be implemented in a follow-up PR. With ATen port support for complex and batched inputs is added. There were no correctness tests, they are added in this PR and I added OpInfo for this operation. We can implement the QR decomposition as a composition of geqrf and orgqr (torch.linalg.householder_product). Also we can implement the least squares solver with geqrf + ormqr + trtrs. So it's useful to have this function renewed at least for the internal code. Resolves pytorch#24705 ghstack-source-id: c663026 Pull Request resolved: pytorch#56249

mruberry · 2021-04-19T08:10:08Z

+static void apply_geqrf(const Tensor& self, const Tensor& tau, int64_t m, int64_t n) {
 #ifndef USE_LAPACK
-  AT_ERROR("qr: LAPACK library not found in compilation");
+  AT_ERROR("geqrf: LAPACK library not found in compilation");


TORCH_CHECK(false, ...)

We have a different error string for LAPACK not being available in the other linalg ops, right?

It's not consistent, some use the message as here and others use a nicer message. I'll change it here to a nicer message.

mruberry · 2021-04-19T08:14:06Z

+    def test_geqrf(self, device, dtype):
+
+        def run_test(shape):
+            # NumPy outputs the result of geqrf operation


This comment is a little confusing as written. It's that NumPy doesn't have a function named geqrf, but np.linalg.qr with mode='raw' computes the same operation, so this test compares against that function.

mruberry · 2021-04-19T08:15:13Z

+        batches = [(), (0, ), (2, ), (2, 1)]
+        ns = [5, 2, 0]
+        for batch, (m, n) in product(batches, product(ns, ns)):
+            # TODO: CUDA path doesn't work with batched or empty inputs


Assert this fails instead of just continuing so when the behavior changes the test demands an update, too

The test is updated in #56251. So can we leave this one as is? In other cases, I agree that it's better to assert failures and not skip the test.

I'll add the assert.

haha, really negotiating against yourself there

mruberry · 2021-04-19T08:16:21Z

-
-Rather, this directly calls the underlying LAPACK function `?geqrf`
+Computes a QR decomposition of :attr:`input`.
+Both `Q` and `R` matrices are stored in the same tensor `a`.


"in the same output tensor `a`."

mruberry · 2021-04-19T08:16:45Z

+The elements of `R` are stored on and above the diagonal.
+Elementary reflectors (or Householder vectors) implicitly defining matrix `Q`
+are stored below the diagonal.
+Result of this function can be used together with :func:`torch.linalg.householder_product`


"Result" -> "The results"

mruberry · 2021-04-19T08:18:27Z

+are stored below the diagonal.
+Result of this function can be used together with :func:`torch.linalg.householder_product`
+to obtain the `Q` matrix or
+with :func:`torch.ormqr` that uses implicit representation of matrix `Q` for efficient matrix-matrix multiplication.


"with :func:`torch.ormqr` that uses implicit representation of matrix `Q` for efficient matrix-matrix multiplication." -> "with :func:`torch.ormqr`, which uses an implicit representation of the `Q` matrix, for an efficient matrix-matrix multiplication."

mruberry · 2021-04-19T08:18:50Z

+to obtain the `Q` matrix or
+with :func:`torch.ormqr` that uses implicit representation of matrix `Q` for efficient matrix-matrix multiplication.
+
+This function directly calls the underlying LAPACK function `geqrf`


This sentence seems redundant - maybe it can be removed?

Yes, we can remove it.

mruberry · 2021-04-19T08:19:10Z

@@ -3368,24 +3368,35 @@ def merge_dicts(*dicts):
 This is a low-level function for calling LAPACK directly. This function


"for calling LAPACK's geqrf directly." ?

mruberry · 2021-04-19T08:20:32Z

+.. note::
+    To obtain explicit Q and R matrices it is recommended to use :func:`torch.linalg.qr`.
+
+.. note::


I would combine this note with the previous to make one "see also" note.

"See also :func:`torch.linalg.qr`, which computes explicit Q and R matrices, and :func:`torch.linalg.lstsq` with the ``driver="gels"`` option for a function that can solve matrix equations using a QR decomposition." ?

mruberry · 2021-04-19T08:22:26Z


+def sample_inputs_geqrf(op_info, device, dtype, requires_grad=False):
+    """
+    This function generates input for torch.geqrf


This comment doesn't hurt anything but I don't think it adds anything to the code, either

I agree, we can remove it.

mruberry · 2021-04-19T08:22:49Z

+    out = []
+    for batch, (m, n) in product(batches, product(ns, ns)):
+        # TODO: CUDA path doesn't work with batched or empty inputs
+        if 'cuda' in device and (batch != () or m == 0 or n == 0):


if torch.device(device).type == 'cuda'

mruberry · 2021-04-19T08:23:08Z

+    """
+    batches = [(), (0, ), (2, ), (1, 1)]
+    ns = [5, 2, 0]
+    out = []


"out" -> "samples"

mruberry

Overall this looks great, as usual; I've suggested a few tweaks and also pinged @ngimel since she's been reviewing most TH->ATen ports.

This PR ports `torch.geqrf` from TH to ATen. CUDA path will be implemented in a follow-up PR. With ATen port support for complex and batched inputs is added. There were no correctness tests, they are added in this PR and I added OpInfo for this operation. We can implement the QR decomposition as a composition of geqrf and orgqr (torch.linalg.householder_product). Also we can implement the least squares solver with geqrf + ormqr + trtrs. So it's useful to have this function renewed at least for the internal code. Resolves #24705 [ghstack-poisoned]

IvanYashchuk · 2021-04-19T12:08:29Z

@mruberry thank you for your suggestions! I've updated this pull request.

This PR ports `torch.geqrf` from TH to ATen. CUDA path will be implemented in a follow-up PR. With ATen port support for complex and batched inputs is added. There were no correctness tests, they are added in this PR and I added OpInfo for this operation. We can implement the QR decomposition as a composition of geqrf and orgqr (torch.linalg.householder_product). Also we can implement the least squares solver with geqrf + ormqr + trtrs. So it's useful to have this function renewed at least for the internal code. Resolves #24705 [ghstack-poisoned]

This PR ports `torch.geqrf` from TH to ATen. CUDA path will be implemented in a follow-up PR. With ATen port support for complex and batched inputs is added. There were no correctness tests, they are added in this PR and I added OpInfo for this operation. We can implement the QR decomposition as a composition of geqrf and orgqr (torch.linalg.householder_product). Also we can implement the least squares solver with geqrf + ormqr + trtrs. So it's useful to have this function renewed at least for the internal code. Resolves pytorch#24705 ghstack-source-id: 956f37f Pull Request resolved: pytorch#56249

This PR ports `torch.geqrf` from TH to ATen. CUDA path will be implemented in a follow-up PR. With ATen port support for complex and batched inputs is added. There were no correctness tests, they are added in this PR and I added OpInfo for this operation. We can implement the QR decomposition as a composition of geqrf and orgqr (torch.linalg.householder_product). Also we can implement the least squares solver with geqrf + ormqr + trtrs. So it's useful to have this function renewed at least for the internal code. Resolves #24705 [ghstack-poisoned]

This PR ports `torch.geqrf` from TH to ATen. CUDA path will be implemented in a follow-up PR. With ATen port support for complex and batched inputs is added. There were no correctness tests, they are added in this PR and I added OpInfo for this operation. We can implement the QR decomposition as a composition of geqrf and orgqr (torch.linalg.householder_product). Also we can implement the least squares solver with geqrf + ormqr + trtrs. So it's useful to have this function renewed at least for the internal code. Resolves pytorch#24705 ghstack-source-id: 3bbcb9f Pull Request resolved: pytorch#56249

ngimel · 2021-04-22T23:47:56Z

 }

+static void geqrf_out_helper(const Tensor& input, const Tensor& QR, const Tensor& tau) {
+  TORCH_INTERNAL_ASSERT(input.dim() >= 2);


Do you want to make the TORCH_INTERNAL_ASSERT_DEBUG_ONLY? You've already checked all these conditions, and this function is not user facing.

mruberry · 2021-04-23T03:21:10Z

@IvanYashchuk just ping me when you're happy with this PR, @IvanYashchuk, and I'll start merging this stack

facebook-github-bot · 2021-04-25T08:18:26Z

@mruberry merged this pull request in 58fcf77.

Summary: Pull Request resolved: pytorch#56249 This PR ports `torch.geqrf` from TH to ATen. CUDA path will be implemented in a follow-up PR. With ATen port support for complex and batched inputs is added. There were no correctness tests, they are added in this PR and I added OpInfo for this operation. We can implement the QR decomposition as a composition of geqrf and orgqr (torch.linalg.householder_product). Also we can implement the least squares solver with geqrf + ormqr + trtrs. So it's useful to have this function renewed at least for the internal code. Resolves pytorch#24705 Test Plan: Imported from OSS Reviewed By: ngimel Differential Revision: D27907357 Pulled By: mruberry fbshipit-source-id: 94e1806078977417e7903db76eab9d578305f585

IvanYashchuk requested a review from ezyang as a code owner April 16, 2021 09:50

facebook-github-bot added the cla signed label Apr 16, 2021

IvanYashchuk added module: linear algebra Issues related to specialized linear algebra operations in PyTorch; includes matrix multiply matmul module: porting Issues related to porting TH/THNN legacy to ATen native labels Apr 16, 2021

IvanYashchuk requested review from mruberry and removed request for ezyang April 16, 2021 09:58

pytorchbot added the open source label Apr 16, 2021

mruberry reviewed Apr 19, 2021

View reviewed changes

mruberry requested a review from ngimel April 19, 2021 08:23

mruberry reviewed Apr 19, 2021

View reviewed changes

ngimel reviewed Apr 22, 2021

View reviewed changes

ngimel approved these changes Apr 22, 2021

View reviewed changes

facebook-github-bot closed this in 58fcf77 Apr 25, 2021

facebook-github-bot added the Merged label Apr 25, 2021

facebook-github-bot deleted the gh/ivanyashchuk/10/head branch April 28, 2021 14:17

ngimel mentioned this pull request May 6, 2021

Roll-up: remaining TH functions #49421

Closed

14 tasks

		@@ -3368,24 +3368,35 @@ def merge_dicts(*dicts):
		This is a low-level function for calling LAPACK directly. This function

Conversation

IvanYashchuk commented Apr 16, 2021 • edited by mruberry Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

facebook-github-bot commented Apr 16, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

💊 CI failures summary and remediations

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mruberry left a comment

Choose a reason for hiding this comment

Uh oh!

IvanYashchuk commented Apr 19, 2021

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mruberry commented Apr 23, 2021

Uh oh!

facebook-github-bot commented Apr 25, 2021

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

IvanYashchuk commented Apr 16, 2021 •

edited by mruberry

Loading

facebook-github-bot commented Apr 16, 2021 •

edited

Loading