-
Notifications
You must be signed in to change notification settings - Fork 54
Open
Labels
gt4py.cartesianIssues concerning the current version with support only for cartesian grids.Issues concerning the current version with support only for cartesian grids.module: frontendFrontend subpackageFrontend subpackagetriage: bugSomething isn't workingSomething isn't working
Description
A slight regression will occur in the next proposed version of the frontend.
The GTIR will be checking that absolute K index are not used for self-assignment. Variable K offset got included as well (check ParAssignment in gtir.py). When computation is in FORWARD/BACKWARD this check is too restrictive, as we have allowed the user to assign to K. gtir doesn't have a concept of computation at this stage so we can:
1/ Trust our frontend checks and remove the guard
2/ Move computation-schedule down enough for GTIR to make a more comprehensive guard
Below is reproducer.
from ndsl.dsl.gt4py import computation, interval, FORWARD
from ndsl.boilerplate import (
get_factories_single_tile,
get_factories_single_tile_orchestrated,
)
from ndsl.constants import X_DIM, Y_DIM, Z_DIM
from ndsl.dsl.typing import FloatField, Float, Int, IntField
from ndsl import StencilFactory, orchestrate
import numpy as np
domain = (2, 2, 5)
ORCHESTRATE = False
if ORCHESTRATE:
stcil_fctry, qty_factry = get_factories_single_tile_orchestrated(
domain[0], domain[1], domain[2], 0, on_cpu=True
)
else:
stcil_fctry, qty_factry = get_factories_single_tile(
domain[0], domain[1], domain[2], 0, backend="dace:cpu"
)
def the_stencil(in_field: FloatField, out_field: FloatField, lev: IntField):
with computation(FORWARD), interval(...):
out_field[0, 0, lev] += 2
class Code:
def __init__(self, stencil_factory: StencilFactory):
orchestrate(obj=self, config=stencil_factory.config.dace_config)
self._the_stencil = stcil_fctry.from_dims_halo(
func=the_stencil,
compute_dims=[X_DIM, Y_DIM, Z_DIM],
)
def __call__(self, in_field: FloatField, out_field: FloatField, lev: IntField):
self._the_stencil(in_field, out_field, lev)
if __name__ == "__main__":
in_arr = qty_factry.zeros(
dims=[X_DIM, Y_DIM, Z_DIM],
units="inputs",
)
sz = domain[0] * domain[1] * domain[2]
in_arr.view[:] = np.arange(sz, dtype=Float).reshape(domain)
out_arr = qty_factry.zeros([X_DIM, Y_DIM, Z_DIM], units="outputs")
lev_arr = qty_factry.zeros([X_DIM, Y_DIM, Z_DIM], units="lev", dtype=Int)
code = Code(stcil_fctry)
code(in_arr, out_arr, lev_arr)
print("Done 🚀")Metadata
Metadata
Assignees
Labels
gt4py.cartesianIssues concerning the current version with support only for cartesian grids.Issues concerning the current version with support only for cartesian grids.module: frontendFrontend subpackageFrontend subpackagetriage: bugSomething isn't workingSomething isn't working