parray_utils.py 文件源码

def get_transpose_function(dtype, conj = False):
    src_type = dtype_to_ctype(dtype)
    name = "trans"
    operation = ""

    if conj:
        if dtype == np.complex128:
            operation = "pycuda::conj"
        elif dtype == np.complex64:
            operation = "pycuda::conj"

    func = SourceModule(
            transpose_template % {
                "name": name,
                "type": src_type,
                "operation": operation
            },
            options=["--ptxas-options=-v"]).get_function(name)
    func.prepare('iiPiPi')#[np.int32, np.int32, np.intp,
    #              np.int32, np.intp, np.int32])
    return func