Fixed Multiple Return Values and Refined Validations

dace/codegen/compiled_sdfg.py

@@ -253,9 +253,17 @@ def __init__(self, sdfg, lib: ReloadableDLL, argnames: Optional[List[str]] = Non
         #  ``False`` either means that a tuple is returned or there are no return values.
         # NOTE: Necessary to handle the case of a tuple with one element.
         self._is_single_value_ret: bool = False
+        self._return_args: Tuple[str] = ()
         if '__return' in self._sdfg.arrays:
             assert not any(aname.startswith('__return_') for aname in self._sdfg.arrays.keys())
             self._is_single_value_ret = True
+            self._return_args = ("__return", )
+        else:
+            # We can not use `sorted` here because it sorts the strings alphabetical order
+            #  and thus `__return_30` ends up before `__return_4`.
+            return_arguments = {aname for aname in self._sdfg.arrays.keys() if aname.startswith('__return_')}
+            self._return_args = tuple(f'__return_{i}' for i in range(len(return_arguments)))
+            assert return_arguments == set(self._return_args)
 
         # Cache SDFG argument properties
         self._typedict = self._sdfg.arglist()
@@ -751,63 +759,63 @@ def _initialize_return_values(self, kwargs):
         self._return_arrays = []
         self._retarray_shapes = []
         self._retarray_is_pyobject = []
-        for arrname, arr in sorted(self.sdfg.arrays.items()):
-            if arrname.startswith('__return'):
-                if arr.transient:
-                    raise ValueError(f'Used the special array name "{arrname}" as transient.')
+        for arrname in self._return_args:
+            arr = self._sdfg.arrays[arrname]
+            if arr.transient:
+                raise ValueError(f'Used the special array name "{arrname}" as transient.')
 
-                elif arrname in kwargs:
-                    # The return value is passed as an argument, in that case store the name in `self._retarray_shapes`.
-                    warnings.warn(f'Return value "{arrname}" is passed as a regular argument.', stacklevel=2)
-                    self._return_arrays.append(kwargs[arrname])
-                    self._retarray_shapes.append((arrname, ))
+            elif arrname in kwargs:
+                # The return value is passed as an argument, in that case store the name in `self._retarray_shapes`.
+                warnings.warn(f'Return value "{arrname}" is passed as a regular argument.', stacklevel=2)
+                self._return_arrays.append(kwargs[arrname])
+                self._retarray_shapes.append((arrname, ))
 
-                elif isinstance(arr, dt.Stream):
-                    raise NotImplementedError('Return streams are unsupported')
+            elif isinstance(arr, dt.Stream):
+                raise NotImplementedError('Return streams are unsupported')
 
+            else:
+                shape = tuple(symbolic.evaluate(s, syms) for s in arr.shape)
+                dtype = arr.dtype.as_numpy_dtype()
+                total_size = int(symbolic.evaluate(arr.total_size, syms))
+                strides = tuple(symbolic.evaluate(s, syms) * arr.dtype.bytes for s in arr.strides)
+                shape_desc = (arrname, dtype, arr.storage, shape, strides, total_size)
+                self._retarray_shapes.append(shape_desc)
+
+                # Create an array with the properties of the SDFG array
+                return_array = self._create_array(*shape_desc)
+                self._return_arrays.append(return_array)
+
+            # BUG COMPATIBILITY(PR#2206):
+            #   In the original version `_retarray_is_pyobject` was named `_retarray_is_scalar`, however
+            #   since scalars could not be returned on an [implementation level](https://github.com/spcl/dace/pull/1609)
+            #   it was only used for `pyobject`s in _some_ cases. Since `pyobject`s are essentially `void` pointers,
+            #   it was possible to return them as scalars.
+            #   However, if the return value was passed as argument, i.e., the first `elif`, then it
+            #   was ignored if `arr` was a `pyobject`. Only if the return value was managed by `self`,
+            #   i.e. the `else` case, then it was considered. The problem is that it was done using the
+            #   following check: `isinstance(arr, dt.Scalar) or isinstance(arr.dtype, dtypes.pyobject)`
+            #   Because of the `or`, _everything_ whose `dtype` is `pyobject` was classified
+            #   as a scalar `pyobject`, i.e., one element.
+            #   The correct behavior would be to change the `or` to an `and` but then several unit
+            #   tests (`test_pyobject_return`, `test_pyobject_return_tuple` and `test_nested_autoparse[False]`
+            #   in `tests/python_frontend/callee_autodetect_test.py`) will fail.
+            #   The following code is bug compatible and also allows to pass a `pyobject` directly, i.e.,
+            #   through `kwargs`.
+            if isinstance(arr.dtype, dtypes.pyobject):
+                if isinstance(arr, dt.Scalar):
+                    # Proper scalar.
+                    self._retarray_is_pyobject.append(True)
+                elif isinstance(arr, dt.Array):
+                    # An array, let's check if it is just a wrapper for a single value.
+                    if not (len(arr.shape) == 1 and arr.shape[0] == 1):
+                        warnings.warn(f'Decay an array of `pyobject`s with shape {arr.shape} to a single one.',
+                                      stacklevel=2)
+                    self._retarray_is_pyobject.append(True)
                 else:
-                    shape = tuple(symbolic.evaluate(s, syms) for s in arr.shape)
-                    dtype = arr.dtype.as_numpy_dtype()
-                    total_size = int(symbolic.evaluate(arr.total_size, syms))
-                    strides = tuple(symbolic.evaluate(s, syms) * arr.dtype.bytes for s in arr.strides)
-                    shape_desc = (arrname, dtype, arr.storage, shape, strides, total_size)
-                    self._retarray_shapes.append(shape_desc)
-
-                    # Create an array with the properties of the SDFG array
-                    return_array = self._create_array(*shape_desc)
-                    self._return_arrays.append(return_array)
-
-                # BUG COMPATIBILITY(PR#2206):
-                #   In the original version `_retarray_is_pyobject` was named `_retarray_is_scalar`, however
-                #   since scalars could not be returned on an [implementation level](https://github.com/spcl/dace/pull/1609)
-                #   it was only used for `pyobject`s in _some_ cases. Since `pyobject`s are essentially `void` pointers,
-                #   it was possible to return them as scalars.
-                #   However, if the return value was passed as argument, i.e., the first `elif`, then it
-                #   was ignored if `arr` was a `pyobject`. Only if the return value was managed by `self`,
-                #   i.e. the `else` case, then it was considered. The problem is that it was done using the
-                #   following check: `isinstance(arr, dt.Scalar) or isinstance(arr.dtype, dtypes.pyobject)`
-                #   Because of the `or`, _everything_ whose `dtype` is `pyobject` was classified
-                #   as a scalar `pyobject`, i.e., one element.
-                #   The correct behavior would be to change the `or` to an `and` but then several unit
-                #   tests (`test_pyobject_return`, `test_pyobject_return_tuple` and `test_nested_autoparse[False]`
-                #   in `tests/python_frontend/callee_autodetect_test.py`) will fail.
-                #   The following code is bug compatible and also allows to pass a `pyobject` directly, i.e.,
-                #   through `kwargs`.
-                if isinstance(arr.dtype, dtypes.pyobject):
-                    if isinstance(arr, dt.Scalar):
-                        # Proper scalar.
-                        self._retarray_is_pyobject.append(True)
-                    elif isinstance(arr, dt.Array):
-                        # An array, let's check if it is just a wrapper for a single value.
-                        if not (len(arr.shape) == 1 and arr.shape[0] == 1):
-                            warnings.warn(f'Decay an array of `pyobject`s with shape {arr.shape} to a single one.',
-                                          stacklevel=2)
-                        self._retarray_is_pyobject.append(True)
-                    else:
-                        raise ValueError(
-                            f'Does not know how to handle "{arrname}", which is a {type(arr).__name__} of `pyobject`.')
-                else:
-                    self._retarray_is_pyobject.append(False)
+                    raise ValueError(
+                        f'Does not know how to handle "{arrname}", which is a {type(arr).__name__} of `pyobject`.')
+            else:
+                self._retarray_is_pyobject.append(False)
 
         assert (not self._is_single_value_ret) or (len(self._return_arrays) == 1)
         assert len(self._return_arrays) == len(self._retarray_shapes) == len(self._retarray_is_pyobject)

dace/sdfg/validation.py

@@ -275,6 +275,33 @@ def validate_sdfg(sdfg: 'dace.sdfg.SDFG', references: Set[int] = None, **context
                     warnings.warn(f'Mismatch between constant and symbol type of "{const_name}", '
                                   f'expected to find "{const_type}" but found "{sdfg.symbols[const_name]}".')
 
+        # Test the return value.
+        tuple_return_args = {n for n in sdfg._arrays if n.startswith('__return_')}
+        if '__return' in sdfg._arrays and tuple_return_args:
+            raise InvalidSDFGError(
+                'Ambiguous return values: an SDFG cannot have both a `__return` (single value) '
+                'and `__return_<i>` (tuple) data descriptor.', sdfg, None)
+        elif '__return' in sdfg._arrays:
+            tuple_return_args = {'__return'}  # This is abuse
+        elif tuple_return_args and tuple_return_args != {f'__return_{i}' for i in range(len(tuple_return_args))}:
+            raise InvalidSDFGError('Tuple return values are not consecutively named')
+        for ret_name_to_check in tuple_return_args:
+            ret_desc = sdfg._arrays[ret_name_to_check]
+            if ret_desc.transient:
+                raise InvalidSDFGError(f'The return value `{ret_name_to_check}` can not be a transient.')
+            if sdfg.parent is None:
+                # These are some top level specific test
+                if isinstance(ret_desc, dt.Scalar):
+                    # This is an implementation level constraint and is thus a separate error.
+                    #  In certain cases the frontend will promote it to a length 1 array.
+                    raise InvalidSDFGError(f'{ret_name_to_check} is a scalar and scalars can not be returned.')
+                if not isinstance(ret_desc, dt.Array):
+                    # This is a limitation of the Python <-> Binary interface, because Python needs to allocate
+                    #  the return value and for that NumPy/CuPy is used.
+                    raise InvalidSDFGError(
+                        f'Only arrays can be returned from SDFG, but `{ret_name_to_check}` is a `{type(desc).__name__}`'
+                    )
+
         # Validate data descriptors
         for name, desc in sdfg._arrays.items():
             if id(desc) in references:
@@ -283,14 +310,6 @@ def validate_sdfg(sdfg: 'dace.sdfg.SDFG', references: Set[int] = None, **context
                     'rather than using multiple references to the same one', sdfg, None)
             references.add(id(desc))
 
-            # Because of how the code generator works Scalars can not be return values.
-            #  TODO: Remove this limitation as the CompiledSDFG contains logic for that.
-            if (sdfg.parent is None and isinstance(desc, dt.Scalar) and name.startswith("__return")
-                    and not desc.transient):
-                raise InvalidSDFGError(
-                    f'Cannot use scalar data descriptor ("{name}") as return value of a top-level function.', sdfg,
-                    None)
-
             # Check for UndefinedSymbol in transient data shape (needed for memory allocation)
             if desc.transient:
                 # Check dimensions

tests/python_frontend/return_value_test.py

@@ -63,6 +63,18 @@ def return_tuple():
     assert res == (5, 6)
 
 
+def test_return_tuple_multi():
+
+    @dace.program
+    def return_tuple_2():
+        return 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12
+
+    res = return_tuple_2()
+    assert isinstance(res, tuple)
+    assert len(res) == 12
+    assert res == (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12)
+
+
 def test_return_array_tuple():
 
     @dace.program