Fix linalg_potri and linalg_potrf operators for large tensor. (#18752)

* Fix linalg_potri operator for large tensor. * Update other variables to support large tensors. * Add to contributors. * Fix whitespace. * Update ZeroTriangular to support large tensors. * Add large tensor unit tests for linalg_potrf and linalg_potri. * Fix crash when accessing already destructed static variables (#18768) (#18778) Co-authored-by: Joe Evans <[email protected]> Co-authored-by: Przemyslaw Tredak <[email protected]>
apache · Jul 24, 2020 · e6de5ae · e6de5ae
1 parent 91d535a
commit e6de5ae
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 5 deletions.
diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md
@@ -251,6 +251,7 @@ List of Contributors
 * [Piljae Chae](https://github.com/IHateMint)
 * [Oliver Kowalke](https://github.com/olk)
 * [Connor Goggins](https://github.com/connorgoggins)
+* [Joe Evans](https://github.com/josephevans)
 
 Label Bot
 ---------

diff --git a/src/operator/tensor/la_op-inl.h b/src/operator/tensor/la_op-inl.h
@@ -36,9 +36,10 @@ using namespace mshadow;
 // Copies lower/upper triangular part to upper/lower, i.e. to the opposite side.
 struct CopyTriangularToOppositeSide {
   template<typename DType>
-  MSHADOW_XINLINE static void Map(int i, int matrix_size, int stride, DType* data, bool to_lower) {
+  MSHADOW_XINLINE static void Map(index_t i, size_t matrix_size, index_t stride,
+                                  DType* data, bool to_lower) {
     // Below computation works even when we are dealing with a batch of matrices.
-    const int row((i % matrix_size) / stride), col(i % stride);
+    const index_t row((i % matrix_size) / stride), col(i % stride);
     if (row > col) {
        if (to_lower) {
          data[i] = data[i + (col - row) * (stride - 1)];
@@ -52,9 +53,9 @@ struct CopyTriangularToOppositeSide {
 // Zero's lower/upper triangular part of a matrix.
 struct ZeroTriangular {
   template<typename DType>
-  MSHADOW_XINLINE static void Map(int i, int matrix_size, int stride, DType* data,
-                                  bool zero_lower) {
-    const int row((i % matrix_size) / stride), col(i % stride);
+  MSHADOW_XINLINE static void Map(index_t i, size_t matrix_size, index_t stride,
+                                  DType* data, bool zero_lower) {
+    const index_t row((i % matrix_size) / stride), col(i % stride);
     if ((!zero_lower && (row < col)) || (zero_lower && (row > col))) data[i] = 0;
   }
 };

diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py
@@ -37,6 +37,7 @@
 LARGE_X = 100000000
 SMALL_X = 100
 SMALL_Y = 50
+LARGE_SQ_X = 80000
 LARGE_SIZE = LARGE_X * SMALL_Y
 LARGE_TENSOR_SHAPE = 2**32
 RNN_LARGE_TENSOR = 2**28
@@ -1167,6 +1168,32 @@ def check_correctness(mxnet_op, numpy_op, atol=1e-3):
     check_gather()
     check_binary_broadcast()
 
+def test_linalg():
+    def check_potrf():
+        # creating an identity matrix input
+        A = nd.zeros((LARGE_SQ_X, LARGE_SQ_X))
+        for i in range(LARGE_SQ_X):
+            A[i,i] = 1
+
+        out = nd.linalg.potrf(A)
+        # output should be an identity matrix
+        for i in range(LARGE_SQ_X):
+            assert out[i,i] == 1
+
+    def check_potri():
+        # creating an identity matrix input
+        A = nd.zeros((LARGE_SQ_X, LARGE_SQ_X))
+        for i in range(LARGE_SQ_X):
+            A[i,i] = 1
+
+        out = nd.linalg.potri(A)
+        # output should be an identity matrix
+        for i in range(LARGE_SQ_X):
+            assert out[i,i] == 1
+
+    check_potrf()
+    check_potri()
+
 
 def test_basic():
     def check_elementwise():