Frontend: add syntactic sugar for schedule and storage types.

dace/data.py

@@ -321,6 +321,21 @@ def set_strides_from_layout(self,
         self.strides = strides
         self.total_size = totalsize
 
+    def __matmul__(self, storage: dtypes.StorageType):
+        """
+        Syntactic sugar for specifying the storage of a data descriptor.
+        This enables controlling the storage location as follows:
+
+        .. code-block:: python
+
+            @dace
+            def add(X: dace.float32[10, 10] @ dace.StorageType.GPU_Global):
+                return X + 1
+        """
+        new_desc = cp.deepcopy(self)
+        new_desc.storage = storage
+        return new_desc
+
 
 @make_properties
 class Scalar(Data):

dace/frontend/python/interface.py

@@ -98,6 +98,7 @@ def method(f: F,
 
     # Create a wrapper class that can bind to the object instance
     class MethodWrapper:
+
         def __init__(self):
             self.wrapped: Dict[int, parser.DaceProgram] = {}
 
@@ -119,8 +120,34 @@ def __get__(self, obj, objtype=None) -> parser.DaceProgram:
 
 
 # Dataflow constructs
+class MapGenerator:
+    """
+    An SDFG map generator class that allows applying operators on it, used
+    for syntactic sugar.
+    """
+
+    def __init__(self, rng: Union[slice, Tuple[slice]]):
+        self.rng = rng
+
+    def __matmul__(self, schedule: dtypes.ScheduleType):
+        """
+        Syntactic sugar for specifying the schedule of the map.
+        This enables controlling the hardware scheduling of the map as follows:
+
+        .. code-block:: python
+
+            for i, j in dace.map[0:N, 0:M] @ ScheduleType.GPU_Global:
+                b[i, j] = a[i, j] + 1
+        """
+        # Ignored in Python mode
+        return self
+
+    def __iter__(self):
+        return ndloop.ndrange(self.rng)
+
 class MapMetaclass(type):
     """ Metaclass for map, to enable ``dace.map[0:N]`` syntax. """
+
     @classmethod
     def __getitem__(cls, rng: Union[slice, Tuple[slice]]) -> Generator[Tuple[int], None, None]:
         """ 
@@ -129,7 +156,7 @@ def __getitem__(cls, rng: Union[slice, Tuple[slice]]) -> Generator[Tuple[int], N
                     N-dimensional range to iterate over.
         :return: Generator of N-dimensional tuples of iterates.
         """
-        yield from ndloop.ndrange(rng)
+        return MapGenerator(rng)
 
 
 class map(metaclass=MapMetaclass):
@@ -142,6 +169,7 @@ class map(metaclass=MapMetaclass):
 
 
 class consume:
+
     def __init__(self, stream: Deque[T], processing_elements: int = 1, condition: Optional[Callable[[], bool]] = None):
         """ 
         Consume is a scope, like ``Map``, that creates parallel execution.
@@ -169,6 +197,7 @@ def __iter__(self) -> Generator[T, None, None]:
 
 class TaskletMetaclass(type):
     """ Metaclass for tasklet, to enable ``with dace.tasklet:`` syntax. """
+
     def __enter__(self):
         # Parse and run tasklet
         frame = inspect.stack()[1][0]
@@ -197,6 +226,7 @@ class tasklet(metaclass=TaskletMetaclass):
 
     The DaCe framework cannot analyze these tasklets for optimization. 
     """
+
     def __init__(self, language: Union[str, dtypes.Language] = dtypes.Language.Python):
         if isinstance(language, str):
             language = dtypes.Language[language]

dace/frontend/python/newast.py

@@ -12,6 +12,7 @@
 import warnings
 from numbers import Number
 from typing import Any, Dict, List, Set, Tuple, Union, Callable, Optional
+import operator
 
 import dace
 from dace import data, dtypes, subsets, symbolic, sdfg as sd
@@ -1620,11 +1621,50 @@ def _parse_for_iterator(self, node: ast.Expr):
             NotImplementedError: If iterator type is not implemented
 
         Returns:
-            Tuple[str, List[str], List[ast.AST]] -- Iterator type, iteration 
+            Tuple[str, List[str], List[ast.AST], Optional[ScheduleType]] --
+                                                    Iterator type, iteration
                                                     ranges, and AST versions of
-                                                    the ranges
+                                                    the ranges. If present, the
+                                                    schedule type is returned.
         """
 
+        if isinstance(node, (ast.BinOp)):
+            # special case:
+            # We allow iterating over binops like:
+            # dace.map[0:N] @ ScheduleType
+            if not isinstance(node.op, ast.MatMult):
+                raise DaceSyntaxError(
+                    self, node, "Binop in for-loop iterator is not supported, "
+                    "except when using the @ operator to specify "
+                    "Schedule types")
+
+            # parse schedule type
+            schedule_name = preprocessing.ModuleResolver(self.modules, True).visit(node.right)
+            schedule_name = rname(schedule_name)
+
+            if schedule_name.startswith("ScheduleType."):
+                # support ScheduleType.<...>
+                schedule_type = schedule_name[len("ScheduleType."):]
+                schedule = getattr(dtypes.ScheduleType, schedule_type)
+            else:
+                # check if it's a module (e.g. dace.ScheduleType or dtypes.ScheduleType)
+                modname = until(schedule_name, '.')
+                if ('.' in schedule_name and modname and modname in self.globals
+                        and dtypes.ismodule(self.globals[modname])):
+                    schedule = operator.attrgetter(schedule_name[len(modname) + 1:])(self.globals[modname])
+                elif schedule_name in self.globals:
+                    schedule = self.globals[schedule_name]
+                else:
+                    schedule = None
+
+                if not isinstance(schedule, dtypes.ScheduleType):
+                    raise DaceSyntaxError(self, node, "RHS of dace.map @ operand must be a ScheduleType")
+
+            node = node.left
+
+        else:
+            schedule = None
+
         if not isinstance(node, (ast.Call, ast.Subscript)):
             raise DaceSyntaxError(self, node, "Iterator of ast.For must be a function or a subscript")
 
@@ -1635,6 +1675,8 @@ def _parse_for_iterator(self, node: ast.Expr):
 
         if iterator not in {'range', 'prange', 'parrange', 'dace.map'}:
             raise DaceSyntaxError(self, node, "Iterator {} is unsupported".format(iterator))
+        if schedule is not None and iterator == "range":
+            raise DaceSyntaxError(self, node, "Cannot specify schedule on range loops")
         elif iterator in ['range', 'prange', 'parrange']:
             # AST nodes for common expressions
             zero = ast.parse('0').body[0]
@@ -1674,7 +1716,7 @@ def visit_ast_or_value(arg):
             else:  # isinstance(node.slice, ast.Index) is True
                 ranges.append(self._parse_index_as_range(node.slice))
 
-        return (iterator, ranges, ast_ranges)
+        return (iterator, ranges, ast_ranges, schedule)
 
     def _parse_map_inputs(self, name: str, params: List[Tuple[str, str]],
                           node: ast.AST) -> Tuple[Dict[str, str], Dict[str, Memlet]]:
@@ -2078,21 +2120,24 @@ def visit_For(self, node: ast.For):
         # 3. `for i,j,k in dace.map[0:M, 0:N, 0:K]`: Creates an ND map
         # print(ast.dump(node))
         indices = self._parse_for_indices(node.target)
-        iterator, ranges, ast_ranges = self._parse_for_iterator(node.iter)
+        iterator, ranges, ast_ranges, schedule = self._parse_for_iterator(node.iter)
 
         if len(indices) != len(ranges):
             raise DaceSyntaxError(self, node, "Number of indices and ranges of for-loop do not match")
 
         if iterator == 'dace.map':
             if node.orelse:
                 raise DaceSyntaxError(self, node, '"else" clause not supported on DaCe maps')
+            if schedule is None:
+                schedule = dtypes.ScheduleType.Default
 
             state = self._add_state('MapState')
             params = [(k, ':'.join([str(t) for t in v])) for k, v in zip(indices, ranges)]
             params, map_inputs = self._parse_map_inputs('map_%d' % node.lineno, params, node)
             me, mx = state.add_map(name='%s_%d' % (self.name, node.lineno),
                                    ndrange=params,
-                                   debuginfo=self.current_lineinfo)
+                                   debuginfo=self.current_lineinfo,
+                                   schedule=schedule)
             # body = SDFG('MapBody')
             body, inputs, outputs, symbols = self._parse_subprogram(
                 self.name,

dace/frontend/python/preprocessing.py

@@ -984,10 +984,13 @@ def __init__(self, globals: Dict[str, Any], filename: str, closure_resolver: Glo
 
     def visit_For(self, node: ast.For) -> Any:
         # Avoid import loops
+        from dace.frontend.python.interface import MapGenerator
+
         EXPLICIT_GENERATORS = [
             range,  # Handled in ProgramVisitor
             dace.map,
             dace.consume,
+            MapGenerator
         ]
 
         node = self.generic_visit(node)
@@ -1033,6 +1036,17 @@ def visit_For(self, node: ast.For) -> Any:
 
         # Find out if unrolling should be done implicitly
         implicit = True
+
+        # Special case for map with @ operator
+        if isinstance(niter, ast.BinOp) and isinstance(niter.op, ast.MatMult):
+            try:
+                geniter = astutils.evalnode(niter.left, self.globals)
+                if isinstance(geniter, MapGenerator):
+                    niter = niter.left
+            except SyntaxError:
+                pass
+        # End of special case
+
         # Anything not a call is implicitly allowed
         if isinstance(niter, (ast.Call, ast.Subscript)):
             if isinstance(niter, ast.Subscript):

tests/python_frontend/device_annotations_test.py

@@ -0,0 +1,64 @@
+# Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved.
+import dace
+import pytest
+
+from dace.dtypes import StorageType, DeviceType, ScheduleType
+from dace import dtypes
+
+cupy = pytest.importorskip("cupy")
+
+
+@pytest.mark.gpu
+def test_storage():
+
+    @dace.program
+    def add(X: dace.float32[32, 32] @ StorageType.GPU_Global):
+        return X + 1
+
+    sdfg = add.to_sdfg()
+    sdfg.apply_gpu_transformations()
+
+    X = cupy.random.random((32, 32)).astype(cupy.float32)
+    Y = sdfg(X=X)
+    assert cupy.allclose(Y, X + 1)
+
+
+@pytest.mark.gpu
+def test_schedule():
+    Seq = ScheduleType.Sequential
+    N = dace.symbol('N')
+
+    @dace.program
+    def add2(X: dace.float32[32, 32] @ StorageType.GPU_Global):
+        for i in dace.map[0:N] @ ScheduleType.GPU_Device:
+            for j in dace.map[0:32] @ dace.ScheduleType.Sequential:
+                X[i, j] = X[i, j] + 1
+        for i in dace.map[0:32] @ dtypes.ScheduleType.GPU_Device:
+            for j in dace.map[0:32] @ Seq:
+                X[i, j] = X[i, j] + 1
+        return X
+
+    X = cupy.random.random((32, 32)).astype(cupy.float32)
+    Y = X.copy()
+    add2(X=X, N=32)
+    assert cupy.allclose(Y + 2, X)
+
+
+@pytest.mark.gpu
+def test_pythonmode():
+
+    def runs_on_gpu(a: dace.float64[20] @ StorageType.GPU_Global, b: dace.float64[20] @ StorageType.GPU_Global):
+        # This map will become a GPU kernel
+        for i in dace.map[0:20] @ ScheduleType.GPU_Device:
+            b[i] = a[i] + 1.0
+
+    gpu_a = cupy.random.rand(20)
+    gpu_b = cupy.random.rand(20)
+    runs_on_gpu(gpu_a, gpu_b)
+    assert cupy.allclose(gpu_b, gpu_a + 1)
+
+
+if __name__ == "__main__":
+    test_storage()
+    test_schedule()
+    test_pythonmode()