[Python parser] Unroll/nounroll transparency over list of objects/functions

[FlorianDeconinck]

Describe the bug

Behavior of loop on constant list of objects/list between unroll and nounroll differs (see example below).

When using nounroll the parser fails during visit_Subscript, raising because it builds a tuple of non symbol/NumConstant. But the unroll behavior is fully capable of laying out the calls properly.

Even if those behavior do not share code, I'd expect a transparency between them.

This is ticket is not urgent as we can go around by using unroll on said loop at the cost of a build time blow up.

Parser trap:

        # Obtain array/tuple
        node_parsed = self._gettype(node.value)

        if len(node_parsed) > 1 or (len(node_parsed) == 1 and isinstance(node_parsed[0], tuple)
                                    and node_parsed[0][1] in ('symbol', 'NumConstant')):
            # If the value is a tuple of constants (e.g., array.shape) and the
            # slice is constant, return the value itself
            nslice = self.visit(node.slice)
            if isinstance(nslice, (Index, Number)):
                v = nslice
                try:
                    value, valtype = node_parsed[int(v)]
                    return value
                except (TypeError, ValueError):
                    pass  # Passthrough to exception

            raise DaceSyntaxError(self, node.value, 'Subscripted object cannot be a tuple')

To Reproduce
DaCe main (0ec62e)

import numpy as np
from dace import program, unroll, nounroll


class ProgramB:
    def __init__(self, parameter):
        self.parameter = parameter

    def __call__(self, A, B):
        A[:] = B[:] + self.parameter


B1 = ProgramB(0.1)
B2 = ProgramB(0.2)
Bs = [B1, B2]
nB = 2


@program
def A(array_a, array_b):
    B1(array_a, array_b)
    B2(array_a, array_b)

    for n in unroll(range(nB)):  # <--- works
        Bs[n](array_a, array_b)

    for n in nounroll(range(nB)):  # <--- fails
        Bs[n](array_a, array_b)


np_a = np.zeros((10,))
np_b = np.ones((10,))

A(np_a, np_b)

[alexnick83]

I believe we can provide a better error message here, as the current one assumes that if the subscripted element is not (DaCe) data, it may (by exception) be an attribute of data that (again, by exception) we allow to be set. This is obviously not the case here, as the reported "tuple" contains two Programs.

If the expectation is that this actually works, then the issue needs further thought, and, currently, I do not see how this would work in a reasonable manner. Here, we need to clarify how the two versions work:

This is roughly the Python program after "preprocessing" and moving into the main Python frontend:

def A(array_a, array_b):
    <dace.frontend.python.parser.DaceProgram object at 0x7fcca57c8560>(array_a, array_b)  # B1(array_a, array_b)
    <dace.frontend.python.parser.DaceProgram object at 0x7fcca57cae70>(array_a, array_b)  # B2(array_a, array_b)
    <dace.frontend.python.parser.DaceProgram object at 0x7fcca37966c0>(array_a, array_b)  # unrolled B1(array_a, array_b)
    <dace.frontend.python.parser.DaceProgram object at 0x7fcca3796c00>(array_a, array_b)  # unrolled B2(array_a, array_b)
    for n in range(2):
        Bs[n](array_a, array_b)

So, when the main frontend parses this program, it will just create a separate SDFGState for each unrolled statement and call the "preprocessed and subsequently parsed" B1 and B2 Programs. There is no indexing actually occurring at the frontend-to-SDFG level, as we cannot use subscripts on a list of Program calls.

As you can also see above, the "nounroll" loop has the preprocessor leave the loop as-is, resulting in indexing a list of Program calls. Such syntax cannot be directly supported as SDFGs do not have Arrays of Program calls. Hence, this is only supported via unrolling. I can only think of this working in a generic manner via an if-clause inside the for-loop, but it also potentially opens a can of worms (e.g., what if nB is a symbol?). Other solutions that, e.g., go deep in the Program calls and find a way to represent the slicing inside the for-loop in an SDFG-compatible way by accessing the underlying data directly, would need significant "preprocessing" anyway, so you may as well just use "unroll."

Do you perhaps have another solution in mind?

[tbennun]

To make things worse, say we did support tuples or python lists in the frontend. A loop that looks like:

    for n in nounroll(range(nB)):
        f(Bs[n], array_a, array_b)

would be more amenable to SDFG creation than calling an arbitrary object. Even if Bs is hinted as tuple[ProgramB, ...] or list[ProgramB], who is to say that the call dunder was not modified on one of the objects? This kind of JIT-OOP analysis is something we do as we parse an unrolled loop.

The approach with f is safer, if possible, but still currently unsupported in the Python frontend until #2227

[FlorianDeconinck]

Thanks gents, good points all around.

So first off, I'll PR an updated error message at some point - which will hint to unroll.

I guess the only thing I could see is force the unroll - but we would need to run some check on the pre-processing steps to catch early that it will be an issue post-parsing. It can also have side-effect where SDFG explodes in size because of the unrolling. We could config it with a flag, and it would at least mirror the callback mechanism of attempting to salvage an un-parsable code.