feat: added economic QR (#903)

* feat: added economic QR

* fix memory alloc in gesvdjBatched

Author: mfzmullen

docs_input/api/linalg/decomp/qr.rst

@@ -8,7 +8,9 @@ Perform a QR decomposition.
 .. doxygenfunction:: qr
 
 .. note::
-   This function is currently not supported with host-based executors (CPU)
+   This function is currently not supported with host-based executors (CPU), and performs a full QR 
+   decomposition of a tensor `A` with shape `... x m x n`, where `Q` is shaped `... x m x m` and `R`
+   is shaped `... x m x n`.
 
 Examples
 ~~~~~~~~
@@ -19,12 +21,27 @@ Examples
    :end-before: example-end qr-test-1
    :dedent:
 
+.. doxygenfunction:: qr_econ
+
+.. note::
+   This function is currently not supported with host-based executors (CPU). It returns an economic 
+   QR decomposition, where `Q/R` are shaped `m x k` and `k x n` respectively, where `k = min(m, n)`. 
+   This is useful when `m >> n` to save memory and computation time.
+
+Examples
+~~~~~~~~
+
+.. literalinclude:: ../../../../test/00_solver/QREcon.cu
+   :language: cpp
+   :start-after: example-begin qr-econ-test-1
+   :end-before: example-end qr-econ-test-1
+   :dedent:
 
 .. doxygenfunction:: qr_solver
 
 .. note::
    This function does not return `Q` explicitly as it only runs :literal:`geqrf` from LAPACK/cuSolver.
-   For full `Q/R`, use :literal:`qr_solver` on a CUDA executor.
+   For full or economic `Q/R`, use :literal:`qr` or :literal:`qr_econ` on a CUDA executor.
 
 Examples
 ~~~~~~~~
@@ -33,4 +50,5 @@ Examples
    :language: cpp
    :start-after: example-begin qr_solver-test-1
    :end-before: example-end qr_solver-test-1
-   :dedent:
+   :dedent:
+

include/matx/operators/qr.h

@@ -184,4 +184,76 @@ __MATX_INLINE__ auto qr_solver(const OpA &a) {
   return detail::SolverQROp(a);
 }
 
+
+namespace detail {
+  template<typename OpA>
+  class EconQROp : public BaseOp<EconQROp<OpA>>
+  {
+    private:
+      typename detail::base_type_t<OpA> a_;
+
+    public:
+      using matxop = bool;
+      using value_type = typename OpA::value_type;
+      using matx_transform_op = bool;
+      using qr_solver_xform_op = bool;
+
+      __MATX_INLINE__ std::string str() const { return "qr_econ()"; }
+      __MATX_INLINE__ EconQROp(const OpA &a) : a_(a) { }    
+
+      // This should never be called
+      template <typename... Is>
+      __MATX_INLINE__ __MATX_DEVICE__ __MATX_HOST__ decltype(auto) operator()(Is... indices) const = delete;
+
+      template <typename Out, typename Executor>
+      void Exec(Out &&out, Executor &&ex) {
+        static_assert(cuda::std::tuple_size_v<remove_cvref_t<Out>> == 3, "Must use mtie with 2 outputs on qr_econ(). ie: (mtie(Q, R) = qr_econ(A))");     
+
+        qr_econ_impl(cuda::std::get<0>(out), cuda::std::get<1>(out), a_, ex);
+      }
+
+      static __MATX_INLINE__ constexpr __MATX_HOST__ __MATX_DEVICE__ int32_t Rank()
+      {
+        return OpA::Rank();
+      }
+
+      template <typename ShapeType, typename Executor>
+      __MATX_INLINE__ void PreRun([[maybe_unused]] ShapeType &&shape, Executor &&ex) noexcept
+      {
+        if constexpr (is_matx_op<OpA>()) {
+          a_.PreRun(std::forward<ShapeType>(shape), std::forward<Executor>(ex));
+        }
+      }
+
+      // Size is not relevant in qr_solver() since there are multiple return values and it
+      // is not allowed to be called in larger expressions
+      constexpr __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ index_t Size(int dim) const
+      {
+        return a_.Size(dim);
+      }
+
+  };
+}
+
+/**
+ * Perform an economic QR decomposition on a matrix using cuSolver.
+ * 
+ * If rank > 2, operations are batched.
+ * 
+ * @tparam OpA
+ *   Data type of input a tensor or operator
+ *
+ * @param a
+ *   Input tensor or operator of shape `... x m x n`
+ * 
+ * @return
+ *   Operator that produces QR outputs.
+ *   - **Q** - Of shape `... x m x min(m, n)`, the reduced orthonormal basis for the span of A.
+ *   - **R** - Upper triangular matrix of shape  `... x min(m, n) x n`.
+ */
+template<typename OpA>
+__MATX_INLINE__ auto qr_econ(const OpA &a) {
+  return detail::EconQROp(a);
 }
+
+}