Repository: GridTools/gt4py
Pull request #1692 (2024-10-17 14:18:37)

refactor[cartesian]: Warn if GT4Py can't find DaCe
<!--
Delete this comment and add a proper description of the changes contained in this PR. The text here will be used in the commit message since the approved PRs are always squash-merged. The preferred format is:

- PR Title: <type>[<scope>]: <one-line-summary>

    <type>:
        - build: Changes that affect the build system or external dependencies
        - ci: Changes to our CI configuration files and scripts
        - docs: Documentation only changes
        - feat: A new feature
        - fix: A bug fix
        - perf: A code change that improves performance
        - refactor: A code change that neither fixes a bug nor adds a feature
        - style: Changes that do not affect the meaning of the code
        - test: Adding missing tests or correcting existing tests

    <scope>: cartesian | eve | next | storage
    # ONLY if changes are limited to a specific subsytem

- PR Description:

    Description of the main changes with links to appropriate issues/documents/references/...
-->

## Description

DaCe backends are optional backends in the cartesian version of GT4Py. Currently, we silently drop support for DaCe backends if DaCe can't be imported. This can can happen because DaCe isn't available or for a couple other reasons (e.g. a circular import in the import path). With this PR we thus add a warning message allowing developers to easily figure out that/why DaCe backends are disabled.

## Requirements

- [ ] All fixes and/or new features come with corresponding tests.
- [ ] Important design decisions have been documented in the approriate ADR inside the [docs/development/ADRs/](docs/development/ADRs/Index.md) folder.

If this PR contains code authored by new contributors please make sure:

- [ ] The PR contains an updated version of the `AUTHORS.md` file adding the names of all the new contributors.


romanc: @egparedes would you have time for a review?

/cc @FlorianDeconinck  as discussed
romanc: > Thanks for the PR. It looks much better now.

Thanks for the quick review. Could you also kick-off the CSCS CI job? Thanks! :pray: 
egparedes: cscs-ci run

Patch:
diff --git a/src/gt4py/cartesian/backend/__init__.py b/src/gt4py/cartesian/backend/__init__.py
index 7a6f8772..e58c7a01 100644
--- a/src/gt4py/cartesian/backend/__init__.py
+++ b/src/gt4py/cartesian/backend/__init__.py
@@ -6,6 +6,8 @@
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
 
