[Broken][WIP] Add Halide support for "for" loops

Disclaimer: this is WIP and does not work yet. In particular
the ScheduleTree generated is incorrect (not yet properly scoped under
for loop). Additionally, the example `For4InvalidHalide` seems to hit
deeper structural issues of the high-level HalideIR that
traditional polyhedral dependence analysis has no issue with (cc @abadams).

This commit tests that the for-loop language construct properly
propagates to the Halide bounds inference.
This commit also explicitly adds checks that only a single
RangeConstraint s`traiindex in min:max`) is used for each index.
This requirement is added because it would be very surprising to mix
explicit ranges coming from for loops with where clauses in a
comprehension.
In particular, the current behavior in Halide inference is to only
keep the last RangeConstraint which systematically discards the
information specified in the loop. `tc2halide` even says:

```// TODO: What if subsequent updates have incompatible bounds
// (e.g. an in-place stencil)?. The .bound directive will use the
// bounds of the last stage for all stages.
```

As a consequence we add an explicit check that multiple bounds
specifications may not coexist.

Author: nicolasvasilache

tc/core/tc2halide.cc

@@ -33,6 +33,21 @@ using std::vector;
 
 namespace {
 
+using FunctionBounds = map<Function, map<string, Interval>, Function::Compare>;
+
+struct TranslationUnit {
+  HalideComponents components;
+  Scope<Interval> enclosingLoopIndices;
+  map<string, Function> funcs;
+  FunctionBounds bounds;
+  bool throwWarnings;
+};
+
+void translateFor(const lang::For& f, TranslationUnit* tu);
+void translateComprehension(
+    const lang::Comprehension& comprehension,
+    TranslationUnit* tu);
+
 Type translateScalarType(int tcType) {
   switch (tcType) {
     case lang::TK_BOOL:
@@ -264,8 +279,6 @@ vector<const Variable*> unboundVariables(const vector<Var>& lhs, Expr rhs) {
   return finder.result;
 }
 
-typedef map<Function, map<string, Interval>, Function::Compare> FunctionBounds;
-
 void forwardBoundsInference(
     const std::vector<Expr>& exprs,
     const FunctionBounds& bounds,
@@ -500,6 +513,29 @@ Expr reductionUpdate(Expr e) {
   return Call::make(e.type(), kReductionUpdate, {e}, Call::Intrinsic);
 }
 
+void translateStatement(const lang::TreeRef& stmt, TranslationUnit* tu) {
+  if (stmt->kind() == lang::TK_COMPREHENSION) {
+    translateComprehension(lang::Comprehension(stmt), tu);
+  } else {
+    CHECK_EQ(stmt->kind(), lang::TK_FOR);
+    translateFor(lang::For(stmt), tu);
+  }
+}
+
+void translateFor(const lang::For& f, TranslationUnit* pTU) {
+  const map<string, Parameter>& params = pTU->components.params;
+  auto constraint = lang::RangeConstraint(f.rangeConstraint());
+  Interval i;
+  const map<string, Expr> lets;
+  i.min = translateExpr(constraint.start(), params, pTU->funcs, lets);
+  i.max = translateExpr(constraint.end(), params, pTU->funcs, lets) - 1;
+  pTU->enclosingLoopIndices.push(f.index().name(), i);
+  for (auto stm : f.statements()) {
+    translateStatement(stm, pTU);
+  }
+  pTU->enclosingLoopIndices.pop(f.index().name());
+}
+
 // Translate a single TC comprehension/statement to Halide components: funcs,
 // bounds, reductions.
 //
@@ -508,10 +544,11 @@ Expr reductionUpdate(Expr e) {
 // in order to be able to apply internal Halide analysis passes on them.
 void translateComprehension(
     const lang::Comprehension& comprehension,
-    const map<string, Parameter>& params,
-    bool throwWarnings,
-    map<string, Function>* funcs,
-    FunctionBounds* bounds) {
+    TranslationUnit* pTU) {
+  const map<string, Parameter>& params = pTU->components.params;
+  bool throwWarnings = pTU->throwWarnings;
+  map<string, Function>* funcs = &pTU->funcs;
+  FunctionBounds* bounds = &pTU->bounds;
   Function f;
   auto it = funcs->find(comprehension.ident().name());
   if (it != funcs->end()) {
@@ -647,6 +684,12 @@ void translateComprehension(
   // demand).
   Scope<Interval> solution;
 
+  // Copy information from enclosing "for" loops
+  for (auto entry = pTU->enclosingLoopIndices.cbegin();
+       entry != pTU->enclosingLoopIndices.cend();
+       ++entry) {
+    solution.push(entry.name(), entry.value());
+  }
   // Put anything explicitly specified with a 'where' class in the solution
   for (auto constraint_ : comprehension.whereClauses()) {
     if (constraint_->kind() != lang::TK_RANGE_CONSTRAINT)
@@ -656,6 +699,11 @@ void translateComprehension(
     i.min = translateExpr(constraint.start(), params, *funcs, lets);
     i.max = translateExpr(constraint.end(), params, *funcs, lets) - 1;
 
+    if (solution.contains(constraint.ident().name())) {
+      throw lang::ErrorReport(constraint_)
+          << "Multiple range constraints per index NYI";
+    }
+
     // TODO: In the future we'll want to make any non-trivial bounds
     // into hidden scalar parameters, and just pass variables to the
     // polyhedral layer instead of potentially complex
@@ -755,25 +803,20 @@ void translateComprehension(
 
 // Translate a semantically checked TC def to HalideComponents struct.
 HalideComponents translateDef(const lang::Def& def, bool throwWarnings) {
-  map<string, Function> funcs;
-  HalideComponents components;
-  components.def = def;
-  FunctionBounds bounds;
+  TranslationUnit tu;
+  tu.components.def = def;
+  tu.throwWarnings = throwWarnings;
 
   for (auto p : def.params()) {
-    translateParam(p, &components.params, &components.inputs);
+    translateParam(p, &tu.components.params, &tu.components.inputs);
   }
-  for (auto c : def.statements()) {
-    translateComprehension(
-        lang::Comprehension(c),
-        components.params,
-        throwWarnings,
-        &funcs,
-        &bounds);
+  // Semantically valid TCs include at most one outer sequential loop for now
+  for (auto stm : def.statements()) {
+    translateStatement(stm, &tu);
   }
   vector<Function> outputs;
   for (auto p : def.returns()) {
-    translateOutput(p, funcs, &outputs);
+    translateOutput(p, tu.funcs, &outputs);
   }
 
   // Now apply an extremely simplified version of Halide lowering
@@ -804,11 +847,12 @@ HalideComponents translateDef(const lang::Def& def, bool throwWarnings) {
   // used in the pipelines we construct here, so just make a host target.
   Target target("host");
   Stmt s = schedule_functions(outputs, fused_groups, env, target, any_memoized);
+  LOG_IF(ERROR, tc::FLAGS_debug_halide) << s;
   // we insert these to allow for inplace mutation of in/out tensors
   s = remove_undef(s);
   // Apply forward bounds inference results. This replaces the usual Halide
   // bounds inference.
-  for (auto p : bounds) {
+  for (auto p : tu.bounds) {
     const Function& f = p.first;
     for (auto b : p.second) {
       const string& var = b.first;
@@ -893,20 +937,20 @@ HalideComponents translateDef(const lang::Def& def, bool throwWarnings) {
   };
   s = SubstituteAllLets().mutate(s);
 
-  components.stmt = s;
+  tu.components.stmt = s;
 
   for (Function f : outputs) {
     OutputImageParam o = Func(f).output_buffers()[0];
     // Apply forward bounds inference results to the output buffers.
-    const auto& b = bounds[f];
+    const auto& b = tu.bounds[f];
     for (int i = 0; i < o.dimensions(); i++) {
       const Interval& bound = b.at(f.args()[i]);
       o.dim(i).set_bounds(bound.min, simplify(bound.max - bound.min + 1));
     }
-    components.outputs.push_back(o);
+    tu.components.outputs.push_back(o);
   }
 
-  return components;
+  return tu.components;
 }
 } // namespace
 

test/test_inference.cc

@@ -24,9 +24,9 @@ using namespace std;
 using namespace lang;
 
 struct InferenceTest : public ::testing::Test {
-  void Check(const string& tc, const string& expected) {
+  void Check(const string& tc, const string& expected, bool warn = true) {
     auto halideComponents =
-        tc2halide::translate(isl::with_exceptions::globalIslCtx(), tc, true);
+        tc2halide::translate(isl::with_exceptions::globalIslCtx(), tc, warn);
 
     stringstream ss;
     // Ordered map for repro
@@ -529,6 +529,94 @@ def fun(float(N) size) -> (O) {
   EXPECT_THROW(Check(tc, {}), ::lang::ErrorReport);
 }
 
+TEST_F(InferenceTest, For1) {
+  string tc = R"TC(
+    def fun(float(M, N) X) -> (R1) {
+        for t in 0:123 {
+            R1(t, m, n) = (t == 0) ? X(m, n) : 0.0
+        }
+    }
+)TC";
+  Check(tc, R"HALIDE(mins:
+X@[0; 0; ]
+R1@[0; 0; 0; ]
+extents:
+X@[M; N; ]
+R1@[123; M; N; ]
+)HALIDE");
+}
+
+TEST_F(InferenceTest, For2) {
+  string tc = R"TC(
+    def fun(float(M, N) X, float(T) Meta) -> (R1, R2) {
+        for t in 0:T {
+            R1(t, m, n) = (t == 0) ? X(m, n) : 0.0
+            R2(t, m, n) = R1(t, m, n)
+        }
+    }
+)TC";
+  Check(
+      tc,
+      R"HALIDE(mins:
+X@[0; 0; ]
+Meta@[0; ]
+R1@[0; 0; 0; ]
+R2@[0; 0; 0; ]
+extents:
+X@[M; N; ]
+Meta@[T; ]
+R1@[T; M; N; ]
+R2@[T; M; N; ]
+)HALIDE");
+}
+
+TEST_F(InferenceTest, For3) {
+  string tc = R"TC(
+    def fun(float(M, N) X, float(T) Meta) -> (R1, R2) {
+        R1(t, m, n) = (t == 0) ? X(m, n) : 0.0 where t in 0:T
+        R2(t, m, n) = 0.0 where t in 0:T, m in 0:M, n in 0:N
+        for t in 1:T {
+            R1(t, m, n) += R1(t-1, m, n)
+            R2(t, m, n) += R1(t, m, n)
+        }
+    }
+)TC";
+  Check(
+      tc,
+      R"HALIDE(mins:
+X@[0; 0; ]
+Meta@[0; ]
+R1@[1; 0; 0; ]
+R2@[1; 0; 0; ]
+extents:
+X@[M; N; ]
+Meta@[T; ]
+R1@[(T + -1); M; N; ]
+R2@[(T + -1); M; N; ]
+)HALIDE",
+      false); // nowarn
+}
+
+TEST_F(InferenceTest, MultipleRangeConstraints) {
+  string tc = R"TC(
+def fun(float(N, M) I) -> (O) {
+  O(i, j) = I(i, j) where i in 3:N-1, i in 1:N
+}
+)TC";
+  EXPECT_THROW(Check(tc, {}), ::lang::ErrorReport);
+}
+
+TEST_F(InferenceTest, MultipleRangeConstraintsFor) {
+  string tc = R"TC(
+    def fun(float(M, N) X) -> (R1) {
+        for t in 0:123 {
+            R1(t, m, n) = (t == 0) ? X(m, n) : 0.0 where t in 0:N
+        }
+    }
+)TC";
+  EXPECT_THROW(Check(tc, {}), ::lang::ErrorReport);
+}
+
 int main(int argc, char** argv) {
   ::testing::InitGoogleTest(&argc, argv);
   ::gflags::ParseCommandLineFlags(&argc, &argv, true);

test/test_tc2halide.cc

@@ -38,6 +38,8 @@ struct TC2Isl : public ::testing::Test {
         isl::with_exceptions::globalIslCtx(), halide);
     auto scheduleHalide = scop->scheduleRoot();
 
+    LOG_IF(ERROR, FLAGS_debug_tc_mapper) << *scheduleHalide;
+
     polyhedral::detail::validateSchedule(scheduleHalide);
   }
 };
@@ -198,6 +200,57 @@ def foo(float(N) A) -> (B) {
   EXPECT_THROW(Check(tc), ::lang::ErrorReport);
 }
 
+TEST_F(TC2Isl, For1) {
+  string tc = R"TC(
+    def fun(float(M, N) X) -> (R1) {
+        for t in 0:123 {
+            R1(t, m, n) = (t == 0) ? X(m, n) : 0.0
+        }
+    }
+)TC";
+  Check(tc);
+}
+
+TEST_F(TC2Isl, For2) {
+  string tc = R"TC(
+    def fun(float(M, N) X, float(T) Meta) -> (R1, R2) {
+        for t in 0:T {
+            R1(t, m, n) = (t == 0) ? X(m, n) : 0.0
+            R2(t, m, n) = R1(t, m, n)
+        }
+    }
+)TC";
+  Check(tc);
+}
+
+TEST_F(TC2Isl, For3) {
+  string tc = R"TC(
+    def fun(float(M, N) X, float(T) Meta) -> (R1, R2) {
+        R1(t, m, n) = (t == 0) ? X(m, n) : 0.0 where t in 0:T
+        R2(t, m, n) = 0.0 where t in 0:T, m in 0:M, n in 0:N
+        for t in 1:T {
+            R1(t, m, n) += R1(t-1, m, n)
+            R2(t, m, n) += R1(t, m, n)
+        }
+    }
+)TC";
+  Check(tc);
+}
+
+TEST_F(TC2Isl, For4InvalidHalide) {
+  string tc = R"TC(
+    def fun(float(M, N) X, float(T) Meta) -> (R1, R2) {
+        R1(t, m, n) = (t == 0) ? X(m, n) : 0.0 where t in 0:T
+        R2(t, m, n) = 0.0 where t in 0:T, m in 0:M, n in 0:N
+        for t in 1:T {
+            R1(t, m, n) += R1(t-1, m, n) + R2(t-1, m, n)
+            R2(t, m, n) += R1(t, m, n)
+        }
+    }
+)TC";
+  EXPECT_THROW(Check(tc), std::exception);
+}
+
 TEST_F(TC2Isl, Types) {
   for (auto type : {"bool",
                     "uint8",