Merge pull request #405 from facebookresearch/pr/space

simplify space manipulations

Author: ftynse

tc/core/halide2isl.cc

@@ -258,17 +258,15 @@ isl::map extractAccess(
 
   isl::space domainSpace = domain.get_space();
   isl::space paramSpace = domainSpace.params();
-
-  isl::space rangeSpace = paramSpace.add_dims(isl::dim_type::set, args.size());
-
-  rangeSpace = rangeSpace.set_tuple_name(isl::dim_type::set, tensor);
+  isl::id tensorID(paramSpace.get_ctx(), tensor);
+  auto rangeSpace = paramSpace.named_set_from_params_id(tensorID, args.size());
 
   // Add a tag to the domain space so that we can maintain a mapping
   // between each access in the IR and the reads/writes maps.
   std::string tag = "__tc_ref_" + std::to_string(accesses->size());
   isl::id tagID(domain.get_ctx(), tag);
   accesses->emplace(op, tagID);
-  isl::space tagSpace = paramSpace.set_tuple_name(isl::dim_type::set, tag);
+  isl::space tagSpace = paramSpace.named_set_from_params_id(tagID, 0);
   domainSpace = domainSpace.product(tagSpace);
 
   // Start with a totally unconstrained relation - every point in

tc/core/polyhedral/cuda/codegen.cc

@@ -481,9 +481,8 @@ isl::multi_aff makeMultiAffAccess(
   CHECK_NE(subscripts.size(), 0u) << "cannot build subscript aff for a scalar";
 
   auto domainSpace = findDomainSpaceById(context);
-  auto tensorSpace = domainSpace.params().set_from_params().add_dims(
-      isl::dim_type::set, subscripts.size());
-  tensorSpace = tensorSpace.set_tuple_id(isl::dim_type::set, tensorId);
+  auto tensorSpace = domainSpace.params().named_set_from_params_id(
+      tensorId, subscripts.size());
   auto space = domainSpace.map_from_domain_and_range(tensorSpace);
 
   auto ma = isl::multi_aff::zero(space);

tc/core/polyhedral/memory_promotion.cc

@@ -98,6 +98,14 @@ ScopedFootprint outputRanges(isl::map access) {
   }
   return footprint;
 }
+
+// Given a set space, construct a map space with the input as domain and
+// a range of the given size.
+isl::space add_range(isl::space space, unsigned dim) {
+  auto range = space.params().unnamed_set_from_params(dim);
+  return space.map_from_domain_and_range(range);
+}
+
 } // namespace
 
 // Access has the shape :: [D -> ref] -> O
@@ -128,8 +136,7 @@ std::unique_ptr<TensorReferenceGroup> TensorReferenceGroup::makeSingleton(
 }
 
 isl::set ScopedFootprint::footprint(isl::set domain) const {
-  auto space = domain.get_space().from_domain();
-  space = space.add_dims(isl::dim_type::out, size());
+  auto space = add_range(domain.get_space(), size());
   auto accessed = isl::map::universe(space).intersect_domain(domain);
   auto lspace = isl::local_space(accessed.get_space().range());
 
@@ -147,8 +154,7 @@ isl::multi_aff ScopedFootprint::lowerBounds() const {
   if (size() == 0) {
     throw promotion::PromotionNYI("promotion for scalars");
   }
-  auto space = at(0).lowerBound.get_space();
-  space = space.add_dims(isl::dim_type::out, size() - 1);
+  auto space = add_range(at(0).lowerBound.get_space().domain(), size());
   auto ma = isl::multi_aff::zero(space);
 
   int i = 0;
@@ -402,8 +408,7 @@ isl::set tensorElementsSet(const Scop& scop, isl::id tensorId) {
   auto halideParameter = scop.findArgument(tensorId).parameter();
   auto space = scop.domain().get_space().params();
   auto nDim = halideParameter.dimensions();
-  space = space.add_dims(isl::dim_type::set, nDim)
-              .set_tuple_id(isl::dim_type::set, tensorId);
+  space = space.named_set_from_params_id(tensorId, nDim);
 
   auto tensorElements = isl::set::universe(space);
   for (int i = 0; i < nDim; ++i) {
@@ -443,7 +448,7 @@ isl::multi_aff dropDummyTensorDimensions(
     }
   }
 
-  space = space.from_domain().add_dims(isl::dim_type::out, list.n());
+  space = add_range(space, list.n());
   return isl::multi_aff(space, list);
 }
 } // namespace

tc/core/polyhedral/schedule_transforms.cc

@@ -559,8 +559,7 @@ detail::ScheduleTree* insertEmptyExtensionAbove(
 isl::map labelExtension(ScheduleTree* root, ScheduleTree* tree, isl::id id) {
   auto prefix = prefixScheduleMupa(root, tree);
   auto scheduleSpace = prefix.get_space();
-  auto space = scheduleSpace.params().set_from_params().set_tuple_id(
-      isl::dim_type::set, id);
+  auto space = scheduleSpace.params().named_set_from_params_id(id, 0);
   auto extensionSpace = scheduleSpace.map_from_domain_and_range(space);
   return isl::map::universe(extensionSpace);
 }

tc/core/polyhedral/scop.cc

@@ -525,15 +525,13 @@ isl::aff Scop::makeIslAffFromStmtExpr(
     const Halide::Expr& e) const {
   auto ctx = stmtId.get_ctx();
   auto iterators = halide.iterators.at(stmtId);
-  auto space = paramSpace.set_from_params();
-  space = space.add_dims(isl::dim_type::set, iterators.size());
+  auto space = paramSpace.named_set_from_params_id(stmtId, iterators.size());
   // Set the names of the set dimensions of "space" for use
   // by halide2isl::makeIslAffFromExpr.
   for (size_t i = 0; i < iterators.size(); ++i) {
     isl::id id(ctx, iterators[i]);
     space = space.set_dim_id(isl::dim_type::set, i, id);
   }
-  space = space.set_tuple_id(isl::dim_type::set, stmtId);
   return halide2isl::makeIslAffFromExpr(space, e);
 }