gt4py.cartesian.gtc.dace.expansion package

Submodules

gt4py.cartesian.gtc.dace.expansion.daceir_builder module

class gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder[source]

Bases: gt4py.eve.visitors.NodeTranslator

class GlobalContext(library_node: "'StencilComputation'", arrays: 'Dict[str, dace.data.Data]')[source]

Bases: object

arrays: Dict[str, dace.data.Data]
get_dcir_decls(access_infos: Dict[gt4py.eve.concepts.SymbolRef, gt4py.cartesian.gtc.daceir.FieldAccessInfo], symbol_collector: gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.SymbolCollector) → List[gt4py.cartesian.gtc.daceir.FieldDecl][source]
library_node: StencilComputation
class IterationContext(grid_subset: 'dcir.GridSubset', parent: "Optional['DaCeIRBuilder.IterationContext']")[source]

Bases: object

grid_subset: gt4py.cartesian.gtc.daceir.GridSubset
classmethod init(*args, **kwargs)[source]
parent: Optional[gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.IterationContext]
pop() → gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.IterationContext[source]
push_axes_extents(axes_extents) → gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.IterationContext[source]
push_expansion_item(item: Union[gt4py.cartesian.gtc.dace.expansion_specification.Map, gt4py.cartesian.gtc.dace.expansion_specification.Loop]) → gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.IterationContext[source]
push_expansion_items(items: Iterable[Union[gt4py.cartesian.gtc.dace.expansion_specification.Map, gt4py.cartesian.gtc.dace.expansion_specification.Loop]]) → gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.IterationContext[source]
push_interval(axis: gt4py.cartesian.gtc.daceir.Axis, interval: Union[gt4py.cartesian.gtc.daceir.DomainInterval, gt4py.cartesian.gtc.oir.Interval]) → gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.IterationContext[source]
class SymbolCollector(symbol_decls: 'Dict[str, dcir.SymbolDecl]' = <factory>)[source]

Bases: object

add_symbol(name: str, dtype: gt4py.cartesian.gtc.common.DataType = DataType.INT32)[source]
remove_symbol(name: gt4py.eve.concepts.SymbolRef)[source]
symbol_decls: Dict[str, gt4py.cartesian.gtc.daceir.SymbolDecl]
to_dataflow(nodes, *, global_ctx: gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.GlobalContext, symbol_collector: gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.SymbolCollector)[source]
to_state(nodes, *, grid_subset: gt4py.cartesian.gtc.daceir.GridSubset)[source]
visit_AssignStmt(node: gt4py.cartesian.gtc.oir.AssignStmt, *, targets, **kwargs: Any) → gt4py.cartesian.gtc.daceir.AssignStmt[source]
visit_BinaryOp(node: gt4py.cartesian.gtc.oir.BinaryOp, **kwargs: Any) → gt4py.cartesian.gtc.daceir.BinaryOp[source]
visit_Cast(node: gt4py.cartesian.gtc.oir.Cast, **kwargs: Any) → gt4py.cartesian.gtc.daceir.Cast[source]
visit_FieldAccess(node: gt4py.cartesian.gtc.oir.FieldAccess, *, is_target: bool, targets: Set[gt4py.eve.concepts.SymbolRef], var_offset_fields: Set[gt4py.eve.concepts.SymbolRef], **kwargs: Any) → Union[gt4py.cartesian.gtc.daceir.IndexAccess, gt4py.cartesian.gtc.daceir.ScalarAccess][source]
visit_HorizontalExecution(node: gt4py.cartesian.gtc.oir.HorizontalExecution, *, global_ctx: gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.GlobalContext, iteration_ctx: gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.IterationContext, symbol_collector: gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.SymbolCollector, loop_order, k_interval, **kwargs)[source]
visit_HorizontalRestriction(node: gt4py.cartesian.gtc.oir.HorizontalRestriction, *, symbol_collector: gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.SymbolCollector, **kwargs: Any) → gt4py.cartesian.gtc.daceir.HorizontalRestriction[source]
visit_Literal(node: gt4py.cartesian.gtc.oir.Literal, **kwargs: Any) → gt4py.cartesian.gtc.daceir.Literal[source]
visit_LocalScalar(node: gt4py.cartesian.gtc.oir.LocalScalar, **kwargs: Any) → gt4py.cartesian.gtc.daceir.LocalScalarDecl[source]
visit_MaskStmt(node: gt4py.cartesian.gtc.oir.MaskStmt, **kwargs: Any) → gt4py.cartesian.gtc.daceir.MaskStmt[source]
visit_NativeFuncCall(node: gt4py.cartesian.gtc.oir.NativeFuncCall, **kwargs: Any) → gt4py.cartesian.gtc.daceir.NativeFuncCall[source]
visit_ScalarAccess(node: gt4py.cartesian.gtc.oir.ScalarAccess, *, global_ctx: gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.GlobalContext, symbol_collector: gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.SymbolCollector, **kwargs: Any) → gt4py.cartesian.gtc.daceir.ScalarAccess[source]
visit_TernaryOp(node: gt4py.cartesian.gtc.oir.TernaryOp, **kwargs: Any) → gt4py.cartesian.gtc.daceir.TernaryOp[source]
visit_UnaryOp(node: gt4py.cartesian.gtc.oir.UnaryOp, **kwargs: Any) → gt4py.cartesian.gtc.daceir.UnaryOp[source]
visit_VariableKOffset(node: gt4py.cartesian.gtc.oir.VariableKOffset, **kwargs)[source]
visit_VerticalLoop(node: gt4py.cartesian.gtc.oir.VerticalLoop, *, global_ctx: gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.GlobalContext, **kwargs)[source]
visit_VerticalLoopSection(node: gt4py.cartesian.gtc.oir.VerticalLoopSection, *, loop_order, iteration_ctx: gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.IterationContext, global_ctx: gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.GlobalContext, symbol_collector: gt4py.cartesian.gtc.dace.expansion.daceir_builder.DaCeIRBuilder.SymbolCollector, **kwargs)[source]
visit_While(node: gt4py.cartesian.gtc.oir.While, **kwargs: Any) → gt4py.cartesian.gtc.daceir.While[source]

