openai_gemm.py

import re
import time
import ctypes
import appdirs
import os.path
import subprocess
import numpy as np
import pycuda.driver as drv
from operator import mul
from struct import unpack_from
from pycuda.tools import context_dependent_memoize
from scikits.cuda import cublas

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



####################################################################################################


# scikits.cuda doesn't expose cublasSgemmEx or cublasHgemm
cublas._libcublas.cublasSgemmEx.restype  = int
cublas._libcublas.cublasSgemmEx.argtypes = [
    cublas._types.handle,
    ctypes.c_int, ctypes.c_int,
    ctypes.c_int, ctypes.c_int, ctypes.c_int,
    ctypes.c_void_p,
    ctypes.c_void_p, ctypes.c_int, ctypes.c_int,
    ctypes.c_void_p, ctypes.c_int, ctypes.c_int,
    ctypes.c_void_p,
    ctypes.c_void_p, ctypes.c_int, ctypes.c_int ]

def cublasSgemmEx(handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc):
    status = cublas._libcublas.cublasSgemmEx(handle,
        cublas._CUBLAS_OP[transa], cublas._CUBLAS_OP[transb],
        m, n, k,
        ctypes.byref(ctypes.c_float(alpha)),
        int(A), 2, lda,
        int(B), 2, ldb,
        ctypes.byref(ctypes.c_float(beta)),
        int(C), 2, ldc)
    cublas.cublasCheckStatus(status)

cublas._libcublas.cublasHgemm.restype  = int
cublas._libcublas.cublasHgemm.argtypes = [
    cublas._types.handle,
    ctypes.c_int, ctypes.c_int,
    ctypes.c_int, ctypes.c_int, ctypes.c_int,
    ctypes.c_void_p,
    ctypes.c_void_p, ctypes.c_int,
    ctypes.c_void_p, ctypes.c_int,
    ctypes.c_void_p,
    ctypes.c_void_p, ctypes.c_int ]

h_dtype = np.dtype(np.float16)

def cublasHgemm(handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc):

    alpha = unpack_from('H', h_dtype.type(alpha))[0]
    beta  = unpack_from('H', h_dtype.type(beta))[0]

    status = cublas._libcublas.cublasHgemm(handle,
        cublas._CUBLAS_OP[transa], cublas._CUBLAS_OP[transb],
        m, n, k,
        ctypes.byref(ctypes.c_uint16(alpha)),
        int(A), lda,
        int(B), ldb,
        ctypes.byref(ctypes.c_uint16(beta)),
        int(C), ldc)
    cublas.cublasCheckStatus(status)

cublasXgemm = {
    "s"  : cublas.cublasSgemm,
    "h"  : cublasSgemmEx,
    "h2" : cublasHgemm,
}


@context_dependent_memoize
def _get_sm_count():
    attributes = drv.Context.get_device().get_attributes()
    return attributes[drv.device_attribute.MULTIPROCESSOR_COUNT]

@context_dependent_memoize
def _get_compute_capability():

    attributes = drv.Context.get_device().get_attributes()
    major = attributes[drv.device_attribute.COMPUTE_CAPABILITY_MAJOR]
    minor = attributes[drv.device_attribute.COMPUTE_CAPABILITY_MINOR]
    return major, minor

@context_dependent_memoize
def _get_events():
    return (drv.Event(), drv.Event())

@context_dependent_memoize
def _get_cublas():
    return cublas.cublasCreate()

@context_dependent_memoize
def _get_bench_data():
    return dict()

