Merge pull request #419 from khaotik/gemmb3d_py

python interface for rgemmBatch_3d

Author: abergeron

pygpu/blas.pyx

@@ -18,6 +18,9 @@ cdef extern from "gpuarray/blas.h":
                        double beta, _GpuArray *C, int nocopy)
     int GpuArray_rger(double alpha, _GpuArray *X, _GpuArray *Y, _GpuArray *A,
                       int nocopy)
+    int GpuArray_rgemmBatch_3d(
+        cb_transpose transA, cb_transpose transB, double alpha,
+        _GpuArray *A, _GpuArray *B, double beta, _GpuArray *C, int nocopy)
 
 cdef api int pygpu_blas_rdot(GpuArray X, GpuArray Y, GpuArray Z, bint nocopy) except -1:
     cdef int err
@@ -52,6 +55,17 @@ cdef api int pygpu_blas_rger(double alpha, GpuArray X, GpuArray Y, GpuArray A,
         raise GpuArrayException(GpuArray_error(&X.ga, err), err)
     return 0
 
+cdef api int pygpu_blas_rgemmBatch_3d(cb_transpose transA, cb_transpose transB,
+                                      double alpha, GpuArray A, GpuArray B,
+                                      double beta, GpuArray C, bint nocopy) except -1:
+    cdef int err
+    err = GpuArray_rgemmBatch_3d(transA, transB,
+                                 alpha, &A.ga, &B.ga,
+                                 beta, &C.ga, nocopy)
+    if err != GA_NO_ERROR:
+        raise GpuArrayException(GpuArray_error(&A.ga, err), err)
+    return 0
+
 
 def dot(GpuArray X, GpuArray Y, GpuArray Z=None, overwrite_z=False):
     """dot(X, Y, Z=None, overwrite_z=False)
@@ -78,14 +92,14 @@ def gemv(double alpha, GpuArray A, GpuArray X, double beta=0.0,
         transA = cb_no_trans
 
     if A.ga.nd != 2:
-        raise TypeError, "A is not a matrix"
+        raise TypeError("A is not a matrix")
     if transA == cb_no_trans:
         Yshp = A.ga.dimensions[0]
     else:
         Yshp = A.ga.dimensions[1]
     if Y is None:
         if beta != 0.0:
-            raise ValueError, "Y not provided and beta != 0"
+            raise ValueError("Y not provided and beta != 0")
         Y = pygpu_empty(1, &Yshp, A.ga.typecode, GA_ANY_ORDER, A.context, None)
         overwrite_y = True
 
@@ -113,9 +127,9 @@ def gemm(double alpha, GpuArray A, GpuArray B, double beta, GpuArray C=None,
         transB = cb_no_trans
 
     if A.ga.nd != 2:
-        raise TypeError, "A is not a matrix"
+        raise TypeError("A is not a matrix")
     if B.ga.nd != 2:
-        raise TypeError, "B is not a matrix"
+        raise TypeError("B is not a matrix")
     if transA == cb_no_trans:
         Cshp[0] = A.ga.dimensions[0]
     else:
@@ -126,7 +140,7 @@ def gemm(double alpha, GpuArray A, GpuArray B, double beta, GpuArray C=None,
         Cshp[1] = B.ga.dimensions[0]
     if C is None:
         if beta != 0.0:
-            raise ValueError, "C not provided and beta != 0"
+            raise ValueError("C not provided and beta != 0")
         C = pygpu_empty(2, Cshp, A.ga.typecode, GA_ANY_ORDER, A.context, None)
         overwrite_c = True
 
@@ -153,3 +167,45 @@ def ger(double alpha, GpuArray X, GpuArray Y, GpuArray A=None,
     pygpu_blas_rger(alpha, X, Y, A, 0)
 
     return A
+
+def gemmBatch_3d(double alpha, GpuArray A, GpuArray B,
+                 double beta, GpuArray C=None,
+                 trans_a=False, trans_b=False, overwrite_c=False):
+    """gemmBatch_3d(alpha, A, B, beta, C=None, trans_a=False, trans_b=False, overwrite_c=False)
+    """
+    cdef cb_transpose transA
+    cdef cb_transpose transB
+    cdef size_t[3] Cshp
+
+    if trans_a:
+        transA = cb_trans
+    else:
+        transA = cb_no_trans
+    if trans_b:
+        transB = cb_trans
+    else:
+        transB = cb_no_trans
+
+    if A.ga.nd != 3:
+        raise TypeError("A is not a batch of matrices")
+    if B.ga.nd != 3:
+        raise TypeError("B is not a batch of matrices")
+
+    Cshp[0] = A.ga.dimensions[0]
+    if transA == cb_no_trans:
+        Cshp[1] = A.ga.dimensions[1]
+    else:
+        Cshp[1] = A.ga.dimensions[2]
+    if transB == cb_no_trans:
+        Cshp[2] = B.ga.dimensions[2]
+    else:
+        Cshp[2] = B.ga.dimensions[1]
+    if C is None:
+        if beta != 0.0:
+            raise ValueError("C not provided and beta != 0")
+        C = pygpu_empty(3, Cshp, A.ga.typecode, GA_ANY_ORDER, A.context, None)
+    elif not overwrite_c:
+        C = pygpu_copy(C, GA_ANY_ORDER)
+    pygpu_blas_rgemmBatch_3d(transA, transB, alpha, A, B, beta, C, 0)
+
+    return C

pygpu/tests/test_blas.py

@@ -167,3 +167,59 @@ def ger(m, n, dtype, order, sliced_x, sliced_y, init_res, overwrite=False):
     gr = gblas.ger(1.0, gX, gY, gA, overwrite_a=overwrite)
 
     numpy.testing.assert_allclose(cr, numpy.asarray(gr), rtol=1e-6)
+
+def test_rgemmBatch_3d():
+    bools = [False, True]
+    for b, (m, n, k), order, trans, offseted_o in product(
+        [1, 17, 31], [(24, 7, 16), (7, 16, 24)], list(product('fc', 'fc', 'c')),
+        list(product(bools, bools)), bools):
+        yield rgemmBatch_3d, b, m, n, k, 'float32', order, trans, \
+            offseted_o, 1, False, False
+    for sliced, overwrite, init_res in product(
+        [1, 2, -1, -2], bools, bools):
+        yield rgemmBatch_3d, 5, 4, 3, 2, 'float32', ('f', 'f', 'c'), \
+            (False, False), False, sliced, overwrite, init_res
+    yield rgemmBatch_3d, 16, 16, 16, 16, 'float64', ('f', 'f', 'c'), (False, False), \
+        False, 1, False, False
+    for alpha, beta, overwrite in product(
+        [0, 1, -1, 0.6], [0, 1, -1, 0.6], bools):
+        yield rgemmBatch_3d, 16, 16, 9, 16, 'float32', ('f', 'f', 'c'), \
+            (False, False), False, 1, overwrite, True, alpha, beta
+
+@guard_devsup
+def rgemmBatch_3d(b, m, n, k, dtype, order, trans, offseted_o, sliced, overwrite,
+         init_res, alpha=1.0, beta=0.0):
+    if trans[0]:
+        shpA = (b,k,m)
+    else:
+        shpA = (b,m,k)
+    if trans[1]:
+        shpB = (b,n,k)
+    else:
+        shpB = (b,k,n)
+
+    cA, gA = gen_gpuarray(shpA, dtype, order=order[0],
+                          offseted_outer=offseted_o,
+                          sliced=sliced, ctx=context)
+    cB, gB = gen_gpuarray(shpB, dtype, order=order[1],
+                          offseted_outer=offseted_o,
+                          sliced=sliced, ctx=context)
+    if init_res:
+        cC, gC = gen_gpuarray((b,m,n), dtype, order=order[2], ctx=context)
+    else:
+        cC, gC = None, None
+
+    cr = numpy.empty((b,m,n), dtype=dtype)
+    if dtype == 'float32':
+        fn_gemm_c = fblas.sgemm
+    else:
+        fn_gemm_c = fblas.dgemm
+    for i in range(b):
+        cCi = cC if cC is None else cC[i]
+        cr[i] = fn_gemm_c(alpha, cA[i], cB[i], beta, cCi, trans_a=trans[0],
+                         trans_b=trans[1], overwrite_c=overwrite)
+
+    gr = gblas.gemmBatch_3d(alpha, gA, gB, beta, gC, trans_a=trans[0],
+                    trans_b=trans[1], overwrite_c=overwrite)
+
+    numpy.testing.assert_allclose(cr, numpy.asarray(gr), rtol=1e-5)