-
Notifications
You must be signed in to change notification settings - Fork 54
Description
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.