Fixed conditions for the early return in cuda qr kernel

chainer · Aug 21, 2019 · 0dbea62 · 0dbea62
1 parent cdbd6be
commit 0dbea62
Showing 1 changed file with 6 additions and 9 deletions.
diff --git a/chainerx_cc/chainerx/cuda/cuda_device/linalg.cu b/chainerx_cc/chainerx/cuda/cuda_device/linalg.cu
@@ -337,15 +337,12 @@ void QrImpl(const Array& a, const Array& q, const Array& r, const Array& tau, Qr
     int64_t k = std::min(m, n);
     int64_t lda = std::max(int64_t{1}, m);
 
-    // Older versions of cuSOLVER (<10.1) might not work well with zero sized arrays
-    // therefore it's better to return earlier
-    if (a.shape().GetTotalSize() == 0 && mode != QrMode::kComplete) {
-        return;
-    }
-
-    // cuSOLVER does not return correct result in this case
-    if (mode == QrMode::kComplete && n == 0) {
-        device.backend().CallKernel<IdentityKernel>(q);
+    // cuSOLVER does not return correct result in this case and older versions of cuSOLVER (<10.1)
+    // might not work well with zero-sized arrays therefore it's better to return earlier
+    if (a.shape().GetTotalSize() == 0) {
+        if  (mode == QrMode::kComplete) {
+            device.backend().CallKernel<IdentityKernel>(q);
+        }
         return;
     }