Repository: numba/numba

Issue:
GPU gufuncs don't broadcast properly
Both calls to `matmult` in the code below should produce the same result, if I am correctly understanding the function of the broadcasting rules; and indeed, they do so for `target='cpu'`. But for `target='cuda'` I get `ValueError: reshape changes the size of the array`.

I've done some digging and I find that it is caused by the `_broadcast` method in `devicearray.py` attempting to reshape `Y` from `(3,3)` to `(10,3,3)`. Of course, it can do that, _if_ it uses a stride of 0 in that first dimension; but it's just attempting a simple `reshape`, and that fails. Earlier code does detect that `Y` should be "pinned" instead.

```
import numba as nb
import numpy as np

@nb.guvectorize('(f8[:,:],f8[:,:],f8[:,:])','(m,n),(n,p)->(m,p)',target='cuda')
def matmul(X,Y,Z):
    for i in range(X.shape[0]):
        for j in range(Y.shape[1]):
            Z[i,j] = 0.0
            for k in range(X.shape[1]):
                Z[i,j] += X[i,k] * Y[k,j]

X = np.random.randn(10,3,3)
Y = np.random.randn(3,3)
print(matmul(X,np.tile(Y,(10,1,1))))
print(matmul(X,Y))
```

GPU gufuncs don't broadcast properly
Both calls to `matmult` in the code below should produce the same result, if I am correctly understanding the function of the broadcasting rules; and indeed, they do so for `target='cpu'`. But for `target='cuda'` I get `ValueError: reshape changes the size of the array`.

I've done some digging and I find that it is caused by the `_broadcast` method in `devicearray.py` attempting to reshape `Y` from `(3,3)` to `(10,3,3)`. Of course, it can do that, _if_ it uses a stride of 0 in that first dimension; but it's just attempting a simple `reshape`, and that fails. Earlier code does detect that `Y` should be "pinned" instead.

```
import numba as nb
import numpy as np

@nb.guvectorize('(f8[:,:],f8[:,:],f8[:,:])','(m,n),(n,p)->(m,p)',target='cuda')
def matmul(X,Y,Z):
    for i in range(X.shape[0]):
        for j in range(Y.shape[1]):
            Z[i,j] = 0.0
            for k in range(X.shape[1]):
                Z[i,j] += X[i,k] * Y[k,j]

X = np.random.randn(10,3,3)
Y = np.random.randn(3,3)
print(matmul(X,np.tile(Y,(10,1,1))))
print(matmul(X,Y))
```



Discussion:

Resolving patch:
diff --git a/numba/cuda/dispatcher.py b/numba/cuda/dispatcher.py
--- a/numba/cuda/dispatcher.py
+++ b/numba/cuda/dispatcher.py
@@ -183,6 +183,15 @@ def _broadcast_scalar_input(self, ary, shape):
                                          dtype=ary.dtype,
                                          gpu_data=ary.gpu_data)
 
+    def _broadcast_add_axis(self, ary, newshape):
+        newax = len(newshape) - len(ary.shape)
+        # Add 0 strides for missing dimension
+        newstrides = (0,) * newax + ary.strides
+        return devicearray.DeviceNDArray(shape=newshape,
+                                         strides=newstrides,
+                                         dtype=ary.dtype,
+                                         gpu_data=ary.gpu_data)
+
 
 class CUDAUFuncMechanism(UFuncMechanism):
     """
diff --git a/numba/npyufunc/deviceufunc.py b/numba/npyufunc/deviceufunc.py
--- a/numba/npyufunc/deviceufunc.py
+++ b/numba/npyufunc/deviceufunc.py
@@ -694,10 +694,29 @@ def _broadcast(self, schedule, params, retval):
                 newparams.append(devary)
             else:
                 # Broadcast vector input
-                newparams.append(p.reshape(odim, *cs))
+                newparams.append(self._broadcast_array(p, odim, cs))
         newretval = retval.reshape(odim, *schedule.oshapes[0])
         return newparams, newretval
 
+    def _broadcast_array(self, ary, newdim, innerdim):
+        newshape = (newdim,) + innerdim
+        # No change in shape
+        if ary.shape == newshape:
+            return ary
+
+        # Creating new dimension
+        elif len(ary.shape) < len(newshape):
+            assert newshape[-len(ary.shape):] == ary.shape, \
+               "cannot add dim and reshape at the same time"
+            return self._broadcast_add_axis(ary, newshape)
+
+        # Collapsing dimension
+        else:
+            return ary.reshape(*newshape)
+
+    def _broadcast_add_axis(self, ary, newshape):
+        raise NotImplementedError("cannot add new axis")
+
     def _broadcast_scalar_input(self, ary, shape):
         raise NotImplementedError
 


Test patch:
diff --git a/numba/cuda/tests/cudapy/test_gufunc.py b/numba/cuda/tests/cudapy/test_gufunc.py
--- a/numba/cuda/tests/cudapy/test_gufunc.py
+++ b/numba/cuda/tests/cudapy/test_gufunc.py
@@ -1,5 +1,5 @@
 from __future__ import print_function, absolute_import
-from numba import void, float32
+from numba import void, float32, float64
 import numpy as np
 import numpy.core.umath_tests as ut
 from numba import guvectorize
@@ -161,6 +161,33 @@ def matmulcore(A, B, C):
 
         self.assertTrue(np.allclose(C, Gold))
 
+    def test_gufunc_new_axis(self):
+
+        @guvectorize([void(float64[:, :], float64[:, :], float64[:, :])],
+                     '(m,n),(n,p)->(m,p)',
+                     target='cuda')
+        def matmulcore(A, B, C):
+            m, n = A.shape
+            n, p = B.shape
+            for i in range(m):
+                for j in range(p):
+                    C[i, j] = 0
+                    for k in range(n):
+                        C[i, j] += A[i, k] * B[k, j]
+
+        gufunc = matmulcore
+
+        X = np.random.randn(10, 3, 3)
+        Y = np.random.randn(3, 3)
+
+        gold = ut.matrix_multiply(X, Y)
+
+        res1 = gufunc(X, Y)
+        np.testing.assert_allclose(gold, res1)
+
+        res2 = gufunc(X, np.tile(Y, (10, 1, 1)))
+        np.testing.assert_allclose(gold, res2)
+
     def test_gufunc_adjust_blocksize(self):
 
         @guvectorize([void(float32[:, :], float32[:, :], float32[:, :])],
@@ -305,7 +332,7 @@ def foo(A, B):
         items = msg[len(head):].strip().split(',')
         items = [i.strip("'\" ") for i in items]
         self.assertEqual(set(['what1', 'ever2']), set(items))
-        
+
 
 if __name__ == '__main__':
     unittest.main()