def is_nonnegative(expr, over_set): space = over_set.get_space() from loopy.symbolic import aff_from_expr try: with isl.SuppressedWarnings(space.get_ctx()): aff = aff_from_expr(space, -expr - 1) except Exception: return None expr_neg_set = isl.BasicSet.universe(space).add_constraint( isl.Constraint.inequality_from_aff(aff)) return over_set.intersect(expr_neg_set).is_empty()
def simplify_using_aff(kernel, expr): inames = get_dependencies(expr) & kernel.all_inames() domain = kernel.get_inames_domain(inames) from pymbolic.mapper.evaluator import UnknownVariableError try: with isl.SuppressedWarnings(kernel.isl_context): aff = aff_from_expr(domain.space, expr) except isl.Error: return expr except TypeError: return expr except UnknownVariableError: return expr # FIXME: Deal with assumptions, too. aff = aff.gist(domain) return aff_to_expr(aff)
def get_iname_length(iname): try: with isl.SuppressedWarnings(kernel.isl_context): return kernel.get_constant_iname_length(iname) except isl.Error: return -1
def assign_axis(recursion_axis, iname, axis=None): """Assign iname to local axis *axis* and start over by calling the surrounding function assign_automatic_axes. If *axis* is None, find a suitable axis automatically. """ try: with isl.SuppressedWarnings(kernel.isl_context): desired_length = kernel.get_constant_iname_length(iname) except isl.Error: # Likely unbounded, automatic assignment is not # going to happen for this iname. new_iname_to_tag = kernel.iname_to_tag.copy() new_iname_to_tag[iname] = None return assign_automatic_axes( kernel.copy(iname_to_tag=new_iname_to_tag), axis=recursion_axis) if axis is None: # {{{ find a suitable axis shorter_possible_axes = [] test_axis = 0 while True: if test_axis >= len(local_size): break if test_axis in assigned_local_axes: test_axis += 1 continue if local_size[test_axis] < desired_length: shorter_possible_axes.append(test_axis) test_axis += 1 continue else: axis = test_axis break # The loop above will find an unassigned local axis # that has enough 'room' for the iname. In the same traversal, # it also finds theoretically assignable axes that are shorter, # in the variable shorter_possible_axes. if axis is None and shorter_possible_axes: # sort as longest first shorter_possible_axes.sort(key=lambda ax: local_size[ax]) axis = shorter_possible_axes[0] # }}} if axis is None: new_tag = None else: new_tag = LocalIndexTag(axis) if desired_length > local_size[axis]: from loopy import split_iname # Don't be tempted to switch the outer tag to unroll--this may # generate tons of code on some examples. return assign_automatic_axes(split_iname( kernel, iname, inner_length=local_size[axis], outer_tag=None, inner_tag=new_tag, do_tagged_check=False), axis=recursion_axis, local_size=local_size) if not isinstance(kernel.iname_to_tag.get(iname), AutoLocalIndexTagBase): raise LoopyError("trying to reassign '%s'" % iname) new_iname_to_tag = kernel.iname_to_tag.copy() new_iname_to_tag[iname] = new_tag return assign_automatic_axes( kernel.copy(iname_to_tag=new_iname_to_tag), axis=recursion_axis, local_size=local_size)