Merge pull request pytorch#320 from yozw/trtrs

Exposing the lapack function trtrs which solves triangular systems of linear equations.

Author: soumith

generic/THLapack.c

@@ -5,6 +5,8 @@
 
 TH_EXTERNC void dgesv_(int *n, int *nrhs, double *a, int *lda, int *ipiv, double *b, int *ldb, int *info);
 TH_EXTERNC void sgesv_(int *n, int *nrhs, float *a, int *lda, int *ipiv, float *b, int *ldb, int *info);
+TH_EXTERNC void dtrtrs_(char *uplo, char *trans, char *diag, int *n, int *nrhs, double *a, int *lda, double *b, int *ldb, int *info);
+TH_EXTERNC void strtrs_(char *uplo, char *trans, char *diag, int *n, int *nrhs, float *a, int *lda, float *b, int *ldb, int *info);
 TH_EXTERNC void dgels_(char *trans, int *m, int *n, int *nrhs, double *a, int *lda, double *b, int *ldb, double *work, int *lwork, int *info);
 TH_EXTERNC void sgels_(char *trans, int *m, int *n, int *nrhs, float *a, int *lda, float *b, int *ldb, float *work, int *lwork, int *info);
 TH_EXTERNC void dsyev_(char *jobz, char *uplo, int *n, double *a, int *lda, double *w, double *work, int *lwork, int *info);
