def are_subsets_contiguous(subset_a: subsets.Subset, subset_b: subsets.Subset, dim: int = None) -> bool:
if dim is not None:
# A version that only checks for contiguity in certain
# dimension (e.g., to prioritize stride-1 range)
if (not isinstance(subset_a, subsets.Range) or not isinstance(subset_b, subsets.Range)):
raise NotImplementedError('Contiguous subset check only ' 'implemented for ranges')
# Other dimensions must be equal
for i, (s1, s2) in enumerate(zip(subset_a.ranges, subset_b.ranges)):
if i == dim:
continue
if s1[0] != s2[0] or s1[1] != s2[1] or s1[2] != s2[2]:
return False
# Set of conditions for contiguous dimension
ab = (subset_a[dim][1] + 1) == subset_b[dim][0]
a_overlap_b = subset_a[dim][1] >= subset_b[dim][0]
ba = (subset_b[dim][1] + 1) == subset_a[dim][0]
b_overlap_a = subset_b[dim][1] >= subset_a[dim][0]
# NOTE: Must check with "==" due to sympy using special types
return (ab == True or a_overlap_b == True or ba == True or b_overlap_a == True)
# General case
bbunion = subsets.bounding_box_union(subset_a, subset_b)
try:
if bbunion.num_elements() == (subset_a.num_elements() + subset_b.num_elements()):
return True
except TypeError:
pass
return False
def cpp_offset_expr(d: data.Data,
subset_in: subsets.Subset,
offset=None,
packed_veclen=1,
indices=None):
""" Creates a C++ expression that can be added to a pointer in order
to offset it to the beginning of the given subset and offset.
:param d: The data structure to use for sizes/strides.
:param subset_in: The subset to offset by.
:param offset: An additional list of offsets or a Subset object
:param packed_veclen: If packed types are targeted, specifies the
vector length that the final offset should be
divided by.
:param indices: A tuple of indices to use for expression.
:return: A string in C++ syntax with the correct offset
"""
# Offset according to parameters, then offset according to array
if offset is not None:
subset = subset_in.offset_new(offset, False)
subset.offset(d.offset, False)
else:
subset = subset_in.offset_new(d.offset, False)
# Obtain start range from offsetted subset
indices = indices or ([0] * len(d.strides))
index = subset.at(indices, d.strides)
if packed_veclen > 1:
index /= packed_veclen
return sym2cpp(index)
def propagate_subset(
memlets: List[Memlet],
arr: data.Data,
params: List[str],
rng: subsets.Subset,
defined_variables: Set[symbolic.SymbolicType] = None) -> Memlet:
""" Tries to propagate a list of memlets through a range (computes the
image of the memlet function applied on an integer set of, e.g., a
map range) and returns a new memlet object.
:param memlets: The memlets to propagate.
:param arr: Array descriptor for memlet (used for obtaining extents).
:param params: A list of variable names.
:param rng: A subset with dimensionality len(params) that contains the
range to propagate with.
:param defined_variables: A set of symbols defined that will remain the
same throughout propagation. If None, assumes
that all symbols outside of `params` have been
defined.
:return: Memlet with propagated subset and volume.
"""
# Argument handling
if defined_variables is None:
# Default defined variables is "everything but params"
defined_variables = set()
defined_variables |= rng.free_symbols
for memlet in memlets:
defined_variables |= memlet.free_symbols
defined_variables -= set(params)
defined_variables = set(
symbolic.pystr_to_symbolic(p) for p in defined_variables)
# Propagate subset
variable_context = [
defined_variables, [symbolic.pystr_to_symbolic(p) for p in params]
]
new_subset = None
for md in memlets:
tmp_subset = None
for pclass in MemletPattern.extensions():
pattern = pclass()
if pattern.can_be_applied([md.subset], variable_context, rng,
[md]):
tmp_subset = pattern.propagate(arr, [md.subset], rng)
break
else:
# No patterns found. Emit a warning and propagate the entire
# array
warnings.warn('Cannot find appropriate memlet pattern to '
'propagate %s through %s' %
(str(md.subset), str(rng)))
tmp_subset = subsets.Range.from_array(arr)
# Union edges as necessary
if new_subset is None:
new_subset = tmp_subset
else:
old_subset = new_subset
new_subset = subsets.union(new_subset, tmp_subset)
if new_subset is None:
warnings.warn('Subset union failed between %s and %s ' %
(old_subset, tmp_subset))
break
# Some unions failed
if new_subset is None:
new_subset = subsets.Range.from_array(arr)
### End of subset propagation
# Create new memlet
new_memlet = copy.copy(memlets[0])
new_memlet.subset = new_subset
new_memlet.other_subset = None
# Propagate volume:
# Number of accesses in the propagated memlet is the sum of the internal
# number of accesses times the size of the map range set (unbounded dynamic)
new_memlet.volume = (sum(m.volume for m in memlets) *
functools.reduce(lambda a, b: a * b, rng.size(), 1))
if any(m.dynamic for m in memlets):
new_memlet.dynamic = True
elif symbolic.issymbolic(new_memlet.volume) and any(
s not in defined_variables
for s in new_memlet.volume.free_symbols):
new_memlet.dynamic = True
new_memlet.volume = 0
return new_memlet