FlorianDeconinck
commented
1 week ago Test file to be dropped in ndsl (when the reference to pyFV3 is removed)
from ndsl.stencils.corners import fill_corners_dgrid_defn
from ndsl.boilerplate import get_factories_single_tile_orchestrated_cpu
from ndsl.constants import X_DIM, Y_DIM, Z_DIM, X_INTERFACE_DIM, Y_INTERFACE_DIM
from ndsl.dsl.typing import Float, FloatField
from ndsl import orchestrate, StencilFactory, DaceConfig
from gt4py.cartesian.gtscript import computation, PARALLEL, interval
class OrchestratedCorner:
def __init__(self, stencil_factory: StencilFactory) -> None:
orchestrate(
obj=self,
config=stencil_factory.config.dace_config
or DaceConfig(communicator=None, backend=stencil_factory.backend),
)
origin, domain = stencil_factory.grid_indexing.get_origin_domain(
dims=[X_DIM, Y_DIM, Z_DIM]
)
axes_offsets = stencil_factory.grid_indexing.axis_offsets(origin, domain)
self.corner_stencil = stencil_factory.from_origin_domain(
fill_corners_dgrid_defn,
externals=axes_offsets,
origin=origin,
domain=domain,
)
def __call__(self, x, y):
self.corner_stencil(x, x, y, y, 1.0)
def test_empty_corners():
stencil_factory, quantity_factory = get_factories_single_tile_orchestrated_cpu(
12, 12, 5, 0
)
# Make the
stencil_factory.grid_indexing.south_edge = False
stencil_factory.grid_indexing.north_edge = False
stencil_factory.grid_indexing.west_edge = False
stencil_factory.grid_indexing.east_edge = False
stencil_factory.grid_indexing.axis_offsets
x = quantity_factory.empty(dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], units="n/a")
y = quantity_factory.empty(dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], units="n/a")
orch_corner = OrchestratedCorner(stencil_factory)
orch_corner(x, y)
def unusued_parameter_stencil(
field: FloatField, # type: ignore
result: FloatField, # type: ignore
weight: Float, # type: ignore
):
with computation(PARALLEL), interval(...):
result = field[1, 0, 0] + field[0, 1, 0] + field[-1, 0, 0] + field[0, -1, 0]
class OrchestratedUnusedParameter:
def __init__(self, stencil_factory: StencilFactory):
orchestrate(
obj=self,
config=stencil_factory.config.dace_config
or DaceConfig(communicator=None, backend=stencil_factory.backend),
)
origin, domain = stencil_factory.grid_indexing.get_origin_domain(
dims=[X_DIM, Y_DIM, Z_DIM]
)
self.unused_stencil = stencil_factory.from_origin_domain(
unusued_parameter_stencil,
origin=origin,
domain=domain,
)
def __call__(self, x, y):
self.unused_stencil(x, y, 1.0)
def test_unused_parameters():
stencil_factory, quantity_factory = get_factories_single_tile_orchestrated_cpu(
12, 12, 5, 2
)
x = quantity_factory.empty(dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], units="n/a")
y = quantity_factory.empty(dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], units="n/a")
orch_unused = OrchestratedUnusedParameter(stencil_factory)
orch_unused(x, y)
def unusued_field_stencil(
field: FloatField, # type: ignore
other_field: FloatField, # type: ignore
result: FloatField, # type: ignore
):
with computation(PARALLEL), interval(...):
result = field[1, 0, 0] + field[0, 1, 0] + field[-1, 0, 0] + field[0, -1, 0]
class OrchestratedunusedField:
def __init__(self, stencil_factory: StencilFactory):
orchestrate(
obj=self,
config=stencil_factory.config.dace_config
or DaceConfig(communicator=None, backend=stencil_factory.backend),
)
origin, domain = stencil_factory.grid_indexing.get_origin_domain(
dims=[X_DIM, Y_DIM, Z_DIM]
)
self.unused_stencil = stencil_factory.from_origin_domain(
unusued_field_stencil,
origin=origin,
domain=domain,
)
def __call__(self, x, unused_field, y):
self.unused_stencil(
x,
unused_field,
y,
)
def test_unused_field():
stencil_factory, quantity_factory = get_factories_single_tile_orchestrated_cpu(
12, 12, 5, 2
)
x = quantity_factory.empty(dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], units="n/a")
x_unused = quantity_factory.empty(dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], units="n/a")
y = quantity_factory.empty(dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], units="n/a")
orch_unused = OrchestratedunusedField(stencil_factory)
orch_unused(x, x_unused, y)
if __name__ == "__main__":
test_unused_parameters()
test_empty_corners()
test_unused_field()
Description
There's a series of bug related to signature at GT level and declared symbols/scalar/fields at DaCe wrapper SDFG. All of those bugs live in the bridge between gt4py & dace.
Those can be classified in 3 groups:
Most of those behavior are linked to the
prune_unused_argumentpass of GT4Py which is called at the very beginning of GTIR. While this is clearly not the design (passes should be pushed down to OIR or backend IR) this was done to deal with some of those issues. Plain removing the prune pass (which could be done considering it gives little to no performance improvement) does not lead to fixing.In the comments below we will put down 3 examples that showcase the issues (either plain
ndslor relying onpyfv3) and some patches that fixes some bugs but creates other.To Reproduce See comment.