Add cuSOLVER path for torch.linalg.qr

IvanYashchuk · IvanYashchuk · commit 62ded9561e94 · 2021-04-29T10:52:07.000Z
Using cuSOLVER path with `pytest test/test_ops.py -k 'linalg_qr' --durations=5` cuts the runtime for these tests by 1 minute locally. Ref. pytorch#51552 ghstack-source-id: e94b357 Pull Request resolved: pytorch#56256
diff --git a/aten/src/ATen/native/BatchLinearAlgebra.cpp b/aten/src/ATen/native/BatchLinearAlgebra.cpp
@@ -1790,7 +1790,7 @@ void linalg_qr_out_helper(const Tensor& input, const Tensor& Q, const Tensor& R,
   orgqr_stub(input.device().type(), const_cast<Tensor&>(Q), tau);
 }
 
-std::tuple<Tensor, Tensor> _linalg_qr_helper_cpu(const Tensor& input, std::string mode) {
+std::tuple<Tensor, Tensor> _linalg_qr_helper_default(const Tensor& input, std::string mode) {
   bool compute_q, reduced_mode;
   std::tie(compute_q, reduced_mode) = _parse_qr_mode(mode);
   auto m = input.size(-2);
diff --git a/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu b/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu
@@ -2133,7 +2133,7 @@ AT_ERROR("qr: MAGMA library not found in "
 #endif
 }
 
-std::tuple<Tensor,Tensor> _linalg_qr_helper_cuda(const Tensor& self, std::string mode) {
+std::tuple<Tensor, Tensor> linalg_qr_helper_magma(const Tensor& self, std::string mode) {
   bool compute_q, reduced;
   std::tie(compute_q, reduced) = _parse_qr_mode(mode);
 
@@ -2178,6 +2178,16 @@ std::tuple<Tensor,Tensor> _linalg_qr_helper_cuda(const Tensor& self, std::string
   return std::make_tuple(q_working_copy, r_working_copy);
 }
 
+std::tuple<Tensor, Tensor> _linalg_qr_helper_cuda(const Tensor& input, std::string mode) {
+#if defined(USE_CUSOLVER)
+  // _linalg_qr_helper_default is a generic function that is implemented using
+  // geqrf_stub and orgqr_stub. It dispatches to cuSOLVER for CUDA inputs if USE_CUSOLVER is defined
+  return _linalg_qr_helper_default(input, mode);
+#else
+  return linalg_qr_helper_magma(input, mode);
+#endif
+}
+
 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ symeig ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 template <typename scalar_t>
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -8902,7 +8902,7 @@
 - func: _linalg_qr_helper(Tensor self, str mode) -> (Tensor, Tensor)
   variants: function
   dispatch:
-    CPU: _linalg_qr_helper_cpu
+    CPU: _linalg_qr_helper_default
     CUDA: _linalg_qr_helper_cuda
 
 - func: linalg_matrix_power(Tensor self, int n) -> Tensor

Original file line number	Diff line number	Diff line change
`@@ -1790,7 +1790,7 @@ void linalg_qr_out_helper(const Tensor& input, const Tensor& Q, const Tensor& R,`
`1790`	`1790`	`orgqr_stub(input.device().type(), const_cast<Tensor&>(Q), tau);`
`1791`	`1791`	`}`
`1792`	`1792`
`1793`		`-std::tuple<Tensor, Tensor> _linalg_qr_helper_cpu(const Tensor& input, std::string mode) {`
	`1793`	`+std::tuple<Tensor, Tensor> _linalg_qr_helper_default(const Tensor& input, std::string mode) {`
`1794`	`1794`	`bool compute_q, reduced_mode;`
`1795`	`1795`	`std::tie(compute_q, reduced_mode) = _parse_qr_mode(mode);`
`1796`	`1796`	`auto m = input.size(-2);`