Merge pull request #265 from facebookresearch/pr/param_order

remove code that depends on order of parameters in isl objects

Author: nicolasvasilache

include/tc/core/polyhedral/scop.h

@@ -118,22 +118,12 @@ struct Scop {
     writes = writes.intersect_params(globalParameterContext);
   }
 
-  // Returns a set that specializes (all) the scop's parameter space to the
-  // integer values passed to the function.
-  // WARNING: this version relies on parameter ordering, be sure you know what
-  // you are doing.
-  template <typename T>
-  isl::set makeContext(const std::vector<T>& sizes = std::vector<T>()) const {
-    auto s = domain().get_space().params();
-    return makeSpecializationSet(s, sizes);
-  }
-
-  // Returns a set that specializes the (positional) scop's subset of
+  // Returns a set that specializes the named scop's subset of
   // parameter space to the integer values passed to the function.
   template <typename T>
   isl::set makeContext(
-      const std::unordered_map<int, T>& sizes =
-          std::unordered_map<int, T>()) const {
+      const std::unordered_map<std::string, T>& sizes =
+          std::unordered_map<std::string, T>()) const {
     auto s = domain().get_space().params();
     return makeSpecializationSet(s, sizes);
   }
@@ -142,8 +132,7 @@ struct Scop {
   // parameter space to the integer values passed to the function.
   template <typename T>
   isl::set makeContext(
-      const std::unordered_map<std::string, T>& sizes =
-          std::unordered_map<std::string, T>()) const {
+      std::initializer_list<std::pair<const std::string, T>> sizes) {
     auto s = domain().get_space().params();
     return makeSpecializationSet(s, sizes);
   }

include/tc/external/detail/islpp.h

@@ -317,17 +317,19 @@ inline isl::set makeParameterContext(
 
 // Given a space and values for parameters, this function creates the set
 // that ties the space parameter to the values.
-// This assumes space.dim(isl::dim_type::param) == paramValues.size()
 //
 template <typename T>
 inline isl::set makeSpecializationSet(
     isl::space space,
-    const std::unordered_map<int, T>& paramValues) {
-  CHECK_GE(space.dim(isl::dim_type::param), paramValues.size());
-  auto lspace = isl::local_space(space);
+    const std::unordered_map<std::string, T>& paramValues) {
+  auto ctx = space.get_ctx();
+  for (auto kvp : paramValues) {
+    space = space.add_param(isl::id(ctx, kvp.first));
+  }
   auto set = isl::set::universe(space);
   for (auto kvp : paramValues) {
-    auto affParam = isl::aff(lspace, isl::dim_type::param, kvp.first);
+    auto id = isl::id(ctx, kvp.first);
+    isl::aff affParam(isl::aff::param_on_domain_space(space, id));
     set = set & (isl::aff_set(affParam) == kvp.second);
   }
   return set;
@@ -336,30 +338,20 @@ inline isl::set makeSpecializationSet(
 template <typename T>
 inline isl::set makeSpecializationSet(
     isl::space space,
-    const std::unordered_map<std::string, T>& paramValues) {
-  CHECK_GE(space.dim(isl::dim_type::param), paramValues.size());
-  std::unordered_map<int, T> aux;
-  for (auto kvp : paramValues) {
-    auto pos = space.find_dim_by_name(isl::dim_type::param, kvp.first);
-    CHECK_LE(0, pos) << "No " << kvp.first << " in: " << space;
-    CHECK_EQ(0, aux.count(pos));
-    aux[pos] = kvp.second;
-  }
-  return makeSpecializationSet(space, aux);
+    std::initializer_list<std::pair<const std::string, T>> paramValues) {
+  std::unordered_map<std::string, T> map(paramValues);
+  return makeSpecializationSet(space, map);
 }
 
-// WARNING: this version relies on parameter ordering, be sure you know what
-// you are doing.
 template <typename T>
 inline isl::set makeSpecializationSet(
     isl::space space,
-    const std::vector<T>& paramValues) {
-  CHECK_EQ(space.dim(isl::dim_type::param), paramValues.size());
-  std::unordered_map<int, T> paramValuesMap;
-  for (int i = 0; i < paramValues.size(); ++i) {
-    paramValuesMap[i] = paramValues[i];
+    std::initializer_list<std::pair<isl::id, T>> paramValues) {
+  std::unordered_map<std::string, T> map;
+  for (auto kvp : paramValues) {
+    map.emplace(kvp.first.get_name(), kvp.second);
   }
-  return makeSpecializationSet(space, paramValuesMap);
+  return makeSpecializationSet(space, map);
 }
 
 namespace detail {

test/test_mapper.cc

@@ -333,7 +333,7 @@ TEST_F(PolyhedralMapperTest, MergedContexts) {
   auto scop = PrepareAndJoinBands(makeMatmulTc());
 
   // Unit test claims to use scop->globalParameterContext properly
-  auto context = scop->makeContext(std::vector<int>{64, 64, 64});
+  auto context = scop->makeContext<int>({{"M", 64}, {"N", 64}, {"K", 64}});
   auto& globalParameterContext =
       const_cast<isl::set&>(scop->globalParameterContext);
   globalParameterContext = globalParameterContext.intersect(context);
@@ -349,7 +349,7 @@ TEST_F(PolyhedralMapperTest, FilterMerge) {
   auto schedule = scop->scheduleRoot();
 
   // Unit test claims to use scop->globalParameterContext properly
-  auto context = scop->makeContext(std::vector<int>{64, 64, 64});
+  auto context = scop->makeContext<int>({{"M", 64}, {"N", 64}, {"K", 64}});
   auto& globalParameterContext =
       const_cast<isl::set&>(scop->globalParameterContext);
   globalParameterContext = globalParameterContext.intersect(context);

test/test_mapper_memory_promotion.cc

@@ -275,7 +275,7 @@ def fun(float(N, M) A, float(N, M) B) -> (C) {
     blockSpace = blockSpace.set_dim_id(isl::dim_type::param, 0, BX)
                      .set_dim_id(isl::dim_type::param, 1, BY);
     isl::set blockZero =
-        makeSpecializationSet(blockSpace, std::vector<int>{{0, 0}});
+        isl::makeSpecializationSet<int>(blockSpace, {{BX, 0}, {BY, 0}});
 
     // Must have groups for these tensors, in arbitrary order.
     unordered_set<string> names{"A", "B", "C"};