GridTools / gt4py

Python library for generating high-performance implementations of stencil kernels for weather and climate modeling from a domain-specific language (DSL).
https://GridTools.github.io/gt4py
Other
106 stars 48 forks source link

[cartesian] Numpy backend lacks non-literal accessor for data dimensions #1578

Open FlorianDeconinck opened 1 month ago

FlorianDeconinck commented 1 month ago

The numpy backend has a limitation on literal accessor not shared by other backends. Consider the following:


import gt4py.cartesian.gtscript as gtscript
from gt4py.cartesian.gtscript import (
    computation,
    interval,
    PARALLEL,
)
import numpy as np

THE_NUMBER_5 = 5
FloatField_NModes = gtscript.Field[gtscript.IJK, (np.float32, (THE_NUMBER_5))]
FloatField = gtscript.Field[gtscript.IJK, np.float32]

@gtscript.stencil(backend="gt:cpu_kfirst")
def ddim_read_in_loop(
        an_array_of_ones: FloatField_NModes,
        out_field: FloatField):
    with computation(PARALLEL), interval(...):
        n_iter = 0
        agg = 0
        while n_iter < THE_NUMBER_5:
            agg = agg + an_array_of_ones[0, 0, 0][n_iter]
            n_iter = n_iter + 1
        out_field = agg

the_ones_field = np.ones((2, 2, 2, THE_NUMBER_5), dtype=np.float32)
out_field = np.zeros((2, 2, 2), dtype=np.float32)

ddim_read_in_loop(the_ones_field, out_field)
print(out_field)

We expect n_iter to be some kind of local scalar variable. But numpy is trying to use as much inner parallelism and masking as possible leading to this code giving a

  File "/home/mad/work/ndsl/external/gt4py/src/gt4py/cartesian/gtc/numpy/npir.py", line 141, in data_indices_are_scalar
    raise ValueError("Data indices must be scalars")
ValueError: Data indices must be scalars

Meanwhile, running this code under a gt:Xor dace:X backend gives the expected results (with n_iter a local scalar`).

PS: this was unearth looking at the Aer Activation physics code for NASA's GEOS model

egparedes commented 1 month ago

Thanks for reporting @FlorianDeconinck . Are you planning to also work in a fix or is it not high priority for you at the moment?

FlorianDeconinck commented 1 month ago

Unclear. We need a fix for sure but we have a mounting amount of physics pattern difficult to write so other feature requests might take precedence. I'd say we are a few weeks away from deciding where we put our efforts