Merge pull request #527 from facebookresearch/pr/single_instance

ScheduleTree::bandTile: drop special casing 0D tile sizes

Author: skimo-openhub

tc/core/polyhedral/cuda/mapped_scop.h

@@ -72,9 +72,16 @@ class MappedScop {
  public:
   static inline std::unique_ptr<MappedScop> makeOneBlockOneThread(
       std::unique_ptr<Scop>&& scop) {
-    return std::unique_ptr<MappedScop>(new MappedScop(
+    auto mscop = std::unique_ptr<MappedScop>(new MappedScop(
         std::move(scop), ::tc::Grid{1, 1, 1}, ::tc::Block{1, 1, 1}, 1, false));
+    auto band = mscop->scop_->obtainOuterBand();
+    mscop->mapBlocksForward(band, 0);
+    mscop->mapThreadsBackward(band);
+    return mscop;
   }
+  // The MappedScop returned by this method does not satisfy the invariant
+  // of having a mapping to blocks and threads.  It is up to the caller
+  // to insert these mappings.
   static inline std::unique_ptr<MappedScop> makeMappedScop(
       std::unique_ptr<Scop>&& scop,
       ::tc::Grid grid,

tc/core/polyhedral/cuda/tighten_launch_bounds.cc

@@ -88,7 +88,8 @@ size_t maxValue(const Scop& scop, const MappingIdType& id) {
   LOG_IF(WARNING, min > 0)
       << "Opportunity for tightening launch bounds with shifting -> min:"
       << min;
-  TC_CHECK(max < sizetMax) << "missing mapping to " << id << *root;
+  TC_CHECK(max < sizetMax) << "missing mapping to " << id << "\n" << *root;
+  TC_CHECK(min < sizetMax) << "missing mapping to " << id << " type\n" << *root;
   // Inclusive range needs + 1 to translate to sizes
   return max + 1;
 }

tc/core/polyhedral/schedule_transforms.cc

@@ -359,9 +359,6 @@ ScheduleTree* bandTile(
   auto eb = st->elemAs<ScheduleTreeElemBand>();
   TC_CHECK(eb) << "Not a band: " << *st;
 
-  if (tileSizes.size() == 0) {
-    return st;
-  }
   auto& band = *eb;
   TC_CHECK(band.permutable_) << "Can't tile a non-permutable band" << band;
 

tc/core/polyhedral/scop.cc

@@ -423,6 +423,8 @@ detail::ScheduleTree* setPermutable(detail::ScheduleTree* tree) {
   return tree;
 }
 
+} // namespace
+
 /*
  * Return the outermost band in the schedule tree with the given root.
  * If there is no single outermost band, then insert a (permutable)
@@ -431,7 +433,8 @@ detail::ScheduleTree* setPermutable(detail::ScheduleTree* tree) {
  * insert the band in the leaf.  If branching is encountered, then
  * insert the band above the branching.
  */
-detail::ScheduleTree* obtainOuterBand(detail::ScheduleTree* root) {
+detail::ScheduleTree* Scop::obtainOuterBand() {
+  auto root = scheduleRoot();
   auto tree = root;
   while (!tree->elemAs<ScheduleTreeElemBand>()) {
     auto n = tree->numChildren();
@@ -449,11 +452,10 @@ detail::ScheduleTree* obtainOuterBand(detail::ScheduleTree* root) {
   }
   return tree;
 }
-} // namespace
 
 detail::ScheduleTree* Scop::tileOuterBand(const TilingView& tileSizes) {
   using namespace tc::polyhedral::detail;
-  auto band = obtainOuterBand(scheduleRoot());
+  auto band = obtainOuterBand();
   auto bandNode = band->elemAs<ScheduleTreeElemBand>();
   std::vector<size_t> sizes = tileSizes.extractVector();
   if (bandNode->nMember() < sizes.size()) {

tc/core/polyhedral/scop.h

@@ -394,6 +394,10 @@ struct Scop {
   static std::unique_ptr<Scop> makeScheduled(
       const Scop& scop,
       const SchedulerOptionsView& schedulerOptions);
+  // Return the outermost band in the schedule tree with the given root.
+  // If there is no single outermost band, then insert a (permutable)
+  // zero-dimensional band and return that.
+  detail::ScheduleTree* obtainOuterBand();
   // Tile the outermost band.
   // Splits the band into tile loop band and point loop band where point loops
   // have fixed trip counts specified in "tiling", and returns a pointer to the

test/test_cuda_mapper.cc

@@ -1151,6 +1151,21 @@ def local_sparse_convolution(float(N, C, H, W) I, float(O, KC, KH, KW) W1) -> (O
   }
 }
 
+/*
+ * Check that a TC with a single instance gets mapped properly.
+ * tightenLaunchBounds (called during codegen) will complain
+ * if it is not.
+ */
+TEST_F(PolyhedralMapperTest, SingleInstance) {
+  string tc = R"TC(
+def f(float(N) I) -> (a)
+{
+    a = 0
+}
+)TC";
+  codegenMapped(tc, DefaultOptions());
+}
+
 int main(int argc, char** argv) {
   ::testing::InitGoogleTest(&argc, argv);
   ::gflags::ParseCommandLineFlags(&argc, &argv, true);