Scoping Issues in ICO Naive Backend

[mroethlin]

There are two very related bugs in the ICO naive backend

First: There is a bug in the cxx-naive-ico code generator in conjunction with an if condition. Each reduction is supposed to span its own scope, to encapsulate the sparse index iterator from the outside world. This fails in the case of an if condition, where a scope over the whole if condition is spanned. This causes problems if both branches or even the conditional contain a reduction or loop fill statement, because the linear / sparse index will not be re-set between the individual iterations. A close to minimal example follows:

dusk code

@stencil
def if_condition(f: Field[Edge], nx: Field[Edge], L: Field[Edge], A: Field[Cell], edge_orientation: Field[Cell > Edge], 
                 f_x: Field[Cell]):
    with levels_downward:
        if sum_over(Cell > Edge, 1) < 4:
            f_x = sum_over(Cell > Edge, L)
        else:
            f_x = sum_over(Cell > Edge, L)

generated code

for (auto const &loc : getCells(LibTag{}, m_mesh)) {
  {
    int sparse_dimension_idx0 = 0;
    if ((reduce(LibTag{}, m_mesh, loc, (::dawn::float_type)0,
                std::vector<::dawn::LocationType>{
                    ::dawn::LocationType::Cells,
                    ::dawn::LocationType::Edges},
                [&](auto &lhs, auto red_loc1) {
                  lhs += (int)1;
                  sparse_dimension_idx0++;
                  return lhs;
                }) < (int)4)) {
      m_f_x(deref(LibTag{}, loc)) =
          reduce(LibTag{}, m_mesh, loc, (::dawn::float_type)0,
                  std::vector<::dawn::LocationType>{
                      ::dawn::LocationType::Cells,
                      ::dawn::LocationType::Edges},
                  [&](auto &lhs, auto red_loc1) {
                    lhs += m_L(deref(LibTag{}, red_loc1));
                    sparse_dimension_idx0++;
                    return lhs;
                  });
    } else {
      m_f_x(deref(LibTag{}, loc)) =
          reduce(LibTag{}, m_mesh, loc, (::dawn::float_type)0,
                  std::vector<::dawn::LocationType>{
                      ::dawn::LocationType::Cells,
                      ::dawn::LocationType::Edges},
                  [&](auto &lhs, auto red_loc1) {
                    lhs += m_L(deref(LibTag{}, red_loc1));
                    sparse_dimension_idx0++;
                    return lhs;
                  });
    }
  }

Second: A very related bug occurs if a local field is demoted to a scalar, written to in one reduction and then consumed in a subsequent reduction. Due to the isolation into different scopes the later reduction can't access the emitted scalar.

dusk code

@stencil
def sph(xn: Field[Vertex], yn: Field[Vertex], xc: Field[Cell], yc: Field[Cell], 
        fn: Field[Vertex], f_intp: Field[Cell], h: Field[Cell], pi: Field[Cell]):
    wij: Field[Cell > Vertex > Edge > Vertex]
    qij: Field[Cell > Vertex > Edge > Vertex]
    Wc: Field[Cell]

    with levels_upward:
        with sparse[Cell > Vertex > Edge > Vertex]:
            qij = sqrt((xc-xn)*(xc-xn) + (yc-yn)*(yc-yn))/h
            wij = 1./(pi*h*h)*exp(-qij*qij)

        Wc = sum_over(Cell > Vertex > Edge > Vertex, wij)    
        f_intp = sum_over(Cell > Vertex > Edge > Vertex, 1/Wc*wij*fn)

generated code

for(auto const& loc : getCells(LibTag{}, m_mesh)) {
  {
    int sparse_dimension_idx0 = 0;
    ::dawn::float_type __local_Wc_2_224 =
        reduce(LibTag{}, m_mesh, loc, (::dawn::float_type)0,
                std::vector<::dawn::LocationType>{
                    ::dawn::LocationType::Cells, ::dawn::LocationType::Vertices,
                    ::dawn::LocationType::Edges, ::dawn::LocationType::Vertices},
                [&](auto& lhs, auto red_loc1) {
                  lhs += m_wij(deref(LibTag{}, loc), sparse_dimension_idx0);
                  sparse_dimension_idx0++;
                  return lhs;
                });
  }
  {
    int sparse_dimension_idx0 = 0;
    m_f_intp(deref(LibTag{}, loc)) =
        reduce(LibTag{}, m_mesh, loc, (::dawn::float_type)0,
                std::vector<::dawn::LocationType>{
                    ::dawn::LocationType::Cells, ::dawn::LocationType::Vertices,
                    ::dawn::LocationType::Edges, ::dawn::LocationType::Vertices},
                [&](auto& lhs, auto red_loc1) {
                  lhs += ((((int)1 / __local_Wc_2_224) *
                          m_wij(deref(LibTag{}, loc), sparse_dimension_idx0)) *
                          m_fn(deref(LibTag{}, red_loc1)));
                  sparse_dimension_idx0++;
                  return lhs;
                });
  }
}

[Stagno]

Some thoughts before we start the pair programming on these.

• The first bug is purely codegen (naive). A trivial solution (and probably a clean one) would be to make the sparse_dimension_idx a static variable local to the lambda function of a reduction. E.g.

(reduce(LibTag{}, m_mesh, loc, (::dawn::float_type)0,
                std::vector<::dawn::LocationType>{
                    ::dawn::LocationType::Cells,
                    ::dawn::LocationType::Edges},
                [&](auto &lhs, auto red_loc1) {
                  static int sparse_dimension_idx0 = 0;
                  lhs += (int)1;
                  sparse_dimension_idx0++;
                  return lhs;
                })

• The second is not a codegen bug at all, as codegen translates correctly the input IIR. The problem is in the lowering which produces such IIR. Probably PassTemporaryType is too greedy and ignores that temporary Wc spans multiple stages, thus it cannot be demoted to a local variable (could be a legacy of structured where multistage->kernel, so local variable across stages is ok).