def map_subscript(self, expr, domain, insn_id): WalkMapper.map_subscript(self, expr, domain, insn_id) from pymbolic.primitives import Variable assert isinstance(expr.aggregate, Variable) shape = None var_name = expr.aggregate.name if var_name in self.kernel.arg_dict: arg = self.kernel.arg_dict[var_name] shape = arg.shape elif var_name in self.kernel.temporary_variables: tv = self.kernel.temporary_variables[var_name] shape = tv.shape if shape is not None: subscript = expr.index if not isinstance(subscript, tuple): subscript = (subscript, ) from loopy.symbolic import get_dependencies available_vars = set(domain.get_var_dict()) shape_deps = set() for shape_axis in shape: if shape_axis is not None: shape_deps.update(get_dependencies(shape_axis)) if not (get_dependencies(subscript) <= available_vars and shape_deps <= available_vars): return if len(subscript) != len(shape): raise LoopyError( "subscript to '%s' in '%s' has the wrong " "number of indices (got: %d, expected: %d)" % (expr.aggregate.name, expr, len(subscript), len(shape))) access_range = self._get_access_range(domain, subscript) if access_range is None: # Likely: index was non-affine, nothing we can do. return shape_domain = isl.BasicSet.universe(access_range.get_space()) for idim in range(len(subscript)): shape_axis = shape[idim] if shape_axis is not None: slab = self._make_slab(shape_domain.get_space(), (dim_type.in_, idim), 0, shape_axis) shape_domain = shape_domain.intersect(slab) if not access_range.is_subset(shape_domain): raise LoopyError( "'%s' in instruction '%s' " "accesses out-of-bounds array element (could not" " establish '%s' is a subset of '%s')." % (expr, insn_id, access_range, shape_domain))
def map_subscript(self, expr): WalkMapper.map_subscript(self, expr) from pymbolic.primitives import Variable assert isinstance(expr.aggregate, Variable) if expr.aggregate.name == self.arg_name: if not isinstance(expr.index, tuple): self.index_ranks.append(1) else: self.index_ranks.append(len(expr.index))
def map_subscript(self, expr): WalkMapper.map_subscript(self, expr) from pymbolic.primitives import Variable assert isinstance(expr.aggregate, Variable) shape = None var_name = expr.aggregate.name if var_name in self.kernel.arg_dict: arg = self.kernel.arg_dict[var_name] shape = arg.shape elif var_name in self.kernel.temporary_variables: tv = self.kernel.temporary_variables[var_name] shape = tv.shape if shape is not None: subscript = expr.index if not isinstance(subscript, tuple): subscript = (subscript,) from loopy.symbolic import get_dependencies, get_access_range available_vars = set(self.domain.get_var_dict()) shape_deps = set() for shape_axis in shape: if shape_axis is not None: shape_deps.update(get_dependencies(shape_axis)) if not (get_dependencies(subscript) <= available_vars and shape_deps <= available_vars): return if len(subscript) != len(shape): raise LoopyError("subscript to '%s' in '%s' has the wrong " "number of indices (got: %d, expected: %d)" % ( expr.aggregate.name, expr, len(subscript), len(shape))) try: access_range = get_access_range(self.domain, subscript, self.kernel.assumptions) except isl.Error: # Likely: index was non-linear, nothing we can do. return except TypeError: # Likely: index was non-linear, nothing we can do. return shape_domain = isl.BasicSet.universe(access_range.get_space()) for idim in range(len(subscript)): shape_axis = shape[idim] if shape_axis is not None: from loopy.isl_helpers import make_slab slab = make_slab( shape_domain.get_space(), (dim_type.in_, idim), 0, shape_axis) shape_domain = shape_domain.intersect(slab) if not access_range.is_subset(shape_domain): raise LoopyError("'%s' in instruction '%s' " "accesses out-of-bounds array element" % (expr, self.insn_id))
def extract_subst(kernel, subst_name, template, parameters=()): """ :arg subst_name: The name of the substitution rule to be created. :arg template: Unification template expression. :arg parameters: An iterable of parameters used in *template*, or a comma-separated string of the same. All targeted subexpressions must match ('unify with') *template* The template may contain '*' wildcards that will have to match exactly across all unifications. """ if isinstance(template, str): from pymbolic import parse template = parse(template) if isinstance(parameters, str): parameters = tuple(s.strip() for s in parameters.split(",")) var_name_gen = kernel.get_var_name_generator() # {{{ replace any wildcards in template with new variables def get_unique_var_name(): based_on = subst_name + "_wc" result = var_name_gen(based_on) return result from loopy.symbolic import WildcardToUniqueVariableMapper wc_map = WildcardToUniqueVariableMapper(get_unique_var_name) template = wc_map(template) # }}} # {{{ gather up expressions expr_descriptors = [] from loopy.symbolic import UnidirectionalUnifier unif = UnidirectionalUnifier(lhs_mapping_candidates=set(parameters)) def gather_exprs(expr, mapper): urecs = unif(template, expr) if urecs: if len(urecs) > 1: raise RuntimeError( "ambiguous unification of '%s' with template '%s'" % (expr, template)) urec, = urecs expr_descriptors.append( ExprDescriptor(insn=insn, expr=expr, unif_var_dict=dict( (lhs.name, rhs) for lhs, rhs in urec.equations))) else: mapper.fallback_mapper(expr) # can't nest, don't recurse from loopy.symbolic import (CallbackMapper, WalkMapper, IdentityMapper) dfmapper = CallbackMapper(gather_exprs, WalkMapper()) for insn in kernel.instructions: dfmapper(insn.assignees) dfmapper(insn.expression) for sr in six.itervalues(kernel.substitutions): dfmapper(sr.expression) # }}} if not expr_descriptors: raise RuntimeError("no expressions matching '%s'" % template) # {{{ substitute rule into instructions def replace_exprs(expr, mapper): found = False for exprd in expr_descriptors: if expr is exprd.expr: found = True break if not found: return mapper.fallback_mapper(expr) args = [exprd.unif_var_dict[arg_name] for arg_name in parameters] result = var(subst_name) if args: result = result(*args) return result # can't nest, don't recurse cbmapper = CallbackMapper(replace_exprs, IdentityMapper()) new_insns = [] for insn in kernel.instructions: new_insns.append(insn.with_transformed_expressions(cbmapper)) from loopy.kernel.data import SubstitutionRule new_substs = { subst_name: SubstitutionRule( name=subst_name, arguments=tuple(parameters), expression=template, ) } for subst in six.itervalues(kernel.substitutions): new_substs[subst.name] = subst.copy( expression=cbmapper(subst.expression)) # }}} return kernel.copy(instructions=new_insns, substitutions=new_substs)
def extract_subst(kernel, subst_name, template, parameters=(), within=None): """ :arg subst_name: The name of the substitution rule to be created. :arg template: Unification template expression. :arg parameters: An iterable of parameters used in *template*, or a comma-separated string of the same. :arg within: An instance of :class:`loopy.match.MatchExpressionBase` or :class:`str` as understood by :func:`loopy.match.parse_match`. All targeted subexpressions must match ('unify with') *template* The template may contain '*' wildcards that will have to match exactly across all unifications. """ if isinstance(kernel, TranslationUnit): kernel_names = [ i for i, clbl in kernel.callables_table.items() if isinstance(clbl, CallableKernel) ] if len(kernel_names) != 1: raise LoopyError() return kernel.with_kernel( extract_subst(kernel[kernel_names[0]], subst_name, template, parameters)) if isinstance(template, str): from pymbolic import parse template = parse(template) if isinstance(parameters, str): parameters = tuple(s.strip() for s in parameters.split(",")) from loopy.match import parse_match within = parse_match(within) var_name_gen = kernel.get_var_name_generator() # {{{ replace any wildcards in template with new variables def get_unique_var_name(): based_on = subst_name + "_wc" result = var_name_gen(based_on) return result from loopy.symbolic import WildcardToUniqueVariableMapper wc_map = WildcardToUniqueVariableMapper(get_unique_var_name) template = wc_map(template) # }}} # {{{ gather up expressions expr_descriptors = [] from loopy.symbolic import UnidirectionalUnifier unif = UnidirectionalUnifier(lhs_mapping_candidates=set(parameters)) def gather_exprs(expr, mapper): urecs = unif(template, expr) if urecs: if len(urecs) > 1: raise RuntimeError( "ambiguous unification of '%s' with template '%s'" % (expr, template)) urec, = urecs expr_descriptors.append( ExprDescriptor(insn=insn, expr=expr, unif_var_dict={ lhs.name: rhs for lhs, rhs in urec.equations })) else: mapper.fallback_mapper(expr) # can't nest, don't recurse from loopy.symbolic import (CallbackMapper, WalkMapper, IdentityMapper) dfmapper = CallbackMapper(gather_exprs, WalkMapper()) from loopy.kernel.instruction import MultiAssignmentBase for insn in kernel.instructions: if isinstance(insn, MultiAssignmentBase) and within(kernel, insn): dfmapper(insn.assignees) dfmapper(insn.expression) for sr in kernel.substitutions.values(): dfmapper(sr.expression) # }}} if not expr_descriptors: raise RuntimeError("no expressions matching '%s'" % template) # {{{ substitute rule into instructions def replace_exprs(expr, mapper): found = False for exprd in expr_descriptors: if expr is exprd.expr: found = True break if not found: return mapper.fallback_mapper(expr) args = [exprd.unif_var_dict[arg_name] for arg_name in parameters] result = var(subst_name) if args: result = result(*args) return result # can't nest, don't recurse cbmapper = CallbackMapper(replace_exprs, IdentityMapper()) new_insns = [] def transform_assignee(expr): # Assignment LHS's cannot be subst rules. Treat them # specially. import pymbolic.primitives as prim if isinstance(expr, tuple): return tuple(transform_assignee(expr_i) for expr_i in expr) elif isinstance(expr, prim.Subscript): return type(expr)(expr.aggregate, cbmapper(expr.index)) elif isinstance(expr, prim.Variable): return expr else: raise ValueError("assignment LHS not understood") for insn in kernel.instructions: if within(kernel, insn): new_insns.append( insn.with_transformed_expressions( cbmapper, assignee_f=transform_assignee)) else: new_insns.append(insn) from loopy.kernel.data import SubstitutionRule new_substs = { subst_name: SubstitutionRule( name=subst_name, arguments=tuple(parameters), expression=template, ) } for subst in kernel.substitutions.values(): new_substs[subst.name] = subst.copy( expression=cbmapper(subst.expression)) # }}} return kernel.copy(instructions=new_insns, substitutions=new_substs)