diff --git a/lappy/core/array.py b/lappy/core/array.py index cce8311583351a5bf022240c2ab7125b60b4d19d..8a20567b41a56134259acfe14ba68a00eee1ac74 100644 --- a/lappy/core/array.py +++ b/lappy/core/array.py @@ -23,6 +23,8 @@ THE SOFTWARE. """ import warnings import logging +from traceback import format_list, extract_stack +from itertools import chain import pyopencl as cl @@ -109,10 +111,9 @@ class LazyObjectBase(object): preconditions are checked at runtime, prior to calling the kernel. It is a list of callables, each of which can take one or more - arguments/bound_arguments. If a precondition fails, a - :class:`PreconditionNotMetError` is thrown (by either the checker - function itself, or by returning False / any negative integer from - the checker). + arguments/bound_arguments. If a precondition fails, an + :class:`Exception` is thrown (by either the checker + function itself, or by returning False from the checker). """ _counter = 0 _name_prefix = '__lappy_object_' @@ -289,9 +290,6 @@ class Array(LazyObjectBase): self._ndim = to_lappy_unsigned(kwargs.pop("ndim", None)) self.domain_expr = domain_expr - if self.ndim == 0: - self.name = self.name.replace('array', 'number') - if self.ndim is None: # infer ndim from given shape # NOTE: use fixed ndim only @@ -315,9 +313,6 @@ class Array(LazyObjectBase): # scalars if self.shape == () or self.ndim == 0: assert self.shape == () and self.ndim == 0 - self.name.replace( - self.__class__._name_prefix, - "lappy_number_") # arrays with concrete ndim try: @@ -351,6 +346,9 @@ class Array(LazyObjectBase): @property def shape(self): + """Returns the shape as a tuple of ints if the value is known, or strings of + names otherwise. + """ return tuple(s.value if s.value is not None else s.name for s in self._shape) @property @@ -489,11 +487,11 @@ class Array(LazyObjectBase): # constant positive dimensions assert self.ndim > 0 - shape = Lookup(Variable(self.name), "shape") + names = ['__' + self.name + '_shape_%d' % d for d in range(self.ndim)] return tuple( Unsigned( - name='__' + self.name + '_shape_%d' % d, - expr=Subscript(shape, d)) + name=names[d], + expr=var(names[d])) for d in range(self.ndim)) def copy(self, stateless=False): @@ -618,28 +616,125 @@ class Array(LazyObjectBase): shape_names = self._shape_names() return ', '.join(shape_names) + @property + def _shape_expr(self): + """Values for the shape (if not None), otherwise expressions. + """ + return tuple( + s.value if s.value is not None else s.expr + for s in self._shape) + def check_preconditions(self, context): """Call checkers in the list of preconditions. + + :param context: a dict which is passed to all checkers. """ + failed_checks = [] for checker in self.preconditions: - res = checker(context) - if res is not None: - if isinstance(res, bool): - if res: - pass + try: + res = checker(context) + if res is not None: + if isinstance(res, bool): + if res: + pass + else: + raise PreconditionNotMetError( + "precondition checker failed: %s" % str(checker)) else: - raise PreconditionNotMetError( - "precondition checker failed: %s" % str(checker)) - elif isinstance(res, int): - if res >= 0: - pass + raise ValueError( + "cannot understand the return value " + "of the precondition checker %s" % str(checker)) + except Exception: # noqa: W0703 + failed_checks.append(checker) + + err_msgs = [] + for fc in failed_checks: + msg = ("Precondition %s not met, which was imposed at\n\n" + % str(checker)) + if hasattr(fc, "frame"): + msg = msg + '\n'.join(format_list(extract_stack( + fc.frame))) + err_msgs.append(msg) + if len(failed_checks) > 0: + precon_err_msg = ( + "%d out of %d preconditions are not met" + % (len(failed_checks), len(self.preconditions))) + for im, msg in enumerate(err_msgs): + precon_err_msg = (precon_err_msg + '\n\n' + + '[%d / %d]: ' % (im + 1, len(failed_checks)) + msg) + raise PreconditionNotMetError(precon_err_msg) + + def with_dtype(self, dtype, new_name=None): + """Returns a copy of self with the specified dtype. + """ + raise NotImplementedError() + + def with_shape_data(self, shape_data, new_name=None): + """Returns a copy of the array with concrete shape. Assuming concrete ndim. + ndim is unchanged by assigning the shape data, meaning that the inputs, when + not None, must be positive. + + :param shape_data: a tuple of shape data. Use None to skip setting certain + axes. + + NOTE: if the shape has nontrivial expression (not just a reference to a + name), lappy will not solve the equation for the "independent variables". + The behavior in such cases is undefined. + """ + if not isinstance(shape_data, (list, tuple)): + raise ValueError() + if self.ndim != len(shape_data): + raise ValueError() + + for s in self._shape: + if False: # TODO: test for nontrivial exprs + warnings.warn( + "the shape variable %s of %s has nontrivial expression, " + " and setting its value may yield undefined behavior" + % (s.name, self.name)) + + if new_name is None: + new_name = 'bound%d_' % self.__class__._counter + self.name + + new_arr = self.with_name(new_name) + + delta_dict = {} + for s, s_val in zip(new_arr._shape, shape_data): + if s_val is None: + continue + else: + assert isinstance(s_val, int) and s_val > 0 + delta_dict[s.name] = (s.value, s_val) + s.value = s_val + + # iterate through arguments and intermediaries and + # propagate the assignments to all symbols of the same name + # + # (they may not refer to the same object, but may have the same name) + # + for arg in new_arr.arguments.values(): + if arg is None: + continue + for sym_s in arg._shape: + if sym_s.name in delta_dict: + if sym_s.value == delta_dict[sym_s.name][0]: # old value + sym_s.value = delta_dict[sym_s.name][1] else: - raise PreconditionNotMetError( - "precondition checker failed: %s" % str(checker)) - else: - raise ValueError( - "cannot understand the return value " - "of the precondition checker %s" % str(checker)) + # the value is updated, since it points to the + # same Unsigned obj + assert sym_s.value == delta_dict[sym_s.name][1] + + for imd in new_arr.intermediaries.values(): + for sym_s in imd._shape: + if sym_s.name in delta_dict: + if sym_s.value == delta_dict[sym_s.name][0]: # old value + sym_s.value = delta_dict[sym_s.name][1] + else: + # the value is updated, since it points to the + # same Unsigned obj + assert sym_s.value == delta_dict[sym_s.name][1] + + return new_arr def with_data(self, new_name=None, **kwargs): """Binds specific data to a slots in the arguments, resulting in a @@ -647,7 +742,7 @@ class Array(LazyObjectBase): environment. """ if new_name is None: - new_name = 'bound_' + self.name + new_name = 'bound%d_' % self.__class__._counter + self.name for key in kwargs.keys(): if key not in self.arguments.keys(): @@ -688,6 +783,102 @@ class Array(LazyObjectBase): return new_arr + def get_data_mapping(self, knl=None): + """Make a data mapping using known data, tailored for giving inputs to + the loopy kernel. Returns all known information if knl is None. + + :param knl: the loopy kernel + """ + data_map = {} + shapeval_expansion_list = [] + + # gather captured data + for arr_name, varr in self.env.items(): + if arr_name in ['isl_ctx', 'cl_ctx', 'cu_ctx']: + continue + + if arr_name == self.name: + # only specify output shape, and let the backend to do the malloc + for out_s in self._shape: + if out_s.value is None: + shapeval_expansion_list.append(out_s) + data_map[out_s.name] = out_s.value + + if isinstance(varr, Array) and varr.value is not None: + data_map[arr_name] = varr.value + elif isinstance(varr, np.ndarray): + data_map[arr_name] = varr + elif varr is None: + pass # self + else: + raise RuntimeError("unrecogonized captured variable %s" % arr_name) + + # try to get as much extra data that loopy wants as possible + # also evaluates the shape expressions + if knl is None: + # gather all known shapes + for s in self._shape: + if s.name in data_map: + continue + if s.value is not None: + data_map[s.name] = s.value + for arg in chain(self.arguments.values(), + self.bound_arguments.values(), self.intermediaries.values()): + if arg is None: + continue # skip self + assert isinstance(arg, Array) + for s in arg._shape: + if s.value is None: + shapeval_expansion_list.append(s) + data_map[s.name] = s.value + else: + for argname, arg in knl.arg_dict.items(): + if argname in data_map: + continue + if argname in self._shape_names(): + shape_val = self._shape[self._shape_names().index(argname)] + if shape_val.value is None: + shapeval_expansion_list.append(shape_val) + data_map[argname] = shape_val.value + for arr_arg in self.bound_arguments.values(): + if arr_arg is None: + continue + if argname in arr_arg._shape_names(): + shape_val = arr_arg._shape[ + arr_arg._shape_names().index(argname)] + if shape_val.value is None: + shapeval_expansion_list.append(shape_val) + data_map[argname] = shape_val.value + for arr_imd in self.intermediaries.values(): + if argname in arr_imd._shape_names(): + shape_val = arr_imd._shape[ + arr_imd._shape_names().index(argname)] + if shape_val.value is None: + shapeval_expansion_list.append(shape_val) + data_map[argname] = shape_val.value + # no need to search arguments + assert len(self.arguments) == 0 + + # evaluate shape expressions + for se in shapeval_expansion_list: + try: + seval = evaluate(se.expr, data_map) + except UnknownVariableError: + warnings.warn( + "cannot get value for %s prior to calling the loopy kernel" + % se.name) + seval = None + data_map[se.name] = seval + + # purge None-valued entries + entries_to_purge = [] + for key, val in data_map.items(): + if val is None: + entries_to_purge.append(key) + for key in entries_to_purge: + data_map.pop(key) + return data_map + def eval(self, shape_dtype=np.int32): """Evaluates the array expression after binding all the arguments. Note that this function has side effects: it sets self.value to the @@ -739,6 +930,9 @@ class Array(LazyObjectBase): # Step 5 # Make and run loopy kernel + # + # FIXME: collect all assumptions and pass to loopy + # kernel_args = [] # array args @@ -810,59 +1004,7 @@ class Array(LazyObjectBase): extra_dtype_dict[argname] = shape_dtype knl = lp.add_and_infer_dtypes(knl, extra_dtype_dict) - - data_map = {} - shapeval_expansion_list = [] - - # gather captured data - for arr_name, varr in self.env.items(): - if arr_name in ['isl_ctx', 'cl_ctx', 'cu_ctx']: - continue - - if arr_name == self.name: - # only specify output shape, and let the backend to do the malloc - for out_s in self._shape: - if out_s.value is None: - shapeval_expansion_list.append(out_s) - data_map[out_s.name] = out_s.value - - if isinstance(varr, Array) and varr.value is not None: - data_map[arr_name] = varr.value - elif isinstance(varr, np.ndarray): - data_map[arr_name] = varr - elif varr is None: - pass # self - else: - raise RuntimeError("unrecogonized captured variable %s" % arr_name) - - # try to get as much extra data that loopy wants as possible - # also evaluates the shape expressions - for argname, arg in knl.arg_dict.items(): - if argname in self.env: - continue - if argname in self._shape_names(): - shape_val = self._shape[self._shape_names().index(argname)] - if shape_val.value is None: - shapeval_expansion_list.append(shape_val) - data_map[argname] = shape_val.value - for arr_arg in self.bound_arguments.values(): - if arr_arg is None: - continue - if argname in arr_arg._shape_names(): - shape_val = arr_arg._shape[arr_arg._shape_names().index(argname)] - if shape_val.value is None: - shapeval_expansion_list.append(shape_val) - data_map[argname] = shape_val.value - - for se in shapeval_expansion_list: - try: - seval = evaluate(se.expr, data_map) - except UnknownVariableError: - warnings.warn( - "cannot get value for %s prior to calling the loopy kernel" - % se.name) - seval = None - data_map[se.name] = seval + data_map = self.get_data_mapping(knl) # TODO: to test everything that the kernel sees, # let loopy infer more context first (e.g. shape of input arrays) @@ -888,6 +1030,9 @@ class Array(LazyObjectBase): def _set_shape_values(self, shape): """Instantiate shape with concrete values. + + This method should not be called from a user. Users should use + Array.with_shape_data() instead. """ for s in tuple(shape): assert int(s) == s @@ -1317,10 +1462,11 @@ def to_lappy_unsigned(unsigned_like, name=None, base_env=None): # {{{ digest shape specifiers -def to_lappy_shape(shape, expr=None): +def to_lappy_shape(shape): """Accepted shape specified in one of the following forms: - - Tuple of nonnegative integers, strings, or Unsigned objects. + - Tuple of nonnegative integers, strings, expressions or + Unsigned objects. - Comma/space-separated string. Nonnegativity assumption is added to the input where the shape is diff --git a/lappy/core/basic_ops.py b/lappy/core/basic_ops.py index 018fa0155418af50a65b2c40c63753fd45958a17..038546a86bad446606526bda249bf5bb0199fef3 100644 --- a/lappy/core/basic_ops.py +++ b/lappy/core/basic_ops.py @@ -25,6 +25,9 @@ from pymbolic.primitives import Variable, Product from numpy import AxisError from lappy.core.array import to_lappy_shape, to_lappy_unsigned +from lappy.core.preconditions import ( + EvaluationPrecondition, + make_size_conservation_condition) def reshape(array, newshape, order='C', name=None, inames=None): @@ -89,6 +92,8 @@ def reshape(array, newshape, order='C', name=None, inames=None): newshape[missing_id] = array.size // Product(shadow_expr) def newaxis_is_integral(context): + """The shape must be an integer along the newaxis. + """ old_size = evaluate(array.size, context) shadow_size = evaluate(Product(shadow_expr), context) @@ -102,14 +107,16 @@ def reshape(array, newshape, order='C', name=None, inames=None): return old_size % shadow_size == 0 # add pre-eval runtime checks on size consistency - new_precond.append(newaxis_is_integral) + new_precond.append( + EvaluationPrecondition(newaxis_is_integral)) if name is None: if len(array.name) > 1 and array.name[:2] == '__': name_prefix = '' else: name_prefix = '__' - name = name_prefix + array.name + '_reshaped' + name = (name_prefix + array.name + + '_reshaped%d' % array.__class__._counter) # noqa: W0212 if inames is None: if len(name) > 1 and name[:2] == '__': @@ -124,21 +131,9 @@ def reshape(array, newshape, order='C', name=None, inames=None): assert all(isinstance(inm, Variable) for inm in inames) if not np.isscalar(array.size): - - def check_size_conservation(context): - old_size = evaluate(array.size, context) - act_new_shape = tuple( - var(s.name) if hasattr(s, 'name') else s - for s in newshape) - new_size = evaluate(Product(act_new_shape), context) - if not old_size == new_size: - raise ValueError("cannot reshape array of size %s into shape %s" - % (str(old_size), str(act_new_shape))) - return old_size == new_size - # add pre-eval runtime checks on size conservation + check_size_conservation = make_size_conservation_condition(array, newshape) new_precond.append(check_size_conservation) - else: # static assersion if not np.prod(newshape) == array.size: @@ -155,9 +150,6 @@ def reshape(array, newshape, order='C', name=None, inames=None): raise ValueError("order must be either 'C' or 'F' (got %s)" % str(order)) # id in the flattened array - print(newshape) - print(inames) - print(array.shape) for iaxis in axis_range: flat_id = flat_id * newshape[iaxis] + inames[iaxis] @@ -230,7 +222,12 @@ def transpose(array, axes=None, name=None): ax_flags[axes[i]] += 1 if name is None: - name = array.name + '_T' + if len(array.name) > 1 and array.name[:2] == '__': + name_prefix = '' + else: + name_prefix = '__' + name = (name_prefix + array.name + + '_T%d' % array.__class__._counter) # noqa: W0212 new_arr = array.with_name(name) diff --git a/lappy/core/broadcast.py b/lappy/core/broadcast.py index ae96d4212e48bff8163ae4c9ec8a904f77b4dba5..ae0012ffa073839cef8ff0f6c049c05da242446b 100644 --- a/lappy/core/broadcast.py +++ b/lappy/core/broadcast.py @@ -20,9 +20,13 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ import numpy as np -from pymbolic import evaluate -from pymbolic.primitives import Product, Variable +from functools import partial +from pymbolic import evaluate, substitute, var +from pymbolic.primitives import Product, Expression, Max + +from lappy.core.array import Unsigned, to_lappy_shape from lappy.core.tools import is_nonnegative_int +from lappy.core.preconditions import EvaluationPrecondition class BroadcastResult(object): @@ -34,68 +38,228 @@ class BroadcastResult(object): list of (lazy) arrays, whose api should support .shape and .ndim queries. + .. attribute:: broadcast_arrays + + arrays that are broadcast. If an array is unchanged, an exact copy + is made; otherwise a new array is created from it. + .. attribute:: ndim ndim of the broadcast result .. attribute:: shape - shape of the broadcast result + shape of the broadcast result. Unlike Array.shape which is a tuple of + :class:`Unsigned`, it is simply a tuple of integers or strings of names, + since broadcasting tracks only the name mappings, not the detailed + expressions therein. .. attribute:: size size (number of elements) of the broadcast result - .. attribute:: reps + .. attribute:: preconditions - list of tuples, representing how each base array is tiled + list of extra preconditions to make the broadcast feasible + + NOTE: only fixed (statically known) ndims are supported """ def __init__(self, array_list): for arr in array_list: if not np.isscalar(arr.ndim): raise ValueError( - "cannot broadcast %s with variable ndim" - % str(arr)) + "cannot broadcast %s with variable ndim" + % str(arr)) self.base_arrays = list(array_list) + self.preconditions = [] self._broadcase_base_arrays() + self._make_broadcast_arrays() def _broadcase_base_arrays(self): + """Compute the ndim, shape and size of the broadcast result. + """ ndims = [arr.ndim for arr in self.base_arrays] self._ndim = max(ndims) + # values when available, names otherwise + shapes = [arr.shape for arr in self.base_arrays] + + # gather shape expressions + expr_map = {} + for arr in self.base_arrays: + for s in arr._shape: + if s.name not in expr_map: + expr_map[var(s.name)] = s.expr + else: + # same name must have the same runtime values + # (may have different expressions) + def check_name_valule_consistency(s, context): + s_val = evaluate(s.expr, context) + another_s_val = evaluate( + expr_map[var(s.name)], context) + return s_val == another_s_val + self.preconditions.append(EvaluationPrecondition( + partial(check_name_valule_consistency, s=s))) + # implicitly reshape to the same ndim # (right-aligned) - shapes = [arr.shape for arr in self.base_arrays] for ishape, shape in enumerate(shapes): assert isinstance(shape, tuple) shapes[ishape] = (1, ) * (self.ndim - len(shape)) + shape # check shape compatibility + # bshape_pre stores value or expressions with array's shape names bshape_pre = {i: 1 for i in range(self.ndim)} for shape in shapes: for iaxis, si in enumerate(shape): if bshape_pre[iaxis] == 1: - bshape_pre[iaxis] = si - else: + if isinstance(si, int): + bshape_pre[iaxis] = si + else: + assert isinstance(si, str) + bshape_pre[iaxis] = var(si) + elif isinstance(bshape_pre[iaxis], int): if bshape_pre[iaxis] == si: pass + elif si == 1: + pass + elif isinstance(si, Expression): + # add precondition + # allow symbol == 1 or symbol == constant at runtime + assert isinstance(si, str) + + def check_broadcast_symbol_val(iaxis, si, context): + si_val = evaluate(var(si), context) + return si_val in (bshape_pre[iaxis], 1) + + self.preconditions.append(EvaluationPrecondition( + partial(check_broadcast_symbol_val, si=si, iaxis=iaxis))) + + bshape_pre[iaxis] = Max((bshape_pre[iaxis], si)) else: raise ValueError( - "operands could not be broadcast " - "together with shapes %s" - % ' '.join(str(s.shape) for s in self.base_arrays)) - bshape = tuple( - int(bshape_pre[i]) if is_nonnegative_int(bshape_pre[i]) - else Variable(bshape_pre[i]) for i in range(self.ndim)) - self._shape = bshape + "operands could not be broadcast " + "together with shapes %s" + % ' '.join(str(s.shape) for s in self.base_arrays)) + else: + # bshape_pre has an symbolic member + assert isinstance(bshape_pre[iaxis], Expression) + if si == 1: + pass + elif isinstance(si, int): + # allow symbol == 1 or symbol == constant at runtime + assert isinstance(bshape_pre[iaxis], Expression) - self._size = Product(self._shape) + def check_broadcast_symbol_val(iaxis, context): + lhs_val = evaluate(bshape_pre[iaxis], context) + return lhs_val in (bshape_pre[iaxis], 1) - self._reps = [] - for shape in shapes: - rep = tuple(bshape[iaxis] if shape[iaxis] == 1 else 1 - for iaxis in range(self.ndim)) - self._reps.append(rep) + self.preconditions.append(EvaluationPrecondition( + partial(check_broadcast_symbol_val, iaxis=iaxis))) + + bshape_pre[iaxis] = Max((bshape_pre[iaxis], si)) + else: + assert isinstance(si, str) + if var(si) == bshape_pre[iaxis]: + pass + else: + # establish 'symbol == symbol' equality + assert isinstance(bshape_pre[iaxis], Expression) + assert isinstance(si, str) + + def check_broadcast_symbol_val(iaxis, si, context): + lhs_val = evaluate(bshape_pre[iaxis], context) + si_val = evaluate(var(si), context) + return (lhs_val == si_val) or ( + lhs_val == 1) or (si_val == 1) + + self.preconditions.append(EvaluationPrecondition( + partial( + check_broadcast_symbol_val, si=si, iaxis=iaxis))) + + bshape_pre[iaxis] = Max((bshape_pre[iaxis], var(si))) + + # expand expressions + for key, bs in bshape_pre.items(): + if isinstance(bs, Expression): + bshape_pre[key] = substitute(bs, expr_map) + elif is_nonnegative_int(bs): + bshape_pre[key] = int(bs) + else: + raise RuntimeError() + + bshape = tuple(bshape_pre[i] for i in range(self.ndim)) + self._shape_exprs = bshape + + self._size_expr = Product(self._shape_exprs) + + def _make_broadcast_arrays(self): + """After knowing ndim, shape and size symbolically, construct the broadcast + copies of the input arrays. + """ + self.broadcast_arrays = [] + + # Unsigneds of broadcast shape + lappy_shape = to_lappy_shape(self._shape_exprs) + + for base_arr in self.base_arrays: + dim_offset = self.ndim - base_arr.ndim + assert dim_offset >= 0 + + # ints and strs + in_shape = base_arr.shape + + if dim_offset == 0: + if in_shape == self.shape: + # unchanged + self.broadcast_arrays.append(base_arr) + continue + + name = base_arr.name + if len(name) > 1 and name[:2] == '__': + name_prefix = '' + else: + name_prefix = '__' + name = (name_prefix + name + + '_broadcast%d' % base_arr.__class__._counter) # noqa: W0212 + + brdc_arr = base_arr.with_name(name) + + assert isinstance(self.ndim, int) and self.ndim >= 0 + brdc_arr._ndim = Unsigned(value=self.ndim) + brdc_arr._shape = lappy_shape + + # make new inames + brdc_arr.inames = brdc_arr._make_default_inames() + + # make a new copy with broadcast expression + # + # account for semi-dynamic broadcasting, i.e., passing shape value 1 + # at runtime, by fetching from (out_iname % in_shape). + # TODO: when it is known that such broadcasting cannot happen, this + # may be optimized out in loopy by imposing assumptions like + # out_iname < in_shape + expr_mapping = {} + for inm, new_inm, in_shape_i in zip( + (None, ) * dim_offset + base_arr.inames, + brdc_arr.inames, + (None, ) * dim_offset + base_arr._shape_expr): + if inm is not None: + expr_mapping[inm] = new_inm % in_shape_i + brdc_arr.expr = substitute(brdc_arr.expr, expr_mapping) + + # update argument list + # if there was a self reference, the metadata needs to be captured + for argname in base_arr.arguments.keys(): + if base_arr.arguments[argname] is None: + brdc_arr.arguments[argname] = base_arr + for argname in base_arr.bound_arguments.keys(): + if base_arr.bound_arguments[argname] is None: + brdc_arr.bound_arguments[argname] = base_arr + + self.broadcast_arrays.append(brdc_arr) + + assert len(self.base_arrays) == len(self.broadcast_arrays) @property def ndim(self): @@ -103,19 +267,16 @@ class BroadcastResult(object): @property def shape(self): - return tuple(s if is_nonnegative_int(s) else s.name for s in self._shape) + return tuple(s if is_nonnegative_int(s) else str(s) + for s in self._shape_exprs) @property def size(self): from pymbolic.mapper.evaluator import UnknownVariableError try: - return evaluate(self._size) + return evaluate(self._size_expr) except UnknownVariableError: - return self._size - - @property - def reps(self): - return self._reps + return self._size_expr def broadcast(*arrays): diff --git a/lappy/core/preconditions.py b/lappy/core/preconditions.py new file mode 100644 index 0000000000000000000000000000000000000000..4e495f619f5c921aeb229398e7f09fba30a1df7b --- /dev/null +++ b/lappy/core/preconditions.py @@ -0,0 +1,154 @@ +from __future__ import division, absolute_import, print_function + +__copyright__ = "Copyright (C) 2019 Xiaoyu Wei" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" +import inspect +from pymbolic import evaluate, var +from pymbolic.primitives import Product + + +class EvaluationPrecondition(object): + """Precondition for kernel evaluations. + Each condition caputures a stack frame to trace back to the place where it is + imposed. + """ + def __init__(self, checker, name=None): + """Captures the stack frame where the condition is constructed for backtrace. + """ + assert callable(checker) + self.checker = checker + + if name is None: + self.name = self._get_checker_name() + else: + self.name = name + + self.frame = self._get_frame_obj() + + # if using Py3, call clear() + if hasattr(self.frame, 'clear'): + self.frame.clear() + + def _get_frame_obj(self): # noqa: W0613,R0201 + """Makes an extra frame so that the checker's frame can be + cleared (to break reference cycles) for better lifetime management. + """ + return inspect.currentframe() + + def _get_checker_name(self, base_checker=None): + """Find a name for the checker. + """ + if base_checker is None: + base_checker = self.checker + + if hasattr(base_checker, '__name__'): + # a function declared with def ... + return base_checker.__name__ + + elif hasattr(base_checker, 'func'): + # a function wrapper, e.g., a partial obj + return self._get_checker_name(base_checker.func) + + else: + # the last resort, e.g. unamed lambda + return 'evaluation_precondition' + + def __str__(self): + return self.name + + def __call__(self, *args, **kwargs): + return self.checker(*args, **kwargs) + + +def logical_or(*checkers): + """A meta-checker, logical OR for multiple checkers. + """ + for chkr in checkers: + assert callable(chkr) + + def meta_checker(*args, **kwargs): + res_list = [] + for chkr in checkers: + try: + res_list.append(chkr(*args, **kwargs)) + except Exception: # noqa: W0703 + res_list.append(False) + return any(res_list) + + return EvaluationPrecondition(meta_checker, 'logical_or') + + +def logical_and(*checkers): + """A meta-checker, logical AND for multiple checkers. + """ + for chkr in checkers: + assert callable(chkr) + + def meta_checker(*args, **kwargs): + res_list = [] + for chkr in checkers: + try: + res_list.append(chkr(*args, **kwargs)) + except Exception: # noqa: W0703 + res_list.append(False) + return all(res_list) + + return EvaluationPrecondition(meta_checker, 'logical_and') + + +def logical_not(checker): + """A meta-checker, logical NOT for multiple checkers. + """ + assert callable(checker) + + def meta_checker(*args, **kwargs): + try: + res = checker(*args, **kwargs) + except Exception: # noqa: W0703 + res = False + return (not res) + + return EvaluationPrecondition(meta_checker, 'logical_not') + + +# {{{ some built-in preconditions + + +def make_size_conservation_condition(array, newshape): + """Check that the array's size is as specified by the newshape. + """ + def check_size_conservation(context): + """ + :param context: a dict of {name: data} to evaluate the expressions in + """ + old_size = evaluate(array.size, context) + act_new_shape = tuple( + var(s.name) if hasattr(s, 'name') else s + for s in newshape) + new_size = evaluate(Product(act_new_shape), context) + if not old_size == new_size: + raise ValueError("cannot reshape array of size %s into shape %s" + % (str(old_size), str(act_new_shape))) + return old_size == new_size + return EvaluationPrecondition(check_size_conservation) + +# }}} End some built-in preconditions diff --git a/lappy/core/ufuncs.py b/lappy/core/ufuncs.py index 8a958f987e9f415dc38633a689f9eabc0640dfcb..7fab3c2bac8a389bdcb31bd51e4285689ae7533b 100644 --- a/lappy/core/ufuncs.py +++ b/lappy/core/ufuncs.py @@ -245,6 +245,7 @@ class BinaryOperation(UFunc): def __call__(self, a, b, name=None): bres = broadcast(a, b) + a, b = bres.broadcast_arrays if self.dtype is None: # FIXME: allow numpy type promotion rules @@ -284,18 +285,21 @@ class BinaryOperation(UFunc): assert len(a.inames) == len(b.inames) b = b.with_inames(a.inames) + # expression handling + new_expr = self.f(a.expr, b.expr) + new_assumptions = None obj = { 'name': name, 'inames': a.inames, - 'expr': self.f(a.expr, b.expr), 'value': None, + 'expr': new_expr, 'value': None, 'domain_expr': a.domain_expr, 'arguments': new_arglist, 'bound_arguments': new_bound_arglist, 'intermediaries': new_interm, 'env': new_env, 'preconditions': a.preconditions + b.preconditions, - 'ndim': bres.ndim, 'shape': bres.shape, + 'ndim': bres.ndim, 'shape': bres._shape_exprs, 'dtype': new_dtype, 'is_integral': all([ a.is_integral, b.is_integral, self.preserve_integral]), diff --git a/test/test_broadcast.py b/test/test_broadcast.py new file mode 100644 index 0000000000000000000000000000000000000000..92ebe9040f3bb7536f06fa8478e3b04fc9346058 --- /dev/null +++ b/test/test_broadcast.py @@ -0,0 +1,97 @@ +from __future__ import division, absolute_import, print_function + +__copyright__ = "Copyright (C) 2019 Xiaoyu Wei" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + +import pytest +import numpy as np +from pymbolic import evaluate +from lappy.core.array import Array +from lappy.core.broadcast import broadcast + + +@pytest.mark.parametrize('shapes', [ + ((12, 12), (12, 1)), + ((12, 12), (1, 12)), + ((12, 12), (1, 1)), + ((12, 12), tuple()), + ((12, 23, 1, 33), (12, 1, 81, 33)), + ((12, 23, 1, 33), (33, )), + ((12, 23, 1, 33), (23, 33)), + ((12, 23, 1, 33), (23, 33), (33, ), tuple(), (1, 1, 23, 1)), + ]) +def test_broadcast_rules(shapes): + + nparrs = [np.zeros(s) for s in shapes] + numpy_res = np.broadcast(*nparrs) + + laarrs = [Array(shape=s) for s in shapes] + lappy_res = broadcast(*laarrs) + + assert numpy_res.ndim == lappy_res.ndim + assert numpy_res.shape == lappy_res.shape + assert numpy_res.size == lappy_res.size + + +@pytest.mark.parametrize('shapes', [ + ((12, 12), (12, 1)), + ((12, 12), (1, 12)), + ((12, 12), (1, 1)), + ((12, 12), tuple()), + ((12, 23, 1, 33), (12, 1, 81, 33)), + ((12, 23, 1, 33), (33, )), + ((12, 23, 1, 33), (23, 33)), + ((12, 23, 1, 33), (23, 33), (33, ), tuple(), (1, 1, 23, 1)), + ]) +def test_symbolic_broadcast_rules(shapes): + + # lappy only knows ndims when doing broadcast + ndims = [len(shape) for shape in shapes] + + nparrs = [np.zeros(s) for s in shapes] + numpy_res = np.broadcast(*nparrs) + + laarrs = [Array(ndim=nd) for nd in ndims] + lappy_res = broadcast(*laarrs) + + assert numpy_res.ndim == lappy_res.ndim + + # inform lappy with actual shapes to evaluate exprs + for i in range(len(laarrs)): + laarrs[i] = laarrs[i].with_shape_data(shapes[i]) + arr_contexts = [arr.get_data_mapping() for arr in laarrs] + broadcast_ctx = {} + for ctx in arr_contexts: + broadcast_ctx.update(ctx) + + # shape and size are now expressions + # evaluate those concretely with a sum + + # check shape + lappy_res_shape = [ + evaluate(lappy_res._shape_exprs[i], broadcast_ctx) + for i in range(lappy_res.ndim)] + assert tuple(lappy_res_shape) == numpy_res.shape + + # check size + lappy_res_size = evaluate(lappy_res.size, broadcast_ctx) + assert lappy_res_size == numpy_res.size diff --git a/test/test_reshape.py b/test/test_reshape.py index 1339adc97d14742dd806a600ee8996b5ce5338dc..66eb8bbfb6eeca8e477965a234529a033dac85ae 100644 --- a/test/test_reshape.py +++ b/test/test_reshape.py @@ -61,10 +61,9 @@ def test_reshape(ctx_factory, in_shape, out_shape, order, dtype=np.float64): C = B.with_data(A=data).with_name('C') # noqa: N806 + print(C.preconditions) + lappy_val = C.eval() numpy_val = data.reshape(out_shape, order=order) - print(lappy_val) - print(numpy_val) - assert np.allclose(lappy_val, numpy_val) diff --git a/test/test_ufunc.py b/test/test_ufunc.py index 63fb49bfd507f97597d74c2684928e5a5c05f70c..296cc8b885c85aa421ed67175a8ae7b4f75ded79 100644 --- a/test/test_ufunc.py +++ b/test/test_ufunc.py @@ -92,3 +92,45 @@ def test_binary_ufunc(ctx_factory, ufunc, shapes, dtype): # compare with numpy assert np.allclose(numpy_res, lappy_res) + + +@pytest.mark.parametrize('shapes', [ + ((256, 256, 3), (3, )), + ((8, 1, 6, 1), (7, 1, 5)), + ((7, 1, 5), (8, 7, 6, 5)), + ((5, 4), (1, ), ), + ((1, ), (5, 4), ), + ((5, 4), (4, ), ), + ((4, ), (5, 4), ), + ]) +@pytest.mark.parametrize('ufunc', ['add', 'subtract', 'multiply', 'divide']) +@pytest.mark.parametrize('dtype', [np.float64, ]) +def test_binary_ufunc_with_broadcast(ctx_factory, ufunc, shapes, dtype): + env = la.core.array.default_env() + env['cl_ctx'] = ctx_factory() + + ndim_a = len(shapes[0]) + ndim_b = len(shapes[1]) + sym_shape_a = tuple('s_a_%d' % i for i in range(ndim_a)) + sym_shape_b = tuple('s_b_%d' % i for i in range(ndim_b)) + + mat_a = la.Array('A', shape=sym_shape_a, dtype=dtype, env=env) + mat_b = la.Array('B', shape=sym_shape_b, dtype=dtype, env=env) + + fmat = getattr(math, ufunc)(mat_a, mat_b).with_name('C') + + shape_a, shape_b = shapes + mat_a_data = np.random.rand(*shape_a).astype(dtype) + mat_b_data = np.random.rand(*shape_b).astype(dtype) + + bound_mat = fmat.with_data( + A=mat_a_data, + B=mat_b_data + ).with_name('C') + + # eval + lappy_res = bound_mat.eval() + numpy_res = getattr(np, ufunc)(mat_a_data, mat_b_data) + + # compare with numpy + assert np.allclose(numpy_res, lappy_res)