Skip to content

[cartesian] Variable K offset self-assignment is over guarded in FORWARD/BACKWARD mode #2220

@FlorianDeconinck

Description

@FlorianDeconinck

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

No one assigned

    Labels

    gt4py.cartesianIssues concerning the current version with support only for cartesian grids.module: frontendFrontend subpackagetriage: bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions