tc2halide: properly tag reduction updates

In particular, only tag reduction updates and only tag them once.
The original code would tag any non-initial write to a tensor
involved in a reduction and it would tag it possibly several times,
once for each reduction on that tensor.
Tag the reduction updates immediately to ensure that
only reduction updates get tagged and that they get tagged exactly once.

Note that the tags are now introduced in the function definitions
rather than only in the lowered statements.  The presence of these
intrinsics may hinder some of the analysis passes in Halide's
standard lowering path.  If these passer ever need to be used,
then the intrinsic introduction mechanism may need to be revisited.

Multiply tagged expressions cause confusion in TC.
In particular, isSupportedReduction would fail to detect
the reduction and the code generation also assumes
that a reduction update is only tagged once.

Tests by Mathieu Fehr.

Author: Sven Verdoolaege

tc/core/tc2halide.cc

@@ -488,13 +488,15 @@ Expr reductionUpdate(Expr e) {
   return Call::make(e.type(), kReductionUpdate, {e}, Call::Intrinsic);
 }
 
+// Note that the function definitions created by translateComprehension may
+// contain kReductionUpdate intrinsics.  These may have to be removed
+// in order to be able to apply internal Halide analysis passes on them.
 void translateComprehension(
     const lang::Comprehension& c,
     const map<string, Parameter>& params,
     bool throwWarnings,
     map<string, Function>* funcs,
-    FunctionBounds* bounds,
-    vector<Function>* reductions) {
+    FunctionBounds* bounds) {
   Function f;
   auto it = funcs->find(c.ident().name());
   if (it != funcs->end()) {
@@ -589,8 +591,9 @@ void translateComprehension(
                                  << c.assignment()->range().text() << "\n";
   }
 
+  // Tag reductions as such
   if (c.assignment()->kind() != '=') {
-    reductions->push_back(f);
+    rhs = reductionUpdate(rhs);
   }
 
   // Bind any scalar params on the rhs to their parameter objects.
@@ -739,13 +742,12 @@ HalideComponents translateDef(const lang::Def& def, bool throwWarnings) {
   components.def = def;
   FunctionBounds bounds;
 
-  vector<Function> reductions;
   for (auto p : def.params()) {
     translateParam(p, &components.params, &components.inputs);
   }
   for (auto c : def.statements()) {
     translateComprehension(
-        c, components.params, throwWarnings, &funcs, &bounds, &reductions);
+        c, components.params, throwWarnings, &funcs, &bounds);
   }
   vector<Function> outputs;
   for (auto p : def.returns()) {
@@ -803,32 +805,6 @@ HalideComponents translateDef(const lang::Def& def, bool throwWarnings) {
   s = uniquify_variable_names(s);
   s = simplify(s);
 
-  // Tag reductions as such
-  for (const Function& f : reductions) {
-    class TagReduction : public IRMutator2 {
-      using IRMutator2::visit;
-      bool found_init = false;
-      Stmt visit(const Provide* op) override {
-        if (op->name == f.name()) {
-          if (found_init) {
-            return Provide::make(
-                op->name, {reductionUpdate(op->values[0])}, op->args);
-          } else {
-            found_init = true;
-            return op;
-          }
-        } else {
-          return op;
-        }
-      }
-      const Function& f;
-
-     public:
-      TagReduction(const Function& f) : f(f) {}
-    } tagReduction(f);
-    s = tagReduction.mutate(s);
-  }
-
   // Trim ProducerConsumer annotations. TC doesn't use them.
   class RemoveProducerConsumer : public IRMutator2 {
     using IRMutator2::visit;

test/test_cuda_mapper.cc

@@ -861,6 +861,77 @@ def fun(float(N) I) -> (O) {
   EXPECT_TRUE(code.find("O[0] = (O") != std::string::npos);
 }
 
+struct ReductionTest : public PolyhedralMapperTest {
+  static CudaMappingOptions reductionTestMappingOptions() {
+    return DefaultOptions()
+        .outerScheduleFusionStrategy(tc::FusionStrategy::Preserve3Coincident)
+        .outerScheduleAllowSkewing(false)
+        .outerSchedulePositiveOrthant(true)
+        .intraTileScheduleFusionStrategy(tc::FusionStrategy::Min)
+        .intraTileScheduleAllowSkewing(false)
+        .intraTileSchedulePositiveOrthant(true)
+        .fixParametersBeforeScheduling(false)
+        .tile(18, 32)
+        .unroll(16)
+        .tileImperfectlyNested(false)
+        .matchLibraryCalls(true)
+        .mapToThreads({512})
+        .mapToBlocks({16384})
+        .useSharedMemory(true)
+        .usePrivateMemory(false)
+        .unrollCopyShared(true);
+  }
+
+  void Check(const string& tc) {
+    auto code = codegenMapped(tc, reductionTestMappingOptions());
+    using tc::code::cuda::kCUBReductionName;
+    EXPECT_TRUE(code.find(kCUBReductionName) != std::string::npos);
+  }
+};
+
+/*
+ * Check that a reduction library call is produced when the reduction
+ * instruction is before an instruction modifying the same tensor.
+ */
+TEST_F(ReductionTest, BeforeInstruction) {
+  Check(R"TC(
+def fun(float(N, K) I) -> (O) {
+    O(n) +=! I(n, r_n)
+    O(n) = O(n) / (K)
+}
+)TC");
+}
+
+/*
+ * Check that a reduction library call is produced when the reduction
+ * instruction is after an instruction modifying the same tensor.
+ */
+TEST_F(ReductionTest, AfterInstruction) {
+  Check(R"TC(
+def fun(float(N, K) I, float(N) O0) -> (O) {
+    O(n) = 0.0 where n in 0:N
+    O(n) += O0(n)
+    O(n) += I(n, r_n)
+}
+)TC");
+}
+
+/*
+ * Check that a reduction library call is produced when the reduction
+ * instruction is placed after an instruction modifying the same tensor and
+ * before an instruction modifying the same tensor.
+ */
+TEST_F(ReductionTest, BetweenInstructions) {
+  Check(R"TC(
+def fun(float(N, K) I, float(N) O0) -> (O) {
+    O(n) = 0.0 where n in 0:N
+    O(n) += O0(n)
+    O(n) += I(n, r_n)
+    O(n) = O(n) / (K)
+}
+)TC");
+}
+
 static const string kTcMM = R"TC(
 def fun(float(M, K) A, float(K, N) B) -> (C) {
     C(m, n) +=! A(m, r_k) * B(r_k, n)