def depends_on(self):
from loopy.kernel.data import auto
from loopy.symbolic import DependencyMapper
if self.stride is auto:
return frozenset()
return DependencyMapper(composite_leaves=auto)(self.stride)
def depends_on(self):
"""
Returns :class:`frozenset` of all the variable names the
:class:`ArrayArgDescriptor` depends on.
"""
from loopy.kernel.data import auto
result = set()
if self.shape:
dep_mapper = DependencyMapper(composite_leaves=False)
for axis_len in self.shape:
if axis_len not in [None, auto]:
result |= dep_mapper(axis_len)
if self.dim_tags:
for dim_tag in self.dim_tags:
result |= dim_tag.depends_on()
return frozenset(var.name for var in result)
def __init__(self, kernel):
# a mapping from parameter names to a list of tuples
# (arg_name, axis_nr, function), where function is a
# unary function of kernel.arg_dict[arg_name].shape[axis_nr]
# returning the desired parameter.
self.param_to_sources = param_to_sources = {}
param_names = kernel.all_params()
from loopy.kernel.data import GlobalArg
from loopy.symbolic import DependencyMapper
from pymbolic import compile
dep_map = DependencyMapper()
from pymbolic import var
for arg in kernel.args:
if isinstance(arg, GlobalArg):
for axis_nr, shape_i in enumerate(arg.shape):
deps = dep_map(shape_i)
if len(deps) == 1:
dep, = deps
if dep.name in param_names:
from pymbolic.algorithm import solve_affine_equations_for
try:
# friggin' overkill :)
param_expr = solve_affine_equations_for(
[dep.name],
[(shape_i, var("shape_i"))])[dep.name]
except:
# went wrong? oh well
pass
else:
param_func = compile(param_expr, ["shape_i"])
param_to_sources.setdefault(
dep.name, []).append(
(arg.name, axis_nr, param_func))
def generate_integer_arg_finding_from_shapes(gen, kernel,
implemented_data_info):
# a mapping from integer argument names to a list of tuples
# (arg_name, expression), where expression is a
# unary function of kernel.arg_dict[arg_name]
# returning the desired integer argument.
iarg_to_sources = {}
from loopy.kernel.data import GlobalArg
from loopy.symbolic import DependencyMapper, StringifyMapper
dep_map = DependencyMapper()
from pymbolic import var
for arg in implemented_data_info:
if arg.arg_class is GlobalArg:
sym_shape = var(arg.name).attr("shape")
for axis_nr, shape_i in enumerate(arg.shape):
if shape_i is None:
continue
deps = dep_map(shape_i)
if len(deps) == 1:
integer_arg_var, = deps
if kernel.arg_dict[
integer_arg_var.name].dtype.is_integral():
from pymbolic.algorithm import solve_affine_equations_for
try:
# friggin' overkill :)
iarg_expr = solve_affine_equations_for(
[integer_arg_var.name], [
(shape_i, sym_shape.index(axis_nr))
])[integer_arg_var]
except Exception as e:
#from traceback import print_exc
#print_exc()
# went wrong? oh well
from warnings import warn
warn(
"Unable to generate code to automatically "
"find '%s' from the shape of '%s':\n%s" %
(integer_arg_var.name, arg.name, str(e)),
ParameterFinderWarning)
else:
iarg_to_sources.setdefault(integer_arg_var.name, []) \
.append((arg.name, iarg_expr))
gen("# {{{ find integer arguments from shapes")
gen("")
for iarg_name, sources in six.iteritems(iarg_to_sources):
gen("if %s is None:" % iarg_name)
with Indentation(gen):
if_stmt = "if"
for arg_name, value_expr in sources:
gen("%s %s is not None:" % (if_stmt, arg_name))
with Indentation(gen):
gen("%s = %s" % (iarg_name, StringifyMapper()(value_expr)))
if_stmt = "elif"
gen("")
gen("# }}}")
gen("")
