def wrap(x):
if isinstance(x, six.integer_types):
return ffi_call(lib.tile_poly_expr_literal, x)
if isinstance(x, TensorDim):
if x.size is None:
raise ValueError('Undefined dimension')
return ffi_call(lib.tile_poly_expr_literal, x.size)
return x.as_ptr()
def wrap(x):
if isinstance(x, six.integer_types):
return ffi_call(lib.tile_expr_int, x)
if isinstance(x, float):
return ffi_call(lib.tile_expr_float, x)
if isinstance(x, Tensor):
return x.as_ptr()
raise TypeError('Unexpected type for call argument: {}. args: {}'.format(x, args))
def __init__(self, dtype=None, sizes=[], strides=None, ptr=None, layout=''):
if ptr:
ffi_obj = ptr
elif dtype is not None:
ffi_obj = ffi_call(lib.tile_shape_alloc, dtype, layout.encode())
if strides is None:
strides = []
if len(strides) != len(sizes):
stride = 1
for i in range(len(sizes) - 1, -1, -1):
strides.insert(0, stride)
stride *= sizes[i]
for (size, stride) in zip(sizes, strides):
ffi_call(lib.tile_shape_add_dimension, ffi_obj, size, stride)
else:
raise ValueError('One of dtype= or ptr= must be specified.')
super(TensorShape, self).__init__(ffi_obj)
def __init__(self, shape=None, dims=None, expr=None, value=None, name=''):
self._name = name
if shape:
self._shape = shape
expr = ffi_call(lib.tile_expr_param, shape.as_ptr(), name.encode())
elif dims is not None:
self._dims = dims
expr = None
elif value is not None:
if isinstance(value, six.integer_types):
expr = ffi_call(lib.tile_expr_int, value)
elif isinstance(value, float):
expr = ffi_call(lib.tile_expr_float, value)
else:
raise TypeError('Invalid type for value={}'.format(value))
elif expr is None:
raise ValueError('One of dims=, shape=, or expr= must be specified.')
super(Tensor, self).__init__(expr)
def __init__(self, agg_op, combo_op, output, inputs, name):
inputs = [x.as_ptr() for x in inputs]
expr = ffi_call(
lib.tile_expr_contraction,
agg_op,
combo_op,
output.as_ptr(),
len(inputs),
inputs,
name.encode(),
)
super(_Contraction, self).__init__(expr)
def call(fn, *args):
def wrap(x):
if isinstance(x, six.integer_types):
return ffi_call(lib.tile_expr_int, x)
if isinstance(x, float):
return ffi_call(lib.tile_expr_float, x)
if isinstance(x, Tensor):
return x.as_ptr()
raise TypeError('Unexpected type for call argument: {}. args: {}'.format(x, args))
args = [wrap(x) for x in args]
return Tensor(expr=ffi_call(lib.tile_expr_call, fn.encode(), len(args), args))
def poly_op(op, *args):
def wrap(x):
if isinstance(x, six.integer_types):
return ffi_call(lib.tile_poly_expr_literal, x)
if isinstance(x, TensorDim):
if x.size is None:
raise ValueError('Undefined dimension')
return ffi_call(lib.tile_poly_expr_literal, x.size)
return x.as_ptr()
args = [wrap(x) for x in args]
return ffi_call(lib.tile_poly_expr_op, op, len(args), args)
def __init__(self, ref, key, dims):
if isinstance(key, tuple) or isinstance(key, list):
idxs = key
else:
idxs = [key]
def wrap_idx(x):
if isinstance(x, six.integer_types):
return ffi_call(lib.tile_poly_expr_literal, x)
return x.as_ptr()
idxs = [wrap_idx(x) for x in idxs]
if dims is None:
dims = ffi.NULL
else:
dims = [_wrap_dim(x) for x in dims]
expr = ffi_call(lib.tile_expr_tensor_spec, ref.as_ptr(), len(idxs), idxs, dims)
super(_TensorSpec, self).__init__(expr)
def type(self):
return ffi_call(lib.tile_shape_get_type, self.as_ptr())
def __init__(self, name, *vars):
exprs = [x.as_ptr() for x in vars]
ffi_obj = ffi_call(lib.tile_program_evaluate, name.encode(), len(exprs), exprs)
super(Program, self).__init__(ffi_obj)
def shape(self):
if self._shape is None:
self._shape = TensorShape(ptr=ffi_call(lib.tile_expr_evaluate_shape, self.as_ptr()))
return self._shape
def use_default(self, rhs):
if not self._is_contraction:
raise TypeError('use_default can only be specified on a contraction.')
ffi_call(lib.tile_expr_contraction_set_use_default, self.as_ptr(), rhs.as_ptr())
return self
def no_defract(self):
if not self._is_contraction:
raise TypeError('no_defract can only be specified on a contraction.')
ffi_call(lib.tile_expr_contraction_set_no_defract, self.as_ptr(), True)
return self
def strides(self):
return [
ffi_call(lib.tile_shape_get_dimension_stride, self.as_ptr(), i)
for i in range(self.rank)
]
def rank(self):
return ffi_call(lib.tile_shape_get_rank, self.as_ptr())
def wrap_idx(x):
if isinstance(x, six.integer_types):
return ffi_call(lib.tile_poly_expr_literal, x)
return x.as_ptr()
def __lt__(self, rhs):
if isinstance(rhs, TensorDim):
rhs = rhs.size
ffi_call(lib.tile_poly_expr_add_constraint, self.as_ptr(), rhs)
return Constraint()
def __init__(self, expr=None, name=''):
if expr is None:
expr = ffi_call(lib.tile_poly_expr_index, name.encode())
super(TensorIndex, self).__init__(expr)
