codegen_type,
register_type_recognizers,
register_type_codegenerators,
)
class CBLAS_LAYOUT:
pass
class CBLAS_TRANSPOSE:
pass
register_type_codegenerators({
CBLAS_LAYOUT: lambda t: "CBLAS_LAYOUT",
CBLAS_TRANSPOSE: lambda t: "CBLAS_TRANSPOSE",
})
class ConvertSGEMMCalls(ast.NodeTransformer):
"""
Converts sgemm(...) calls into a valid cblas_sgemm call for MKL by
prepending 101 for Row-Major and converting trans_A, trans_B in to the
proper constant.
"""
def visit_Call(self, node):
if node.func.id == "sgemm":
node.func.id = "cblas_sgemm"
for i in range(2):
node.args[i] = C.Cast(
CBLAS_TRANSPOSE(),
int: lambda t: ctypes.c_long(t),
bool: lambda t: ctypes.c_bool(t),
float: lambda t: ctypes.c_double(t),
str: lambda t: ctypes.c_char(str.encode(t)) if len(t) == 1 else
ctypes.c_char_p(str.encode(t)),
type(None): lambda t: None,
})
register_type_codegenerators({
ctypes.c_int: lambda t: "int",
ctypes.c_long: lambda t: "long",
ctypes.c_float: lambda t: "float",
ctypes.c_double: lambda t: "double",
ctypes.c_char: lambda t: "char",
ctypes.c_char_p: lambda t: "char*",
ctypes.c_void_p: lambda t: "void*",
ctypes.c_bool: lambda t: "bool",
ctypes.c_ulong: lambda t: "size_t",
type(None): lambda n: "void",
_ctypes.Array: lambda ct: "%s*" % codegen_type(ct._type_()),
_ctypes._Pointer: lambda ct: "%s*" % codegen_type(ct._type_()),
})
else:
register_type_recognizers({
types.IntType: lambda t: ctypes.c_long(t),
types.LongType: lambda t: ctypes.c_long(t),
types.BooleanType: lambda t: ctypes.c_bool(t),
types.FloatType: lambda t: ctypes.c_double(t),
types.StringType: lambda t: ctypes.c_char(t) if len(t) == 1 else
ctypes.c_char_p(t),
register_type_codegenerators({
# pointers
np.ctypeslib._ndptr: codegen_ndptr,
# boolean types
np.bool8: lambda t: "bool",
# signed integer types
np.byte: lambda t: "char",
np.short: lambda t: "short",
np.intc: lambda t: "int",
np.longlong: lambda t: "long long",
# technically not universally compatible
np.int8: lambda t: "char",
np.int16: lambda t: "short",
np.int32: lambda t: "int",
np.int64: lambda t: "long",
# unsigned integer types
np.ubyte: lambda t: "unsigned byte",
np.ushort: lambda t: "unsigned short",
np.uintc: lambda t: "unsigned int",
np.ulonglong: lambda t: "unsigned long long",
# floating point types
np.single: lambda t: "float",
np.float32: lambda t: "float",
np.double: lambda t: "double",
np.float64: lambda t: "double",
})
if sys.maxsize > 2 ** 32:
X64_BIT = True
else:
# Python alias c_int to c_long on 32 bit platforms
X64_BIT = False
register_type_codegenerators({
ctypes.c_short: lambda t: "short",
ctypes.c_int: lambda t: "int",
ctypes.c_long: lambda t: "long" if X64_BIT else "int",
ctypes.c_float: lambda t: "float",
ctypes.c_double: lambda t: "double",
ctypes.c_char: lambda t: "char",
ctypes.c_char_p: lambda t: "char*",
ctypes.c_void_p: lambda t: "void*",
ctypes.c_bool: lambda t: "bool",
ctypes.c_ulong: lambda t: "size_t",
ctypes.c_uint32: lambda t: "uint32_t",
type(None): lambda n: "void",
_ctypes.Array: lambda ct: "%s*" % codegen_type(ct._type_()),
_ctypes._Pointer: lambda ct: "%s*" % codegen_type(ct._type_()),
})
#register version specific nodes
if sys.version_info >= (3, 0):
pass
else:
register_type_codegenerators({
# pointers
np.ctypeslib._ndptr:
codegen_ndptr,
# boolean types
np.bool8:
lambda t: "bool",
# signed integer types
np.byte:
lambda t: "char",
np.short:
lambda t: "short",
np.intc:
lambda t: "int",
np.longlong:
lambda t: "long long",
# technically not universally compatible
np.int8:
lambda t: "char",
np.int16:
lambda t: "short",
np.int32:
lambda t: "int",
np.int64:
lambda t: "long",
# unsigned integer types
np.ubyte:
lambda t: "unsigned byte",
np.ushort:
lambda t: "unsigned short",
np.uintc:
lambda t: "unsigned int",
np.ulonglong:
lambda t: "unsigned long long",
# floating point types
np.single:
lambda t: "float",
np.float32:
lambda t: "float",
np.double:
lambda t: "double",
np.float64:
lambda t: "double",
})
except:
log.warn("Failed to load OpenCL runtime.")
import pycl
from ctree.types import (
codegen_type,
register_type_recognizers,
register_type_codegenerators,
)
register_type_recognizers({})
register_type_codegenerators({
pycl.cl_context: lambda t: "cl_context",
pycl.cl_command_queue: lambda t: "cl_command_queue",
pycl.cl_kernel: lambda t: "cl_kernel",
pycl.cl_mem: lambda t: "cl_mem",
})
devices_context_queue_map = {}
def get_context_and_queue_from_devices(devices):
key = tuple(device.vendor_id for device in devices)
try:
return devices_context_queue_map[key]
except KeyError:
context = pycl.clCreateContext(devices)
queue = pycl.clCreateCommandQueue(context)
devices_context_queue_map[key] = (context, queue)
return devices_context_queue_map[key]
if ctree.OCL_ENABLED:
import pycl
from ctree.types import (
codegen_type,
register_type_recognizers,
register_type_codegenerators,
)
register_type_recognizers({
})
register_type_codegenerators({
pycl.cl_context: lambda t: "cl_context",
pycl.cl_command_queue: lambda t: "cl_command_queue",
pycl.cl_kernel: lambda t: "cl_kernel",
pycl.cl_mem: lambda t: "cl_mem",
})
devices_context_queue_map = {}
def get_context_and_queue_from_devices(devices):
key = tuple(device.vendor_id for device in devices)
try:
return devices_context_queue_map[key]
except KeyError:
context = pycl.clCreateContext(devices)
queue = pycl.clCreateCommandQueue(context)
devices_context_queue_map[key] = (context, queue)
from ctree.simd.types import m256d, m256, m512
from ctree.types import register_type_codegenerators
register_type_codegenerators({
m256d: lambda t: "__m256d",
m256: lambda t: "__m256",
m512: lambda t: "__m512"
})
from ctypes import Structure
from ctree.c.nodes import Statement, Block, SymbolRef
from ctree.types import register_type_codegenerators
register_type_codegenerators({
Structure: lambda t: "struct " + type(t).__name__})
class CodegenableStruct(Structure):
def codegen(self):
return "struct " + type(self).__name__
def __str__(self):
return self.codegen()
class StructDef(Statement):
def __init__(self, struct_type, initializer=None):
super(StructDef, self).__init__()
self.struct_type = struct_type
fields = map(lambda (n, t): SymbolRef(n, t()), struct_type._fields_)
self.block = Block(fields)
self.struct_name = struct_type.__name__
self.initializer = initializer
def codegen(self, indent=0):
return "struct " + self.struct_name + self.block.codegen(indent)
