def copy_assign():
yield ParAssignStmt(
loc=SourceLocation(line=3, column=2, source="copy_gtir"),
left=FieldAccess.centered(
name="a", loc=SourceLocation(line=3, column=1, source="copy_gtir")
),
right=FieldAccess.centered(
name="b", loc=SourceLocation(line=3, column=3, source="copy_gtir")
),
)
def copy_v_loop(copy_assign, interval):
yield VerticalLoop(
loc=SourceLocation(line=2, column=1, source="copy_gtir"),
loop_order=LoopOrder.FORWARD,
interval=interval,
body=[copy_assign],
temporaries=[],
)
def copy_computation(copy_v_loop):
yield Stencil(
name="copy_gtir",
loc=SourceLocation(line=1, column=1, source="copy_gtir"),
params=[
FieldDecl(name="a", dtype=DataType.FLOAT32),
FieldDecl(name="b", dtype=DataType.FLOAT32),
],
vertical_loops=[copy_v_loop],
)
def copy_computation(copy_v_loop):
yield Stencil(
name="copy_gtir",
loc=SourceLocation(line=1, column=1, source="copy_gtir"),
params=[
FieldDecl(
name="foo",
dtype=DataType.FLOAT32,
dimensions=(True, True, True),
),
FieldDecl(
name="bar",
dtype=DataType.FLOAT32,
dimensions=(True, True, True),
),
],
vertical_loops=[copy_v_loop],
)
def interval(copy_assign):
yield Interval(
loc=SourceLocation(line=2, column=11, source="copy_gtir"),
start=AxisBound(level=LevelMarker.START, offset=0),
end=AxisBound(level=LevelMarker.END, offset=0),
)
def location_to_source_location(loc: Optional["Location"]) -> Optional[SourceLocation]:
if loc is None or loc.line <= 0 or loc.column <= 0:
return None
return SourceLocation(loc.line, loc.column, loc.scope)
