def parse_memlet_subset(array: data.Data, node: Union[ast.Name, ast.Subscript],
das: Dict[str, Any],
parsed_slice: Any = None) -> Tuple[subsets.Range, List[int]]:
"""
Parses an AST subset and returns access range, as well as new dimensions to
add.
:param array: Accessed data descriptor (used for filling in missing data,
e.g., negative indices or empty shapes).
:param node: AST node representing whole array or subset thereof.
:param das: Dictionary of defined arrays and symbols mapped to their values.
:return: A 2-tuple of (subset, list of new axis indices).
"""
# Get memlet range
ndslice = [(0, s - 1, 1) for s in array.shape]
extra_dims = []
arrdims: Dict[int, str] = {}
if isinstance(node, ast.Subscript):
# Parse and evaluate ND slice(s) (possibly nested)
if parsed_slice:
cnode = copy.deepcopy(node)
cnode.slice = parsed_slice
else:
cnode = node
ast_ndslices = astutils.subscript_to_ast_slice_recursive(cnode)
offsets = list(range(len(array.shape)))
# Loop over nd-slices (A[i][j][k]...)
subset_array = []
for idx, ast_ndslice in enumerate(ast_ndslices):
# Cut out dimensions that were indexed in the previous slice
narray = copy.deepcopy(array)
narray.shape = [
s for i, s in enumerate(array.shape) if i in offsets
]
# Loop over the N dimensions
ndslice, offsets, new_extra_dims, arrdims = _fill_missing_slices(
das, ast_ndslice, narray, offsets)
if new_extra_dims and idx != (len(ast_ndslices) - 1):
raise NotImplementedError('New axes only implemented for last '
'slice')
if arrdims and len(ast_ndslices) != 1:
raise NotImplementedError('Array dimensions not implemented '
'for consecutive subscripts')
extra_dims = new_extra_dims
subset_array.append(_ndslice_to_subset(ndslice))
subset = subset_array[0]
# Compose nested indices, e.g., of the form "A[i,:,j,:][k,l]"
for i in range(1, len(subset_array)):
subset = subset.compose(subset_array[i])
else: # Use entire range
subset = _ndslice_to_subset(ndslice)
if isinstance(subset, subsets.Indices):
subset = subsets.Range([(i, i, 1) for i in subset])
return subset, extra_dims, arrdims
def parse_memlet_subset(array: data.Data, node: Union[ast.Name, ast.Subscript],
das: Dict[str, Any]):
array_dependencies = {}
# Get memlet range
ndslice = [(0, s - 1, 1) for s in array.shape]
if isinstance(node, ast.Subscript):
# Parse and evaluate ND slice(s) (possibly nested)
ast_ndslices = astutils.subscript_to_ast_slice_recursive(node)
offsets = list(range(len(array.shape)))
# Loop over nd-slices (A[i][j][k]...)
subset_array = []
for ast_ndslice in ast_ndslices:
# Cut out dimensions that were indexed in the previous slice
narray = copy.deepcopy(array)
narray.shape = [
s for i, s in enumerate(array.shape) if i in offsets
]
# Loop over the N dimensions
ndslice, offsets = _fill_missing_slices(das, ast_ndslice, narray,
offsets)
subset_array.append(_ndslice_to_subset(ndslice))
subset = subset_array[0]
# Compose nested indices, e.g., of the form "A[i,:,j,:][k,l]"
for i in range(1, len(subset_array)):
subset = subset.compose(subset_array[i])
# Compute additional array dependencies (as a result of
# indirection)
# for dim in subset:
# if not isinstance(dim, tuple): dim = [dim]
# for r in dim:
# for expr in symbolic.swalk(r):
# if symbolic.is_sympy_userfunction(expr):
# arr = expr.func.__name__
# array_dependencies[arr] = self.curnode.globals[arr]
else: # Use entire range
subset = _ndslice_to_subset(ndslice)
if isinstance(subset, subsets.Indices):
subset = subsets.Range([(i, i, 1) for i in subset])
return subset