Preliminary implementation for globals

dusk/grammar.py

@@ -463,8 +463,12 @@ def constant(self, value):
             _type,
         )
 
-    @transform(Name(id=Capture(str).to("name"), ctx=AnyContext))
-    def var(self, name: str, index: expr = None):
+    @transform(Capture(Name(id=Capture(str).to("name"), ctx=AnyContext)).to("node"))
+    def var(self, node: Name, name: str, index: expr = None):
+
+        if hasattr(node, "sir"):
+            # FIXME: remove once symbol resolution has been properly unified
+            return node.sir
 
         if not self.ctx.scope.current_scope.contains(name):
             raise SemanticError(f"Undeclared variable '{name}'!", name)

dusk/integration.py

@@ -10,6 +10,7 @@ class StencilObject:
     callable: Callable
     filename: str
     pyast: Optional[ast.FunctionDef] = None
+    globals: Optional[sir.GlobalVariableMap]
     sir_node: Optional[sir.SIR] = None
 
     def __init__(self, callable: Callable, filename: str = "<unknown>"):

dusk/passes/pipeline.py

@@ -7,13 +7,15 @@
 from dusk.integration import StencilObject
 from dusk.grammar import Grammar
 from dusk.passes.symbol_resolution import resolve_symbols
+from dusk.passes.resolve_globals import resolve_globals
 
 
 def stencil_object_to_sir(stencil_object: StencilObject) -> sir.Stencil:
 
     add_filename(stencil_object)
     add_pyast(stencil_object)
     resolve_symbols(stencil_object)
+    resolve_globals(stencil_object)
     add_sir(stencil_object)
 
     return stencil_object.sir_node
@@ -42,5 +44,8 @@ def add_sir(stencil_object: StencilObject) -> None:
 
     sir_stencil = Grammar().stencil(stencil_object.pyast)
     stencil_object.sir_node = ser.make_sir(
-        stencil_object.filename, sir.GridType.Value("Unstructured"), [sir_stencil]
+        stencil_object.filename,
+        sir.GridType.Value("Unstructured"),
+        [sir_stencil],
+        global_variables=stencil_object.globals,
     )

dusk/passes/resolve_globals.py

@@ -0,0 +1,24 @@
+import ast
+
+import dawn4py.serialization as ser
+import dawn4py.serialization.SIR as sir
+
+import dusk.script as script
+from dusk.integration import StencilObject
+
+
+def resolve_globals(stencil_object: StencilObject) -> None:
+
+    # TODO: check preconditions & postconditions
+
+    stencil_object.globals = sir.GlobalVariableMap()
+
+    for node in ast.walk(stencil_object.pyast):
+        if (
+            isinstance(node, ast.Name)
+            and hasattr(node, "decl")
+            and isinstance(node.decl, script.Global)
+        ):
+            name = node.decl.name
+            stencil_object.globals.map[name].double_value = 0
+            node.sir = ser.make_var_access_expr(name, is_external=True)

dusk/script/__init__.py

@@ -9,6 +9,7 @@
 
 __all__ = [
     "stencil",
+    "Global",
     "Edge",
     "Cell",
     "Vertex",
@@ -38,6 +39,13 @@ def stencil(stencil: typing.Callable) -> typing.Callable:
     return stencil
 
 
+class Global:
+    name: str
+
+    def __init__(self, name: str):
+        self.name = name
+
+
 class Edge(metaclass=internal.LocationType):
     pass
 

tests/stencils/test_globals.py

@@ -0,0 +1,39 @@
+from dusk.script import *
+from dusk.test import stencil_test
+
+
+dt = Global("dt")
+
+
+@stencil_test()
+def test_simple_global1(a: Field[Edge]):
+    with levels_downward:
+        a = dt
+
+
+g = Global("some_global")
+
+
+@stencil_test()
+def test_simple_global2(a: Field[Edge], b: Field[Cell]):
+    with levels_downward:
+        a = sum_over(Edge > Cell, g * b)
+
+
+one = Global("one")
+two = Global("two")
+
+
+@stencil_test()
+def test_two_globals(a: Field[Edge]):
+    with levels_downward:
+        a = one + two
+
+
+aliased = dt
+
+
+@stencil_test()
+def test_aliased_global(a: Field[Cell]):
+    with levels_upward:
+        a = aliased / (one - two)