Skip to content

Commit

Permalink
Improve memlet label and string initialization
Browse files Browse the repository at this point in the history
Follow up on the discussion in #1678.

Supports `src[expr] -> dst[expr]`, `src[expr] -> [expr]`, and `[expr] -> dst[expr]` initializations for memlets. Also improves memlet label printouts.

@philip-paul-mueller @phschaad the expression mentioned in the other PR will now be printed as `[0, 0] -> B[0]` for clarity and can be reparsed.
  • Loading branch information
tbennun authored Oct 11, 2024
1 parent 6525bc5 commit 1733237
Showing 1 changed file with 32 additions and 13 deletions.
45 changes: 32 additions & 13 deletions dace/memlet.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,9 @@ def __init__(self,
of use API. Must follow one of the following forms:
1. ``ARRAY``,
2. ``ARRAY[SUBSET]``,
3. ``ARRAY[SUBSET] -> OTHER_SUBSET``.
3. ``ARRAY[SUBSET] -> [OTHER_SUBSET]``,
4. ``[OTHER_SUBSET] -> ARRAY[SUBSET]``,
5. ``SRC_ARRAY[SRC_SUBSET] -> DST_ARRAY[DST_SUBSET]``.
:param data: Data descriptor name attached to this memlet.
:param subset: The subset to take from the data attached to the edge,
represented either as a string or a Subset object.
Expand Down Expand Up @@ -330,6 +332,10 @@ def _parse_from_subexpr(self, expr: str):
raise SyntaxError('Invalid memlet syntax "%s"' % expr)
return expr, None

# [subset] syntax
if expr.startswith('['):
return None, SubsetProperty.from_string(expr[1:-1])

# array[subset] syntax
arrname, subset_str = expr[:-1].split('[')
if not dtypes.validate_name(arrname):
Expand All @@ -342,27 +348,40 @@ def _parse_memlet_from_str(self, expr: str):
or the _data,_subset fields.
:param expr: A string expression of the this memlet, given as an ease
of use API. Must follow one of the following forms:
1. ``ARRAY``,
2. ``ARRAY[SUBSET]``,
3. ``ARRAY[SUBSET] -> OTHER_SUBSET``.
Note that modes 2 and 3 are deprecated and will leave
the memlet uninitialized until inserted into an SDFG.
of use API. Must follow one of the following forms:
1. ``ARRAY``,
2. ``ARRAY[SUBSET]``,
3. ``ARRAY[SUBSET] -> [OTHER_SUBSET]``,
4. ``[OTHER_SUBSET] -> ARRAY[SUBSET]``,
5. ``SRC_ARRAY[SRC_SUBSET] -> DST_ARRAY[DST_SUBSET]``.
Note that options 1-2 will leave the memlet uninitialized
until added into an SDFG.
"""
expr = expr.strip()
if '->' not in expr: # Options 1 and 2
self.data, self.subset = self._parse_from_subexpr(expr)
return

# Option 3
# Options 3-5
src_expr, dst_expr = expr.split('->')
src_expr = src_expr.strip()
dst_expr = dst_expr.strip()
if '[' not in src_expr and not dtypes.validate_name(src_expr):
raise SyntaxError('Expression without data name not yet allowed')

self.data, self.subset = self._parse_from_subexpr(src_expr)
self.other_subset = SubsetProperty.from_string(dst_expr)
src_data, src_subset = self._parse_from_subexpr(src_expr)
dst_data, dst_subset = self._parse_from_subexpr(dst_expr)
if src_data is None and dst_data is None:
raise SyntaxError('At least one data name needs to be given')

if src_data is not None: # Prefer src[subset] -> [other_subset]
self.data = src_data
self.subset = src_subset
self.other_subset = dst_subset
self._is_data_src = True
else:
self.data = dst_data
self.subset = src_subset
self.other_subset = dst_subset
self._is_data_src = False

def try_initialize(self, sdfg: 'dace.sdfg.SDFG', state: 'dace.sdfg.SDFGState',
edge: 'dace.sdfg.graph.MultiConnectorEdge'):
Expand Down Expand Up @@ -660,7 +679,7 @@ def _label(self, shape):

if self.other_subset is not None:
if self._is_data_src is False:
result += ' <- [%s]' % str(self.other_subset)
result = f'[{self.other_subset}] -> {result}'
else:
result += ' -> [%s]' % str(self.other_subset)
return result
Expand Down

0 comments on commit 1733237

Please sign in to comment.