Added early return for empty arrays in cuda qr kernel

chainer · Aug 20, 2019 · cdbd6be · cdbd6be
1 parent 1874f05
commit cdbd6be
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/chainerx_cc/chainerx/cuda/cuda_device/linalg.cu b/chainerx_cc/chainerx/cuda/cuda_device/linalg.cu
@@ -337,6 +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);