gt4py.cartesian.gtc.dace.expansion.expansion module

class gt4py.cartesian.gtc.dace.expansion.expansion.StencilComputationExpansion(*args, **kwargs)[source]

Bases: dace.transformation.transformation.ExpandTransformation

static _fix_context(nsdfg, node: StencilComputation, parent_state: dace.sdfg.state.SDFGState, daceir: gt4py.cartesian.gtc.daceir.NestedSDFG)[source]

Apply changes to StencilComputation and the SDFG it is embedded in to satisfy post-expansion constraints.

  • change connector names to match inner array name (before expansion prefixed to satisfy uniqueness)

  • change in- and out-edges’ subsets so that they have the same shape as the corresponding array inside

  • determine the domain size based on edges to StencilComputation

environments: List = []
static expansion(node: StencilComputation, parent_state: dace.sdfg.state.SDFGState, parent_sdfg: dace.sdfg.sdfg.SDFG) → dace.sdfg.nodes.NestedSDFG[source]

Expand the coarse SDFG in parent_sdfg to a NestedSDFG with all the states.

gt4py.cartesian.gtc.dace.expansion.sdfg_builder module

class gt4py.cartesian.gtc.dace.expansion.sdfg_builder.StencilComputationSDFGBuilder[source]

Bases: gt4py.eve.traits.VisitorWithSymbolTableTrait

class NodeContext(input_node_and_conns: Dict[Optional[str], Tuple[dace.sdfg.nodes.Node, Optional[str]]], output_node_and_conns: Dict[Optional[str], Tuple[dace.sdfg.nodes.Node, Optional[str]]])[source]

Bases: object

input_node_and_conns: Dict[Optional[str], Tuple[dace.sdfg.nodes.Node, Optional[str]]]
output_node_and_conns: Dict[Optional[str], Tuple[dace.sdfg.nodes.Node, Optional[str]]]
class SDFGContext(sdfg: dace.sdfg.sdfg.SDFG, state: dace.sdfg.state.SDFGState, state_stack: List[dace.sdfg.state.SDFGState] = <factory>)[source]

Bases: object

