Merge pull request #434 from nicolasvasilache/pr/flag-unroll-factor

Add a flag to control max unrolling during autotuning

Author: ftynse

tc/autotuner/autotuner-inl.h

@@ -399,7 +399,8 @@ void setupTuningParameters(
   configuration.tilingParams.setRange(nTilesDim, tileRange);
   configuration.blockParams.setRange(range, "b");
   configuration.gridParams.setRange(range, "g");
-  configuration.unrollFactor = RangeParameter({1, 2, 4, 8, 16, 32}, "unroll");
+  configuration.unrollFactor =
+      RangeParameter(powers2(FLAGS_tuner_max_unroll_size), "unroll");
 }
 } // namespace
 

tc/autotuner/parameters.cc

@@ -281,7 +281,11 @@ void TuningConfiguration::fromMappingOptions(
       options.proto.fix_parameters_before_scheduling());
   tilingParams.fromMappingOptions(options.tiling);
   unrollFactor.selectFromValue(
-      (options.proto.has_unroll() ? options.proto.unroll() : 1));
+      (options.proto.has_unroll()
+           ? std::min(
+                 static_cast<uint32_t>(options.proto.unroll()),
+                 FLAGS_tuner_max_unroll_size)
+           : 1));
   tileImperfectlyNested.selectValue(options.proto.tile_imperfectly_nested());
   matchLibraryCalls.selectValue(options.proto.match_library_calls());
 }

tc/autotuner/utils.cc

@@ -27,17 +27,21 @@
 namespace tc {
 namespace autotune {
 
-std::vector<std::size_t> powers2andCeilDivisors(std::size_t val) {
+std::vector<std::size_t> powers2(std::size_t val) {
   auto numPowers = static_cast<std::size_t>(std::ceil(std::log2(val)));
-  // 1. generate `numPowers' powers of 2
   std::vector<std::size_t> res(numPowers + 1);
   std::size_t p = 1;
   std::generate(res.begin(), res.end(), [p]() mutable {
     auto old_p = p;
     p *= 2;
     return old_p;
   });
-  // 2. additionally insert ceil(val / powers2)
+  return res;
+}
+
+std::vector<std::size_t> powers2andCeilDivisors(std::size_t val) {
+  std::vector<std::size_t> res = powers2(val);
+  // Additionally insert ceil(val / powers2)
   res.reserve(res.size() * 2);
   for (std::size_t i = 0, s = res.size(); i < s; ++i) {
     if (res[i] > val) {

tc/autotuner/utils.h

@@ -29,6 +29,9 @@
 namespace tc {
 namespace autotune {
 
+/// Returns all the powers of 2 up to the first one that is larger than val
+std::vector<std::size_t> powers2(std::size_t val);
+
 /// Returns all the powers of 2 up to the first one that is larger than val
 /// and the result of ceil(val/pow2) for each of those powers of 2 (except for
 /// the larger one)

tc/core/flags.cc

@@ -55,6 +55,10 @@ DEFINE_bool(
     "Print debug spew for experimental schedule_tree");
 
 // Autotuner flags
+DEFINE_uint32(
+    tuner_max_unroll_size,
+    32,
+    "Polyhedral unrolling is expensive, limit to 32 by default");
 DEFINE_uint32(
     tuner_gen_pop_size,
     100,

tc/core/flags.h

@@ -39,6 +39,7 @@ DECLARE_uint32(benchmark_warmup);
 DECLARE_uint32(benchmark_iterations);
 
 // Used in autotuning
+DECLARE_uint32(tuner_max_unroll_size);
 DECLARE_uint32(tuner_gen_pop_size);
 DECLARE_uint32(tuner_gen_crossover_rate);
 DECLARE_uint32(tuner_gen_mutation_rate);