Vertical indirection (#1023)

## Technical Description

This PR introduces vertical indirections. In Pseudocode:
```
for all levels k:
  for all cells c:
    out[c,k] = in[c,vert_nbh[k]]
    out[c,k] = in[c,vert_nbh[k]+1] //additional shift
```
More complicated constructs are disallowed, e.g. `out[c,k] = in[c,vert_nbh[k+1]]`, `out[c,k] = in[c,vert_nbh[vert_nbh_inner[k]]]` and `out[c,k] = in[c,vert_nbh1[k+1]+vert_nbh2[k+1]]` would all be illegal.

### Design

A `VerticalOffset` class was introduced. What was formerly known as `verticalOffset` (i.e. the integer to shift k) was renamed to `verticalShift` and encapsulated into that `VerticalOffset` class, along with a `FieldAccessExpr` known as `verticalIndirection`. The interface prohibits direct construction of the `verticalIndirection`, the underlying constructor is wrapped by the `VerticalOffset` s.t. only `FieldAccessExpr`s without Offsets can be constructed. 

Unfortunately, some limitations of the Visitor infrastructure drove some design decisions, e.g.
* the `verticalIndirection` is actually stored as `std::shared_ptr<Expr>` because `getChildren()` returns `ArrayRef<std::shared_ptr<Expr>>`, i.e. access to the original `shared_ptr` is required due to the `ArrayRef`
* getting a _mutable_ reference `std::shared_ptr<Expr>& getIndirectionFieldAsExpr();` is required since there is no proper `const` visitor. 
These observations highlight the importance of [this issue](#617)

The presence of the vertical Indirection has deep consequences with regards to the `Extents` and `Intervals`, since they can both be constructed or modified with the `Offsets`. The presence of a vertical indirections implies that an **extent or an interval can become undefined**. This undefined state is either propagated, e.g. merging an offset with a an undefined offset makes the offset undefined, or simply leads to assertions, e.g. computing the gap intervals of a vector of intervals with some intervals being undefined is illegal. It was carefully evaluated with operation should propagate and which operations should be prohibited, in parts by the requirements imposed by the passes. However, there may still be edge cases we didn't catch. 

## Testing

* The usual slew of de/serialization tests for the new `Offsets` as well as the undefined value of the `Extent` (not required for `Interval` since only `SIR::Interval` is serialized, which does not carry an undefined state)
* Tests ensuring the propagation / prohibiting of `Extent` / `Interval` ops in undefined state.
* `Dawn4Py` tests exercising the newly introduced vertical indirections, also in presence of field versioning etc.
* In a preliminary step, a utility called `testMutator` was introduced to mutate IIR in such a way that each read becomes indirected. This was used to get a good idea about the affected parts of the code, but is currently unused. It remains in the PR in the hopes it may be useful for possible future debugging.

Unfortunately, `dusk` is not ready for vertical indirections yet. It is expected that going end to end might uncover some currently missed edge cases

## Dependencies

This PR is independent, but relies on a [previous refactoring](#1022) to work. 

## Resolves / Enhances

Fixes #979
Fixes #980

dawn/src/dawn/AST/ASTExpr.cpp

@@ -400,6 +400,13 @@ void FieldAccessExpr::setPureOffset(const Offsets& offset) {
   argumentOffset_ = Array3i{{0, 0, 0}};
 }
 
+ArrayRef<std::shared_ptr<Expr>> FieldAccessExpr::getChildren() {
+  if(offset_.hasVerticalIndirection()) {
+    return ExprRangeType(offset_.getVerticalIndirectionFieldAsExpr());
+  }
+  return ExprRangeType();
+}
+
 std::shared_ptr<Expr> FieldAccessExpr::clone() const {
   return std::make_shared<FieldAccessExpr>(*this);
 }

dawn/src/dawn/AST/ASTExpr.h

@@ -541,6 +541,8 @@ class FieldAccessExpr : public Expr {
     return (argumentMap_[0] != -1 || argumentMap_[1] != -1 || argumentMap_[2] != -1);
   }
 
+  ExprRangeType getChildren() override;
+
   /// @brief Set the `offset` and reset the argument and argument-offset maps
   ///
   /// This function is used during the inlining when we now all the offsets.
@@ -564,6 +566,7 @@ class FieldAccessExpr : public Expr {
   void setName(std::string name) { name_ = name; }
 
   const Offsets& getOffset() const { return offset_; }
+  Offsets& getOffset() { return offset_; }
 
   const Array3i& getArgumentMap() const { return argumentMap_; }
   Array3i& getArgumentMap() { return argumentMap_; }

dawn/src/dawn/AST/ASTStringifier.cpp

@@ -279,7 +279,7 @@ class StringVisitor : public ASTVisitor {
       auto hOffset = ast::offset_cast<CartesianOffset const&>(offset.horizontalOffset());
 
       std::array<int, 3> offsetArray = {hOffset.offsetI(), hOffset.offsetJ(),
-                                        offset.verticalOffset()};
+                                        offset.verticalShift()};
       ss_ << expr->getName() << "[";
 
       const auto& argMap = expr->getArgumentMap();

dawn/src/dawn/AST/Offsets.cpp

@@ -13,7 +13,12 @@
 //===------------------------------------------------------------------------------------------===//
 
 #include "dawn/AST/Offsets.h"
+#include "dawn/AST/ASTExpr.h"
+#include "dawn/IIR/ASTExpr.h"
 #include "dawn/Support/Assert.h"
+#include "dawn/Support/Logger.h"
+#include <memory>
+#include <optional>
 
 namespace dawn {
 namespace ast {
@@ -125,22 +130,116 @@ GridType HorizontalOffset::getGridType() const {
   }
 }
 
+// Vertical Offset
+
+bool VerticalOffset::operator==(VerticalOffset const& other) const {
+  bool offsetEqual = verticalShift_ == other.verticalShift_;
+  if(!verticalIndirection_ && !other.verticalIndirection_) {
+    return offsetEqual;
+  } else if(verticalIndirection_ && other.verticalIndirection_) {
+    return offsetEqual && (*verticalIndirection_ == *other.verticalIndirection_);
+  } else {
+    return false;
+  }
+}
+
+std::string VerticalOffset::getIndirectionFieldName() const {
+  DAWN_ASSERT(hasIndirection());
+  return std::dynamic_pointer_cast<FieldAccessExpr>(verticalIndirection_)->getName();
+}
+std::shared_ptr<const FieldAccessExpr> VerticalOffset::getIndirectionField() const {
+  DAWN_ASSERT(hasIndirection());
+  return std::dynamic_pointer_cast<const FieldAccessExpr>(verticalIndirection_);
+}
+void VerticalOffset::setIndirectionAccessID(int accessID) {
+  DAWN_ASSERT(hasIndirection());
+  auto field = std::dynamic_pointer_cast<FieldAccessExpr>(verticalIndirection_);
+  field->getData<iir::IIRAccessExprData>().AccessID = std::make_optional(accessID);
+}
+std::optional<int> VerticalOffset::getIndirectionAccessID() const {
+  DAWN_ASSERT(hasIndirection());
+  auto field = std::dynamic_pointer_cast<FieldAccessExpr>(verticalIndirection_);
+  auto data = field->getData<iir::IIRAccessExprData>();
+  return data.AccessID;
+}
+std::shared_ptr<Expr>& VerticalOffset::getIndirectionFieldAsExpr() {
+  DAWN_ASSERT(hasIndirection());
+  return verticalIndirection_;
+}
+
+VerticalOffset VerticalOffset::operator+=(VerticalOffset const& other) {
+  DAWN_ASSERT_MSG(!verticalIndirection_ && !other.verticalIndirection_,
+                  "operator += not well defined for vertical offsets with indirection");
+  verticalShift_ += other.verticalShift_;
+  return *this;
+}
+
+VerticalOffset::VerticalOffset(int shift, const std::string& fieldName)
+    : verticalShift_(shift), verticalIndirection_(std::make_shared<FieldAccessExpr>(fieldName)) {}
+
+VerticalOffset::VerticalOffset(const VerticalOffset& other) { *this = other; }
+
+VerticalOffset& VerticalOffset::operator=(VerticalOffset const& other) {
+  verticalShift_ = other.verticalShift_;
+  if(other.verticalIndirection_) {
+    verticalIndirection_ = std::make_shared<FieldAccessExpr>(*other.getIndirectionField());
+  } else {
+    verticalIndirection_ = nullptr;
+  }
+  return *this;
+}
+
 // Offsets
 
 Offsets::Offsets(HorizontalOffset const& hOffset, int vOffset)
-    : horizontalOffset_(hOffset), verticalOffset_(vOffset) {}
+    : horizontalOffset_(hOffset), verticalOffset_(VerticalOffset(vOffset)) {}
+Offsets::Offsets(HorizontalOffset const& hOffset, int vOffset, const std::string& vIndirection)
+    : horizontalOffset_(hOffset), verticalOffset_(VerticalOffset(vOffset, vIndirection)) {}
 
 Offsets::Offsets(cartesian_, int i, int j, int k)
-    : horizontalOffset_(cartesian, i, j), verticalOffset_(k) {}
+    : horizontalOffset_(cartesian, i, j), verticalOffset_(VerticalOffset(k)) {}
 Offsets::Offsets(cartesian_, std::array<int, 3> const& structuredOffsets)
     : Offsets(cartesian, structuredOffsets[0], structuredOffsets[1], structuredOffsets[2]) {}
+Offsets::Offsets(cartesian_, int i, int j, int k, const std::string& fieldName)
+    : horizontalOffset_(cartesian, i, j), verticalOffset_(VerticalOffset(k, fieldName)) {}
+Offsets::Offsets(cartesian_, std::array<int, 3> const& structuredOffsets,
+                 const std::string& fieldName)
+    : Offsets(cartesian, structuredOffsets[0], structuredOffsets[1], structuredOffsets[2],
+              fieldName) {}
 Offsets::Offsets(cartesian_) : horizontalOffset_(cartesian) {}
 
 Offsets::Offsets(unstructured_, bool hasOffset, int k)
     : horizontalOffset_(unstructured, hasOffset), verticalOffset_(k) {}
+Offsets::Offsets(unstructured_, bool hasOffset, int k, const std::string& fieldName)
+    : horizontalOffset_(unstructured, hasOffset), verticalOffset_(k, fieldName) {}
 Offsets::Offsets(unstructured_) : horizontalOffset_(unstructured) {}
 
-int Offsets::verticalOffset() const { return verticalOffset_; }
+std::string Offsets::getVerticalIndirectionFieldName() const {
+  DAWN_ASSERT(hasVerticalIndirection());
+  return verticalOffset_.getIndirectionFieldName();
+}
+
+std::shared_ptr<const FieldAccessExpr> Offsets::getVerticalIndirectionField() const {
+  DAWN_ASSERT(hasVerticalIndirection());
+  return verticalOffset_.getIndirectionField();
+}
+void Offsets::setVerticalIndirectionAccessID(int accessID) {
+  DAWN_ASSERT(hasVerticalIndirection());
+  verticalOffset_.setIndirectionAccessID(accessID);
+}
+std::optional<int> Offsets::getVerticalIndirectionAccessID() const {
+  DAWN_ASSERT(hasVerticalIndirection());
+  return verticalOffset_.getIndirectionAccessID();
+}
+std::shared_ptr<Expr>& Offsets::getVerticalIndirectionFieldAsExpr() {
+  DAWN_ASSERT(hasVerticalIndirection());
+  return verticalOffset_.getIndirectionFieldAsExpr();
+}
+
+void Offsets::setVerticalIndirection(const std::string& fieldName) {
+  verticalOffset_ = VerticalOffset(verticalOffset_.getShift(), fieldName);
+}
+
 HorizontalOffset const& Offsets::horizontalOffset() const { return horizontalOffset_; }
 
 bool Offsets::operator==(Offsets const& other) const {
@@ -158,11 +257,13 @@ bool Offsets::isZero() const { return verticalOffset_ == 0 && horizontalOffset_.
 
 std::string to_string(unstructured_, Offsets const& offset) {
   auto const& hoffset = offset_cast<UnstructuredOffset const&>(offset.horizontalOffset());
-  auto const& voffset = offset.verticalOffset();
+  auto const& vshift = offset.verticalShift();
 
   using namespace std::string_literals;
   return (hoffset.hasOffset() ? "<has_horizontal_offset>"s : "<no_horizontal_offset>"s) + "," +
-         std::to_string(voffset);
+         (offset.hasVerticalIndirection()
+              ? offset.getVerticalIndirectionFieldName() + "[" + std::to_string(vshift) + "]"
+              : std::to_string(vshift));
 }
 
 std::string to_string(cartesian_, Offsets const& offsets, std::string const& sep) {
@@ -177,7 +278,10 @@ std::string to_string(Offsets const& offset) {
       [&](UnstructuredOffset const&) { return to_string(unstructured, offset); },
       [&]() {
         using namespace std::string_literals;
-        return "<no_horizontal_offset>,"s + std::to_string(offset.verticalOffset());
+        return "<no_horizontal_offset>,"s + (offset.hasVerticalIndirection()
+                                                 ? offset.getVerticalIndirectionFieldName() + "[" +
+                                                       std::to_string(offset.verticalShift()) + "]"
+                                                 : std::to_string(offset.verticalShift()));
       });
 }
 

dawn/src/dawn/AST/Offsets.h

@@ -20,12 +20,15 @@
 
 #include <array>
 #include <memory>
+#include <optional>
 #include <string>
 
 namespace dawn {
 namespace ast {
 
 class Offsets;
+class FieldAccessExpr;
+class Expr;
 
 static constexpr cartesian_ cartesian;
 static constexpr unstructured_ unstructured;
@@ -162,19 +165,57 @@ auto offset_dispatch(HorizontalOffset const& hOffset, CartFn const& cartFn,
   }
 }
 
+class VerticalOffset {
+
+  int verticalShift_ = 0;
+  // shared ptr required for visitors
+  std::shared_ptr<Expr> verticalIndirection_ = nullptr;
+
+public:
+  VerticalOffset() = default;
+  VerticalOffset(int shift) : verticalShift_(shift){};
+  VerticalOffset(int shift, const std::string& fieldName);
+  VerticalOffset(const std::string& fieldName) : VerticalOffset(0, fieldName){};
+  VerticalOffset(const VerticalOffset&);
+  int getShift() const { return verticalShift_; }
+  bool hasIndirection() const { return bool(verticalIndirection_); }
+  std::string getIndirectionFieldName() const;
+  std::shared_ptr<const FieldAccessExpr> getIndirectionField() const;
+  // unfortunately we need this to be compatible with the visitor infrastructure
+  std::shared_ptr<Expr>& getIndirectionFieldAsExpr();
+  void setIndirectionAccessID(int accessID);
+  std::optional<int> getIndirectionAccessID() const;
+  bool operator==(VerticalOffset const& other) const;
+  VerticalOffset operator+=(VerticalOffset const& other);
+  VerticalOffset& operator=(VerticalOffset const& other);
+};
+
 class Offsets {
 public:
   Offsets() = default;
   Offsets(HorizontalOffset const& hOffset, int vOffset);
+  Offsets(HorizontalOffset const& hOffset, int vOffset, const std::string& vIndirection);
 
   Offsets(cartesian_, int i, int j, int k);
   Offsets(cartesian_, std::array<int, 3> const& structuredOffsets);
+  Offsets(cartesian_, int i, int j, int k, const std::string& fieldName);
+  Offsets(cartesian_, std::array<int, 3> const& structuredOffsets, const std::string& fieldName);
   explicit Offsets(cartesian_);
 
   Offsets(unstructured_, bool hasOffset, int k);
+  Offsets(unstructured_, bool hasOffset, int k, const std::string& fieldName);
   explicit Offsets(unstructured_);
 
-  int verticalOffset() const;
+  int verticalShift() const { return verticalOffset_.getShift(); }
+  bool hasVerticalIndirection() const { return verticalOffset_.hasIndirection(); }
+  std::string getVerticalIndirectionFieldName() const;
+  std::shared_ptr<const FieldAccessExpr> getVerticalIndirectionField() const;
+  void setVerticalIndirectionAccessID(int accessID);
+  std::optional<int> getVerticalIndirectionAccessID() const;
+  // unfortunately we need this to be compatible with the visitor infrastructure
+  std::shared_ptr<Expr>& getVerticalIndirectionFieldAsExpr();
+  void setVerticalIndirection(const std::string& fieldName);
+
   HorizontalOffset const& horizontalOffset() const;
 
   bool operator==(Offsets const& other) const;
@@ -185,7 +226,7 @@ class Offsets {
 
 private:
   HorizontalOffset horizontalOffset_;
-  int verticalOffset_ = 0;
+  VerticalOffset verticalOffset_;
 };
 Offsets operator+(Offsets o1, Offsets const& o2);
 
@@ -197,7 +238,7 @@ template <typename F>
 std::string to_string(cartesian_, Offsets const& offset, std::string const& sep,
                       F const& offset_to_string) {
   auto const& hoffset = offset_cast<CartesianOffset const&>(offset.horizontalOffset());
-  auto const& voffset = offset.verticalOffset();
+  auto const& voffset = offset.verticalShift();
   std::string s;
   std::string csep = "";
   if(std::string ret = offset_to_string("i", hoffset.offsetI()); ret != "") {

dawn/src/dawn/CodeGen/CXXNaive-ico/ASTStencilBody.cpp

@@ -256,8 +256,14 @@ void ASTStencilBody::visit(const std::shared_ptr<iir::FieldAccessExpr>& expr) {
     }
   } else {
     if(metadata_.getFieldDimensions(iir::getAccessID(expr)).isVertical()) {
-      ss_ << "m_" << getName(expr) << "("
-          << "k+" << expr->getOffset().verticalOffset() << ")";
+      if(expr->getOffset().hasVerticalIndirection()) {
+        ss_ << "m_" << getName(expr) << "("
+            << "m_" << expr->getOffset().getVerticalIndirectionFieldName() + "("
+            << "k)+" << expr->getOffset().verticalShift() << ")";
+      } else {
+        ss_ << "m_" << getName(expr) << "("
+            << "k+" << expr->getOffset().verticalShift() << ")";
+      }
       return;
     }
 
@@ -278,7 +284,14 @@ void ASTStencilBody::visit(const std::shared_ptr<iir::FieldAccessExpr>& expr) {
       if(isHorizontal) {
         ss_ << ")";
       } else {
-        ss_ << ",k+" << expr->getOffset().verticalOffset() << ")";
+        if(expr->getOffset().hasVerticalIndirection()) {
+          ss_ << ","
+              << "m_" << expr->getOffset().getVerticalIndirectionFieldName() << "(deref(LibTag{}, "
+              << resArgName << "), k)+" << expr->getOffset().verticalShift() << ")";
+
+        } else {
+          ss_ << ",k+" << expr->getOffset().verticalShift() << ")";
+        }
       }
     } else {
       std::string sparseIdx = parentIsReduction_
@@ -289,7 +302,13 @@ void ASTStencilBody::visit(const std::shared_ptr<iir::FieldAccessExpr>& expr) {
       if(isHorizontal) {
         ss_ << ")";
       } else {
-        ss_ << ",k+" << expr->getOffset().verticalOffset() << ")";
+        if(expr->getOffset().hasVerticalIndirection()) {
+          ss_ << ","
+              << "m_" << expr->getOffset().getVerticalIndirectionFieldName() << "(deref(LibTag{}, "
+              << sparseArgName_ << "), k)+" << expr->getOffset().verticalShift() << ")";
+        } else {
+          ss_ << ",k+" << expr->getOffset().verticalShift() << ")";
+        }
       }
     }
   }

dawn/src/dawn/CodeGen/CodeGen.cpp

@@ -498,7 +498,11 @@ std::string extentToString(iir::Extents const& extents, ast::GridType const& gri
       result += "false";
   }
   auto const& vExtents = extents.verticalExtent();
-  result += ", " + std::to_string(vExtents.minus()) + "," + std::to_string(vExtents.plus());
+  if(!vExtents.isUndefined()) {
+    result += ", " + std::to_string(vExtents.minus()) + "," + std::to_string(vExtents.plus());
+  } else {
+    result += ", std::numeric_limits<int>::min() , std::numeric_limits<int>::max()";
+  }
   return result;
 }
 } // namespace