Skip to content

cartesian: Extend 2D temporaries to gt:X backend #2322

@FlorianDeconinck

Description

@FlorianDeconinck

Following this PR we have introduced 2D temporaries as described in the ADR.

gt:X backends were excluded from this PR because GridTools pre-suppose that all computation of temporaries are done on the 3D grid at least - and won't allocate below that dimensionality. This is a fair limitation since GridTools doesn't provide the guardrails against race condition that GT4Py does.

We can circumvent this by creating a pool of buffers passed as arguments. In details a stencils written like this

@stencil(backend=..., dtype={"My2DType": Field[IJ, np.float64]})
def the_stencil(A: Field[IJK, np.float64], B: Field[IJK, np.float64]):
    with computation(FORWARD), inteval(...):
        tmp_2D: Field[IJ, np.float64] = 0
        ...

would become during OIR -> GTCPP

@stencil(backend=..., dtype={"My2DType": Field[IJ, np.float64]})
def the_stencil(A: Field[IJK, np.float64], B: Field[IJK, np.float64], tmp_2D: Field[IJ, np.float64]):
    with computation(FORWARD), inteval(...):
        tmp_2D = 0
        ...

And the tmp_2D would be sourced from a pooled memory of 2D temporaries properly scoped to the IJ required - with all information generated and stored on the stencil compiled during codegen of GTCPP

Code would inject it right before the pybind11 bindings.

A pass would need to be done to make sure the temporaries don't re-use the name.

Metadata

Metadata

Assignees

No one assigned

    Labels

    gt4py.cartesianIssues concerning the current version with support only for cartesian grids.triage: enhancementNew feature or request

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions