Merge pull request #457 from facebookresearch/pr/update

tc2halide: properly tag reduction updates

Author: skimo-openhub

tc/core/halide2isl.cc

@@ -482,14 +482,6 @@ std::vector<Reduction> findReductions(const Stmt& s) {
   class FindReductions : public IRVisitor {
     using IRVisitor::visit;
 
-    bool isReductionInit(const Provide* op) {
-      if (const Call* call = op->values[0].as<Call>()) {
-        return call->is_intrinsic(tc2halide::kReductionInit);
-      } else {
-        return false;
-      }
-    }
-
     bool isReductionUpdate(const Provide* op) {
       if (const Call* call = op->values[0].as<Call>()) {
         return call->is_intrinsic(tc2halide::kReductionUpdate);

tc/core/halide_utils.cc

@@ -123,8 +123,7 @@ std::string halideCodegenC(const Stmt& stmt) {
     using IRPrinter::visit;
 
     void visit(const Call* op) override {
-      if (op->is_intrinsic(tc2halide::kReductionInit) ||
-          op->is_intrinsic(tc2halide::kReductionUpdate)) {
+      if (op->is_intrinsic(tc2halide::kReductionUpdate)) {
         op->args[0].accept(this);
       } else if (
           op->call_type == Call::Halide || op->call_type == Call::Image) {

tc/core/polyhedral/codegen_llvm.cc

@@ -234,9 +234,7 @@ class CodeGen_TC : public Halide::Internal::CodeGen_X86 {
       auto addr = builder->CreateInBoundsGEP(baseAddr, args);
       value = builder->CreateLoad(addr);
       return;
-    } else if (
-        call->is_intrinsic(tc2halide::kReductionInit) ||
-        call->is_intrinsic(tc2halide::kReductionUpdate)) {
+    } else if (call->is_intrinsic(tc2halide::kReductionUpdate)) {
       call->args[0].accept(this);
       return;
     } else {

tc/core/polyhedral/cuda/codegen.cc

@@ -550,9 +550,7 @@ void emitHalideExpr(
           op->call_type == Halide::Internal::Call::CallType::Image) {
         tc::polyhedral::detail::emitMappedTensorAccess(
             op->name, op, op->args, context);
-      } else if (
-          op->is_intrinsic(tc2halide::kReductionInit) ||
-          op->is_intrinsic(tc2halide::kReductionUpdate)) {
+      } else if (op->is_intrinsic(tc2halide::kReductionUpdate)) {
         op->args[0].accept(this);
       } else {
         IRPrinter::visit(op);

tc/core/tc2halide.cc

@@ -484,21 +484,19 @@ void forwardBoundsInference(
   }
 }
 
-Expr reductionInit(Expr e) {
-  return Call::make(e.type(), kReductionInit, {e}, Call::Intrinsic);
-}
-
 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()) {
@@ -593,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.
@@ -743,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()) {
@@ -807,33 +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 Provide::make(
-                op->name, {reductionInit(op->values[0])}, op->args);
-          }
-        } 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;

tc/core/tc2halide.h

@@ -40,8 +40,7 @@ struct HalideComponents {
 };
 
 // For TC reductions, the right-hand-sides of the corresponding
-// Provide nodes are tagged with intrinsics with the following names.
-Halide::Internal::Call::ConstString kReductionInit = "ReductionInit";
+// Provide nodes are tagged with intrinsics with the following name.
 Halide::Internal::Call::ConstString kReductionUpdate = "ReductionUpdate";
 
 // Translate a TC parse tree into equivalent Halide imperative IR with

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)