[NASA:Update] Update GT4Py and DaCe + submoduling of DaCe (#21)

.gitmodules

@@ -1,6 +1,6 @@
 [submodule "external/gt4py"]
 	path = external/gt4py
 	url = https://github.com/gridtools/gt4py.git
-[submodule "buildenv"]
-	path = buildenv
-	url = https://github.com/ai2cm/buildenv.git
+[submodule "external/dace"]
+	path = external/dace
+	url = https://github.com/spcl/dace.git

constraints.txt

@@ -2,7 +2,7 @@
 # This file is autogenerated by pip-compile with Python 3.8
 # by the following command:
 #
-#    pip-compile --output-file=constraints.txt driver/setup.py dsl/setup.py external/gt4py/setup.cfg fv3core/setup.py physics/setup.py requirements_dev.txt requirements_docs.txt requirements_lint.txt stencils/setup.py util/requirements.txt util/setup.py
+#    pip-compile --output-file=constraints.txt driver/setup.py dsl/setup.py fv3core/setup.py physics/setup.py requirements_dev.txt requirements_docs.txt requirements_lint.txt stencils/setup.py util/requirements.txt util/setup.py
 #
 aenum==3.1.11
     # via dace
@@ -21,36 +21,31 @@ asttokens==2.0.5
     #   devtools
     #   stack-data
 astunparse==1.6.3
-    # via dace
+    # via
+    #   dace
+    #   gt4py
 async-timeout==3.0.1
     # via aiohttp
 attrs==22.1.0
     # via
     #   aiohttp
     #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
     #   jsonschema
     #   pytest
 babel==2.9.1
     # via sphinx
 backcall==0.2.0
     # via ipython
-backports.entry-points-selectable==1.1.1
+backports-entry-points-selectable==1.1.1
     # via virtualenv
 black==22.3.0
-    # via
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
+    # via gt4py
 boltons==21.0.0
-    # via
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
+    # via gt4py
 bump2version==1.0.1
     # via -r util/requirements.txt
 cached-property==1.5.2
-    # via
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
+    # via gt4py
 cachetools==4.2.2
     # via google-auth
 certifi==2021.5.30
@@ -74,20 +69,19 @@ click==8.0.1
     #   black
     #   flask
     #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
 cloudpickle==2.0.0
     # via dask
+cmake==3.26.4
+    # via gt4py
 commonmark==0.9.1
     # via recommonmark
 coverage==5.5
     # via
     #   -r util/requirements.txt
     #   pytest-cov
-cytoolz==0.11.2
-    # via
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
-dace==0.14.0
+cytoolz==0.12.1
+    # via gt4py
+dace==0.14.4
     # via
     #   -r requirements_dev.txt
     #   pace-dsl
@@ -109,13 +103,9 @@ decorator==5.0.9
     #   gcsfs
     #   ipython
 deepdiff==6.2.1
-    # via
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
+    # via gt4py
 devtools==0.8.0
-    # via
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
+    # via gt4py
 dill==0.3.5.1
     # via dace
 distlib==0.3.2
@@ -155,9 +145,7 @@ flake8==3.8.4
 flask==2.1.2
     # via dace
 frozendict==2.3.4
-    # via
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
+    # via gt4py
 fsspec==2021.7.0
     # via
     #   dask
@@ -196,10 +184,8 @@ googleapis-common-protos==1.53.0
     # via google-api-core
 gprof2dot==2021.2.21
     # via pytest-profiling
-gridtools-cpp==2.2.2
-    # via
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
+gridtools-cpp==2.3.0
+    # via gt4py
 h5netcdf==0.11.0
     # via -r util/requirements.txt
 h5py==2.10.0
@@ -217,7 +203,9 @@ imagesize==1.2.0
 importlib-metadata==4.11.3
     # via flask
 importlib-resources==5.10.0
-    # via jsonschema
+    # via
+    #   gt4py
+    #   jsonschema
 iniconfig==1.1.1
     # via pytest
 ipykernel==6.16.2
@@ -232,7 +220,6 @@ jinja2==3.0.1
     # via
     #   flask
     #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
     #   sphinx
 jsonschema==4.16.0
     # via nbformat
@@ -244,12 +231,12 @@ jupyter-core==4.11.2
     # via
     #   jupyter-client
     #   nbformat
+lark==1.1.5
+    # via gt4py
 locket==0.2.1
     # via partd
 mako==1.1.6
-    # via
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
+    # via gt4py
 markupsafe==2.0.1
     # via
     #   jinja2
@@ -260,6 +247,11 @@ matplotlib-inline==0.1.6
     #   ipython
 mccabe==0.6.1
     # via flake8
+mpi4py==3.1.4
+    # via
+    #   -r requirements_dev.txt
+    #   pace-driver
+    #   pace-driver (driver/setup.py)
 mpmath==1.2.1
     # via sympy
 multidict==5.1.0
@@ -272,7 +264,6 @@ mypy-extensions==0.4.3
     # via
     #   black
     #   mypy
-    #   typing-inspect
 nbclient==0.6.8
     # via nbmake
 nbformat==5.7.0
@@ -292,10 +283,9 @@ netcdf4==1.5.7
     #   pace-driver
     #   pace-driver (driver/setup.py)
 networkx==2.6.3
-    # via
-    #   dace
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
+    # via dace
+ninja==1.11.1
+    # via gt4py
 nodeenv==1.6.0
     # via pre-commit
 numcodecs==0.7.2
@@ -309,7 +299,6 @@ numpy==1.21.2
     #   cftime
     #   dace
     #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
     #   h5py
     #   netcdf4
     #   numcodecs
@@ -322,7 +311,6 @@ numpy==1.21.2
     #   pace-util
     #   pace-util (util/setup.py)
     #   pandas
-    #   scipy
     #   xarray
     #   zarr
 oauthlib==3.1.1
@@ -333,7 +321,6 @@ packaging==21.0
     # via
     #   dask
     #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
     #   ipykernel
     #   pytest
     #   sphinx
@@ -387,9 +374,7 @@ pyasn1==0.4.8
 pyasn1-modules==0.2.8
     # via google-auth
 pybind11==2.8.1
-    # via
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
+    # via gt4py
 pycodestyle==2.6.0
     # via flake8
 pycparser==2.20
@@ -465,10 +450,6 @@ requests-oauthlib==1.3.0
     # via google-auth-oauthlib
 rsa==4.7.2
     # via google-auth
-scipy==1.7.1
-    # via
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
 six==1.16.0
     # via
     #   asttokens
@@ -513,10 +494,8 @@ stack-data==0.5.1
     # via ipython
 sympy==1.9
     # via dace
-tabulate==0.8.9
-    # via
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
+tabulate==0.9.0
+    # via gt4py
 toml==0.10.2
     # via
     #   pre-commit
@@ -530,8 +509,6 @@ toolz==0.11.1
     #   -r util/requirements.txt
     #   cytoolz
     #   dask
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
     #   partd
 tornado==6.2
     # via
@@ -555,15 +532,9 @@ typing-extensions==4.3.0
     #   aiohttp
     #   black
     #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
     #   mypy
     #   pace-util
     #   pace-util (util/setup.py)
-    #   typing-inspect
-typing-inspect==0.7.1
-    # via
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
 urllib3==1.26.6
     # via requests
 virtualenv==20.7.2
@@ -590,9 +561,7 @@ xarray==0.19.0
     #   pace-physics
     #   pace-physics (physics/setup.py)
 xxhash==2.0.2
-    # via
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
+    # via gt4py
 yarl==1.6.3
     # via aiohttp
 zarr==2.9.2
@@ -601,7 +570,9 @@ zarr==2.9.2
     #   pace-driver
     #   pace-driver (driver/setup.py)
 zipp==3.8.0
-    # via importlib-metadata
+    # via
+    #   importlib-metadata
+    #   importlib-resources
 
 # The following packages are considered to be unsafe in a requirements file:
 # setuptools

dsl/pace/dsl/dace/build.py

@@ -135,7 +135,7 @@ def set_distributed_caches(config: "DaceConfig"):
         verb = "reading"
 
     gt_config.cache_settings["dir_name"] = get_cache_directory(config.code_path)
-    pace.util.pace_log.critical(
+    pace.util.pace_log.info(
         f"[{orchestration_mode}] Rank {config.my_rank} "
         f"{verb} cache {gt_config.cache_settings['dir_name']}"
     )

dsl/pace/dsl/dace/orchestration.py

@@ -10,6 +10,7 @@
 from dace.frontend.python.parser import DaceProgram
 from dace.transformation.auto.auto_optimize import make_transients_persistent
 from dace.transformation.helpers import get_parent_map
+from dace.transformation.passes.simplify import SimplifyPass
 
 from pace.dsl.dace.build import get_sdfg_path, write_build_info
 from pace.dsl.dace.dace_config import (
@@ -65,17 +66,13 @@ def _download_results_from_dace(
             gt4py_results = [
                 gt4py.storage.from_array(
                     r,
-                    default_origin=(0, 0, 0),
                     backend=config.get_backend(),
-                    managed_memory=True,
                 )
                 for r in dace_result
             ]
         else:
             gt4py_results = [
-                gt4py.storage.from_array(
-                    r, default_origin=(0, 0, 0), backend=config.get_backend()
-                )
+                gt4py.storage.from_array(r, backend=config.get_backend())
                 for r in dace_result
             ]
     return gt4py_results
@@ -111,6 +108,17 @@ def _to_gpu(sdfg: dace.SDFG):
         sd.openmp_sections = False
 
 
+def _simplify(sdfg: dace.SDFG, validate=True, verbose=False):
+    """Override of sdfg.simplify to skip failing transformation
+    per https://github.com/spcl/dace/issues/1328
+    """
+    return SimplifyPass(
+        validate=validate,
+        verbose=verbose,
+        skip=["ConstantPropagation"],
+    ).apply_pass(sdfg, {})
+
+
 def _build_sdfg(
     daceprog: DaceProgram, sdfg: dace.SDFG, config: DaceConfig, args, kwargs
 ):
@@ -144,7 +152,7 @@ def _build_sdfg(
                 del sdfg_kwargs[k]
 
         with DaCeProgress(config, "Simplify (1/2)"):
-            sdfg.simplify(validate=False, verbose=True)
+            _simplify(sdfg, validate=False, verbose=True)
 
         # Perform pre-expansion fine tuning
         with DaCeProgress(config, "Split regions"):
@@ -155,7 +163,7 @@ def _build_sdfg(
             sdfg.expand_library_nodes()
 
         with DaCeProgress(config, "Simplify (2/2)"):
-            sdfg.simplify(validate=False, verbose=True)
+            _simplify(sdfg, validate=False, verbose=True)
 
         # Move all memory that can be into a pool to lower memory pressure.
         # Change Persistent memory (sub-SDFG) into Scope and flag it.

dsl/pace/dsl/dace/sdfg_opt_passes.py

@@ -21,4 +21,4 @@ def splittable_region_expansion(sdfg: dace.SDFG, verbose: bool = False):
                     "K",
                 ]
                 if verbose:
-                    pace_log.info(f"Reordered schedule for {node.label}")
+                    pace_log.debug(f"Reordered schedule for {node.label}")