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)
