Iterator view prototype

.gitignore

@@ -107,6 +107,7 @@ venv.bak/
 /site
 
 # DaCe
+.dacecache
 *.sdfg
 .dace.conf
 

setup.cfg

@@ -183,7 +183,7 @@ lines_after_imports = 2
 default_section = THIRDPARTY
 sections = FUTURE,STDLIB,THIRDPARTY,FIRSTPARTY,LOCALFOLDER
 known_first_party = eve,gtc,gt4py,__externals__,__gtscript__
-known_third_party = attr,black,boltons,cached_property,click,dace,dawn4py,devtools,factory,hypothesis,jinja2,mako,networkx,numpy,packaging,pkg_resources,pybind11,pydantic,pytest,pytest_factoryboy,setuptools,tabulate,tests,typing_extensions,typing_inspect,xxhash
+known_third_party = atlas4py,attr,black,boltons,cached_property,click,dace,dawn4py,devtools,factory,hypothesis,jinja2,mako,networkx,numpy,packaging,pkg_resources,pybind11,pydantic,pytest,pytest_factoryboy,setuptools,tabulate,tests,typing_extensions,typing_inspect,xxhash
 
 #-- mypy --
 [mypy]

src/iterator/ARCHITECTURE.md

@@ -0,0 +1,67 @@
+# Architecture
+
+Implements the iterator view as described [here](https://github.com/GridTools/concepts/wiki/Iterator-View).
+
+## Iterator view program in Python
+
+A program for the iterator view consists of Python functions decorated with `@fundef` and an entry point, *fencil*, which is a Python function decorated with `@fendef`. The *fencil* must only contain calls to the `closure(...)` function.
+
+Legal functions much not have side-effects, however, e.g., for debugging purposes, side-effects can be used in embedded execution.
+
+There are 2 modes of execution: *embedded* (direct execution in Python) and *tracing* (trace function calls -> Eve IR representation -> code generation).
+The implementations of *embedded* and *tracing* are decoupled by registering themselves (dependency inversion) in the functions defined in `builtins.py` (contains dispatch functions for all builtins of the model) and `runtime.py` (contains dispatch mechanism for the `fendef` and `fundef` decorators and the `closure(...)` function).
+
+The builtins dispatcher is implemented in `dispatcher.py`. Implementations are registered with a key (`str`) (currently `tracing` and `embedded`). The active implementation is selected by pushing a key to the dispatcher stack.
+
+`fundef` returns a wrapper around the function, which dispatches `__call__` to a hook if a predicate is met (used for *tracing*). By default the original function is called (used in *embedded* mode).
+
+`fendef` return a wrapper that dispatches to a registered function. Should be simplified. *Embedded* registers itself as default, *tracing* registers itself such that it's used if the fencil is called with `backend` keyword argument.
+
+## Embedded execution
+
+Embedded execution is implemented in the file `embedded.py`.
+
+Sketch:
+- fields are np arrays with named axes; names are instances of `CartesianAxis`
+- in `closure()`, the stencil is executed for each point in the domain, the fields are wrapped in an iterator pointing to the current point of execution.
+- `shift()` is lazy (to allow lift implementation), offsets are accumulated in the iterator and only executed when `deref()` is called.
+- as described in the design, offsets are abstract; on fencil execution the `offset_provider` keyword argument needs to be specified, which is a dict of `str` to either `CartesianAxis` or `NeighborTableOffsetProvider`
+- if `column_axis` keyword argument is specified on fencil execution (or in the fencil decorator), all operations will be done column wise in the give axis; `column_axis` needs to be specified if `scan` is used
+
+## Tracing
+
+An iterator view program is traced (implemented in `tracing.py`) and represented in a tree structure defined by the nodes in (`ir.py`).
+
+Sketch:
+- Each builtin returns a `FunctionCall` node representing the builtin.
+- Foreach `fundef`, the signature of the wrapped function is extracted, then it is invoked with `Sym` nodes as arguments.
+- Expressions involving an `Expr` node (e.g. `Sym`) are converted to appropriate builtin calls, e.g. `4. + Sym(id='foo')` is converted to `FunCall(fun=SymRef(id='plus'), args=...)`
+- In appropriate places values are converted to nodes, see `make_node()`.
+- Finally the IR tree will be passed to `execute_program()` in `backend_executor.py` which will generator code for the program (and execute, if appropriate).
+
+## Backends
+
+See directory `backends/`.
+
+### Cpptoy
+
+Generates C++ code in the spirit of https://github.com/GridTools/gridtools/pull/1643. Incomplete, will be adapted to the full C++ prototype. (only code generation)
+
+### Lisp
+
+Incomplete. Example for the grammar used in the model design document. (not executable)
+
+### Embedded
+
+Generates from the IR an aquivalent Python iterator view program which is then executed in embedded mode (round trip).
+
+### Double roundtrip
+
+Generates the Python iterator view program, traces it again, generates again and executes. Ensures that the generated Python code can still be traced. While the original program might be different from the generated program (e.g. `+` will be converted to `plus()` builtin). The programs from the embedded and double roundtrip backends should be identical.
+
+## Adding a new builtin
+
+Currently there are 4 places where a new builtin needs to be added
+- `builtin.py`: for dispatching to an actual implementation
+- `embedded.py` and `tracing.py`: for the respective implementation
+- `ir.py`: we check for consistent use of symbols, therefore if a `FunCall` to the new builtin is used, it needs to be available in the symbol table.

src/iterator/README.md

@@ -0,0 +1 @@
+# README

src/iterator/__init__.py

@@ -0,0 +1,19 @@
+from typing import Optional, Union
+
+from . import builtins, runtime, tracing
+
+
+__all__ = ["builtins", "runtime", "tracing"]
+
+from packaging.version import LegacyVersion, Version, parse
+from pkg_resources import DistributionNotFound, get_distribution
+
+
+try:
+    __version__: str = get_distribution("gt4py").version
+except DistributionNotFound:
+    __version__ = "X.X.X.unknown"
+
+__versioninfo__: Optional[Union[LegacyVersion, Version]] = parse(__version__)
+
+del DistributionNotFound, LegacyVersion, Version, get_distribution, parse

src/iterator/atlas_utils.py

@@ -0,0 +1,33 @@
+# GT4Py New Semantic Model - GridTools Framework
+#
+# Copyright (c) 2014-2021, ETH Zurich All rights reserved.
+#
+# This file is part of the GT4Py project and the GridTools framework.  GT4Py
+# New Semantic Model is free software: you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the Free
+# Software Foundation, either version 3 of the License, or any later version.
+# See the LICENSE.txt file at the top-level directory of this distribution for
+# a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from atlas4py import IrregularConnectivity
+
+
+class AtlasTable:
+    def __init__(self, atlas_connectivity) -> None:
+        self.atlas_connectivity = atlas_connectivity
+
+    def __getitem__(self, indices):
+        primary_index = indices[0]
+        neigh_index = indices[1]
+        if isinstance(self.atlas_connectivity, IrregularConnectivity):
+            if neigh_index < self.atlas_connectivity.cols(primary_index):
+                return self.atlas_connectivity[primary_index, neigh_index]
+            else:
+                return None
+        else:
+            if neigh_index < 2:
+                return self.atlas_connectivity[primary_index, neigh_index]
+            else:
+                raise AssertionError()

src/iterator/backend_executor.py

@@ -0,0 +1,21 @@
+from devtools import debug
+
+from iterator.backends import backend
+from iterator.ir import Program
+
+
+def execute_program(prog: Program, *args, **kwargs):
+    assert "backend" in kwargs
+    assert len(prog.fencil_definitions) == 1
+
+    if "debug" in kwargs and kwargs["debug"]:
+        debug(prog)
+
+    if not len(args) == len(prog.fencil_definitions[0].params):
+        raise RuntimeError("Incorrect number of arguments")
+
+    if kwargs["backend"] in backend._BACKENDS:
+        b = backend.get_backend(kwargs["backend"])
+        b(prog, *args, **kwargs)
+    else:
+        raise RuntimeError(f"Backend {kwargs['backend']} is not registered.")

src/iterator/backends/__init__.py

@@ -0,0 +1 @@
+from . import cpptoy, double_roundtrip, embedded, lisp

src/iterator/backends/backend.py

@@ -0,0 +1,9 @@
+_BACKENDS = {}
+
+
+def register_backend(name, backend):
+    _BACKENDS[name] = backend
+
+
+def get_backend(name):
+    return _BACKENDS[name]

src/iterator/backends/cpptoy.py

@@ -0,0 +1,59 @@
+from typing import Any
+
+from eve import codegen
+from eve.codegen import FormatTemplate as as_fmt
+from eve.codegen import MakoTemplate as as_mako
+from iterator.backends import backend
+from iterator.ir import OffsetLiteral
+from iterator.transforms import apply_common_transforms
+
+
+class ToyCpp(codegen.TemplatedGenerator):
+    Sym = as_fmt("{id}")
+    SymRef = as_fmt("{id}")
+    IntLiteral = as_fmt("{value}")
+    FloatLiteral = as_fmt("{value}")
+    AxisLiteral = as_fmt("{value}")
+
+    def visit_OffsetLiteral(self, node: OffsetLiteral, **kwargs):
+        return node.value if isinstance(node.value, str) else f"{node.value}_c"
+
+    StringLiteral = as_fmt("{value}")
+    FunCall = as_fmt("{fun}({','.join(args)})")
+    Lambda = as_mako(
+        "[=](${','.join('auto ' + p for p in params)}){return ${expr};}"
+    )  # TODO capture
+    StencilClosure = as_mako(
+        "closure(${domain}, ${stencil}, out(${','.join(outputs)}), ${','.join(inputs)})"
+    )
+    FencilDefinition = as_mako(
+        """
+    auto ${id} = [](${','.join('auto&& ' + p for p in params)}){
+        fencil(${'\\n'.join(closures)});
+    };
+    """
+    )
+    FunctionDefinition = as_mako(
+        """
+    inline constexpr auto ${id} = [](${','.join('auto ' + p for p in params)}){
+        return ${expr};
+        };
+    """
+    )
+    Program = as_fmt("{''.join(function_definitions)} {''.join(fencil_definitions)}")
+
+    @classmethod
+    def apply(cls, root, **kwargs: Any) -> str:
+        transformed = apply_common_transforms(
+            root,
+            use_tmps=kwargs.get("use_tmps", False),
+            offset_provider=kwargs.get("offset_provider", None),
+        )
+        generated_code = super().apply(transformed, **kwargs)
+        formatted_code = codegen.format_source("cpp", generated_code, style="LLVM")
+        return formatted_code
+
+
+backend.register_backend(
+    "cpptoy", lambda prog, *args, **kwargs: print(ToyCpp.apply(prog, **kwargs))
+)

src/iterator/backends/double_roundtrip.py

@@ -0,0 +1,9 @@
+from eve.concepts import Node
+from iterator.backends import backend, embedded
+
+
+def executor(ir: Node, *args, **kwargs):
+    embedded.executor(ir, *args, dispatch_backend=embedded._BACKEND_NAME, **kwargs)
+
+
+backend.register_backend("double_roundtrip", executor)

src/iterator/backends/embedded.py

@@ -0,0 +1,154 @@
+import importlib.util
+import tempfile
+
+import iterator
+from eve import codegen
+from eve.codegen import FormatTemplate as as_fmt
+from eve.codegen import MakoTemplate as as_mako
+from eve.concepts import Node
+from iterator.backends import backend
+from iterator.ir import AxisLiteral, FencilDefinition, OffsetLiteral
+from iterator.transforms import apply_common_transforms
+
+
+class EmbeddedDSL(codegen.TemplatedGenerator):
+    Sym = as_fmt("{id}")
+    SymRef = as_fmt("{id}")
+    BoolLiteral = as_fmt("{value}")
+    IntLiteral = as_fmt("{value}")
+    FloatLiteral = as_fmt("{value}")
+    NoneLiteral = as_fmt("None")
+    OffsetLiteral = as_fmt("{value}")
+    AxisLiteral = as_fmt("{value}")
+    StringLiteral = as_fmt("{value}")
+    FunCall = as_fmt("{fun}({','.join(args)})")
+    Lambda = as_mako("(lambda ${','.join(params)}: ${expr})")
+    StencilClosure = as_mako(
+        "closure(${domain}, ${stencil}, [${','.join(outputs)}], [${','.join(inputs)}])"
+    )
+    FencilDefinition = as_mako(
+        """
+@fendef
+def ${id}(${','.join(params)}):
+    ${'\\n    '.join(closures)}
+    """
+    )
+    FunctionDefinition = as_mako(
+        """
+@fundef
+def ${id}(${','.join(params)}):
+    return ${expr}
+    """
+    )
+    Program = as_fmt(
+        """
+{''.join(function_definitions)} {''.join(fencil_definitions)}"""
+    )
+
+
+# TODO this wrapper should be replaced by an extension of the IR
+class WrapperGenerator(EmbeddedDSL):
+    def visit_FencilDefinition(self, node: FencilDefinition, *, tmps):
+        params = self.visit(node.params)
+        non_tmp_params = [param for param in params if param not in tmps]
+
+        body = []
+        for tmp, domain in tmps.items():
+            axis_literals = [named_range.args[0].value for named_range in domain.args]
+            origin = (
+                "{"
+                + ", ".join(
+                    f"{named_range.args[0].value}: -{self.visit(named_range.args[1])}"
+                    for named_range in domain.args
+                )
+                + "}"
+            )
+            shape = (
+                "("
+                + ", ".join(
+                    f"{self.visit(named_range.args[2])}-{self.visit(named_range.args[1])}"
+                    for named_range in domain.args
+                )
+                + ")"
+            )
+            body.append(
+                f"{tmp} = np_as_located_field({','.join(axis_literals)}, origin={origin})(np.full({shape}, np.nan))"
+            )
+
+        body.append(f"{node.id}({','.join(params)}, **kwargs)")
+        body = "\n    ".join(body)
+        return f"\ndef {node.id}_wrapper({','.join(non_tmp_params)}, **kwargs):\n    {body}\n"
+
+
+_BACKEND_NAME = "embedded"
+
+
+def executor(ir: Node, *args, **kwargs):
+    debug = "debug" in kwargs and kwargs["debug"] is True
+    use_tmps = "use_tmps" in kwargs and kwargs["use_tmps"] is True
+
+    tmps = dict()
+
+    def register_tmp(tmp, domain):
+        tmps[tmp] = domain
+
+    ir = apply_common_transforms(
+        ir, use_tmps=use_tmps, offset_provider=kwargs["offset_provider"], register_tmp=register_tmp
+    )
+
+    program = EmbeddedDSL.apply(ir)
+    wrapper = WrapperGenerator.apply(ir, tmps=tmps)
+    offset_literals = (
+        ir.iter_tree().if_isinstance(OffsetLiteral).getattr("value").if_isinstance(str).to_set()
+    )
+    axis_literals = ir.iter_tree().if_isinstance(AxisLiteral).getattr("value").to_set()
+    with tempfile.NamedTemporaryFile(
+        mode="w",
+        suffix=".py",
+        delete=not debug,
+    ) as tmp:
+        if debug:
+            print(tmp.name)
+        header = """
+import numpy as np
+from iterator.builtins import *
+from iterator.runtime import *
+from iterator.embedded import np_as_located_field
+"""
+        offset_literals = [f'{o} = offset("{o}")' for o in offset_literals]
+        axis_literals = [f'{o} = CartesianAxis("{o}")' for o in axis_literals]
+        tmp.write(header)
+        tmp.write("\n".join(offset_literals))
+        tmp.write("\n")
+        tmp.write("\n".join(axis_literals))
+        tmp.write("\n")
+        tmp.write(program)
+        tmp.write(wrapper)
+        tmp.flush()
+
+        spec = importlib.util.spec_from_file_location("module.name", tmp.name)
+        foo = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(foo)  # type: ignore
+
+        fencil_name = ir.fencil_definitions[0].id
+        fencil = getattr(foo, fencil_name + "_wrapper")
+        assert "offset_provider" in kwargs
+
+        new_kwargs = {}
+        new_kwargs["offset_provider"] = kwargs["offset_provider"]
+        if "column_axis" in kwargs:
+            new_kwargs["column_axis"] = kwargs["column_axis"]
+
+        if "dispatch_backend" not in kwargs:
+            iterator.builtins.builtin_dispatch.push_key("embedded")
+            fencil(*args, **new_kwargs)
+            iterator.builtins.builtin_dispatch.pop_key()
+        else:
+            fencil(
+                *args,
+                **new_kwargs,
+                backend=kwargs["dispatch_backend"],
+            )
+
+
+backend.register_backend(_BACKEND_NAME, executor)