def make_copy_kernel(new_dim_tags, old_dim_tags=None):
"""Returns a :class:`LoopKernel` that changes the data layout
of a variable (called "input") to the new layout specified by
*new_dim_tags* from the one specified by *old_dim_tags*.
*old_dim_tags* defaults to an all-C layout of the same rank
as the one given by *new_dim_tags*.
"""
from loopy.kernel.array import (parse_array_dim_tags,
SeparateArrayArrayDimTag, VectorArrayDimTag)
new_dim_tags = parse_array_dim_tags(new_dim_tags, n_axes=None)
rank = len(new_dim_tags)
if old_dim_tags is None:
old_dim_tags = parse_array_dim_tags(
",".join(rank * ["c"]), n_axes=None)
elif isinstance(old_dim_tags, str):
old_dim_tags = parse_array_dim_tags(
old_dim_tags, n_axes=None)
indices = ["i%d" % i for i in range(rank)]
shape = ["n%d" % i for i in range(rank)]
commad_indices = ", ".join(indices)
bounds = " and ".join(
f"0<={ind}<{shape_i}"
for ind, shape_i in zip(indices, shape))
set_str = "{{[{}]: {}}}".format(
commad_indices,
bounds
)
result = make_kernel(set_str,
"output[%s] = input[%s]"
% (commad_indices, commad_indices),
lang_version=MOST_RECENT_LANGUAGE_VERSION)
result = tag_array_axes(result, "input", old_dim_tags)
result = tag_array_axes(result, "output", new_dim_tags)
unrolled_tags = (SeparateArrayArrayDimTag, VectorArrayDimTag)
for i in range(rank):
if (isinstance(new_dim_tags[i], unrolled_tags)
or isinstance(old_dim_tags[i], unrolled_tags)):
result = tag_inames(result, {indices[i]: "unr"})
return result
def make_copy_kernel(new_dim_tags, old_dim_tags=None):
"""Returns a :class:`LoopKernel` that changes the data layout
of a variable (called "input") to the new layout specified by
*new_dim_tags* from the one specified by *old_dim_tags*.
*old_dim_tags* defaults to an all-C layout of the same rank
as the one given by *new_dim_tags*.
"""
from loopy.kernel.array import (parse_array_dim_tags,
SeparateArrayArrayDimTag, VectorArrayDimTag)
new_dim_tags = parse_array_dim_tags(new_dim_tags, n_axes=None)
rank = len(new_dim_tags)
if old_dim_tags is None:
old_dim_tags = parse_array_dim_tags(
",".join(rank * ["c"]), n_axes=None)
elif isinstance(old_dim_tags, str):
old_dim_tags = parse_array_dim_tags(
old_dim_tags, n_axes=None)
indices = ["i%d" % i for i in range(rank)]
shape = ["n%d" % i for i in range(rank)]
commad_indices = ", ".join(indices)
bounds = " and ".join(
"0<=%s<%s" % (ind, shape_i)
for ind, shape_i in zip(indices, shape))
set_str = "{[%s]: %s}" % (
commad_indices,
bounds
)
result = make_kernel(set_str,
"output[%s] = input[%s]"
% (commad_indices, commad_indices))
result = tag_array_axes(result, "input", old_dim_tags)
result = tag_array_axes(result, "output", new_dim_tags)
unrolled_tags = (SeparateArrayArrayDimTag, VectorArrayDimTag)
for i in range(rank):
if (isinstance(new_dim_tags[i], unrolled_tags)
or isinstance(old_dim_tags[i], unrolled_tags)):
result = tag_inames(result, {indices[i]: "unr"})
return result