def _ceil_div(x, y):
    return -(-x // y)

def _closest_divisor(val, div):
    divisors = sorted([(abs(i - div), i) for i in range(2, 8) if val % i == 0])
    if len(divisors):
        return (divisors[0][1], val // divisors[0][1])
    else:
        return (1, val)


# Tile sizes:      m,   n,  k, vA,vB,vC  div,  op (dynamic shared options)
k128x128x8    = (128, 128,  8,  4, 4, 1,   2,  0, (0,))
k32x32x32     = ( 32,  32, 32,  4, 4, 1,   4,  0, (0, 2**14))
k32x64x32_NN  = ( 32,  64, 32,  8, 4, 4,   4,  1, (0, 2**13))
k32x32x64_NT  = ( 32,  32, 64,  8, 8, 4,   4,  1, (0,))
k16x64x64_NN  = ( 16,  64, 64,  8, 4, 4,   4,  1, (0,))
k16x64x64_NT  = ( 16,  64, 64,  8, 8, 4,   4,  1, (0,))

selections = {
    "s" : {
        "TN" : (k128x128x8, k32x32x32),
        "NN" : (k128x128x8, k32x32x32),
        "NT" : (k128x128x8, k32x32x32),
        "TT" : (k128x128x8, k32x32x32),
    },
    "h" : {
        "TN" : (k128x128x8, k32x32x32),
        "NN" : (k128x128x8, k32x32x32, k32x64x32_NN, k16x64x64_NN),
        "NT" : (k128x128x8, k32x32x32, k32x32x64_NT, k16x64x64_NT),
        "TT" : (k128x128x8, k32x32x32),
    },
}

# Autotune kernel selection
@context_dependent_memoize
def _get_gemm_kernel(prefix, op, cda, cdb, cdc, m, n, k):

    if op[0] == 'T':
         vec4A = (cda & 3) == 0 and (m & 3) == 0
         vec8A = (cda & 7) == 0 and (m & 7) == 0
         dimA  = (k,cda)
    else:
         vec4A = (cda & 3) == 0 and (k & 3) == 0
         vec8A = (cda & 7) == 0 and (k & 7) == 0
         dimA  = (m,cda)

    if op[1] == 'T':
         vec4B = (cdb & 3) == 0 and (k & 3) == 0
         vec8B = (cdb & 7) == 0 and (k & 7) == 0
         dimB  = (n,cdb)
    else:
         vec4B = (cdb & 3) == 0 and (n & 3) == 0
         vec8B = (cdb & 7) == 0 and (n & 7) == 0
         dimB  = (k,cdb)

    vec4C = (cdc & 3) == 0 and (n & 3) == 0
    dimC  = (m,cdc)

    dtype = np.dtype(np.float32 if prefix == 's' else np.float16)

    A = drv.mem_alloc(mul(*dimA) * dtype.itemsize)
    B = drv.mem_alloc(mul(*dimB) * dtype.itemsize)
    C = drv.mem_alloc(mul(*dimC) * dtype.itemsize)

    # TODO: use curand
    dataA = np.random.uniform(-1.0, 1.0, dimA).astype(dtype)
    dataB = np.random.uniform(-1.0, 1.0, dimB).astype(dtype)
    drv.memcpy_htod(int(A), dataA)
    drv.memcpy_htod(int(B), dataB)

    # Using random data gets you more accurate autotune results
    # drv.memset_d8(int(A), 0, mul(*dimA) * dtype.itemsize)
    # drv.memset_d8(int(B), 0, mul(*dimB) * dtype.itemsize)

    timings = []
    cache   = []

    # scale the repeat count to amount of work
    repeat = min(max(int(5e11 * 28 / (m*n*k * 2.0 * _get_sm_count()) ), 10), 5000)
    warmup = repeat
    #print repeat

    start, end = _get_events()
    flops = m * n * k * 2.0

    for tileM, tileN, tileK, vecA, vecB, vecC, div, base_op, dyn_shared in selections[prefix][op]:

        vecA = (vecA == 4 and vec4A) or (vecA == 8 and vec8A)
        vecB = (vecB == 4 and vec4B) or (vecB == 8 and vec8B)
        vecC =  vecC == 1 or  vec4C
        vec  = vecA and vecB and vecC

        if base_op:
            # The op is part of the base kernel name
            base = "%sgemm_%dx%dx%d_%s" % (prefix, tileM, tileN, tileK, op)
            opts = ( "vec", ) if vec else ()
        else:
            # The op is an option passed to a more generic kernel
            base = "%sgemm_%dx%dx%d" % (prefix, tileM, tileN, tileK)
            opts = ( op, "vec" ) if vec else (op,)

        kernel = get_kernel(base, opts)

        blk_A = _ceil_div(m, tileM)
        blk_B = _ceil_div(n, tileN)

        # TODO: perhaps autotune all possible small divisors
        blk_a, blk_A = _closest_divisor(blk_A, div)
        blk_b, blk_B = _closest_divisor(blk_B, div)
        if blk_a == 1:
            blk_a, blk_A = (blk_A, 1)

        for dynamic_shared in dyn_shared:

            params = [
                (blk_a * blk_b, blk_B, blk_A), (kernel.threads, 1, 1), None,
                C, A, B, 1.0, 0.0,
                cda, cdb, cdc, m, n, k, blk_a, blk_b ]

            #print kernel.name, params, dynamic_shared

            # Warmup (once per config)
            for r in range(warmup):
                kernel.prepared_async_call(*params)
            warmup = 0

            # Benchmark
            start.record()
            for r in range(repeat):
                kernel.prepared_async_call(*params, shared_size=dynamic_shared)
            end.record()
            end.synchronize()
            msecs = end.time_since(start) / float(repeat)
            gflops = flops / (msecs * 1000000.0)

            params[3:8] = (None,) * 5

            timings.append((msecs, gflops, kernel, params, dynamic_shared))
            cache.append((msecs, gflops, kernel.name, dynamic_shared))

    major, minor = _get_compute_capability()
    if prefix == "h" and major == 6 and minor == 0:
        cublas_gemm = cublasXgemm["h2"]
    else:
        cublas_gemm = cublasXgemm[prefix]

    # record a cublas time for reference
    cublas_handle = _get_cublas()
    start.record()
    for r in range(repeat):
        # convert row order to col order
        cublas_gemm(cublas_handle, op[1], op[0], n, m, k, 1.0, B, cdb, A, cda, 0.0, C, cdc)
    end.record()
    end.synchronize()
    msecs = end.time_since(start) / float(repeat)
    gflops = flops / (msecs * 1000000.0)
    cache.append( (msecs, gflops, "cuBLAS", 0) )

    # cache complete timing data for benchmark comparisons
    # this data could be cached to disk for quicker autotuning on future runs
    _get_bench_data()[(prefix, op, cda, cdb, cdc, m, n, k)] = cache

    # return the fastest kernel
    return tuple(sorted(timings)[0][2:5])



# Utility function to test all tiles for the given dimensions and dtype
def matmul_test(ng, dtype, op, m, n, k, ones=False, out=False):

    prefix = "s" if dtype is np.float32 else "h"

    if op[0] == 'T':
         vec4A = (m & 3) == 0
         vec8A = (m & 7) == 0
         dimA  = (k,m)
         cda   = m
    else:
         vec4A = (k & 3) == 0
         vec8A = (k & 7) == 0
         dimA  = (m,k)
         cda   = k

    if op[1] == 'T':
         vec4B = (k & 3) == 0
         vec8B = (k & 7) == 0
         dimB  = (n,k)
         cdb   = k
    else:
         vec4B = (n & 3) == 0
         vec8B = (n & 7) == 0
         dimB  = (k,n)
         cdb   = n

    vec4C = (n & 3) == 0
    dimC  = (m,n)
    cdc   = n

    A1 = ng.empty(dimA, dtype=dtype)
    B1 = ng.empty(dimB, dtype=dtype)
    C1 = ng.empty(dimC, dtype=dtype)
    C2 = ng.empty(dimC, dtype=dtype)

    if ones:
        A1[:] = 1.0
        B1[:] = 1.0
    else:
        # fill with uniform randoms from -1 to 1
        A1[:] = 2 * (.5 - ng.rand())
        B1[:] = 2 * (.5 - ng.rand())

    # for reducing outputs
    partial1 = ng.empty((C1.shape[0],1), dtype=np.float32)
    partial2 = partial1[0:1,0:1]

    cublas_handle = _get_cublas()

    for tileM, tileN, tileK, vecA, vecB, vecC, div, base_op, dyn_shared in selections[prefix][op]:

        vecA = (vecA == 4 and vec4A) or (vecA == 8 and vec8A)
        vecB = (vecB == 4 and vec4B) or (vecB == 8 and vec8B)
        vecC =  vecC == 1 or  vec4C
        vec  = vecA and vecB and vecC

        if base_op:
            # The op is part of the base kernel name
            base = "%sgemm_%dx%dx%d_%s" % (prefix, tileM, tileN, tileK, op)
            opts = ( "vec", ) if vec else ()
        else:
            # The op is an option passed to a more generic kernel
            base = "%sgemm_%dx%dx%d" % (prefix, tileM, tileN, tileK)
            opts = ( op, "vec" ) if vec else (op,)

        kernel = get_kernel(base, opts)

        blk_A = _ceil_div(m, tileM)
        blk_B = _ceil_div(n, tileN)

        blk_a, blk_A = _closest_divisor(blk_A, div)
        blk_b, blk_B = _closest_divisor(blk_B, div)
        if blk_a == 1:
            blk_a, blk_A = (blk_A, 1)

        for alpha, beta in ( (1.0,0.0), (0.5,0.5) ):

            try:
                if ones:
                    C1[:] = 1.0
                else:
                    C1[:] = 2 * (.5 - ng.rand())
                C2[:] = C1

                params = [
                    (blk_a * blk_b, blk_B, blk_A), (kernel.threads, 1, 1), None,
                    C1.gpudata, A1.gpudata, B1.gpudata, alpha, beta,
                    cda, cdb, cdc, m, n, k, blk_a, blk_b ]

                kernel.prepared_async_call(*params)

                # convert row order to col order
                cublasXgemm[prefix](cublas_handle, op[1], op[0], n, m, k, alpha, B1.gpudata, cdb, A1.gpudata, cda, beta, C2.gpudata, cdc)

                # Check for NaNs
                partial1[:] = ng.min(ng.finite(C1), axis=1)
                partial2[:] = ng.min(partial1, axis=0)
                if partial2.get()[0,0] == 0.0:
                    print "Error: NaN kernel: %s mnk: (%d,%d,%d) ab: (%f,%f)" % (kernel.name, m,n,k, alpha,beta)
                    exit()

                # Get Max Diff
                partial1[:] = ng.max(abs(C2 - C1), axis=1)
                partial2[:] = ng.max(partial1, axis=0)
                diff = partial2.get()[0,0]

                # Get Mean
                partial1[:] = ng.sum(abs(C2), axis=1)
                partial2[:] = ng.sum(partial1, axis=0)
                mean = partial2.get()[0,0] / C2.size

                # Scale diff by the mean
                pctErr = 100 * diff / mean

                #print "Error: %.3f %s" % (pctErr, kernel.name)

                maxerr = .005 if dtype is np.float32 else 0.7

                if pctErr > maxerr:
                    print "Error: %.3f%% diff: %.5f mean %.5f kernel: %s mnk: (%d,%d,%d) ab: (%f,%f)" % (pctErr, diff, mean, kernel.name, m,n,k, alpha,beta)
                    print params
                    if out:
                        C1  = C1.get()
                        C2  = C2.get()
                        D  = abs(C2 - C1)
                        np.savetxt("out_diff.txt",    D,  fmt='%3.1f')
                        np.savetxt("out_correct.txt", C2, fmt='%5.1f')
                        np.savetxt("out_error",       C1, fmt='%5.1f')
                    exit()

            except drv.Error as e:
                print "kernel: %s mnk: (%d,%d,%d) ab: (%f,%f)" % (kernel.name, m,n,k, alpha,beta)
                print e
                exit()

### below code adapted from Nervana Neon: kernel_specs.py

def _get_cache_dir(subdir=None):

    cache_dir = appdirs.user_cache_dir("openai-gemm")

    if subdir:
        subdir = subdir if isinstance(subdir, list) else [subdir]
        cache_dir = os.path.join(cache_dir, *subdir)

    if not os.path.exists(cache_dir):
        os.makedirs(cache_dir)

    return cache_dir

# helpful for kernel development
debug = 0

base_dir  = os.path.dirname(__file__)
maxas_dir = os.path.join(base_dir, "maxas")
sass_dir  = os.path.join(base_dir, "sass")

kernels = {
    # Generic gemm tiles
    "sgemm_128x128x8":   {"threads": 256, "sass": "xgemm_128x128x8",   "params": "xgemm", "share": "(128*8 + 32)*4 + 4", "args": {"type": "s"} },
    "hgemm_128x128x8":   {"threads": 256, "sass": "xgemm_128x128x8",   "params": "xgemm", "share": "(128*8 + 32)*4 + 4", "args": {"type": "h"} },
    "sgemm_32x32x32":    {"threads": 128, "sass": "xgemm_32x32x32",    "params": "xgemm", "share": "(32*33)*4 + 4",      "args": {"type": "s"} },
    "hgemm_32x32x32":    {"threads": 128, "sass": "xgemm_32x32x32",    "params": "xgemm", "share": "(32*33)*4 + 4",      "args": {"type": "h"} },

    # Custom hgemm tiles designed for small minibatch RNNs
    "hgemm_32x64x32_NN": {"threads": 128, "sass": "hgemm_32x64x32_NN", "params": "xgemm", "share": "32*33*2 + 64*32*2 + 4"               },
    "hgemm_32x32x64_NT": {"threads": 128, "sass": "hgemm_32x32x64_NT", "params": "xgemm", "share": "32*65*4 + 4"                         },
    "hgemm_16x64x64_NN": {"threads": 128, "sass": "hgemm_16x64x64_NN", "params": "xgemm", "share": "(16*64 + 32)*2 + 64*64*2 + 4"        },
    "hgemm_16x64x64_NT": {"threads": 128, "sass": "hgemm_16x64x64_NT", "params": "xgemm", "share": "(16*64 + 32)*2 + (64*64 + 32)*2 + 4" },
}

_params = {
    "xgemm": [
        "float* param_C",
        "float* param_A",
        "float* param_B",
        "float param_alpha",
        "float param_beta",
        "unsigned param_cda",
        "unsigned param_cdb",
        "unsigned param_cdc",
        "unsigned param_m",
        "unsigned param_n",
        "unsigned param_k",
        "unsigned param_blk_a",
        "unsigned param_blk_b",
    ],
}

_space_re = re.compile(r"\s+")

_share_template = r"""
    .shared .align 4 .b32 share[{0}];
"""

_kernel_template = r"""
.version {6}
.target {0}
.address_size 64

// args: {5}

.visible .entry  {1}(
{2}
)
.reqntid {3}
{{
{4}
    ret;
}}
"""

def get_ptx_file(kernel_spec, kernel_name, arch, ptx_ver):

    ptx_dir = _get_cache_dir([arch, 'ptx'])

    thread_spec = kernel_spec["threads"]
    args_spec   = str(kernel_spec.get("args",""))
    param_spec  = _params[kernel_spec["params"]]

    kernel_params = []
    for p in param_spec:
        ptype, pname = _space_re.split(p)

        if ptype[-1] == '*':
            ptype = '.u64'
        elif ptype == 'float':
            ptype = '.f32'
        else:
            ptype = '.u32'

        kernel_params.append("    .param %s %s" % (ptype, pname))

    kernel_params = ",\n".join(kernel_params)

    if "share" in kernel_spec:
        share = _share_template.format(eval(kernel_spec["share"]))
    else:
        share = ""

    kernel_text = _kernel_template.format(arch, kernel_name, kernel_params, thread_spec, share, args_spec, ptx_ver)
    kernel_ptx  = os.path.join(ptx_dir, kernel_name + ".ptx")

    current_text = ""
    if os.path.exists(kernel_ptx):
        f = open(kernel_ptx, "r")
        current_text = f.read()
        f.close()
    # only write out the kernel if text has changed.
    if kernel_text != current_text:
        f = open(kernel_ptx, "w")
        f.write(kernel_text)
        f.close()

    return kernel_ptx


include_re = re.compile(r'^<INCLUDE\s+file="([^"]+)"\s*/>')

def extract_includes(name, includes=None):
    if not includes:
        includes = list()
    sass_file = os.path.join(sass_dir, name)
    includes.append((sass_file, os.path.getmtime(sass_file)))
    for line in open(sass_file, "r"):
        match = include_re.search(line)
        if match:
            extract_includes(match.group(1), includes)
    return includes

def run_command(cmdlist):
    cmd  = " ".join(cmdlist)
    proc = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    out, err = proc.communicate()
    if proc.returncode:
        raise RuntimeError("Error(%d):\n%s\n%s" % (proc.returncode, cmd, err))
    if debug:
        print cmd
        if out: print out
        if err: print err


@context_dependent_memoize
def get_kernel(base_name, options=None):

    major, minor = _get_compute_capability()
    if major < 5:
        raise RuntimeError("sass kernels require Maxwell or greater class hardware")

    arch  = "sm_%d%d" % (major, minor)

    libprefix = "PERL5LIB=%s" % maxas_dir
    maxas_i = [libprefix, os.path.join(maxas_dir, "maxas.pl") + " -i -w"]
    maxas_p = [libprefix, os.path.join(maxas_dir, "maxas.pl") + " -p"]

    kernel_spec = kernels[base_name]
    kernel_name = base_name

    # static options
    if "args" in kernel_spec:
        for pair in kernel_spec["args"].items():
            maxas_i.append("-D%s %s" % pair)
            maxas_p.append("-D%s %s" % pair)

    # dynamic options
    if options is not None:
        for opt in options:
            if type(opt) is tuple:
                maxas_i.append("-D%s %s" % opt)
                maxas_p.append("-D%s %s" % opt)
                kernel_name += "_%s%s" % opt
            else:
                maxas_i.append("-D%s 1" % opt)
                maxas_p.append("-D%s 1" % opt)
                kernel_name += "_%s" % opt

    maxas_i.insert(2, "-k " + kernel_name)

    sass_name  = kernel_spec["sass"] + ".sass"
    cubin_name = kernel_name + ".cubin"
    cubin_dir  = _get_cache_dir([arch, 'cubin'])

    ptx_version = "4.2" if major < 6 else "5.0"
    ptx_file   = get_ptx_file(kernel_spec, kernel_name, arch, ptx_version)
    sass_file  = os.path.join(sass_dir, sass_name)
    cubin_file = os.path.join(cubin_dir, cubin_name)

    if not os.path.exists(sass_file):
        raise RuntimeError("Missing: %s for kernel: %s" % (sass_name, kernel_name))

    ptx_mtime   = os.path.getmtime(ptx_file)
    cubin_mtime = os.path.getmtime(cubin_file) if os.path.exists(cubin_file) else 0

    build_cubin = False
    if ptx_mtime > cubin_mtime:
        build_cubin = True

    includes = extract_includes(sass_name)
    for include, include_mtime in includes:
        if include_mtime > cubin_mtime:
            build_cubin = True
            break

    if build_cubin:
        # build the cubin and run maxas in the same command
        # we don't want the chance of a generated cubin not processed by maxas (in case user hits ^C in between these steps)
        run_command([ "ptxas -v -arch", arch, "-o", cubin_file, ptx_file, ";" ] + maxas_i + [sass_file, cubin_file])
        cubin_mtime = time.time()

    # output preprocessed and disassembled versions in debug mode
    if debug:
        pre_dir  = _get_cache_dir([arch, 'pre'])
        dump_dir = _get_cache_dir([arch, 'dump'])

        pre_file   = os.path.join(pre_dir,  kernel_name + "_pre.sass")
        dump_file  = os.path.join(dump_dir, kernel_name + "_dump.sass")
        pre_mtime  = os.path.getmtime(pre_file)  if os.path.exists(pre_file)  else 0
        dump_mtime = os.path.getmtime(dump_file) if os.path.exists(dump_file) else 0

        for include, include_mtime in includes:
            if include_mtime > pre_mtime:
                run_command(maxas_p + [sass_file, pre_file])
                break

        if cubin_mtime > dump_mtime:
            run_command(["nvdisasm -c", cubin_file, ">", dump_file])

    # generate the function signature for pycuda
    params  = _params[kernel_spec["params"]]
    sig = ""
    for p in params:
        ptype, pname = _space_re.split(p)
        if ptype[-1] == '*':
            sig += "Q"
        elif ptype == 'float':
            sig += "f"
        elif ptype == 'unsigned':
            sig += "I"
        else:
            sig += "i"

    module = drv.module_from_file(cubin_file)
    func   = module.get_function(kernel_name)
    func.prepare(sig)
    func.threads = kernel_spec["threads"]
    func.name = kernel_name
    func.static_shared = eval(kernel_spec["share"])

    return func