class NativeFuncCall(Expr):
func = attribute(of=NativeFunction)
args = attribute(of=ListOf[Expr])
data_type = attribute(of=DataType)
loc = attribute(of=Location, optional=True)
class ParameterAccessor(Accessor):
symbol = attribute(of=str)
class StencilID:
qualified_name = attribute(of=str)
version = attribute(of=str)
def __iter__(self):
return iter([self.qualified_name, self.version])
class While(Statement):
condition = attribute(of=Expr)
body = attribute(of=BlockStmt)
loc = attribute(of=Location, optional=True)
class ComputationBlock(Node):
interval = attribute(of=AxisInterval)
iteration_order = attribute(of=IterationOrder)
body = attribute(of=BlockStmt)
loc = attribute(of=Location, optional=True)
class TernaryOpExpr(CompositeExpr):
condition = attribute(of=Expr)
then_expr = attribute(of=Expr)
else_expr = attribute(of=Expr)
loc = attribute(of=Location, optional=True)
class Assign(Statement):
target = attribute(of=Ref)
value = attribute(of=Expr)
loc = attribute(of=Location, optional=True)
class StageGroup(IIRNode):
stages = attribute(of=ListOf[Stage])
class MultiStage(IIRNode):
name = attribute(of=str)
iteration_order = attribute(of=IterationOrder)
# caches = attribute(of=ListOf[Cache])
groups = attribute(of=ListOf[StageGroup])
class ApplyBlock(Node):
interval = attribute(of=AxisInterval)
local_symbols = attribute(of=DictOf[str, VarDecl])
body = attribute(of=BlockStmt)
class Stage(IIRNode):
name = attribute(of=str)
accessors = attribute(of=ListOf[Accessor])
apply_blocks = attribute(of=ListOf[ApplyBlock])
compute_extent = attribute(of=Extent)
class FieldRef(Ref):
name = attribute(of=str)
offset = attribute(of=DictOf[str, int])
loc = attribute(of=Location, optional=True)
class ShapedExpr(gt_ir.Expr):
axes = attribute(of=SetOf[str])
expr = attribute(of=gt_ir.Expr)
class StencilImplementation(IIRNode):
name = attribute(of=str)
api_signature = attribute(of=ListOf[ArgumentInfo])
domain = attribute(of=Domain)
fields = attribute(of=DictOf[str, FieldDecl]) # All fields, including temporaries
parameters = attribute(of=DictOf[str, VarDecl])
multi_stages = attribute(of=ListOf[MultiStage])
fields_extents = attribute(of=DictOf[str, Extent])
unreferenced = attribute(of=ListOf[str], factory=list)
axis_splitters_var = attribute(of=str, optional=True)
externals = attribute(of=DictOf[str, Any], optional=True)
sources = attribute(of=DictOf[str, str], optional=True)
docstring = attribute(of=str)
@property
def has_effect(self):
"""
Determine whether the stencil modifies any of its arguments.
Note that the only guarantee of this function is that the stencil has no effect if it returns ``false``. It
might however return true in cases where the optimization passes were not able to deduce this.
"""
return self.multi_stages and not all(
arg_field in self.unreferenced for arg_field in self.arg_fields
)
@property
def arg_fields(self):
result = [f.name for f in self.fields.values() if f.is_api]
return result
@property
def temporary_fields(self):
result = [f.name for f in self.fields.values() if not f.is_api]
return result
class UnaryOpExpr(CompositeExpr):
op = attribute(of=UnaryOperator)
arg = attribute(of=Expr)
loc = attribute(of=Location, optional=True)
class StencilImplementation(IIRNode):
name = attribute(of=str)
api_signature = attribute(of=ListOf[ArgumentInfo])
domain = attribute(of=Domain)
fields = attribute(
of=DictOf[str, FieldDecl]) # All fields, including temporaries
parameters = attribute(of=DictOf[str, VarDecl])
multi_stages = attribute(of=ListOf[MultiStage])
fields_extents = attribute(of=DictOf[str, Extent])
unreferenced = attribute(of=ListOf[str], factory=list)
axis_splitters_var = attribute(of=str, optional=True)
externals = attribute(of=DictOf[str, Any], optional=True)
sources = attribute(of=DictOf[str, str], optional=True)
@property
def arg_fields(self):
result = [f.name for f in self.fields.values() if f.is_api]
return result
@property
def temporary_fields(self):
result = [f.name for f in self.fields.values() if not f.is_api]
return result
class BinOpExpr(CompositeExpr):
op = attribute(of=BinaryOperator)
lhs = attribute(of=Expr)
rhs = attribute(of=Expr)
loc = attribute(of=Location, optional=True)
class Axis(Node):
name = attribute(of=str)
class BlockStmt(Statement):
stmts = attribute(of=ListOf[Statement])
loc = attribute(of=Location, optional=True)
class ScalarLiteral(Literal):
value = attribute(of=Any) # Potentially an array of numeric structs
data_type = attribute(of=DataType)
loc = attribute(of=Location, optional=True)
class If(Statement):
condition = attribute(of=Expr)
main_body = attribute(of=BlockStmt)
else_body = attribute(of=BlockStmt, optional=True)
loc = attribute(of=Location, optional=True)
class BuiltinLiteral(Literal):
value = attribute(of=Builtin)
loc = attribute(of=Location, optional=True)
class AxisBound(Node):
level = attribute(of=UnionOf[LevelMarker, VarRef])
offset = attribute(of=int, default=0)
loc = attribute(of=Location, optional=True)
class VarRef(Ref):
name = attribute(of=str)
index = attribute(of=int, optional=True)
loc = attribute(of=Location, optional=True)
class ArgumentInfo(Node):
name = attribute(of=str)
is_keyword = attribute(of=bool, default=False)
default = attribute(of=Any, default=Empty)
class Cast(Expr):
data_type = attribute(of=DataType)
expr = attribute(of=Expr)
loc = attribute(of=Location, optional=True)
class FieldAccessor(Accessor):
symbol = attribute(of=str)
intent = attribute(of=AccessKind)
extent = attribute(of=Extent, default=Extent.zeros())
class AugAssign(Statement):
target = attribute(of=Ref)
value = attribute(of=Expr)
op = attribute(of=BinaryOperator)
loc = attribute(of=Location, optional=True)
