Merge pull request #537 from facebookresearch/shared-promotion-anywhere

Tunable shared promotion depth

Author: ftynse

tc/autotuner/autotuner-inl.h

@@ -374,6 +374,8 @@ void setupTuningParameters(
       RangeParameter(powers2(FLAGS_tuner_max_unroll_size), "unroll");
   configuration.privateDepth =
       RangeParameter({0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, "pdepth");
+  configuration.sharedDepth =
+      RangeParameter({0, 1, 2, 3, 4, 5, 6, 7}, "sdepth");
 }
 } // namespace
 

tc/autotuner/parameters.cc

@@ -240,6 +240,7 @@ void TuningConfiguration::applyToParameters(
   useReadOnlyCache.apply(f);
   matchLibraryCalls.apply(f);
   privateDepth.apply(f);
+  sharedDepth.apply(f);
 }
 
 bool TuningConfiguration::isValid() const {
@@ -275,6 +276,7 @@ std::vector<ParameterView> TuningConfiguration::collectParameters() {
   params.emplace_back(useReadOnlyCache);
   params.emplace_back(matchLibraryCalls);
   params.emplace_back(privateDepth);
+  params.emplace_back(sharedDepth);
 
   return params;
 }
@@ -306,6 +308,7 @@ void TuningConfiguration::fromCudaMappingOptions(
   unrollCopyShared.selectValue(options.proto().unroll_copy_shared());
   useReadOnlyCache.selectValue(options.proto().use_readonly_cache());
   privateDepth.selectFromValue(options.proto().private_depth());
+  sharedDepth.selectFromValue(options.proto().shared_depth());
 }
 
 void TuningConfiguration::fromCpuMappingOptions(
@@ -335,6 +338,7 @@ void TuningConfiguration::applyToCudaMappingOptions(
   options.unrollCopyShared(unrollCopyShared.value());
   options.useReadOnlyCache(useReadOnlyCache.value());
   options.privateDepth(privateDepth.value());
+  options.sharedDepth(sharedDepth.value());
 }
 
 void TuningConfiguration::applyToCpuMappingOptions(

tc/autotuner/parameters.h

@@ -191,6 +191,7 @@ class TuningConfiguration {
   BoolParameter useReadOnlyCache;
   BoolParameter matchLibraryCalls;
   RangeParameter privateDepth;
+  RangeParameter sharedDepth;
 
  private:
   std::vector<std::function<bool(const TuningConfiguration&)>> validators_;
@@ -229,6 +230,7 @@ class TuningParameterFixer {
   llvm::Optional<bool> useReadOnlyCache;
   llvm::Optional<bool> matchLibraryCalls;
   llvm::Optional<uint32_t> privateDepth;
+  llvm::Optional<uint32_t> sharedDepth;
 
   friend class TuningConfiguration;
 };

tc/core/cuda/cuda_mapping_options.cc

@@ -294,6 +294,11 @@ CudaMappingOptions& CudaMappingOptions::privateDepth(uint32_t depth) {
   return *this;
 }
 
+CudaMappingOptions& CudaMappingOptions::sharedDepth(uint32_t depth) {
+  ownedProto_.set_shared_depth(depth);
+  return *this;
+}
+
 CudaMappingOptions& CudaMappingOptions::mapToThreads(
     const std::string& commaSeparatedSizes) {
   auto sizes = parseCommaSeparatedIntegers<uint64_t>(commaSeparatedSizes);

tc/core/cuda/cuda_mapping_options.h

@@ -196,6 +196,7 @@ class CudaMappingOptions {
   CudaMappingOptions& unrollCopyShared(bool b);
   CudaMappingOptions& useReadOnlyCache(bool b);
   CudaMappingOptions& privateDepth(uint32_t depth);
+  CudaMappingOptions& sharedDepth(uint32_t depth);
   ///@}
 
   /// Static constructors for predefined strategies.

tc/core/cuda/cuda_mapping_options_cpp_printer.cc

@@ -39,6 +39,7 @@ CudaMappingOptionsCppPrinter& operator<<(
         "maxSharedMemory", cudaOptions.proto().max_shared_memory());
   }
   prn.printValueOption("privateDepth", cudaOptions.proto().private_depth());
+  prn.printValueOption("sharedDepth", cudaOptions.proto().shared_depth());
   prn.endStmt();
   return prn;
 }

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -1049,8 +1049,9 @@ std::unique_ptr<MappedScop> MappedScop::makeWithOuterBlockInnerThreadStrategy(
   LOG_IF(INFO, FLAGS_debug_tc_mapper) << "After mapping to blocks:" << std::endl
                                       << *mappedScop->schedule();
 
-  // 8. Promote to shared memory below the loops mapped to blocks.
-  // This may split the outer band, so find the new outer band after promotion.
+  // 8. Promote to shared memory.
+  // If shared promotion depth is specified in the mapping options, use the
+  // specified value.  Otherwise, promote below the loops mapped to blocks.
   if (cudaOptions.proto().use_shared_memory()) {
     size_t sharedMemorySize = cudaOptions.proto().has_max_shared_memory()
         ? cudaOptions.proto().max_shared_memory()
@@ -1069,29 +1070,24 @@ std::unique_ptr<MappedScop> MappedScop::makeWithOuterBlockInnerThreadStrategy(
       sharedMemorySize -= reductionMemoryRequirement;
     }
 
-    auto band = outerBand->as<ScheduleTreeBand>();
-    LOG_IF(WARNING, FLAGS_debug_tc_mapper && band->nMember() == 0)
-        << "Aborting memory promotion because outer band has 0 members (NYI)";
-    if (band->nMember() > 0 && sharedMemorySize > 0) {
+    if (sharedMemorySize > 0) {
       LOG_IF(
           WARNING,
           cudaOptions.proto().unroll_copy_shared() &&
               !generic.proto.has_unroll())
           << "requested to unroll copies to shared memory without providing the unroll size";
 
-      promoteGreedilyAtDepth(
+      auto depth = cudaOptions.proto().has_shared_depth()
+          ? cudaOptions.proto().shared_depth()
+          : std::min(
+                outerBand->as<ScheduleTreeBand>()->nOuterCoincident(),
+                mappedScop->numBlocks.view.size());
+      promoteToSharedAtDepth(
           *mappedScop,
-          std::min(band->nOuterCoincident(), mappedScop->numBlocks.view.size()),
+          depth,
           sharedMemorySize,
           cudaOptions.proto().unroll_copy_shared() &&
               generic.proto.has_unroll());
-
-      auto bands = ScheduleTree::collectDFSPreorder(
-          scop->scheduleRoot(), ScheduleTreeType::Band);
-      if (bands.size() == 0) { // Sanity check.
-        throw NoBandsException("no bands after promotion");
-      }
-      outerBand = bands[0];
     }
   }