FlorianDeconinck
commented
1 month ago Example showing the state of things in SDFG pass and a trivial application of the stree
from gt4py.cartesian.gtscript import (
computation,
interval,
PARALLEL,
)
from ndsl.boilerplate import get_factories_single_tile_orchestrated_cpu
from ndsl.constants import X_DIM, Y_DIM, Z_DIM
from ndsl.dsl.typing import FloatField
from ndsl import StencilFactory, orchestrate
import numpy as np
import dace
import dace.sdfg.analysis.schedule_tree.treenodes as dace_stree
from dace.transformation.dataflow import MapFusion
from dace.transformation.interstate import StateFusion
domain = (3, 3, 4)
stcil_fctry, ijk_qty_fctry = get_factories_single_tile_orchestrated_cpu(
domain[0], domain[1], domain[2], 0
)
def double_map(in_field: FloatField, out_field: FloatField):
with computation(PARALLEL), interval(...):
out_field = in_field
with computation(PARALLEL), interval(...):
out_field = in_field * 3
class DaCeGT4Py_Bridge:
def __init__(self, stencil_factory: StencilFactory):
orchestrate(obj=self, config=stencil_factory.config.dace_config)
self.stencil = stcil_fctry.from_dims_halo(
func=double_map,
compute_dims=[X_DIM, Y_DIM, Z_DIM],
)
def __call__(self, in_field: FloatField, out_field: FloatField):
self.stencil(in_field, out_field)
if __name__ == "__main__":
I = np.arange(domain[0] * domain[1] * domain[2], dtype=np.float64).reshape(domain)
O = np.zeros(domain)
# # Trigger NDSL orchestration pipeline & grab cached SDFG
bridge = DaCeGT4Py_Bridge(stcil_fctry)
bridge(I, O)
sdfg: dace.sdfg.SDFG = bridge.__sdfg__(I, O).csdfg.sdfg
sdfg.save("orig.sdfg")
stree = sdfg.as_schedule_tree()
#### Classic SDFG transform
# Strategy: fuse states out of the way, then merge maps
# State fusion
# No fusion occurs because of a potential write race since
# it would be putting both maps under one input
r = sdfg.apply_transformations_repeated(
StateFusion,
print_report=True,
validate_all=True,
)
print(f"Fused {r} states")
sdfg.save("state_fusion.sdfg")
# No fusion occurs because maps are in different states
# (previous failure to merge)
r = sdfg.apply_transformations_repeated(
MapFusion,
print_report=True,
validate_all=True,
)
print(f"Fused {r} maps")
sdfg.save("map_fusion.sdfg")
#### Schedule Tree transform
# We demonstrate here a very basic usage of Schedule Tree to merge
# maps that have the same range, which should be the first pass
# we write
print("Before merge")
print(stree.as_string())
first_map: dace_stree.MapScope = stree.children[2]
second_map: dace_stree.MapScope = stree.children[3]
if first_map.node.range == second_map.node.range:
first_map.children.extend(second_map.children)
first_map.parent.children.remove(second_map)
print("After merge")
print(stree.as_string())
romanc
Cache miss occurs in the code, even the orchestrated, because we have IJK loops around all stencils instead of having K-loop shared by the stencils. Even after applying "map fusion", it's nowhere near enough to match the K-block/loop of Fortran.
For instance the entire code of
D_SWis within a single K-loop in Fortran.Using Tal's new
streeor directSDFGmanipulation, we shall implement an aggressive transform ondaceto move the K-loop as far as possible, including by introducingifto block out indices that should be computed.Another code, called
partial expansionwhich work onSDFGalone and develop by Linus, exists inGT4Py.cartesianas a branch.streebranch: https://github.com/spcl/dace/tree/stree-to-sdfgFvTp2d(code) or other)C_SWand/orD_SW