add_loop(index_range: gt4py.cartesian.gtc.daceir.Range)[source]
add_state()[source]
pop_loop()[source]
sdfg: dace.sdfg.sdfg.SDFG
state: dace.sdfg.state.SDFGState
state_stack: List[dace.sdfg.state.SDFGState]
visit_ComputationState(node: gt4py.cartesian.gtc.daceir.ComputationState, *, sdfg_ctx: gt4py.cartesian.gtc.dace.expansion.sdfg_builder.StencilComputationSDFGBuilder.SDFGContext, **kwargs) → None[source]
visit_DomainLoop(node: gt4py.cartesian.gtc.daceir.DomainLoop, *, sdfg_ctx: gt4py.cartesian.gtc.dace.expansion.sdfg_builder.StencilComputationSDFGBuilder.SDFGContext, **kwargs) → None[source]
visit_DomainMap(node: gt4py.cartesian.gtc.daceir.DomainMap, *, node_ctx: gt4py.cartesian.gtc.dace.expansion.sdfg_builder.StencilComputationSDFGBuilder.NodeContext, sdfg_ctx: gt4py.cartesian.gtc.dace.expansion.sdfg_builder.StencilComputationSDFGBuilder.SDFGContext, **kwargs) → None[source]
visit_FieldDecl(node: gt4py.cartesian.gtc.daceir.FieldDecl, *, sdfg_ctx: gt4py.cartesian.gtc.dace.expansion.sdfg_builder.StencilComputationSDFGBuilder.SDFGContext, non_transients: Set[gt4py.eve.concepts.SymbolRef], **kwargs) → None[source]
visit_Memlet(node: gt4py.cartesian.gtc.daceir.Memlet, *, scope_node: gt4py.cartesian.gtc.daceir.ComputationNode, sdfg_ctx: gt4py.cartesian.gtc.dace.expansion.sdfg_builder.StencilComputationSDFGBuilder.SDFGContext, node_ctx: gt4py.cartesian.gtc.dace.expansion.sdfg_builder.StencilComputationSDFGBuilder.NodeContext, connector_prefix='', symtable: ChainMap[gt4py.eve.concepts.SymbolRef, gt4py.cartesian.gtc.daceir.Decl]) → None[source]
visit_NestedSDFG(node: gt4py.cartesian.gtc.daceir.NestedSDFG, *, sdfg_ctx: Optional[gt4py.cartesian.gtc.dace.expansion.sdfg_builder.StencilComputationSDFGBuilder.SDFGContext] = None, node_ctx: Optional[gt4py.cartesian.gtc.dace.expansion.sdfg_builder.StencilComputationSDFGBuilder.NodeContext] = None, symtable: ChainMap[gt4py.eve.concepts.SymbolRef, Any], **kwargs) → dace.sdfg.nodes.NestedSDFG[source]
visit_Range(node: gt4py.cartesian.gtc.daceir.Range, **kwargs) → Dict[str, str][source]
visit_SymbolDecl(node: gt4py.cartesian.gtc.daceir.SymbolDecl, *, sdfg_ctx: gt4py.cartesian.gtc.dace.expansion.sdfg_builder.StencilComputationSDFGBuilder.SDFGContext, **kwargs) → None[source]
visit_Tasklet(node: gt4py.cartesian.gtc.daceir.Tasklet, *, sdfg_ctx: gt4py.cartesian.gtc.dace.expansion.sdfg_builder.StencilComputationSDFGBuilder.SDFGContext, node_ctx: gt4py.cartesian.gtc.dace.expansion.sdfg_builder.StencilComputationSDFGBuilder.NodeContext, symtable: ChainMap[gt4py.eve.concepts.SymbolRef, gt4py.cartesian.gtc.daceir.Decl], **kwargs) → None[source]

gt4py.cartesian.gtc.dace.expansion.tasklet_codegen module

class gt4py.cartesian.gtc.dace.expansion.tasklet_codegen.TaskletCodegen[source]

Bases: gt4py.eve.codegen.TemplatedGenerator, gt4py.eve.traits.VisitorWithSymbolTableTrait

