Merge pull request #483 from facebookresearch/pr/fix_462

MappedScop::detectReductions: avoid inserting filters that turn out to be empty

Author: ftynse

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -189,7 +189,7 @@ void fixThreadsBelow(
 }
 
 /*
- * Try and order the other statements in "domain" (if any)
+ * Try and order the other active statements at "tree" (if any)
  * away from the "updates" statements, returning true is the operation succeeds.
  * In particular, only do this if it doesn't violate any dependences.
  * Anything that depends on an update statement is ordered after
@@ -198,8 +198,8 @@ void fixThreadsBelow(
 bool separatedOut(
     Scop& scop,
     detail::ScheduleTree* tree,
-    isl::union_set domain,
     isl::union_set updates) {
+  auto domain = activeDomainPoints(scop.scheduleRoot(), tree);
   auto other = domain.subtract(updates);
   if (other.is_empty()) {
     return true;
@@ -271,7 +271,7 @@ bool MappedScop::detectReductions(detail::ScheduleTree* tree) {
   // Order the other statements (if any) before the update statements
   // to ensure the band from which the reduction band has been split off
   // only contains update statements.
-  if (!separatedOut(scop(), tree, domain, updates)) {
+  if (!separatedOut(scop(), tree, updates)) {
     return false;
   }
   reductionBandUpdates_.emplace(tree, Reduction(updateIds));

test/test_cuda_mapper.cc

@@ -1022,6 +1022,29 @@ def fun(float(N) I) -> (O) {
   ASSERT_TRUE(code.find("__ldg(&I") != std::string::npos) << code; // yes
 }
 
+/*
+ * Check that isolating the update statements does not introduce
+ * an empty mapping filter.
+ */
+TEST_F(PolyhedralMapperTest, EmptyMappingFilter) {
+  constexpr static auto tc = R"TC(
+  def var_2D_1D(float(N, K) I, float(N) mean) -> (var)
+  {
+       var(n) +=! I(n, r_k) * I(n, r_k)
+       var(n)  =  var(n) / (K) - mean(n) * mean(n)
+  }
+)TC";
+  auto mappingOptions = DefaultOptions()
+                            .fixParametersBeforeScheduling(false)
+                            .matchLibraryCalls(true)
+                            .mapToThreads(256);
+  auto scop = Prepare(tc);
+  scop->fixParameters<int>({{"N", 1024}, {"K", 36864}});
+  auto mscop = MappedScop::makeWithOuterBlockInnerThreadStrategy(
+      std::move(scop), mappingOptions);
+  mscop->codegen(specializedName);
+}
+
 int main(int argc, char** argv) {
   ::testing::InitGoogleTest(&argc, argv);
   ::gflags::ParseCommandLineFlags(&argc, &argv, true);