Merge pull request #281 from abergeron/err_blas

Check for overflow in cublas wrapper

Author: nouiz

pygpu/gpuarray.pyx

@@ -1978,7 +1978,10 @@ cdef class GpuArray:
         return str(numpy.asarray(self))
 
     def __repr__(self):
-        return 'gpuarray.' + repr(numpy.asarray(self))
+        try:
+            return 'gpuarray.' + repr(numpy.asarray(self))
+        except Exception:
+            return 'gpuarray.array(<content not available>)'
 
 
 

src/gpuarray/error.h

@@ -34,6 +34,7 @@ enum ga_error {
   GA_NODEV_ERROR,
   GA_MISC_ERROR,
   GA_COMM_ERROR,
+  GA_XLARGE_ERROR,
   /* Add more error types if needed, but at the end */
   /* Don't forget to sync with Gpu_error() */
 };

src/gpuarray_blas_cuda_cublas.c

@@ -5,7 +5,9 @@
 #include "gpuarray/kernel.h"
 #include "gpuarray/error.h"
 
-#include "cublas_v2.h"
+#include <limits.h>
+
+#include <cublas_v2.h>
 
 extern const gpuarray_buffer_ops cuda_ops;
 
@@ -33,6 +35,8 @@ typedef struct _blas_handle {
   cublasStatus_t err;
 } blas_handle;
 
+#define LARGE_VAL(v) (v >= INT_MAX)
+
 static const char *code_sgemvBH_N_a1_b1_small =                         \
   "extern \"C\"__global__ void sgemv(const float *A[], size_t lda, "    \
   "                                  const float *x[], size_t incx, "   \
@@ -326,6 +330,11 @@ static int sgemm(cb_order order, cb_transpose transA, cb_transpose transB,
   ASSERT_BUF(B);
   ASSERT_BUF(C);
 
+  if (LARGE_VAL(M) || LARGE_VAL(N) || LARGE_VAL(K) ||
+      LARGE_VAL(lda) || LARGE_VAL(ldb) || LARGE_VAL(ldc) ||
+      LARGE_VAL(M * N) || LARGE_VAL(M * K) || LARGE_VAL(K * N))
+    return GA_XLARGE_ERROR;
+
   if (order == cb_c) {
     /* swap A and B */
     t = N;
@@ -386,6 +395,11 @@ static int dgemm(cb_order order, cb_transpose transA, cb_transpose transB,
   ASSERT_BUF(B);
   ASSERT_BUF(C);
 
+  if (LARGE_VAL(M) || LARGE_VAL(N) || LARGE_VAL(K) ||
+      LARGE_VAL(lda) || LARGE_VAL(ldb) || LARGE_VAL(ldc) ||
+      LARGE_VAL(M * N) || LARGE_VAL(M * K) || LARGE_VAL(K * N))
+    return GA_XLARGE_ERROR;
+
   if (order == cb_c) {
     /* swap A and B */
     t = N;
@@ -450,6 +464,11 @@ static int hgemm(cb_order order, cb_transpose transA, cb_transpose transB,
   ASSERT_BUF(B);
   ASSERT_BUF(C);
 
+  if (LARGE_VAL(M) || LARGE_VAL(N) || LARGE_VAL(K) ||
+      LARGE_VAL(lda) || LARGE_VAL(ldb) || LARGE_VAL(ldc) ||
+      LARGE_VAL(M * N) || LARGE_VAL(M * K) || LARGE_VAL(K * N))
+    return GA_XLARGE_ERROR;
+
   if (order == cb_c) {
     /* swap A and B */
     t = N;
@@ -539,6 +558,11 @@ static int sgemmBatch(cb_order order, cb_transpose transA, cb_transpose transB,
 
   if (batchCount == 0) return GA_NO_ERROR;
 
+  if (LARGE_VAL(M) || LARGE_VAL(N) || LARGE_VAL(K) ||
+      LARGE_VAL(lda) || LARGE_VAL(ldb) || LARGE_VAL(ldc) ||
+      LARGE_VAL(M * N) || LARGE_VAL(M * K) || LARGE_VAL(K * N))
+    return GA_XLARGE_ERROR;
+
   ASSERT_BUF(A[0]);
   ctx = A[0]->ctx;
   h = (blas_handle *)ctx->blas_handle;
@@ -659,6 +683,11 @@ static int dgemmBatch(cb_order order, cb_transpose transA, cb_transpose transB,
 
   if (batchCount == 0) return GA_NO_ERROR;
 
+  if (LARGE_VAL(M) || LARGE_VAL(N) || LARGE_VAL(K) ||
+      LARGE_VAL(lda) || LARGE_VAL(ldb) || LARGE_VAL(ldc) ||
+      LARGE_VAL(M * N) || LARGE_VAL(M * K) || LARGE_VAL(K * N))
+    return GA_XLARGE_ERROR;
+
   ASSERT_BUF(A[0]);
   ctx = A[0]->ctx;
   h = (blas_handle *)ctx->blas_handle;
@@ -782,6 +811,10 @@ static int sgemv(cb_order order, cb_transpose transA, size_t M, size_t N,
   ASSERT_BUF(X);
   ASSERT_BUF(Y);
 
+  if (LARGE_VAL(M) || LARGE_VAL(N) || LARGE_VAL(M * N) ||
+      LARGE_VAL(lda) || LARGE_VAL(incX) || LARGE_VAL(incY))
+    return GA_XLARGE_ERROR;
+
   if (order == cb_c) {
     t = N;
     N = M;
@@ -833,6 +866,10 @@ static int dgemv(cb_order order, cb_transpose transA, size_t M, size_t N,
   ASSERT_BUF(X);
   ASSERT_BUF(Y);
 
+  if (LARGE_VAL(M) || LARGE_VAL(N) || LARGE_VAL(M * N) ||
+      LARGE_VAL(lda) || LARGE_VAL(incX) || LARGE_VAL(incY))
+    return GA_XLARGE_ERROR;
+
   if (order == cb_c) {
     t = N;
     N = M;
@@ -1149,6 +1186,10 @@ static int sger(cb_order order, size_t M, size_t N, float alpha, gpudata *X,
   ASSERT_BUF(Y);
   ASSERT_BUF(A);
 
+  if (LARGE_VAL(M) || LARGE_VAL(N) || LARGE_VAL(M * N) ||
+      LARGE_VAL(lda) || LARGE_VAL(incX) || LARGE_VAL(incY))
+    return GA_XLARGE_ERROR;
+
   if (order == cb_c) {
     t = M;
     M = N;
@@ -1202,6 +1243,10 @@ static int dger(cb_order order, size_t M, size_t N, double alpha, gpudata *X,
   ASSERT_BUF(Y);
   ASSERT_BUF(A);
 
+  if (LARGE_VAL(M) || LARGE_VAL(N) || LARGE_VAL(M * N) ||
+      LARGE_VAL(lda) || LARGE_VAL(incX) || LARGE_VAL(incY))
+    return GA_XLARGE_ERROR;
+
   if (order == cb_c) {
     t = M;
     M = N;

src/gpuarray_error.c

@@ -23,6 +23,7 @@ const char *gpuarray_error_str(int err) {
   case GA_NODEV_ERROR:       return "No devices are available";
   case GA_MISC_ERROR:        return "Undeterminate error";
   case GA_COMM_ERROR:        return "Error in collectives call";
+  case GA_XLARGE_ERROR:      return "Input size too large for operation";
   default: return "Unknown GA error";
   }
 }