[feature] Add 2D temporary support for NDSL's BoolFieldIJ

ndsl/dsl/stencil.py

@@ -25,6 +25,7 @@
 from ndsl.dsl.dace.orchestration import SDFGConvertible
 from ndsl.dsl.stencil_config import CompilationConfig, RunMode, StencilConfig
 from ndsl.dsl.typing import (
+    BoolFieldIJ,
     Float,
     FloatFieldIJ,
     FloatFieldIJ32,
@@ -329,6 +330,7 @@ def __init__(
             "IntFieldIJ": IntFieldIJ,
             "IntFieldIJ32": IntFieldIJ32,
             "IntFieldIJ64": IntFieldIJ64,
+            "BoolFieldIJ": BoolFieldIJ,
         }
 
         # Keep compilation at __init__ if we are not orchestrated.

tests/dsl/test_stencil.py

@@ -13,6 +13,7 @@
 )
 from ndsl.dsl.gt4py import FORWARD, PARALLEL, Field, computation, interval
 from ndsl.dsl.typing import (
+    BoolFieldIJ,
     FloatField,
     FloatFieldIJ,
     FloatFieldIJ32,
@@ -130,16 +131,18 @@ def two_dim_temporaries_stencil(q_out: FloatField) -> None:
         tmp_3d_iij: IntFieldIJ = 4
         tmp_3d_iij32: IntFieldIJ32 = 5
         tmp_3d_iij64: IntFieldIJ64 = 6
+        mask: BoolFieldIJ = q_out >= 0
 
     with computation(PARALLEL), interval(...):
-        q_out = (
-            tmp_2d_fij
-            + tmp_2d_fij32
-            + tmp_3d_fij64
-            + tmp_3d_iij
-            + tmp_3d_iij32
-            + tmp_3d_iij64
-        )
+        if mask:
+            q_out = (
+                tmp_2d_fij
+                + tmp_2d_fij32
+                + tmp_3d_fij64
+                + tmp_3d_iij
+                + tmp_3d_iij32
+                + tmp_3d_iij64
+            )
 
 
 def test_stencil_2D_temporaries() -> None:

tests/dsl/test_stencil_wrapper.py

@@ -16,6 +16,7 @@
 from ndsl.dsl.gt4py_utils import make_storage_from_shape
 from ndsl.dsl.stencil import _convert_quantities_to_storage
 from ndsl.dsl.typing import (
+    BoolFieldIJ,
     Float,
     FloatField,
     FloatFieldIJ,
@@ -285,6 +286,7 @@ def test_frozen_stencil_kwargs_passed_to_init(
             "IntFieldIJ": IntFieldIJ,
             "IntFieldIJ32": IntFieldIJ32,
             "IntFieldIJ64": IntFieldIJ64,
+            "BoolFieldIJ": BoolFieldIJ,
         },
     )