def register_scalar_type(ParakeetClass, dtype, equiv_python_types=[]):
parakeet_type = ParakeetClass(dtype)
_dtype_to_parakeet_type[dtype] = parakeet_type
python_types = [dtype.type] + equiv_python_types
for python_type in python_types:
type_conv.register(python_type, parakeet_type)
return parakeet_type
def register_scalar_type(ParakeetClass, dtype, equiv_python_types = []):
parakeet_type = ParakeetClass(dtype)
_dtype_to_parakeet_type[dtype] = parakeet_type
python_types = [dtype.type] + equiv_python_types
for python_type in python_types:
type_conv.register(python_type, parakeet_type)
return parakeet_type
def make_slice_type(start_t, stop_t, step_t):
key = (start_t, stop_t, step_t)
if key in _slice_type_cache:
return _slice_type_cache[key]
else:
t = SliceT(start_t, stop_t, step_t)
_slice_type_cache[key] = t
return t
def typeof_slice(s):
start_type = type_conv.typeof(s.start)
stop_type = type_conv.typeof(s.stop)
step_type = type_conv.typeof(s.step)
return make_slice_type(start_type, stop_type, step_type)
type_conv.register(slice, SliceT, typeof_slice)
class ArrayT(StructT):
_members = ['elt_type', 'rank']
def node_init(self):
tuple_t = repeat_tuple(Int64, self.rank)
self.shape_t = tuple_t
self.strides_t = tuple_t
self.ptr_t = ptr_type(self.elt_type)
self._fields_ = [
('data', self.ptr_t),
('shape', tuple_t),
('strides', tuple_t),
('offset', Int64),
def make_slice_type(start_t, stop_t, step_t):
key = (start_t, stop_t, step_t)
if key in _slice_type_cache:
return _slice_type_cache[key]
else:
t = SliceT(start_t, stop_t, step_t)
_slice_type_cache[key] = t
return t
def typeof_slice(s):
start_type = type_conv.typeof(s.start)
stop_type = type_conv.typeof(s.stop)
step_type = type_conv.typeof(s.step)
return make_slice_type(start_type, stop_type, step_type)
type_conv.register(slice, SliceT, typeof_slice)
class ArrayT(StructT):
_members = ['elt_type', 'rank']
def node_init(self):
assert isinstance(self.elt_type, core_types.ScalarT), \
"Can't create array with element type %s, currently only scalar elements supported" % \
(self.elt_type,)
tuple_t = repeat_tuple(Int64, self.rank)
self.shape_t = tuple_t
self.strides_t = tuple_t
self.ptr_t = ptr_type(self.elt_type)
_tuple_types = {}
def repeat_tuple(t, n):
"""Given the base type t, construct the n-tuple t*t*...*t"""
elt_types = tuple([t] * n)
if elt_types in _tuple_types:
return _tuple_types[elt_types]
else:
tuple_t = TupleT(elt_types)
_tuple_types[elt_types] = tuple_t
return tuple_t
def make_tuple_type(elt_types):
"""
Use this memoized construct to avoid constructing too many distinct tuple type
objects and speeding up equality checks
"""
key = tuple(elt_types)
if key in _tuple_types:
return _tuple_types[key]
else:
t = TupleT(key)
_tuple_types[key] = t
return t
def typeof(python_tuple):
return make_tuple_type(map(type_conv.typeof, python_tuple))
type_conv.register(tuple, TupleT, typeof)
def combine(self, other):
if self == other: return self
else:raise IncompatibleTypes(self, other)
def __str__(self):
return "SliceT(%s, %s, %s)" % (self.start_type,
self.stop_type,
self.step_type)
def __repr__(self):
return str(self)
_slice_type_cache = {}
def make_slice_type(start_t, stop_t, step_t):
key = (start_t, stop_t, step_t)
if key in _slice_type_cache:
return _slice_type_cache[key]
else:
t = SliceT(start_t, stop_t, step_t)
_slice_type_cache[key] = t
return t
def typeof_slice(s):
start_type = type_conv.typeof(s.start)
stop_type = type_conv.typeof(s.stop)
step_type = type_conv.typeof(s.step)
return make_slice_type(start_type, stop_type, step_type)
type_conv.register(slice, SliceT, typeof_slice)
for (arg_types, typed_fn) in sorted(specializations):
print " -- %s ==> %s" % (arg_types, typed_fn.name)
count += 1
print
print "Total: %d function specializations" % count
import atexit
atexit.register(print_specializations)
def typeof_fn(f):
import ast_conversion
untyped_fn = ast_conversion.translate_function_value(f)
closure_args = untyped_fn.python_nonlocals()
closure_arg_types = map(type_conv.typeof, closure_args)
return make_closure_type(untyped_fn, closure_arg_types)
type_conv.register(FunctionType, ClosureT, typeof_fn)
import prims
def typeof_prim(p):
untyped_fn = prims.prim_wrapper(p)
return make_closure_type(untyped_fn, [])
type_conv.register(prims.class_list, ClosureT, typeof_prim)
"""
Map each (untyped fn id, fixed arg) types to a distinct integer so that the
runtime representation of closures just need to carry this ID
"""
closure_type_to_id = {}
def __str__(self):
return "NoneT"
def __hash__(self):
return 0
def __eq__(self, other):
return isinstance(other, NoneT)
def __repr__(self):
return str(self)
NoneType = NoneT()
def typeof_none(_):
return NoneType
type_conv.register(type(None), NoneT, typeof_none)
def is_struct(c_repr):
return type(c_repr) == type(ctypes.Structure)
class FieldNotFound(Exception):
def __init__(self, struct_t, field_name):
self.struct_t = struct_t
self.field_name = field_name
class StructT(Type):
"""All concrete types excluding scalars and pointers"""
# expect each child class to fill this list
_fields_ = []
import numpy as np import types import scalar_types import type_conv from array_type import make_array_type, ArrayT from tuple_type import make_tuple_type, TupleT from core_types import NoneT, NoneType, TypeValueT type_conv.register(type(None), NoneT, lambda _: NoneType) def typeof_dtype(dt): return TypeValueT(scalar_types.from_dtype(dt)) type_conv.register([np.dtype], TypeValueT, typeof_dtype) def typeof_type(t): assert hasattr(t, 'dtype'), "Can only convert numpy types" dt = t(0).dtype pt = scalar_types.from_dtype(dt) return TypeValueT(pt) type_conv.register(types.TypeType, TypeValueT, typeof_type) def typeof_tuple(python_tuple): return make_tuple_type(map(type_conv.typeof, python_tuple)) type_conv.register(types.TupleType, TupleT, typeof_tuple) def typeof_array(x): x = np.asarray(x) elt_t = scalar_types.from_dtype(x.dtype)
