def matmul()

in openai_gemm.py [0:0]

• 40 lines of code
• 11 McCabe index (conditional complexity)

def matmul(A, B, C, alpha=1.0, beta=0.0, stream=None, bench=False):
    """
        C = alpha * A   . B   + beta * C
        C = alpha * A.T . B   + beta * C
        C = alpha * A   . B.T + beta * C
        C = alpha * A.T . B.T + beta * C

        bench: return benchmark data for all available tiles + cublas
    """

    # this could be relaxed, kernels are capable of mixed precision (with minor tweaks)
    # the s/h prefix would then go away and each type would be specified with kernel build option
    assert A.dtype.type == B.dtype.type == C.dtype.type

    if   C.dtype.type is np.float32:
        prefix = "s"
    elif C.dtype.type is np.float16:
        prefix = "h"
    else:
        raise TypeError("Only floating point dot currently supported.")

    # (m,n) = (m,k) . (k,n)
    m = A.shape[0]
    n = B.shape[1]
    k = A.shape[1]
    assert m == C.shape[0]
    assert n == C.shape[1]
    assert k == B.shape[0]

    # Extract the operations and contiguous dimension sizes (cda, cdb, cdc).
    # Note that these can be the same as from the shape unless the non-contiguous dimension is sliced.
    # One dimension must be contiguous (DRAM efficiency demands this).
    # Note that the strides here do not include the datatype size as they would in numpy.
    # A transpose op (.T) on a GPUTensor reverses the shape and strides then flags the tensor as transposed (is_trans=True) -
    #    The underlying data is unchanged.
    if A.is_trans:
         opA  = 'T'
         cda  = A.strides[1]
         assert A.strides[0] == 1
    else:
         opA  = 'N'
         cda  = A.strides[0]
         assert A.strides[1] == 1

    if B.is_trans:
         opB  = 'T'
         cdb  = B.strides[1]
         assert B.strides[0] == 1
    else:
         opB  = 'N'
         cdb  = B.strides[0]
         assert B.strides[1] == 1

    cdc  = C.strides[0]
    assert C.strides[1] == 1

    op = opA + opB

    # get and autotune the kernel selection
    kernel, params, dynamic_shared = _get_gemm_kernel(prefix, op, cda, cdb, cdc, m, n, k)

    # bind dynamic params
    params[2:8] = (stream, C.gpudata, A.gpudata, B.gpudata, alpha, beta)

    # call the kernel
    kernel.prepared_async_call(*params, shared_size=dynamic_shared)

    # unbind dynamic params
    params[2:8] = (None,) * 6

    # return benchmark data if requested
    if bench:
        return _get_bench_data()[(prefix, op, cda, cdb, cdc, m, n, k)]

    return C