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.
Following this PR we have introduced 2D temporaries as described in the ADR.
gt:Xbackends 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
would become during OIR -> GTCPP
And the
tmp_2Dwould 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 GTCPPCode would inject it right before the
pybind11bindings.A pass would need to be done to make sure the temporaries don't re-use the name.