from loopy.symbolic import get_dependencies

for insn in kernel.instructions:

pred_vars = (

frozenset.union(

*(get_dependencies(pred) for pred in insn.predicates))

if insn.predicates else frozenset())

written_pred_vars = frozenset(insn.assignee_var_names()) & pred_vars

if written_pred_vars:

raise LoopyError("In instruction '%s': may not write to "

"variable(s) '%s' involved in the instruction's predicates"

iname_to_reduction_count = {}

iname_to_nonsimultaneous_reduction_count = {}

def map_reduction(expr, rec):

rec(expr.expr)

for iname in expr.inames:

iname_to_reduction_count[iname] = (

iname_to_reduction_count.get(iname, 0) + 1)

if not expr.allow_simultaneous:

iname_to_nonsimultaneous_reduction_count[iname] = (

iname_to_nonsimultaneous_reduction_count.get(iname, 0) + 1)

return expr

from loopy.symbolic import ReductionCallbackMapper

cb_mapper = ReductionCallbackMapper(map_reduction)

for insn in kernel.instructions:

insn.with_transformed_expressions(cb_mapper)

for iname, count in iname_to_reduction_count.items():

nonsimul_count = iname_to_nonsimultaneous_reduction_count.get(iname, 0)

if nonsimul_count and count > 1:

raise LoopyError("iname '%s' used in more than one reduction. "

"(%d of them, to be precise.) "

"Since this usage can easily cause loop scheduling "

"problems, this is prohibited by default. "

"Use loopy.make_reduction_inames_unique() to fix this. "

"If you are sure that this is OK, write the reduction "

"as 'simul_reduce(...)' instead of 'reduce(...)'"

indices: frozenset[int], values: Iterable[T] | None

) -> tuple[T, ...] | None:

"""

if values is None:

return values

from loopy.kernel.array import SeparateArrayArrayDimTag

new_args = []

vng = kernel.get_var_name_generator()

made_changes = False

# {{{ rewrite arguments

for arg in kernel.args:

if not isinstance(arg, ArrayArg) or arg.dim_tags is None:

new_args.append(arg)

continue

sep_axis_indices = [

i for i, dim_tag in enumerate(arg.dim_tags)

if isinstance(dim_tag, SeparateArrayArrayDimTag)]

if not sep_axis_indices or arg._separation_info:

new_args.append(arg)

continue

made_changes = True

sep_axis_indices_set = frozenset(sep_axis_indices)

assert isinstance(arg.shape, tuple)

new_shape: tuple[Expression, ...] | None = \

_remove_at_indices(sep_axis_indices_set, arg.shape)

new_dim_tags: tuple[ArrayDimImplementationTag, ...] | None = \

_remove_at_indices(sep_axis_indices_set, arg.dim_tags)

new_dim_names: tuple[str | None, ...] | None = \

_remove_at_indices(sep_axis_indices_set, arg.dim_names)

sep_shape: list[Expression] = [arg.shape[i] for i in sep_axis_indices]

for i, sep_shape_i in enumerate(sep_shape):

if not isinstance(sep_shape_i, (int, np.integer)):

raise LoopyError(

f"Axis {sep_axis_indices[i]+1} (1-based) of "

f"argument '{arg.name}' is tagged 'sep', but "

"does not have constant length.")

sep_info = _ArraySeparationInfo(

sep_axis_indices_set=sep_axis_indices_set,

subarray_names=constantdict({

ind: vng(f"{arg.name}_s{'_'.join(str(i) for i in ind)}")

for ind in np.ndindex(*cast("list[int]", sep_shape))}))

new_args.append(arg.copy(_separation_info=sep_info))

for _ind, san in sorted(sep_info.subarray_names.items()):

new_args.append(

arg.copy(

name=san,

shape=new_shape,

dim_tags=new_dim_tags,

dim_names=new_dim_names))

# }}}

if not made_changes:

return kernel

additional_args: list[KernelArgument] = []

vng = kernel.get_var_name_generator()

from pymbolic.primitives import ExpressionNode, Variable

from loopy.kernel.array import FixedStrideArrayDimTag

# {{{ process arguments

new_args = []

for arg in kernel.args:

if isinstance(arg, ArrayArg) and not arg._separation_info:

what = f"offset for argument '{arg.name}'"

if arg.offset is auto:

offset_name = vng(arg.name+"_offset")

additional_args.append(ValueArg(

offset_name, kernel.index_dtype))

arg = arg.copy(offset=offset_name)

elif isinstance(arg.offset, (int, np.integer, ExpressionNode, str)):

pass

else:

raise LoopyError(f"invalid value of {what}")

if arg.dim_tags is None:

new_dim_tags: tuple[ArrayDimImplementationTag, ...] | None \

= arg.dim_tags

else:

new_dim_tags = ()

for iaxis, dim_tag in enumerate(arg.dim_tags):

if isinstance(dim_tag, FixedStrideArrayDimTag):

what = ("axis stride for axis "

f"{iaxis+1} (1-based) of '{arg.name}'")

if dim_tag.stride is auto:

stride_name = vng(f"{arg.name}_stride{iaxis}")

dim_tag = dim_tag.copy(stride=Variable(stride_name))

additional_args.append(ValueArg(

stride_name, kernel.index_dtype))

elif isinstance(

dim_tag.stride, (int, np.integer, ExpressionNode)):

pass

else:

raise LoopyError(f"invalid value of {what}")

new_dim_tags = (*new_dim_tags, dim_tag)

arg = arg.copy(dim_tags=new_dim_tags)

new_args.append(arg)

# }}}

if not additional_args:

return kernel

else:

made_changes = False

from pymbolic.primitives import ExpressionNode

# {{{ process arguments

new_args = []

for arg in kernel.args:

if isinstance(arg, ArrayArg):

if arg.offset is auto:

made_changes = True

arg = arg.copy(offset=0)

elif isinstance(arg.offset, (int, np.integer, ExpressionNode, str)):

from pymbolic.primitives import is_zero

if not is_zero(arg.offset):

raise LoopyError(

f"Non-zero offset on argument '{arg.name}' "

f"of callable kernel '{kernel.name}. This is not allowed.")

else:

raise LoopyError(f"invalid value of offset for '{arg.name}'")

new_args.append(arg)

# }}}

if not made_changes:

return kernel

else:

return {iname

for aname, adeps in zip(

insn.assignee_var_names(),

insn.assignee_subscript_deps(), strict=True)

for iname in adeps & kernel.all_inames()

if aname in tv_names

logger.debug("%s: find temporary address space" % kernel.name)

new_temp_vars = {}

import loopy as lp

from loopy.kernel.data import AddressSpace, GroupInameTag, LocalInameTagBase

writers = kernel.writer_map()

base_storage_to_aliases = {}

for temp_var in kernel.temporary_variables.values():

if temp_var.base_storage is not None:

base_storage_to_aliases.setdefault(

temp_var.base_storage, []).append(temp_var.name)

for temp_var in kernel.temporary_variables.values():

# Only fill out for variables that do not yet know if they're

# local. (I.e. those generated by implicit temporary generation.)

if temp_var.address_space is not lp.auto:

new_temp_vars[temp_var.name] = temp_var

continue

tv_names = (frozenset([temp_var.name])

| frozenset(base_storage_to_aliases.get(temp_var.base_storage, [])))

my_writers = writers.get(temp_var.name, frozenset())

if temp_var.base_storage is not None:

for alias in base_storage_to_aliases.get(temp_var.base_storage, []):

my_writers = my_writers | writers.get(alias, frozenset())

desired_aspace_per_insn = []

for insn_id in my_writers:

insn = kernel.id_to_insn[insn_id]

# A write race will emerge if:

#

# - the variable is local

# and

# - the instruction is run across more inames (locally) parallel

# than are reflected in the assignee indices.

locparallel_compute_inames = _get_compute_inames_tagged(

kernel, insn, LocalInameTagBase)

locparallel_assignee_inames = _get_assignee_inames_tagged(

kernel, insn, LocalInameTagBase, tv_names)

grpparallel_compute_inames = _get_compute_inames_tagged(

kernel, insn, GroupInameTag)

grpparallel_assignee_inames = _get_assignee_inames_tagged(

kernel, insn, GroupInameTag, temp_var.name)

assert locparallel_assignee_inames <= locparallel_compute_inames

assert grpparallel_assignee_inames <= grpparallel_compute_inames

desired_aspace = AddressSpace.PRIVATE

for iname_descr, aspace_descr, apin, cpin, aspace in [

("local", "local", locparallel_assignee_inames,

locparallel_compute_inames, AddressSpace.LOCAL),

("group", "global", grpparallel_assignee_inames,

grpparallel_compute_inames, AddressSpace.GLOBAL),

]:

if (apin != cpin and bool(apin)):

warn_with_kernel(

kernel,

f"write_race_{aspace_descr}({insn_id})",

"instruction '%s' looks invalid: "

"it assigns to indices based on %s IDs, but "

"its temporary '%s' cannot be made %s because "

"a write race across the iname(s) '%s' would emerge. "

"(Do you need to add an extra iname to your prefetch?)"

% (insn_id, iname_descr, temp_var.name, aspace_descr,

", ".join(cpin - apin)),

WriteRaceConditionWarning)

if (apin == cpin

# doesn't want to be in this address space if there

# aren't any parallel inames of that kind

and bool(cpin)):

desired_aspace = max(desired_aspace, aspace)

desired_aspace_per_insn.append(desired_aspace)

if not desired_aspace_per_insn:

warn_with_kernel(kernel, "temp_to_write(%s)" % temp_var.name,

"cannot automatically determine address space of '%s'"

% temp_var.name, LoopyAdvisory)

new_temp_vars[temp_var.name] = temp_var

continue

overall_aspace = max(desired_aspace_per_insn)

if not all(iaspace == overall_aspace for iaspace in desired_aspace_per_insn):

raise LoopyError("not all instructions agree on the "

"the desired address space (private/local/global) of the "

"temporary '%s'" % temp_var.name)

new_temp_vars[temp_var.name] = temp_var.copy(address_space=overall_aspace)

logger.debug("%s: add axes to temporaries for ilp" % kernel.name)

from loopy.kernel.data import IlpBaseTag, VectorizeTag, filter_iname_tags_by_type

privatizing_inames = frozenset(

name for name, iname in kernel.inames.items()

if filter_iname_tags_by_type(iname.tags, (IlpBaseTag, VectorizeTag))

)

if not privatizing_inames:

return kernel

from loopy.transform.privatize import privatize_temporaries_with_inames

"""Find instances of AtomicInit or AtomicUpdate with use of other atomic

logger.debug("%s: check atomic loads" % kernel.name)

from loopy.kernel.array import ArrayBase

from loopy.kernel.instruction import Assignment, AtomicLoad

from loopy.types import AtomicType

# find atomic variables

atomicity_candidates = (

{v.name for v in kernel.temporary_variables.values()

if isinstance(v.dtype, AtomicType)}

|

{v.name for v in kernel.args

if isinstance(v, ArrayBase)

and isinstance(v.dtype, AtomicType)})

new_insns = []

for insn in kernel.instructions:

if isinstance(insn, Assignment):

# look for atomic variables

atomic_accesses = {a.var_name for a in insn.atomicity}

accessed_atomic_vars = (insn.dependency_names() & atomicity_candidates)\

- {insn.assignee_var_names()[0]}

if not accessed_atomic_vars <= atomic_accesses:

# if we're missing some

missed = accessed_atomic_vars - atomic_accesses

for x in missed:

if {x} & atomicity_candidates:

insn = insn.copy(

atomicity=(*insn.atomicity, AtomicLoad(x)))

new_insns.append(insn)

"""

def __init__(self, rule_mapping_context, caller_kernel, clbl_inf_ctx):

super().__init__(rule_mapping_context)

self.caller_kernel = caller_kernel

self.clbl_inf_ctx = clbl_inf_ctx

def map_call(self, expr, expn_state, assignees=None):

from pymbolic.mapper.substitutor import make_subst_func

from pymbolic.primitives import Call, Variable

from loopy.kernel.array import ArrayBase

from loopy.kernel.data import ValueArg

from loopy.kernel.function_interface import (

ValueArgDescriptor,

get_arg_descriptor_for_expression,

)

from loopy.symbolic import ResolvedFunction, SubstitutionMapper

if not isinstance(expr.function, ResolvedFunction):

# ignore if the call is not to a ResolvedFunction

return super().map_call(expr, expn_state)

arg_id_to_arg = dict(enumerate(expr.parameters))

if assignees is not None:

# If supplied with assignees then this is a CallInstruction

for i, arg in enumerate(assignees):

arg_id_to_arg[-i-1] = arg

arg_id_to_descr = {

arg_id: get_arg_descriptor_for_expression(self.caller_kernel, arg)

for arg_id, arg in arg_id_to_arg.items()}

clbl = self.clbl_inf_ctx[expr.function.name]

# {{{ translating descriptor expressions to the callable's namespace

deps_as_params = []

subst_map = {}

deps = frozenset().union(*(descr.depends_on()

for descr in arg_id_to_descr.values()))

assert deps <= self.caller_kernel.all_variable_names()

for dep in deps:

caller_arg = self.caller_kernel.arg_dict.get(dep, (self.caller_kernel

.temporary_variables

.get(dep)))

if not (isinstance(caller_arg, ValueArg)

or (isinstance(caller_arg, ArrayBase)

and caller_arg.shape == ())):

raise NotImplementedError(f"Obtained '{dep}' as a dependency for"

f" call '{expr.function.name}' which is not a scalar.")

clbl, callee_name = clbl.with_added_arg(caller_arg.dtype,

ValueArgDescriptor())

subst_map[dep] = Variable(callee_name)

deps_as_params.append(Variable(dep))

mapper = SubstitutionMapper(make_subst_func(subst_map))

arg_id_to_descr = {id_: descr.map_expr(mapper)

for id_, descr in arg_id_to_descr.items()}

# }}}

# specializing the function according to the parameter description

new_clbl, self.clbl_inf_ctx = clbl.with_descrs(arg_id_to_descr,

self.clbl_inf_ctx)

self.clbl_inf_ctx, new_func_id = (self.clbl_inf_ctx

.with_callable(expr.function.function,

new_clbl))

return Call(ResolvedFunction(new_func_id),

tuple(self.rec(child, expn_state)

for child in expr.parameters)

+ tuple(deps_as_params))

def map_call_with_kwargs(self, expr):

# See https://github.com/inducer/loopy/pull/323

raise NotImplementedError

@override

def __call__(self,

expr: Expression,

kernel: LoopKernel,

insn: InstructionBase,

assignees=None):

from loopy.symbolic import ExpansionState, UncachedIdentityMapper

return UncachedIdentityMapper.__call__(self, expr,

ExpansionState(

kernel=kernel,

instruction=insn,

stack=(),

arg_context=constantdict()), assignees=assignees)

@override

def map_kernel(self, kernel: LoopKernel):

new_insns = []

for insn in kernel.instructions:

if isinstance(insn, CallInstruction):

# In call instructions the assignees play an important in

# determining the arg_id_to_descr

mapper = partial(self, kernel=kernel, insn=insn,

assignees=insn.assignees)

new_insns.append(insn.with_transformed_expressions(mapper))

elif isinstance(insn, MultiAssignmentBase):

mapper = partial(self, kernel=kernel, insn=insn)

new_insns.append(insn.with_transformed_expressions(mapper))

elif isinstance(insn, (_DataObliviousInstruction, CInstruction)):

new_insns.append(insn)

else:

raise NotImplementedError("arg_descr_inference for %s instruction" %

type(insn))

"""

from loopy.symbolic import SubstitutionRuleMappingContext

rule_mapping_context = SubstitutionRuleMappingContext(

kernel.substitutions, kernel.get_var_name_generator())

arg_descr_inf_mapper = ArgDescrInferenceMapper(rule_mapping_context,

kernel, callables_table)

descr_inferred_kernel = rule_mapping_context.finish_kernel(

arg_descr_inf_mapper.map_kernel(kernel))

"""

from loopy import ValueArg, auto

from loopy.kernel.array import ArrayBase

from loopy.kernel.function_interface import ArrayArgDescriptor, ValueArgDescriptor

from loopy.translation_unit import make_clbl_inf_ctx, resolve_callables

t_unit = resolve_callables(t_unit)

clbl_inf_ctx = make_clbl_inf_ctx(t_unit.callables_table,

t_unit.entrypoints)

for e in t_unit.entrypoints:

def _tuple_or_none(s):

if isinstance(s, tuple) or s in [None, auto]:

return s

else:

return (s,)

arg_id_to_descr: dict[str, ArgDescriptor] = {}

for arg in t_unit[e].args:

if isinstance(arg, ArrayBase):

if arg.shape not in (None, auto):

arg_id_to_descr[arg.name] = ArrayArgDescriptor(

_tuple_or_none(arg.shape), arg.address_space,

arg.dim_tags)

elif isinstance(arg, ValueArg):

arg_id_to_descr[arg.name] = ValueArgDescriptor()

else:

raise NotImplementedError()

new_callable, clbl_inf_ctx = t_unit.callables_table[e].with_descrs(

# FIXME: Keyword args are at best half-implemented

arg_id_to_descr, clbl_inf_ctx) # pyright: ignore[reportArgumentType]

clbl_inf_ctx, _new_name = clbl_inf_ctx.with_callable(e, new_callable,

is_entrypoint=True)

from pytools.graph import compute_topological_order

from loopy.kernel.instruction import BarrierInstruction

from loopy.kernel.tools import get_call_graph

from loopy.transform.callable import inline_callable_kernel

def has_gbarrier(knl):

return any((isinstance(insn, BarrierInstruction)

and insn.synchronization_kind == "global")

for insn in knl.instructions)

call_graph = get_call_graph(program, only_kernel_callables=True)

# traverse the kernel calls in a reverse topological sort so that barriers

# are rightly passed to the entrypoints.

toposort = compute_topological_order(call_graph,

# pass key to have deterministic codegen

key=lambda x: x

)

for name in toposort[::-1]:

if has_gbarrier(program[name]):

program = inline_callable_kernel(program, name)

from loopy.translation_unit import get_reachable_resolved_callable_ids

reachable_function_ids = get_reachable_resolved_callable_ids(t_unit

.callables_table,

t_unit.entrypoints)

new_callables = {name: clbl for name, clbl in t_unit.callables_table.items()

if name in (reachable_function_ids | t_unit.entrypoints)}

from loopy.kernel import KernelState

prepro_logger = ProcessLogger(logger, "%s: preprocess" % kernel.name)

kernel = make_arrays_for_sep_arrays(kernel)

if is_entrypoint:

kernel = make_args_for_offsets_and_strides(kernel)

else:

# No need for offsets internally, we can pass arbitrary pointers.

kernel = zero_offsets_and_strides(kernel)

from loopy.check import check_identifiers_in_subst_rules

check_identifiers_in_subst_rules(kernel)

# {{{ check that there are no l.auto-tagged inames

from loopy.kernel.data import AutoLocalInameTagBase

for name, iname in kernel.inames.items():

if (filter_iname_tags_by_type(iname.tags, AutoLocalInameTagBase)

and name in kernel.all_inames()):

raise LoopyError("kernel with automatically-assigned "

"local axes passed to preprocessing")

# }}}

# Ordering restriction:

# Type inference and reduction iname uniqueness don't handle substitutions.

# Get them out of the way.

check_for_writes_to_predicates(kernel)

check_reduction_iname_uniqueness(kernel)

# Ordering restriction:

# add_axes_to_temporaries_for_ilp because reduction accumulators

# need to be duplicated by this.

kernel = realize_ilp(kernel)

kernel = find_temporary_address_space(kernel)

# check for atomic loads, much easier to do here now that the dependencies

# have been established

kernel = check_atomic_loads(kernel)

kernel = kernel.target.preprocess(kernel)

kernel = kernel.copy(

state=KernelState.PREPROCESSED)

prepro_logger.done()

from loopy.kernel import KernelState

if t_unit.state >= KernelState.PREPROCESSED:

return t_unit

if len([clbl for clbl in t_unit.callables_table.values() if

isinstance(clbl, CallableKernel)]) == 1:

t_unit = t_unit.with_entrypoints(",".join(clbl.name for clbl in

t_unit.callables_table.values() if isinstance(clbl,

CallableKernel)))

if not t_unit.entrypoints:

raise LoopyError("Translation unit did not receive any entrypoints")

from loopy.translation_unit import resolve_callables

t_unit = resolve_callables(t_unit)

t_unit = filter_reachable_callables(t_unit)

t_unit = infer_unknown_types(t_unit, expect_completion=False)

from loopy.transform.subst import expand_subst

t_unit = expand_subst(t_unit)

from loopy.kernel.creation import apply_single_writer_dependency_heuristic

t_unit = apply_single_writer_dependency_heuristic(t_unit)

# Ordering restrictions:

#

# - realize_reduction must happen after type inference because it needs

# to be able to determine the types of the reduced expressions.

#

# - realize_reduction must happen after default dependencies are added

# because it manipulates the depends_on field, which could prevent

# defaults from being applied.

from loopy.transform.realize_reduction import realize_reduction

t_unit = realize_reduction(t_unit, unknown_types_ok=False)

# {{{ preprocess callable kernels

# Callable editing restrictions:

#

# - should not edit callables_table in :meth:`preprocess_single_kernel`

# as we are iterating over it.[1]

#

# [1] https://docs.python.org/3/library/stdtypes.html#dictionary-view-objects

new_callables = {}

for func_id, in_knl_callable in t_unit.callables_table.items():

if isinstance(in_knl_callable, CallableKernel):

new_subkernel = _preprocess_single_kernel(

in_knl_callable.subkernel,

is_entrypoint=func_id in t_unit.entrypoints)

in_knl_callable = in_knl_callable.copy(

subkernel=new_subkernel)

elif isinstance(in_knl_callable, ScalarCallable):

pass

else:

raise NotImplementedError("Unknown callable type %s." % (

type(in_knl_callable).__name__))

new_callables[func_id] = in_knl_callable

t_unit = t_unit.copy(callables_table=constantdict(new_callables))

# }}}

# infer arg descrs of the callables

t_unit = infer_arg_descr(t_unit)

# Ordering restriction:

# callees with gbarrier in them must be inlined after inferrring arg_descr.

return inline_kernels_with_gbarriers(t_unit)

# {{{ prepare for caching

# PicklableDtype instances for example need to know the target they're working

# towards in order to pickle and unpickle them. This is the first pass that

# uses caching, so we need to be ready to pickle. This means propagating

# this target information.