Arg = <gt4py.eve.codegen.FormatTemplate object>
BinaryOp = <gt4py.eve.codegen.FormatTemplate object>
Cast = <gt4py.eve.codegen.FormatTemplate object>
LocalScalarDecl = <gt4py.eve.codegen.FormatTemplate object>
Param = <gt4py.eve.codegen.FormatTemplate object>
ScalarAccess = <gt4py.eve.codegen.FormatTemplate object>
TernaryOp = <gt4py.eve.codegen.FormatTemplate object>
UnaryOp = <gt4py.eve.codegen.FormatTemplate object>
classmethod apply_codegen(node: gt4py.cartesian.gtc.daceir.Tasklet, **kwargs: Any) → str[source]
visit_AssignStmt(node: gt4py.cartesian.gtc.daceir.AssignStmt, **kwargs)[source]
visit_BuiltInLiteral(builtin: gt4py.cartesian.gtc.common.BuiltInLiteral, **kwargs: Any) → str[source]
visit_CartesianOffset(node: gt4py.cartesian.gtc.common.CartesianOffset, **kwargs)[source]
visit_DataType(dtype: gt4py.cartesian.gtc.common.DataType, **kwargs: Any) → str[source]
visit_HorizontalMask(node: gt4py.cartesian.gtc.common.HorizontalMask, **kwargs)[source]
visit_HorizontalRestriction(node: gt4py.cartesian.gtc.daceir.HorizontalRestriction, **kwargs)[source]
visit_IndexAccess(node: gt4py.cartesian.gtc.daceir.IndexAccess, *, is_target, sdfg_ctx, symtable: ChainMap[gt4py.eve.concepts.SymbolRef, gt4py.cartesian.gtc.daceir.Decl], **kwargs)[source]
visit_Literal(literal: gt4py.cartesian.gtc.daceir.Literal, *, in_idx=False, **kwargs)[source]
visit_MaskStmt(node: gt4py.cartesian.gtc.daceir.MaskStmt, **kwargs)[source]
visit_NativeFuncCall(call: gt4py.cartesian.gtc.common.NativeFuncCall, **kwargs: Any) → str[source]
visit_NativeFunction(func: gt4py.cartesian.gtc.common.NativeFunction, **kwargs: Any) → str[source]
visit_Tasklet(node: gt4py.cartesian.gtc.daceir.Tasklet, **kwargs)[source]
visit_UnaryOperator(op: gt4py.cartesian.gtc.common.UnaryOperator, **kwargs: Any) → str[source]
visit_VariableKOffset(node: gt4py.cartesian.gtc.common.CartesianOffset, **kwargs)[source]
visit_While(node: gt4py.cartesian.gtc.daceir.While, **kwargs)[source]

gt4py.cartesian.gtc.dace.expansion.utils module

class gt4py.cartesian.gtc.dace.expansion.utils.HorizontalExecutionSplitter[source]

Bases: gt4py.eve.visitors.NodeTranslator

static is_horizontal_execution_splittable(he: gt4py.cartesian.gtc.oir.HorizontalExecution)[source]
visit_HorizontalExecution(node: gt4py.cartesian.gtc.oir.HorizontalExecution, *, extents, library_node)[source]
visit_VerticalLoopSection(node: gt4py.cartesian.gtc.oir.VerticalLoopSection, **kwargs)[source]
class gt4py.cartesian.gtc.dace.expansion.utils.HorizontalIntervalRemover[source]

Bases: gt4py.eve.visitors.NodeTranslator

visit_HorizontalInterval(node: gt4py.cartesian.gtc.common.HorizontalInterval)[source]
visit_HorizontalMask(node: gt4py.cartesian.gtc.common.HorizontalMask, *, axis: gt4py.cartesian.gtc.daceir.Axis)[source]
class gt4py.cartesian.gtc.dace.expansion.utils.HorizontalMaskRemover[source]

Bases: gt4py.eve.visitors.NodeTranslator

visit_MaskStmt(node: gt4py.cartesian.gtc.oir.MaskStmt)[source]
visit_Tasklet(node: gt4py.cartesian.gtc.daceir.Tasklet)[source]
gt4py.cartesian.gtc.dace.expansion.utils.get_dace_debuginfo(node: gt4py.cartesian.gtc.common.LocNode)[source]
gt4py.cartesian.gtc.dace.expansion.utils.mask_includes_inner_domain(mask: gt4py.cartesian.gtc.common.HorizontalMask)[source]
gt4py.cartesian.gtc.dace.expansion.utils.remove_horizontal_region(node, axis)[source]
gt4py.cartesian.gtc.dace.expansion.utils.split_horizontal_executions_regions(node: StencilComputation)[source]

Module contents