Frontend: add syntactic sugar for schedule and storage types.

dace/data.py

@@ -321,6 +321,20 @@ 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 lets us write code like:
+
+        .. 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/newast.py

@@ -11,6 +11,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
@@ -1618,11 +1619,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")
 
@@ -1633,6 +1673,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]
@@ -1672,7 +1714,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]]:
@@ -2076,21 +2118,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,

tests/python_frontend/device_annotations_test.py

@@ -0,0 +1,47 @@
+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
+
+    @dace.program
+    def add2(X: dace.float32[32, 32] @ StorageType.GPU_Global):
+        for i in dace.map[0:32] @ 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)
+    assert cupy.allclose(Y + 2, X)
+
+
+if __name__ == "__main__":
+    test_storage()
+    test_schedule()