GridTools Stage Extents

[eddie-c-davis]

Description

This PR modifies the GridTools backends to use the make_stage_with_extent method to define stages rather than make_stage as in the Dawn backends. For example, the qx_edge_west stencil in the FV3 dycore,

def qx_edge_west2(qin: sd, dxa: sd, qx: sd):
    with computation(PARALLEL), interval(...):
        g_in = dxa / dxa[-1, 0, 0]
        qx0 = qx
        qx = (
            3.0 * (g_in * qin[-1, 0, 0] + qin) - (g_in * qx0[-1, 0, 0] + qx0[1, 0, 0])
        ) / (2.0 + 2.0 * g_in)

the qx0 temporary is introduced to prevent a nonzero horizontal extents static assertion error. The generated code with make_stage compiles but produces incorrect data. The proposed change generates the following code which produces correct results.

computation_t gt_computation = gt::make_computation<backend_t>(
        make_grid(domain),
        gt::make_multistage(gt::execute::parallel(),
            gt::make_stage_with_extent<stage__10_func, gt::extent<0, 0, 0, 0>>(
                p_dxa(), p_g_in()
            ),
            gt::make_stage_with_extent<stage__13_func, gt::extent<-1, 1, 0, 0>>(
                p_qx(), p_qx0()
            ),
            gt::make_stage_with_extent<stage__16_func, gt::extent<0, 0, 0, 0>>(
                p_g_in(), p_qin(), p_qx0(), p_qx()
            )
        )
    );

This change is required because the GridTools dependence analysis cannot always compute the necessary extents when temporaries are reused. The is documented in issue #920 in the GridTools repository.

[havogt]

Finally, I found the root cause of the issue, see #143. Did this PR fix the issue for all GT backends (especially GPU). There should be a data race even with this PR. The stage_with_extent probably "fixed" the dependency analysis, but not the race condition.