diff --git a/src/gt4py/cartesian/frontend/defir_to_gtir.py b/src/gt4py/cartesian/frontend/defir_to_gtir.py index 0c2274ce78..5059aa3b27 100644 --- a/src/gt4py/cartesian/frontend/defir_to_gtir.py +++ b/src/gt4py/cartesian/frontend/defir_to_gtir.py @@ -39,6 +39,8 @@ Expr, FieldDecl, FieldRef, + For, + ForIndex, HorizontalIf, If, IterationOrder, @@ -582,6 +584,19 @@ def visit_While(self, node: While) -> gtir.While: loc=location_to_source_location(node.loc), ) + def visit_ForIndex(self, node: ForIndex) -> gtir.ForIndex: + return gtir.ForIndex(name=node.name, dtype=common.DataType(node.data_type.value)) + + def visit_For(self, node: For) -> gtir.For: + return gtir.For( + index_name=node.index_name, + iter_start=node.iter_start, + iter_stop=node.iter_stop, + iter_step=node.iter_step, + body=self.visit(node.body), + loc=location_to_source_location(node.loc), + ) + def visit_VarRef(self, node: VarRef, **kwargs) -> gtir.ScalarAccess: return gtir.ScalarAccess(name=node.name, loc=location_to_source_location(node.loc)) diff --git a/src/gt4py/cartesian/frontend/gtscript_frontend.py b/src/gt4py/cartesian/frontend/gtscript_frontend.py index 6d8ff204e0..246fed7559 100644 --- a/src/gt4py/cartesian/frontend/gtscript_frontend.py +++ b/src/gt4py/cartesian/frontend/gtscript_frontend.py @@ -18,6 +18,7 @@ from typing import ( Any, Callable, + ClassVar, Dict, Final, List, @@ -46,6 +47,7 @@ GTScriptValueError, ) from gt4py.cartesian.utils import meta as gt_meta, warn_experimental_feature +from gt4py.eve import datamodels as gt_datamodels PYTHON_AST_VERSION: Final = (3, 10) @@ -454,11 +456,15 @@ def visit_FunctionDef(self, node: ast.FunctionDef): class ReturnReplacer(gt_meta.ASTTransformPass): @classmethod - def apply(cls, ast_object: ast.AST, target_node: ast.AST) -> None: + def apply(cls, ast_object: ast.AST, target_node: Optional[ast.AST]) -> None: """Ensure that there is only a single return statement (can still return a tuple).""" ret_count = sum(isinstance(node, ast.Return) for node in ast.walk(ast_object)) - if ret_count != 1: - raise GTScriptSyntaxError("GTScript Functions should have a single return statement") + if ret_count > 1: + raise GTScriptSyntaxError("GTScript Functions cannot have multiple return statements") + elif ret_count == 0 and target_node is not None: + raise GTScriptSyntaxError( + "Attempting to assign the return value of a GTScript function that does not return anything." + ) cls().visit(ast_object, target_node=target_node) @staticmethod @@ -485,6 +491,98 @@ def _filter_absolute_K_index_method(node: ast.Call) -> bool: return isinstance(node.func, ast.Attribute) and node.func.attr == "at" +@gt_datamodels.datamodel(frozen=True) +class ForIndex(ast.AST): + name: str + + def __deepcopy__(self, memo: dict) -> "ForIndex": + return self + + +@gt_datamodels.datamodel(frozen=True) +class ForIndexTransformer(ast.NodeTransformer): + name: str + + def visit_Name(self, node: ast.Name) -> Union[ForIndex, ast.Name]: + super().generic_visit(node) + return ForIndex(self.name) if node.id == self.name else node + + +@gt_datamodels.datamodel +class For(ast.AST): + index_name: str + iter_start: int + iter_stop: int + iter_step: int + body: list[ast.AST] + lineno: Optional[int] = None + col_offset: Optional[int] = None + _fields: ClassVar[tuple[str, ...]] = ("body",) + + def __post_init__(self) -> None: + transformer = ForIndexTransformer(self.index_name) + self.body = [transformer.visit(stmt) for stmt in self.body] + + def __deepcopy__(self, memo: dict) -> "For": + return For( + index_name=self.index_name, + iter_start=self.iter_start, + iter_stop=self.iter_stop, + iter_step=self.iter_step, + body=copy.deepcopy(self.body), + lineno=self.lineno, + col_offset=self.col_offset, + ) + + +@gt_datamodels.datamodel(frozen=True) +class ForTransformer(ast.NodeTransformer): + context: dict + + @classmethod + def apply(cls, func_node: ast.FunctionDef, context: dict): + unroller = cls(context) + unroller(func_node) + + def __call__(self, func_node: ast.FunctionDef) -> None: + self.visit(func_node) + + def visit_For(self, node: Union[ast.For, For]) -> For: + super().generic_visit(node) + if isinstance(node, ast.For): + assert isinstance(node.target, ast.Name) + + if ( + isinstance(node.iter, ast.Call) + and isinstance(node.iter.func, ast.Name) + and node.iter.func.id == "range" + ): + range_args = [eval(ast.unparse(arg), self.context) for arg in node.iter.args] + assert 1 <= len(range_args) <= 3 + if len(range_args) == 1: + start, stop, step = 0, *range_args, 1 + elif len(range_args) == 2: + start, stop, step = *range_args, 1 + else: + start, stop, step = range_args + else: + raise GTScriptSyntaxError( + "For-loop index values can only be specified using range()." + ) + + return For( + index_name=node.target.id, + iter_start=start, + iter_stop=stop, + iter_step=step, + body=[self.visit(item) for item in node.body], + lineno=node.lineno, + col_offset=node.col_offset, + ) + else: + return node + + class CallInliner(ast.NodeTransformer): """Inlines calls to gtscript.function calls. @@ -554,6 +652,10 @@ def visit_While(self, node: ast.While): node.body = self._process_stmts(node.body) return node + def visit_For(self, node: Union[ast.For, For]): + node.body = self._process_stmts(node.body) + return node + def visit_Assign(self, node: ast.Assign): if isinstance(node.value, ast.Call): if _filter_absolute_K_index_method(node.value): @@ -570,6 +672,12 @@ def visit_Assign(self, node: ast.Assign): return self.generic_visit(node) + def _get_sliced_symbol(self, node): + if isinstance(node, ast.Name): + return node.id + elif isinstance(node, ast.Subscript): + return self._get_sliced_symbol(node.value) + def visit_Call(self, node: ast.Call, *, target_node=None): # Cyclomatic complexity too high if _filter_absolute_K_index_method(node): return node @@ -589,17 +697,8 @@ def visit_Call(self, node: ast.Call, *, target_node=None): # Cyclomatic complex if call_name not in self.context or not hasattr(self.context[call_name], "_gtscript_"): raise GTScriptSyntaxError("Unknown call", loc=nodes.Location.from_ast_node(node)) - # Recursively inline any possible nested subroutine call - call_info = self.context[call_name]._gtscript_ - call_ast = copy.deepcopy(call_info["ast"]) - self.current_name = call_name - CallInliner.apply( - call_ast, - call_info["local_context"], - call_stack={*self.call_stack, call_name}, - ) - # Extract call arguments + call_info = self.context[call_name]._gtscript_ call_signature = call_info["api_signature"] arg_infos = {arg.name: arg.default for arg in call_signature} try: @@ -623,6 +722,21 @@ def visit_Call(self, node: ast.Call, *, target_node=None): # Cyclomatic complex loc=nodes.Location.from_ast_node(node), ) from ex + # Inline constant arguments + call_ast = copy.deepcopy(call_info["ast"]) + local_context = { + name: arg_node.value + for name, arg_node in call_args.items() + if isinstance(arg_node, ast.Constant) + } + ValueInliner.apply(call_ast, local_context) + + # Recursively inline any possible nested subroutine call + self.current_name = call_name + CallInliner.apply( + call_ast, call_info["local_context"], call_stack={*self.call_stack, call_name} + ) + # Rename local names in subroutine to avoid conflicts with caller context names try: assign_targets = gt_meta.collect_assign_targets(call_ast, allow_multiple_targets=False) @@ -633,10 +747,15 @@ def visit_Call(self, node: ast.Call, *, target_node=None): # Cyclomatic complex assigned_symbols = set() for target in assign_targets: - if not isinstance(target, ast.Name): - raise GTScriptSyntaxError(message="Unsupported assignment target.", loc=target) + if isinstance(target, ast.Subscript): + sliced_symbol = self._get_sliced_symbol(target) + if sliced_symbol not in call_args: + raise GTScriptSyntaxError(message="Unsupported assignment target.", loc=target) + else: + if not isinstance(target, ast.Name): + raise GTScriptSyntaxError(message="Unsupported assignment target.", loc=target) - assigned_symbols.add(target.id) + assigned_symbols.add(target.id) name_mapping = { name: value.id @@ -657,32 +776,32 @@ def visit_Call(self, node: ast.Call, *, target_node=None): # Cyclomatic complex # Replace returns by assignments in subroutine if target_node is None: - if any( - isinstance(nd.value, ast.Tuple) - for nd in ast.walk(call_ast) - if isinstance(nd, ast.Return) - ): + return_nodes = [nd for nd in ast.walk(call_ast) if isinstance(nd, ast.Return)] + if any(isinstance(nd.value, ast.Tuple) for nd in return_nodes): raise GTScriptSyntaxError( "Only functions with a single return value can be used in expressions, including as call arguments. " "Please assign the function results to symbols first." ) - target_node = ast.Name( - ctx=ast.Store(), - lineno=node.lineno, - col_offset=node.col_offset, - id=template_fmt.format(name="RETURN_VALUE"), - ) - assert isinstance(target_node, (ast.Name, ast.Tuple, ast.Subscript)) and isinstance( - target_node.ctx, ast.Store - ) + if len(return_nodes) > 0: + target_node = ast.Name( + ctx=ast.Store(), + lineno=node.lineno, + col_offset=node.col_offset, + id=template_fmt.format(name="RETURN_VALUE"), + ) + assert target_node is None or ( + isinstance(target_node, (ast.Name, ast.Tuple, ast.Subscript)) + and isinstance(target_node.ctx, ast.Store) + ) ReturnReplacer.apply(call_ast, target_node) # Add subroutine sources prepending the required arg assignments inlined_stmts = [] for arg_name, arg_value in call_args.items(): - if arg_name not in name_mapping: + # note(stubbiali): filter out constant arguments (which have been previously inlined) + if arg_name not in name_mapping and not isinstance(arg_value, ast.Constant): inlined_stmts.append( ast.Assign( lineno=node.lineno, @@ -716,7 +835,7 @@ def visit_Call(self, node: ast.Call, *, target_node=None): # Cyclomatic complex col_offset=target_node.col_offset, elts=target_node.elts, ) - else: + elif isinstance(target_node, ast.Subscript): result_node = ast.Subscript( ctx=ast.Load(), lineno=target_node.lineno, @@ -724,6 +843,8 @@ def visit_Call(self, node: ast.Call, *, target_node=None): # Cyclomatic complex value=target_node.value, slice=target_node.slice, ) + else: # target_node is None + result_node = call_ast.body[0] # Add the temp_annotations and temp_init_values to the parent current_info = self.context[self.current_name]._gtscript_ @@ -738,9 +859,17 @@ def visit_Call(self, node: ast.Call, *, target_node=None): # Cyclomatic complex return result_node def visit_Expr(self, node: ast.Expr): - """Ignore pure string statements in callee.""" - pure_str_types = (ast.Constant,) + ((ast.Str,) if hasattr(ast, "Str") else ()) - if not isinstance(node.value, pure_str_types): + if ( + isinstance(node.value, ast.Call) + and gt_meta.get_qualified_name_from_node(node.value.func) not in gtscript.MATH_BUILTINS + ): + # Inline a function with no return value and then remove the current node + self.visit(node.value, target_node=None) + return None + elif not isinstance( + node.value, (ast.Constant,) + ((ast.Str,) if hasattr(ast, "Str") else ()) + ): + # Ignore pure string statements in callee return super().visit(node.value) @@ -1658,6 +1787,42 @@ def visit_While(self, node: ast.While) -> list: return result + def visit_ForIndex(self, node: ForIndex) -> nodes.ForIndex: + return nodes.ForIndex( + name=node.name, + data_type=nodes.DataType.from_dtype( + self.dtypes[int] if self.dtypes and int in self.dtypes else int + ), + ) + + def visit_For(self, node: For) -> list: + assert isinstance(node, For) + + loc = nodes.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, nodes.Statement) for item in stmts) + + result = [ + nodes.For( + index_name=node.index_name, + iter_start=node.iter_start, + iter_stop=node.iter_stop, + iter_step=node.iter_step, + body=nodes.BlockStmt(stmts=stmts, loc=loc), + loc=nodes.Location.from_ast_node(node), + ) + ] + + 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 _absolute_K_index_method(self, node: ast.Call): # Dev note: we enforce .at(K=..., ddim=[...]) # A better version of this code would look through the keywords @@ -2508,9 +2673,17 @@ def run(self, backend_name: str): ValueInliner.apply(main_func_node, context=local_context) + # Insert custom nodes for for-loops + # note(stubbiali): address the case of a function called within a for-loop + ForTransformer.apply(main_func_node, context=local_context) + # Inline function calls CallInliner.apply(main_func_node, context=local_context) + # Insert custom nodes for for-loops + # note(stubbiali): address the case of a for-loop inside a function + ForTransformer.apply(main_func_node, context=local_context) + # Evaluate and inline compile-time conditionals CompiledIfInliner.apply(main_func_node, context=local_context) diff --git a/src/gt4py/cartesian/frontend/nodes.py b/src/gt4py/cartesian/frontend/nodes.py index f863243e37..4cb3b6635d 100644 --- a/src/gt4py/cartesian/frontend/nodes.py +++ b/src/gt4py/cartesian/frontend/nodes.py @@ -416,6 +416,12 @@ class AxisIndex(Expr): data_type = attribute(of=DataType, default=DataType.INT32) +@attribclass +class ForIndex(Expr): + name = attribute(of=str) + data_type = attribute(of=DataType) + + @enum.unique class NativeFunction(enum.Enum): ABS = enum.auto() @@ -703,6 +709,16 @@ class While(Statement): loc = attribute(of=Location, optional=True) +@attribclass +class For(Statement): + index_name = attribute(of=str) + iter_start = attribute(of=int) + iter_stop = attribute(of=int) + iter_step = attribute(of=int) + body = attribute(of=BlockStmt) + loc = attribute(of=Location, optional=None) + + # ---- IR: computations ---- @enum.unique class IterationOrder(enum.Enum): diff --git a/src/gt4py/cartesian/gtc/common.py b/src/gt4py/cartesian/gtc/common.py index 63c03c92c9..26a8feb118 100644 --- a/src/gt4py/cartesian/gtc/common.py +++ b/src/gt4py/cartesian/gtc/common.py @@ -438,6 +438,19 @@ def condition_is_boolean(self, attribute: datamodels.Attribute, value: Expr) -> verify_condition_is_boolean(self, value) +class ForIndex(eve.Node): + name: str + dtype: DataType + + +class For(eve.GenericNode, Generic[StmtT]): + index_name: str + iter_start: int + iter_stop: int + iter_step: int + body: List[StmtT] + + class AssignStmt(eve.GenericNode, Generic[TargetT, ExprT]): left: TargetT right: ExprT diff --git a/src/gt4py/cartesian/gtc/dace/oir_to_tasklet.py b/src/gt4py/cartesian/gtc/dace/oir_to_tasklet.py index 2329128d70..09ec8be198 100644 --- a/src/gt4py/cartesian/gtc/dace/oir_to_tasklet.py +++ b/src/gt4py/cartesian/gtc/dace/oir_to_tasklet.py @@ -290,6 +290,9 @@ def visit_IteratorAccess(self, node: oir.IteratorAccess, ctx: Context, **kwargs: return tir.Axis(node.name).iteration_symbol() + def visit_ForIndex(self, node: oir.ForIndex, **kwargs: Any) -> str: + return node.name + # Not (yet) supported section def visit_CacheDesc(self, node: oir.CacheDesc, **kwargs: Any) -> None: raise NotImplementedError("To be implemented: Caches") @@ -316,6 +319,9 @@ def visit_MaskStmt(self, node: oir.MaskStmt, **kwargs: Any) -> None: def visit_While(self, node: oir.While, **kwargs: Any) -> None: raise RuntimeError("visit_While should not be called") + def visit_For(self, node: oir.For, **kwargs: Any) -> None: + raise RuntimeError("visit_For should not be called") + def visit_HorizontalRestriction(self, node: oir.HorizontalRestriction, **kwargs: Any) -> None: raise RuntimeError("visit_HorizontalRestriction: should be dealt with in TreeIR") diff --git a/src/gt4py/cartesian/gtc/dace/oir_to_treeir.py b/src/gt4py/cartesian/gtc/dace/oir_to_treeir.py index c578fea24e..5e18c438bb 100644 --- a/src/gt4py/cartesian/gtc/dace/oir_to_treeir.py +++ b/src/gt4py/cartesian/gtc/dace/oir_to_treeir.py @@ -20,7 +20,7 @@ ControlFlow: TypeAlias = ( - oir.HorizontalExecution | oir.While | oir.MaskStmt | oir.HorizontalRestriction + oir.HorizontalExecution | oir.While | oir.For | oir.MaskStmt | oir.HorizontalRestriction ) """All control flow OIR nodes""" @@ -236,6 +236,20 @@ def visit_While(self, node: oir.While, ctx: tir.Context) -> None: groups = self._group_statements(node) self.visit(groups, ctx=ctx) + def visit_For(self, node: oir.For, ctx: tir.Context) -> None: + for_ = tir.For( + index_name=node.index_name, + iter_start=node.iter_start, + iter_stop=node.iter_stop, + iter_step=node.iter_step, + children=[], + parent=ctx.current_scope, + ) + + with for_.scope(ctx): + groups = self._group_statements(node) + self.visit(groups, ctx=ctx) + def visit_AxisBound(self, node: common.AxisBound, axis_start: str, axis_end: str) -> str: if node.level == common.LevelMarker.START: return f"({axis_start}) + ({node.offset})" @@ -523,6 +537,9 @@ def visit_Decl(self, node: oir.Decl, **kwargs: Any) -> None: def visit_FieldDecl(self, node: oir.FieldDecl, **kwargs: Any) -> None: raise RuntimeError("visit_FieldDecl should not be called") + def visit_ForIndex(self, node: oir.ForIndex, **kwargs: Any) -> None: + raise RuntimeError("visit_ForIndex should not be called") + def visit_LocalScalar(self, node: oir.LocalScalar, **kwargs: Any) -> None: raise RuntimeError("visit_LocalScalar should not be called") diff --git a/src/gt4py/cartesian/gtc/dace/treeir.py b/src/gt4py/cartesian/gtc/dace/treeir.py index 668ce866bf..e31f28270d 100644 --- a/src/gt4py/cartesian/gtc/dace/treeir.py +++ b/src/gt4py/cartesian/gtc/dace/treeir.py @@ -119,6 +119,17 @@ class While(TreeScope): """Condition as ScheduleTree worthy code""" +class ForIndex(TreeNode): + name: str + + +class For(TreeScope): + index_name: str + iter_start: int + iter_stop: int + iter_step: int + + class HorizontalLoop(TreeScope): bounds_i: Bounds bounds_j: Bounds diff --git a/src/gt4py/cartesian/gtc/dace/treeir_to_stree.py b/src/gt4py/cartesian/gtc/dace/treeir_to_stree.py index 0988c593d3..46e766d355 100644 --- a/src/gt4py/cartesian/gtc/dace/treeir_to_stree.py +++ b/src/gt4py/cartesian/gtc/dace/treeir_to_stree.py @@ -124,6 +124,12 @@ def visit_While(self, node: tir.While, ctx: Context) -> None: with ContextPushPop(ctx, while_scope): self.visit(node.children, ctx=ctx) + def visit_For(self, node: tir.For, ctx: Context) -> None: + for_scope = tn.ForScope(loop=_loop_region_for_over_data_dim(node), children=[]) + + with ContextPushPop(ctx, for_scope): + self.visit(node.children, ctx=ctx) + def visit_TreeRoot(self, node: tir.TreeRoot) -> tn.ScheduleTreeRoot: """Construct a schedule tree from TreeIR.""" tree = tn.ScheduleTreeRoot( @@ -168,3 +174,21 @@ def _loop_region_while(node: tir.While) -> LoopRegion: :return: DaCe LoopRegion to use in `tn.WhileScope` """ return LoopRegion(label=f"while_loop_{id(node)}", condition_expr=CodeBlock(node.condition_code)) + + +def _loop_region_for_over_data_dim(node: tir.For) -> LoopRegion: + """ + Translates a for-loop over a data dimension into a Dace LoopRegion to be used in `tn.ForScope`. + + :param node: For loop to translate + :return: DaCe LoopRegion to use in `tn.ForScope` + """ + plus_minus = "+" if (iter_step := node.iter_step) > 0 else "-" + comparison = "<" if iter_step > 0 else ">" + return LoopRegion( + label=f"for_loop_data_dim_{id(node)}", + loop_var=(index_name := node.index_name), + initialize_expr=CodeBlock(f"{index_name} = {node.iter_start}"), + condition_expr=CodeBlock(f"{index_name} {comparison} {node.iter_stop}"), + update_expr=CodeBlock(f"{index_name} = {index_name} {plus_minus} {abs(iter_step)}"), + ) diff --git a/src/gt4py/cartesian/gtc/gtcpp/gtcpp.py b/src/gt4py/cartesian/gtc/gtcpp/gtcpp.py index f1401a776a..8000c6c342 100644 --- a/src/gt4py/cartesian/gtc/gtcpp/gtcpp.py +++ b/src/gt4py/cartesian/gtc/gtcpp/gtcpp.py @@ -72,6 +72,14 @@ class While(common.While[Stmt, Expr], Stmt): pass +class ForIndex(common.ForIndex, Expr): + pass + + +class For(common.For[Stmt], Stmt): + pass + + class UnaryOp(common.UnaryOp[Expr], Expr): pass diff --git a/src/gt4py/cartesian/gtc/gtcpp/gtcpp_codegen.py b/src/gt4py/cartesian/gtc/gtcpp/gtcpp_codegen.py index 62eae238a9..a0ffc77a70 100644 --- a/src/gt4py/cartesian/gtc/gtcpp/gtcpp_codegen.py +++ b/src/gt4py/cartesian/gtc/gtcpp/gtcpp_codegen.py @@ -256,6 +256,14 @@ def visit_Temporary(self, node: gtcpp.Temporary, **kwargs: Any) -> str: While = as_mako("while(${cond}) {${''.join(body)}}") + ForIndex = as_mako("${name}") + For = as_mako( + "for(std::size_t ${index_name}=${iter_start}; " + "${index_name}${'<' if _this_node.iter_step > 0 else '>'}${iter_stop}; " + "${index_name}+=(${iter_step})) " + "{${''.join(body)}}" + ) + BlockStmt = as_mako("{${''.join(body)}}") def visit_GTComputationCall( diff --git a/src/gt4py/cartesian/gtc/gtcpp/oir_to_gtcpp.py b/src/gt4py/cartesian/gtc/gtcpp/oir_to_gtcpp.py index be07b81e08..c3aaeacbb1 100644 --- a/src/gt4py/cartesian/gtc/gtcpp/oir_to_gtcpp.py +++ b/src/gt4py/cartesian/gtc/gtcpp/oir_to_gtcpp.py @@ -301,6 +301,18 @@ def visit_While(self, node: oir.While, **kwargs: Any) -> gtcpp.While: cond=self.visit(node.cond, **kwargs), body=self.visit(node.body, **kwargs) ) + def visit_ForIndex(self, node: oir.ForIndex, **kwargs: Any) -> gtcpp.ForIndex: + return gtcpp.ForIndex(name=node.name, dtype=node.dtype) + + def visit_For(self, node: common.For, **kwargs: Any) -> gtcpp.For: + return gtcpp.For( + index_name=node.index_name, + iter_start=node.iter_start, + iter_stop=node.iter_stop, + iter_step=node.iter_step, + body=self.visit(node.body, **kwargs), + ) + def visit_HorizontalExecution( self, node: oir.HorizontalExecution, diff --git a/src/gt4py/cartesian/gtc/gtir.py b/src/gt4py/cartesian/gtc/gtir.py index 8b7900a4c5..5ec2ccda52 100644 --- a/src/gt4py/cartesian/gtc/gtir.py +++ b/src/gt4py/cartesian/gtc/gtir.py @@ -165,6 +165,14 @@ def _no_write_and_read_with_horizontal_offset_all( raise ValueError(f"Illegal write and read with horizontal offset detected for {names}.") +class ForIndex(common.ForIndex, Expr): + pass + + +class For(common.For[Stmt], Stmt): + pass + + class UnaryOp(common.UnaryOp[Expr], Expr): pass diff --git a/src/gt4py/cartesian/gtc/gtir_to_oir.py b/src/gt4py/cartesian/gtc/gtir_to_oir.py index 249201040e..df9021063a 100644 --- a/src/gt4py/cartesian/gtc/gtir_to_oir.py +++ b/src/gt4py/cartesian/gtc/gtir_to_oir.py @@ -143,6 +143,24 @@ def visit_While(self, node: gtir.While, **kwargs: Any) -> oir.While: condition: oir.Expr = self.visit(node.cond) return oir.While(cond=condition, body=body, loc=node.loc) + def visit_ForIndex(self, node: gtir.ForIndex, **kwargs: Any) -> oir.ForIndex: + return oir.ForIndex(name=node.name, dtype=node.dtype) + + def visit_For(self, node: gtir.For, **kwargs: Any) -> oir.For: + body: List[oir.Stmt] = [] + for statement in node.body: + oir_statement = self.visit(statement, **kwargs) + body.extend(utils.flatten(utils.listify(oir_statement))) + + return oir.For( + index_name=node.index_name, + iter_start=node.iter_start, + iter_stop=node.iter_stop, + iter_step=node.iter_step, + body=body, + loc=node.loc, + ) + def visit_FieldIfStmt( self, node: gtir.FieldIfStmt, diff --git a/src/gt4py/cartesian/gtc/numpy/npir.py b/src/gt4py/cartesian/gtc/numpy/npir.py index ae5be3903a..3f4b77b592 100644 --- a/src/gt4py/cartesian/gtc/numpy/npir.py +++ b/src/gt4py/cartesian/gtc/numpy/npir.py @@ -118,6 +118,10 @@ class VarKOffset(common.VariableKOffset[Expr]): pass +class ForIndex(common.ForIndex, Expr): + pass + + class KMaskFieldAccess(Expr): dtype = common.DataType.INT64 @@ -130,14 +134,6 @@ class FieldSlice(VectorLValue): data_index: List[Expr] = eve.field(default_factory=list) kind: common.ExprKind = common.ExprKind.FIELD - @datamodels.validator("data_index") - def data_indices_are_scalar( - self, attribute: datamodels.Attribute, data_index: List[Expr] - ) -> None: - for index in data_index: - if index.kind != common.ExprKind.SCALAR: - raise ValueError("Data indices must be scalars") - class ParamAccess(Expr): name: eve.Coerced[eve.SymbolRef] @@ -193,6 +189,10 @@ class While(common.While[Stmt, Expr], Stmt): pass +class For(common.For[Stmt], Stmt): + pass + + # --- Control Flow --- class HorizontalBlock(common.LocNode, eve.SymbolTableTrait): body: List[Stmt] diff --git a/src/gt4py/cartesian/gtc/numpy/npir_codegen.py b/src/gt4py/cartesian/gtc/numpy/npir_codegen.py index 8f9224e62b..015a396e25 100644 --- a/src/gt4py/cartesian/gtc/numpy/npir_codegen.py +++ b/src/gt4py/cartesian/gtc/numpy/npir_codegen.py @@ -109,6 +109,8 @@ def visit_TemporaryDecl( VarKOffset = as_fmt("lk + {k}") + ForIndex = as_fmt("{name}") + def visit_KMaskFieldAccess(self, node: npir.KMaskFieldAccess, **kwargs: Any) -> str: return "__k_mask[i:I, j:J, k:K]" @@ -139,7 +141,8 @@ def visit_FieldSlice(self, node: npir.FieldSlice, **kwargs: Any) -> Union[str, C ) args = _make_slice_access(offsets, kwargs["is_serial"], kwargs.get("horizontal_mask")) - data_index = self.visit(node.data_index, inside_slice=True, **kwargs) + kwargs["inside_slice"] = True + data_index = self.visit(node.data_index, **kwargs) access_slice = ", ".join(args + list(data_index)) @@ -266,6 +269,28 @@ def visit_While(self, node: npir.While, **kwargs: Any) -> str: body.extend(stmt.split("\n")) return self.While.render(cond=cond, body=body) + For = as_jinja( + textwrap.dedent( + """\ + for {{ index_name }} in range({{ iter_start }}, {{ iter_stop }}, {{ iter_step }}): + {% for stmt in body %}{{ stmt }} + {% endfor %} + """ + ) + ) + + def visit_For(self, node: npir.For, **kwargs: Any) -> str: + body = [] + for stmt in self.visit(node.body, **kwargs): + body.extend(stmt.split("\n")) + return self.For.render( + index_name=self.visit(node.index_name, **kwargs), + iter_start=self.visit(node.iter_start, **kwargs), + iter_stop=self.visit(node.iter_stop, **kwargs), + iter_step=self.visit(node.iter_step, **kwargs), + body=body, + ) + def visit_VerticalPass(self, node: npir.VerticalPass, **kwargs): is_serial = node.direction != common.LoopOrder.PARALLEL has_variable_k = bool(node.walk_values().if_isinstance(npir.VarKOffset).to_list()) diff --git a/src/gt4py/cartesian/gtc/numpy/oir_to_npir.py b/src/gt4py/cartesian/gtc/numpy/oir_to_npir.py index 39d94505f4..02e8d4344c 100644 --- a/src/gt4py/cartesian/gtc/numpy/oir_to_npir.py +++ b/src/gt4py/cartesian/gtc/numpy/oir_to_npir.py @@ -79,6 +79,9 @@ def visit_VariableKOffset( ) -> Tuple[int, int, eve.Node]: return 0, 0, npir.VarKOffset(k=self.visit(node.k, **kwargs)) + def visit_ForIndex(self, node: oir.ForIndex, **kwargs: Any) -> npir.ForIndex: + return npir.ForIndex(name=node.name, dtype=node.dtype) + def visit_AbsoluteKIndex(self, node: oir.AbsoluteKIndex, **kwargs: Any) -> None: raise NotImplementedError( "Absolute K indexation (e.g. `field.at(...)`) is an experimental feature and not yet implemented for the `numpy` backend." @@ -116,7 +119,9 @@ def visit_BinaryOp( self, node: oir.BinaryOp, **kwargs: Any ) -> Union[npir.VectorArithmetic, npir.VectorLogic]: args = dict( - op=node.op, left=self.visit(node.left, **kwargs), right=self.visit(node.right, **kwargs) + op=node.op, + left=self.visit(node.left, **kwargs), + right=self.visit(node.right, **kwargs), ) if isinstance(node.op, common.LogicalOperator): return npir.VectorLogic(**args) @@ -184,6 +189,15 @@ def visit_While( cond=cond_expr, body=utils.flatten(self.visit(node.body, mask=cond_expr, **kwargs)) ) + def visit_For(self, node: oir.For, **kwargs: Any) -> npir.For: + return npir.For( + index_name=node.index_name, + iter_start=node.iter_start, + iter_stop=node.iter_stop, + iter_step=node.iter_step, + body=utils.flatten(self.visit(node.body, **kwargs)), + ) + def visit_HorizontalRestriction( self, node: oir.HorizontalRestriction, *, extent: Extent, **kwargs: Any ) -> Any: @@ -210,11 +224,17 @@ def visit_HorizontalExecution( stmts = utils.flatten(self.visit(node.body, extent=extent, **kwargs)) return npir.HorizontalBlock( - body=stmts, extent=extent, declarations=self.visit(node.declarations, **kwargs) + body=stmts, + extent=extent, + declarations=self.visit(node.declarations, **kwargs), ) def visit_VerticalLoopSection( - self, node: oir.VerticalLoopSection, *, loop_order: common.LoopOrder, **kwargs: Any + self, + node: oir.VerticalLoopSection, + *, + loop_order: common.LoopOrder, + **kwargs: Any, ) -> npir.VerticalPass: return npir.VerticalPass( body=self.visit(node.horizontal_executions, **kwargs), diff --git a/src/gt4py/cartesian/gtc/oir.py b/src/gt4py/cartesian/gtc/oir.py index b73bd1b253..b170e34c9a 100644 --- a/src/gt4py/cartesian/gtc/oir.py +++ b/src/gt4py/cartesian/gtc/oir.py @@ -119,6 +119,14 @@ class While(common.While[Stmt, Expr], Stmt): pass +class ForIndex(common.ForIndex, Expr): + pass + + +class For(common.For[Stmt], Stmt): + pass + + class Decl(LocNode): name: eve.Coerced[eve.SymbolName] dtype: common.DataType diff --git a/src/gt4py/cartesian/gtscript.py b/src/gt4py/cartesian/gtscript.py index bcf3b91e8e..1273186429 100644 --- a/src/gt4py/cartesian/gtscript.py +++ b/src/gt4py/cartesian/gtscript.py @@ -96,6 +96,7 @@ "__externals__", "__INLINED", "compile_assert", + "range", *MATH_BUILTINS, *TYPE_HINT_AND_CAST_BUILTINS, } @@ -311,7 +312,13 @@ def stencil( # Setup build_info timings if build_info is not None: - time_keys = ("parse_time", "module_time", "codegen_time", "build_time", "load_time") + time_keys = ( + "parse_time", + "module_time", + "codegen_time", + "build_time", + "load_time", + ) build_info.update({time_key: 0.0 for time_key in time_keys}) build_options = gt_definitions.BuildOptions( diff --git a/src/gt4py/cartesian/utils/meta.py b/src/gt4py/cartesian/utils/meta.py index dff436d436..22c4ec9415 100644 --- a/src/gt4py/cartesian/utils/meta.py +++ b/src/gt4py/cartesian/utils/meta.py @@ -254,7 +254,7 @@ def apply(cls, ast_root, context, default=None): return result def __init__(self, context: dict): - self.context = copy.deepcopy(context) + self.context = {**context} def visit_Name(self, node): return self.context[node.id] diff --git a/src/gt4py/storage/allocators.py b/src/gt4py/storage/allocators.py index 394374c2a4..5cdc904e54 100644 --- a/src/gt4py/storage/allocators.py +++ b/src/gt4py/storage/allocators.py @@ -212,7 +212,7 @@ def allocate( # Compute the padding required in the contiguous dimension to get aligned blocks dims_layout = [layout_map.index(i) for i in range(len(shape))] # Convert shape size to same data type (note that `np.int16` can overflow) - padded_shape_lst = [np.int32(x) for x in shape] + padded_shape_lst = [np.int64(x) for x in shape] if ndim > 0: padded_shape_lst[dims_layout[-1]] = ( # type: ignore[call-overload] math.ceil(shape[dims_layout[-1]] / items_per_aligned_block)