While loops in GTC toolchain

.pre-commit-config.yaml

@@ -12,6 +12,7 @@ repos:
   - id: check-case-conflict
 #  - id: check-json
   - id: check-merge-conflict
+    exclude_types: [rst]
   - id: check-toml
   - id: check-yaml
   - id: debug-statements

src/gt4py/backend/gtc_backend/defir_to_gtir.py

@@ -15,7 +15,7 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 import numbers
-from typing import Any, Dict, List, Union, cast
+from typing import Any, Dict, List, Tuple, Union, cast
 
 from gt4py.ir import IRNodeVisitor
 from gt4py.ir.nodes import (
@@ -45,6 +45,7 @@
     UnaryOpExpr,
     VarDecl,
     VarRef,
+    While,
 )
 from gtc import common, gtir
 from gtc.common import ExprKind
@@ -245,15 +246,15 @@ def visit_NativeFuncCall(self, node: NativeFuncCall) -> gtir.NativeFuncCall:
             loc=common.location_to_source_location(node.loc),
         )
 
-    def visit_FieldRef(self, node: FieldRef):
+    def visit_FieldRef(self, node: FieldRef) -> gtir.FieldAccess:
         return gtir.FieldAccess(
             name=node.name,
             offset=self.transform_offset(node.offset),
             data_index=[self.visit(index) for index in node.data_index],
             loc=common.location_to_source_location(node.loc),
         )
 
-    def visit_If(self, node: If):
+    def visit_If(self, node: If) -> Union[gtir.FieldIfStmt, gtir.ScalarIfStmt]:
         cond = self.visit(node.condition)
         if cond.kind == ExprKind.FIELD:
             return gtir.FieldIfStmt(
@@ -274,19 +275,26 @@ def visit_If(self, node: If):
                 loc=common.location_to_source_location(node.loc),
             )
 
+    def visit_While(self, node: While) -> gtir.While:
+        return gtir.While(
+            cond=self.visit(node.condition),
+            body=self.visit(node.body),
+            loc=common.location_to_source_location(node.loc),
+        )
+
     def visit_VarRef(self, node: VarRef, **kwargs):
         return gtir.ScalarAccess(name=node.name, loc=common.location_to_source_location(node.loc))
 
-    def visit_AxisInterval(self, node: AxisInterval):
+    def visit_AxisInterval(self, node: AxisInterval) -> Tuple[gtir.AxisBound, gtir.AxisBound]:
         return self.visit(node.start), self.visit(node.end)
 
-    def visit_AxisBound(self, node: AxisBound):
+    def visit_AxisBound(self, node: AxisBound) -> gtir.AxisBound:
         # TODO(havogt) add support VarRef
         return gtir.AxisBound(
             level=self.GT4PY_LEVELMARKER_TO_GTIR_LEVELMARKER[node.level], offset=node.offset
         )
 
-    def visit_FieldDecl(self, node: FieldDecl):
+    def visit_FieldDecl(self, node: FieldDecl) -> gtir.FieldDecl:
         dimension_names = ["I", "J", "K"]
         dimensions = [dim in node.axes for dim in dimension_names]
         # datatype conversion works via same ID
@@ -298,7 +306,7 @@ def visit_FieldDecl(self, node: FieldDecl):
             loc=common.location_to_source_location(node.loc),
         )
 
-    def visit_VarDecl(self, node: VarDecl):
+    def visit_VarDecl(self, node: VarDecl) -> gtir.ScalarDecl:
         # datatype conversion works via same ID
         return gtir.ScalarDecl(
             name=node.name,
@@ -308,10 +316,10 @@ def visit_VarDecl(self, node: VarDecl):
 
     def transform_offset(
         self, offset: Dict[str, Union[int, Expr]], **kwargs: Any
-    ) -> Union[gtir.CartesianOffset, gtir.VariableKOffset]:
+    ) -> Union[common.CartesianOffset, gtir.VariableKOffset]:
         k_val = offset.get("K", 0)
         if isinstance(k_val, numbers.Integral):
-            return gtir.CartesianOffset(i=offset.get("I", 0), j=offset.get("J", 0), k=k_val)
+            return common.CartesianOffset(i=offset.get("I", 0), j=offset.get("J", 0), k=k_val)
         elif isinstance(k_val, Expr):
             return gtir.VariableKOffset(k=self.visit(k_val, **kwargs))
         else:

src/gt4py/frontend/gtscript_frontend.py