+from warnings import warn
+
 from .base import (
     REGISTRY,
     Backend,
@@ -16,13 +18,6 @@ from .base import (
     from_name,
     register,
 )
-
-
-try:
-    from .dace_backend import DaceCPUBackend, DaceGPUBackend
-except ImportError:
-    pass
-
 from .cuda_backend import CudaBackend
 from .gtcpp_backend import GTCpuIfirstBackend, GTCpuKfirstBackend, GTGpuBackend
 from .module_generator import BaseModuleGenerator
@@ -47,5 +42,12 @@ __all__ = [
 ]
 
 
-if "DaceCPUBackend" in globals():
+try:
+    from .dace_backend import DaceCPUBackend, DaceGPUBackend
+
     __all__ += ["DaceCPUBackend", "DaceGPUBackend"]
+except ImportError:
+    warn(
+        "GT4Py was unable to load DaCe. DaCe backends (`dace:cpu` and `dace:gpu`) will not be available.",
+        stacklevel=2,
+    )
diff --git a/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py b/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py
index 8fb1451e..a8ae1cc0 100644
--- a/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py
+++ b/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py
@@ -209,8 +209,8 @@ def translate_as_fieldop(
     The dataflow can be as simple as a single tasklet, or implement a local computation
     as a composition of tasklets and even include a map to range on local dimensions (e.g.
     neighbors and map builtins).
-    The stencil dataflow is instantiated inside a map scope, which apply the stencil over
-    the field domain.
+    The stencil dataflow is instantiated inside a map scope, which applies the stencil
+    over the field domain.
     """
     assert isinstance(node, gtir.FunCall)
     assert cpm.is_call_to(node.fun, "as_fieldop")
@@ -219,11 +219,24 @@ def translate_as_fieldop(
     fun_node = node.fun
     assert len(fun_node.args) == 2
     stencil_expr, domain_expr = fun_node.args
-    assert isinstance(stencil_expr, gtir.Lambda)
-    assert isinstance(domain_expr, gtir.FunCall)
+
+    if isinstance(stencil_expr, gtir.Lambda):
+        # Default case, handled below: the argument expression is a lambda function
+        # representing the stencil operation to be computed over the field domain.
+        pass
+    elif cpm.is_ref_to(stencil_expr, "deref"):
+        # Special usage of 'deref' as argument to fieldop expression, to pass a scalar
+        # value to 'as_fieldop' function. It results in broadcasting the scalar value
+        # over the field domain.
+        return translate_broadcast_scalar(node, sdfg, state, sdfg_builder, reduce_identity)
+    else:
+        raise NotImplementedError(
+            f"Expression type '{type(stencil_expr)}' not supported as argument to 'as_fieldop' node."
+        )
 
     # parse the domain of the field operator
     domain = extract_domain(domain_expr)
+    domain_indices = sbs.Indices([dace_gtir_utils.get_map_variable(dim) for dim, _, _ in domain])
 
     # The reduction identity value is used in place of skip values when building
     # a list of neighbor values in the unstructured domain.
@@ -273,16 +286,15 @@ def translate_as_fieldop(
     input_edges, output = taskgen.visit(stencil_expr, args=stencil_args)
     output_desc = output.result.node.desc(sdfg)
 
-    domain_index = sbs.Indices([dace_gtir_utils.get_map_variable(dim) for dim, _, _ in domain])
     if isinstance(node.type.dtype, itir_ts.ListType):
         assert isinstance(output_desc, dace.data.Array)
         assert set(output_desc.offset) == {0}
         # additional local dimension for neighbors
         # TODO(phimuell): Investigate if we should swap the two.
-        output_subset = sbs.Range.from_indices(domain_index) + sbs.Range.from_array(output_desc)
+        output_subset = sbs.Range.from_indices(domain_indices) + sbs.Range.from_array(output_desc)
     else:
         assert isinstance(output_desc, dace.data.Scalar)
-        output_subset = sbs.Range.from_indices(domain_index)
+        output_subset = sbs.Range.from_indices(domain_indices)
 
     # create map range corresponding to the field operator domain
     me, mx = sdfg_builder.add_map(
@@ -307,6 +319,62 @@ def translate_as_fieldop(
     return result_field
 
 
+def translate_broadcast_scalar(
+    node: gtir.Node,
+    sdfg: dace.SDFG,
+    state: dace.SDFGState,
+    sdfg_builder: gtir_sdfg.SDFGBuilder,
+    reduce_identity: Optional[gtir_dataflow.SymbolExpr],
+) -> FieldopResult:
+    """
+    Generates the dataflow subgraph for the 'as_fieldop' builtin function for the
+    special case where the argument to 'as_fieldop' is a 'deref' scalar expression,
+    rather than a lambda function. This case corresponds to broadcasting the scalar
+    value over the field domain. Therefore, it is lowered to a mapped tasklet that
+    just writes the scalar value out to all elements of the result field.
+    """
+    assert isinstance(node, gtir.FunCall)
+    assert cpm.is_call_to(node.fun, "as_fieldop")
+    assert isinstance(node.type, ts.FieldType)
+
+    fun_node = node.fun
+    assert len(fun_node.args) == 2
+    stencil_expr, domain_expr = fun_node.args
+    assert cpm.is_ref_to(stencil_expr, "deref")
+
+    domain = extract_domain(domain_expr)
+    domain_indices = sbs.Indices([dace_gtir_utils.get_map_variable(dim) for dim, _, _ in domain])
+
+    assert len(node.args) == 1
+    assert isinstance(node.args[0].type, ts.ScalarType)
+    scalar_expr = _parse_fieldop_arg(
+        node.args[0], sdfg, state, sdfg_builder, domain, reduce_identity=None
+    )
+    assert isinstance(scalar_expr, gtir_dataflow.MemletExpr)
+    assert scalar_expr.subset == sbs.Indices.from_string("0")
+    result = gtir_dataflow.DataflowOutputEdge(
+        state, gtir_dataflow.DataExpr(scalar_expr.node, node.args[0].type)
+    )
+    result_field = _create_temporary_field(sdfg, state, domain, node.type, dataflow_output=result)
+
+    sdfg_builder.add_mapped_tasklet(
+        "broadcast",
+        state,
+        map_ranges={
+            dace_gtir_utils.get_map_variable(dim): f"{lower_bound}:{upper_bound}"
+            for dim, lower_bound, upper_bound in domain
+        },
+        inputs={"__inp": dace.Memlet(data=scalar_expr.node.data, subset="0")},
+        code="__val = __inp",
+        outputs={"__val": dace.Memlet(data=result_field.data_node.data, subset=domain_indices)},
+        input_nodes={scalar_expr.node.data: scalar_expr.node},
+        output_nodes={result_field.data_node.data: result_field.data_node},
+        external_edges=True,
+    )
+
+    return result_field
+
+
 def translate_if(
     node: gtir.Node,
     sdfg: dace.SDFG,
@@ -638,6 +706,7 @@ if TYPE_CHECKING:
     # Use type-checking to assert that all translator functions implement the `PrimitiveTranslator` protocol
     __primitive_translators: list[PrimitiveTranslator] = [
         translate_as_fieldop,
+        translate_broadcast_scalar,
         translate_if,
         translate_literal,
         translate_make_tuple,
diff --git a/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_sdfg.py b/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_sdfg.py
index d79d8873..3697609d 100644
--- a/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_sdfg.py
+++ b/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_sdfg.py
@@ -32,7 +32,6 @@ from gt4py.next.program_processors.runners.dace_common import utility as dace_ut
 from gt4py.next.program_processors.runners.dace_fieldview import (
     gtir_builtin_translators,
     gtir_dataflow,
-    transformations as gtx_transformations,
     utility as dace_gtir_utils,
 )
 from gt4py.next.type_system import type_specifications as ts, type_translation as tt
@@ -665,7 +664,6 @@ def build_sdfg_from_gtir(
 
     The lowering to SDFG requires that the program node is type-annotated, therefore this function
     runs type ineference as first step.
-    As a final step, it runs the `simplify` pass to ensure that the SDFG is in the DaCe canonical form.
 
     Arguments:
         ir: The GTIR program node to be lowered to SDFG
@@ -677,10 +675,8 @@ def build_sdfg_from_gtir(
 
     ir = gtir_type_inference.infer(ir, offset_provider=offset_provider)
     ir = ir_prune_casts.PruneCasts().visit(ir)
-    ir = dace_gtir_utils.patch_gtir(ir)
     sdfg_genenerator = GTIRToSDFG(offset_provider)
     sdfg = sdfg_genenerator.visit(ir)
     assert isinstance(sdfg, dace.SDFG)
 
-    gtx_transformations.gt_simplify(sdfg)
     return sdfg
diff --git a/src/gt4py/next/program_processors/runners/dace_fieldview/utility.py b/src/gt4py/next/program_processors/runners/dace_fieldview/utility.py
index 855dc9c9..b5c447a1 100644
--- a/src/gt4py/next/program_processors/runners/dace_fieldview/utility.py
+++ b/src/gt4py/next/program_processors/runners/dace_fieldview/utility.py
@@ -11,10 +11,7 @@ from __future__ import annotations
 import itertools
 from typing import Any
 
-from gt4py import eve
 from gt4py.next import common as gtx_common
-from gt4py.next.iterator import ir as gtir
-from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm
 from gt4py.next.type_system import type_specifications as ts
 
 
@@ -65,41 +62,3 @@ def get_tuple_type(data: tuple[Any, ...]) -> ts.TupleType:
     return ts.TupleType(
         types=[get_tuple_type(d) if isinstance(d, tuple) else d.data_type for d in data]
     )
-
-
-def patch_gtir(ir: gtir.Program) -> gtir.Program:
-    """
-    Make the IR compliant with the requirements of lowering to SDFG.
-
-    Applies canonicalization of as_fieldop expressions as well as some temporary workarounds.
-    This allows to lower the IR to SDFG for some special cases.
-    """
-
-    class PatchGTIR(eve.PreserveLocationVisitor, eve.NodeTranslator):
-        def visit_FunCall(self, node: gtir.FunCall) -> gtir.Node:
-            if cpm.is_applied_as_fieldop(node):
-                assert isinstance(node.fun, gtir.FunCall)
-                assert isinstance(node.type, ts.FieldType)
-
-                # Handle the case of fieldop without domain. This case should never happen, but domain
-                # inference currently produces this kind of nodes for unreferenced tuple fields.
-                # TODO(tehrengruber): remove this workaround once domain ineference supports this case
-                if len(node.fun.args) == 1:
-                    return gtir.Literal(value="0", type=node.type.dtype)
-
-                assert len(node.fun.args) == 2
-                stencil = node.fun.args[0]
-
-                # Canonicalize as_fieldop: always expect a lambda expression.
-                # Here we replace the call to deref with a lambda expression and empty arguments list.
-                if cpm.is_ref_to(stencil, "deref"):
-                    node.fun.args[0] = gtir.Lambda(
-                        expr=gtir.FunCall(fun=stencil, args=node.args), params=[]
-                    )
-                    node.args = []
-
-            node.args = self.visit(node.args)
-            node.fun = self.visit(node.fun)
-            return node
-
-    return PatchGTIR().visit(ir)


Test patch:
diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py
index 98e15dac..728b4b02 100644
--- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py
+++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py
@@ -40,10 +40,11 @@ dace_backend = pytest.importorskip("gt4py.next.program_processors.runners.dace_f
 
 
 N = 10
-IFTYPE = ts.FieldType(dims=[IDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64))
-CFTYPE = ts.FieldType(dims=[Cell], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64))
-EFTYPE = ts.FieldType(dims=[Edge], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64))
-VFTYPE = ts.FieldType(dims=[Vertex], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64))
+FLOAT_TYPE = ts.ScalarType(kind=ts.ScalarKind.FLOAT64)
+IFTYPE = ts.FieldType(dims=[IDim], dtype=FLOAT_TYPE)
+CFTYPE = ts.FieldType(dims=[Cell], dtype=FLOAT_TYPE)
+EFTYPE = ts.FieldType(dims=[Edge], dtype=FLOAT_TYPE)
+VFTYPE = ts.FieldType(dims=[Vertex], dtype=FLOAT_TYPE)
 V2E_FTYPE = ts.FieldType(dims=[Vertex, V2EDim], dtype=EFTYPE.dtype)
 CARTESIAN_OFFSETS = {
     "IDim": IDim,
@@ -315,6 +316,62 @@ def test_gtir_tuple_expr():
     assert np.allclose(c, a * 2 + b)
 
 
+def test_gtir_tuple_broadcast_scalar():
+    domain = im.domain(gtx_common.GridType.CARTESIAN, ranges={IDim: (0, "size")})
+    testee = gtir.Program(
+        id="gtir_tuple_broadcast_scalar",
+        function_definitions=[],
+        params=[
+            gtir.Sym(
+                id="x",
+                type=ts.TupleType(types=[FLOAT_TYPE, ts.TupleType(types=[FLOAT_TYPE, FLOAT_TYPE])]),
+            ),
+            gtir.Sym(id="y", type=IFTYPE),
+            gtir.Sym(id="size", type=SIZE_TYPE),
+        ],
+        declarations=[],
+        body=[
+            gtir.SetAt(
+                expr=im.as_fieldop("deref", domain)(
+                    im.plus(
+                        im.tuple_get(0, "x"),
+                        im.plus(
+                            im.multiplies_(
+                                im.tuple_get(
+                                    0,
+                                    im.tuple_get(1, "x"),
+                                ),
+                                2.0,
+                            ),
+                            im.multiplies_(
+                                im.tuple_get(
+                                    1,
+                                    im.tuple_get(1, "x"),
+                                ),
+                                3.0,
+                            ),
+                        ),
+                    )
+                ),
+                domain=domain,
+                target=gtir.SymRef(id="y"),
+            )
+        ],
+    )
+
+    a = np.random.rand()
+    b = np.random.rand()
+    c = np.random.rand()
+    d = np.empty(N, dtype=type(a))
+
+    sdfg = dace_backend.build_sdfg_from_gtir(testee, CARTESIAN_OFFSETS)
+
+    x_fields = (a, b, c)
+
+    sdfg(*x_fields, d, **FSYMBOLS)
+    assert np.allclose(d, a + 2 * b + 3 * c)
+
+
 def test_gtir_tuple_return():
     domain = im.domain(gtx_common.GridType.CARTESIAN, ranges={IDim: (0, "size")})
     testee = gtir.Program(
@@ -954,8 +1011,8 @@ def test_gtir_connectivity_shift():
         im.op_as_fieldop("plus", edge_domain)("e2v_offset", 0),
     )
 
-    CE_FTYPE = ts.FieldType(dims=[Cell, Edge], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64))
-    EV_FTYPE = ts.FieldType(dims=[Edge, Vertex], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64))
+    CE_FTYPE = ts.FieldType(dims=[Cell, Edge], dtype=FLOAT_TYPE)
+    EV_FTYPE = ts.FieldType(dims=[Edge, Vertex], dtype=FLOAT_TYPE)
     CELL_OFFSET_FTYPE = ts.FieldType(dims=[Cell], dtype=SIZE_TYPE)
     EDGE_OFFSET_FTYPE = ts.FieldType(dims=[Edge], dtype=SIZE_TYPE)