@@ -29,6 +31,7 @@ TH_EXTERNC void sorgqr_(int *m, int *n, int *k, float *a, int *lda, float *tau,
 TH_EXTERNC void dorgqr_(int *m, int *n, int *k, double *a, int *lda, double *tau, double *work, int *lwork, int *info);
 
 
+/* Compute the solution to a real system of linear equations A * X = B */
 void THLapack_(gesv)(int n, int nrhs, real *a, int lda, int *ipiv, real *b, int ldb, int* info)
 {
 #ifdef USE_LAPACK
@@ -43,6 +46,23 @@ void THLapack_(gesv)(int n, int nrhs, real *a, int lda, int *ipiv, real *b, int
  return;
 }
 
+/* Solve a triangular system of the form A * X = B or A^T * X = B */
+void THLapack_(trtrs)(char uplo, char trans, char diag, int n, int nrhs, real *a, int lda, real *b, int ldb, int* info)
+{
+#ifdef USE_LAPACK
+#if defined(TH_REAL_IS_DOUBLE)
+ dtrtrs_(&uplo, &trans, &diag, &n, &nrhs, a, &lda, b, &ldb, info);
+#else
+ strtrs_(&uplo, &trans, &diag, &n, &nrhs, a, &lda, b, &ldb, info);
+#endif
+#else
+ THError("trtrs : Lapack library not found in compile time\n");
+#endif
+ return;
+}
+
+/* Solve overdetermined or underdetermined real linear systems involving an
+M-by-N matrix A, or its transpose, using a QR or LQ factorization of A */
 void THLapack_(gels)(char trans, int m, int n, int nrhs, real *a, int lda, real *b, int ldb, real *work, int lwork, int *info)
 {
 #ifdef USE_LAPACK
@@ -56,6 +76,8 @@ void THLapack_(gels)(char trans, int m, int n, int nrhs, real *a, int lda, real
 #endif
 }
 
+/* Compute all eigenvalues and, optionally, eigenvectors of a real symmetric
+matrix A */
 void THLapack_(syev)(char jobz, char uplo, int n, real *a, int lda, real *w, real *work, int lwork, int *info)
 {
 #ifdef USE_LAPACK
@@ -69,6 +91,8 @@ void THLapack_(syev)(char jobz, char uplo, int n, real *a, int lda, real *w, rea
 #endif
 }
 
+/* Compute for an N-by-N real nonsymmetric matrix A, the eigenvalues and,
+optionally, the left and/or right eigenvectors */
 void THLapack_(geev)(char jobvl, char jobvr, int n, real *a, int lda, real *wr, real *wi, real* vl, int ldvl, real *vr, int ldvr, real *work, int lwork, int *info)
 {
 #ifdef USE_LAPACK
@@ -82,6 +106,8 @@ void THLapack_(geev)(char jobvl, char jobvr, int n, real *a, int lda, real *wr,
 #endif
 }
 
+/* Compute the singular value decomposition (SVD) of a real M-by-N matrix A,
+optionally computing the left and/or right singular vectors */
 void THLapack_(gesvd)(char jobu, char jobvt, int m, int n, real *a, int lda, real *s, real *u, int ldu, real *vt, int ldvt, real *work, int lwork, int *info)
 {
 #ifdef USE_LAPACK

generic/THLapack.h

@@ -4,6 +4,8 @@
 
 /* AX=B */
 TH_API void THLapack_(gesv)(int n, int nrhs, real *a, int lda, int *ipiv, real *b, int ldb, int* info);
+/* Solve a triangular system of the form A * X = B or A^T * X = B */
+TH_API void THLapack_(trtrs)(char uplo, char trans, char diag, int n, int nrhs, real *a, int lda, real *b, int ldb, int* info);
 /* ||AX-B|| */
 TH_API void THLapack_(gels)(char trans, int m, int n, int nrhs, real *a, int lda, real *b, int ldb, real *work, int lwork, int *info);
 /* Eigenvals */

generic/THTensorLapack.c

@@ -164,6 +164,90 @@ void THTensor_(gesv)(THTensor *rb_, THTensor *ra_, THTensor *b, THTensor *a)
  THIntTensor_free(ipiv);
 }
 
+void THTensor_(trtrs)(THTensor *rb_, THTensor *ra_, THTensor *b, THTensor *a,
+ const char *uplo, const char *trans, const char *diag)
+{
+ int n, nrhs, lda, ldb, info;
+ THTensor *ra__; // working version of A matrix to be passed into lapack TRTRS
+ THTensor *rb__; // working version of B matrix to be passed into lapack TRTRS
+
+ int clonea; // set to 1 if ra__ should be copied into ra_ at return
+ int cloneb; // set to 1 if rb__ should be copied into rb_ at return
+ int destroya; // set to 1 if ra__ needs to be destroyed at return
+ int destroyb; // set to 1 if rb__ needs to be destroyed at return
+
+
+ if (a == NULL || ra_ == a) /* possibly destroy the inputs */
+ {
+ THArgCheck(ra_->nDimension == 2, 1, "A should be 2 dimensional");
+ ra__ = THTensor_(new)();
+ clonea = THTensor_(lapackClone)(ra__,ra_,0);
+ destroya = 1;
+ }
+ else /*we want to clone and use ra_ as computational space*/
+ {
+ THArgCheck(a->nDimension == 2, 1, "A should be 2 dimensional");
+ clonea = THTensor_(lapackClone)(ra_,a,1);
+ ra__ = ra_;
+ destroya = 0;
+ }
+ if (b == NULL || rb_ == b) /* possibly destroy the inputs */
+ {
+ THArgCheck(rb_->nDimension == 2, 2, "B should be 2 dimensional");
+ rb__ = THTensor_(new)();
+ cloneb = THTensor_(lapackClone)(rb__,rb_,0);
+ destroyb = 1;
+ }
+ else /*we want to clone and use rb_ as computational space*/
+ {
+ THArgCheck(b->nDimension == 2, 2, "B should be 2 dimensional");
+ cloneb = THTensor_(lapackClone)(rb_,b,1);
+ rb__ = rb_;
+ destroyb = 0;
+ }
+
+ THArgCheck(ra__->nDimension == 2, 1, "A should be 2 dimensional");
+ THArgCheck(rb__->nDimension == 2, 2, "b should be 2 dimensional");
+ THArgCheck(ra__->size[0] == ra__->size[1], 1, "A should be square");
+ THArgCheck(rb__->size[0] == ra__->size[0], 2, "A,b size incompatible");
+
+ n = (int)ra__->size[0];
+ nrhs = (int)rb__->size[1];
+ lda = n;
+ ldb = n;
+
+ THLapack_(trtrs)(uplo[0], trans[0], diag[0], n, nrhs,
+ THTensor_(data)(ra__), lda,
+ THTensor_(data)(rb__), ldb, &info);
+
+ /* clean up */
+ if (destroya)
+ {
+ if (clonea)
+ {
+ THTensor_(copy)(ra_,ra__);
+ }
+ THTensor_(free)(ra__);
+ }
+ if (destroyb)
+ {
+ if (cloneb)
+ {
+ THTensor_(copy)(rb_,rb__);
+ }
+ THTensor_(free)(rb__);
+ }
+
+ if (info < 0)
+ {
+ THError("Lapack trtrs : Argument %d : illegal value", -info);
+ }
+ else if (info > 0)
+ {
+ THError("Lapack trtrs : A(%d,%d) is zero, singular A.", info,info);
+ }
+}
+
 void THTensor_(gels)(THTensor *rb_, THTensor *ra_, THTensor *b, THTensor *a)
 {
  // Note that a = NULL is interpreted as a = ra_, and b = NULL as b = rb_.

generic/THTensorLapack.h

@@ -3,6 +3,7 @@
 #else
 
 TH_API void THTensor_(gesv)(THTensor *rb_, THTensor *ra_, THTensor *b_, THTensor *a_);
+TH_API void THTensor_(trtrs)(THTensor *rb_, THTensor *ra_, THTensor *b_, THTensor *a_, const char *uplo, const char *trans, const char *diag);
 TH_API void THTensor_(gels)(THTensor *rb_, THTensor *ra_, THTensor *b_, THTensor *a_);
 TH_API void THTensor_(syev)(THTensor *re_, THTensor *rv_, THTensor *a_, const char *jobz, const char *uplo);
 TH_API void THTensor_(geev)(THTensor *re_, THTensor *rv_, THTensor *a_, const char *jobvr);