@@ -737,7 +737,7 @@ def __init__(
         self.extra_temp_decls = extra_temp_decls or {}
         self.parsing_context = None
         self.iteration_order = None
-        self.if_decls_stack = []
+        self.decls_stack = []
         gt_ir.NativeFunction.PYTHON_SYMBOL_TO_IR_OP = {
             "abs": gt_ir.NativeFunction.ABS,
             "min": gt_ir.NativeFunction.MIN,
@@ -1237,7 +1237,7 @@ def visit_IfExp(self, node: ast.IfExp) -> gt_ir.TernaryOpExpr:
         return result
 
     def visit_If(self, node: ast.If) -> list:
-        self.if_decls_stack.append([])
+        self.decls_stack.append([])
 
         main_stmts = []
         for stmt in node.body:
@@ -1251,11 +1251,11 @@ def visit_If(self, node: ast.If) -> list:
             assert all(isinstance(item, gt_ir.Statement) for item in else_stmts)
 
         result = []
-        if len(self.if_decls_stack) == 1:
-            result.extend(self.if_decls_stack.pop())
-        elif len(self.if_decls_stack) > 1:
-            self.if_decls_stack[-2].extend(self.if_decls_stack[-1])
-            self.if_decls_stack.pop()
+        if len(self.decls_stack) == 1:
+            result.extend(self.decls_stack.pop())
+        elif len(self.decls_stack) > 1:
+            self.decls_stack[-2].extend(self.decls_stack[-1])
+            self.decls_stack.pop()
 
         result.append(
             gt_ir.If(
@@ -1270,17 +1270,28 @@ def visit_If(self, node: ast.If) -> list:
 
         return result
 
-    def visit_While(self, node: ast.While) -> gt_ir.While:
-        if node.orelse:
-            raise GTScriptSyntaxError("orelse is not supported on while loops")
-        stmts = []
-        for stmt in node.body:
-            stmts.extend(self.visit(stmt))
-        return gt_ir.While(
-            condition=self.visit(node.test),
-            loc=gt_ir.Location.from_ast_node(node),
-            body=gt_ir.BlockStmt(stmts=stmts, loc=gt_ir.Location.from_ast_node(node)),
-        )
+    def visit_While(self, node: ast.While) -> list:
+        loc = gt_ir.Location.from_ast_node(node)
+
+        self.decls_stack.append([])
+        stmts = gt_utils.flatten([self.visit(stmt) for stmt in node.body])
+        assert all(isinstance(item, gt_ir.Statement) for item in stmts)
+
+        result = [
+            gt_ir.While(
+                condition=self.visit(node.test),
+                loc=gt_ir.Location.from_ast_node(node),
+                body=gt_ir.BlockStmt(stmts=stmts, loc=loc),
+            )
+        ]
+
+        if len(self.decls_stack) == 1:
+            result.extend(self.decls_stack.pop())
+        elif len(self.decls_stack) > 1:
+            self.decls_stack[-2].extend(self.decls_stack[-1])
+            self.decls_stack.pop()
+
+        return result
 
     def visit_Call(self, node: ast.Call):
         native_fcn = gt_ir.NativeFunction.PYTHON_SYMBOL_TO_IR_OP[node.func.id]
@@ -1371,8 +1382,8 @@ def visit_Assign(self, node: ast.Assign) -> list:
                     # layout_id=t.id,
                     is_api=False,
                 )
-                if len(self.if_decls_stack):
-                    self.if_decls_stack[-1].append(field_decl)
+                if len(self.decls_stack):
+                    self.decls_stack[-1].append(field_decl)
                 else:
                     result.append(field_decl)
                 self.fields[field_decl.name] = field_decl

src/gtc/common.py

@@ -394,6 +394,21 @@ def condition_is_boolean(cls, cond: Expr) -> Expr:
         return verify_condition_is_boolean(cls, cond)
 
 
+class While(GenericNode, Generic[StmtT, ExprT]):
+    """
+    Generic while loop.
+
+    Verifies that `cond` is a boolean expr (if `dtype` is set).
+    """
+
+    cond: ExprT
+    body: List[StmtT]
+
+    @validator("cond")
+    def condition_is_boolean(cls, cond: Expr) -> Expr:
+        return verify_condition_is_boolean(cls, cond)
+
+
 class AssignStmt(GenericNode, Generic[TargetT, ExprT]):
     left: TargetT
     right: ExprT

src/gtc/cuir/cuir.py

@@ -69,11 +69,17 @@ class KCacheAccess(common.FieldAccess[Expr, VariableKOffset], Expr):
     k_cache_is_different_from_field_access = True
 
     @validator("offset")
-    def zero_ij_offset(cls, v: CartesianOffset) -> CartesianOffset:
+    def has_no_ij_offset(cls, v: Union[CartesianOffset, VariableKOffset]) -> CartesianOffset:
         if not v.i == v.j == 0:
             raise ValueError("No ij-offset allowed")
         return v
 
+    @validator("offset")
+    def not_variable_offset(cls, v: Union[CartesianOffset, VariableKOffset]) -> CartesianOffset:
+        if isinstance(v, VariableKOffset):
+            raise ValueError("Cannot k-cache a variable k offset")
+        return v
+
     @validator("data_index")
     def no_additional_dimensions(cls, v: List[int]) -> List[int]:
         if v:
@@ -92,6 +98,10 @@ class MaskStmt(Stmt):
     body: List[Stmt]
 
 
+class While(common.While[Stmt, Expr], Stmt):
+    pass
+
+
 class UnaryOp(common.UnaryOp[Expr], Expr):
     pass
 

src/gtc/cuir/cuir_codegen.py

@@ -46,6 +46,14 @@ class CUIRCodegen(codegen.TemplatedGenerator):
         """
     )
 
+    While = as_mako(
+        """
+        while (${cond}) {
+            ${'\\n'.join(body)}
+        }
+        """
+    )
+
     def visit_FieldAccess(self, node: cuir.FieldAccess, **kwargs: Any):
         def maybe_const(s):
             try:

src/gtc/cuir/oir_to_cuir.py

@@ -119,6 +119,11 @@ def visit_MaskStmt(self, node: oir.MaskStmt, **kwargs: Any) -> cuir.MaskStmt:
             mask=self.visit(node.mask, **kwargs), body=self.visit(node.body, **kwargs)
         )
 
+    def visit_While(self, node: oir.While, **kwargs: Any) -> cuir.While:
+        return cuir.While(
+            cond=self.visit(node.cond, **kwargs), body=self.visit(node.body, **kwargs)
+        )
+
     def visit_Cast(self, node: oir.Cast, **kwargs: Any) -> cuir.Cast:
         return cuir.Cast(dtype=node.dtype, expr=self.visit(node.expr, **kwargs))
 

src/gtc/dace/expansion.py

@@ -171,6 +171,13 @@ def visit_MaskStmt(self, node: oir.MaskStmt, **kwargs):
         body_code = [indent + b for b in body_code]
         return "\n".join([mask_str] + body_code)
 
+    def visit_While(self, node: oir.While, **kwargs):
+        body = self.visit(node.body, **kwargs)
+        cond = self.visit(node.cond, is_target=False, **kwargs)
+        indent = " " * 4
+        delim = f"\n{indent}"
+        return f"while {cond}:\n{indent}{delim.join(body)}"
+
     class RemoveCastInIndexVisitor(eve.NodeTranslator):
         def visit_FieldAccess(self, node: oir.FieldAccess):
             if node.data_index:

src/gtc/dace/utils.py

@@ -232,16 +232,19 @@ def domain() -> "CartesianIterationSpace":
         )
 
     @staticmethod
-    def from_offset(offset: CartesianOffset) -> "CartesianIterationSpace":
+    def from_offset(
+        offset: Union[CartesianOffset, oir.VariableKOffset]
+    ) -> "CartesianIterationSpace":
 
+        dict_offsets = offset.to_dict()
         return CartesianIterationSpace(
             i_interval=oir.Interval(
-                start=oir.AxisBound.from_start(min(0, offset.i)),
-                end=oir.AxisBound.from_end(max(0, offset.i)),
+                start=oir.AxisBound.from_start(min(0, dict_offsets["i"])),
+                end=oir.AxisBound.from_end(max(0, dict_offsets["i"])),
             ),
             j_interval=oir.Interval(
-                start=oir.AxisBound.from_start(min(0, offset.j)),
-                end=oir.AxisBound.from_end(max(0, offset.j)),
+                start=oir.AxisBound.from_start(min(0, dict_offsets["j"])),
+                end=oir.AxisBound.from_end(max(0, dict_offsets["j"])),
             ),
         )
 

src/gtc/gtcpp/gtcpp.py

@@ -77,6 +77,10 @@ class IfStmt(common.IfStmt[Stmt, Expr], Stmt):
     pass
 
 
+class While(common.While[Stmt, Expr], Stmt):
+    pass
+
+
 class UnaryOp(common.UnaryOp[Expr], Expr):
     pass
 

src/gtc/gtcpp/gtcpp_codegen.py

@@ -203,6 +203,8 @@ def visit_LoopOrder(self, looporder: LoopOrder, **kwargs: Any) -> str:
         """
     )
 
+    While = as_mako("while(${cond}) {${''.join(body)}}")
+
     BlockStmt = as_mako("{${''.join(body)}}")
 
     def visit_GTComputationCall(

src/gtc/gtcpp/oir_to_gtcpp.py

@@ -180,6 +180,11 @@ def visit_MaskStmt(self, node: oir.MaskStmt, **kwargs: Any) -> gtcpp.IfStmt:
             true_branch=gtcpp.BlockStmt(body=self.visit(node.body, **kwargs)),
         )
 
+    def visit_While(self, node: oir.While, **kwargs: Any) -> gtcpp.While:
+        return gtcpp.While(
+            cond=self.visit(node.cond, **kwargs), body=self.visit(node.body, **kwargs)
+        )
+
     def visit_HorizontalExecution(
         self,
         node: oir.HorizontalExecution,