Merge pull request #570 from facebookresearch/pr/tuple

use tuple_id related convenience functions

Author: skimo-openhub

tc/core/polyhedral/cuda/codegen.cc

@@ -438,7 +438,7 @@ void emitRegisterAccess(
 void emitGlobalAccess(
     isl::multi_pw_aff access,
     const CodegenStatementContext& context) {
-  LdgWrapper ldgWrapper(context, access.get_tuple_id(isl::dim_type::out));
+  LdgWrapper ldgWrapper(context, access.get_range_tuple_id());
   emitAccess(access, context);
 }
 } // namespace
@@ -674,7 +674,7 @@ void emitMappedTensorAccess(
   auto access =
       makeMultiAffAccess(tensorId, subscripts, context); // MA :: D -> O
   auto promotion = promotionInfo.group->promotion(); // MA :: [S -> O] -> P
-  promotion = promotion.set_tuple_id(isl::dim_type::out, promotionInfo.groupId);
+  promotion = promotion.set_range_tuple_id(promotionInfo.groupId);
   auto iteratorMap = context.iteratorMap(); // PMA :: A -> D
   auto schedule =
       isl::map::from_union_map(promotionInfo.outerSchedule.intersect_domain(

tc/core/polyhedral/cuda/codegen.h

@@ -114,7 +114,7 @@ struct CodegenStatementContext : CodegenContext {
     return this->nodeInfoMap.at(astNodeId).build;
   }
   isl::id statementId() const {
-    return this->iteratorMap().get_tuple_id(isl::dim_type::out);
+    return this->iteratorMap().get_range_tuple_id();
   }
   isl::set domain() const {
     return isl::map::from(this->iteratorMap()).range();

tc/core/polyhedral/memory_promotion.cc

@@ -87,8 +87,8 @@ std::unique_ptr<TensorReferenceGroup> TensorReferenceGroup::makeSingleton(
     isl::map scopedAccess,
     AccessType type) {
   auto ref = std::unique_ptr<TensorReference>(new TensorReference);
-  auto refId = scopedAccess.get_space().domain().unwrap().get_tuple_id(
-      isl::dim_type::out);
+  auto refId =
+      scopedAccess.get_space().domain().unwrap().get_map_range_tuple_id();
   scopedAccess = scopedAccess.domain_factor_domain();
   ref->originalAccess = originalAccess.domain_factor_domain();
   ref->scopedAccess = scopedAccess;
@@ -306,7 +306,7 @@ void addSingletonReferenceGroups(
       continue;
     }
 
-    auto tensorId = a.get_tuple_id(isl::dim_type::out);
+    auto tensorId = a.get_range_tuple_id();
     if (unapproximatable.count(tensorId) != 0) {
       continue;
     }
@@ -474,21 +474,20 @@ ScheduleTree* insertCopiesUnder(
   // Take the set of all tensor elements.
   auto tensorElements = tensorElementsSet(scop, tensorId);
 
-  auto promotion =
-      isl::map(group.promotion()).set_tuple_id(isl::dim_type::out, groupId);
+  auto promotion = isl::map(group.promotion()).set_range_tuple_id(groupId);
   auto promotionSpace = promotion.get_space();
 
   auto identityCopySchedule =
       isl::multi_aff::identity(promotionSpace.range().map_from_set());
-  identityCopySchedule =
-      identityCopySchedule.pullback(isl::multi_aff::range_map(promotionSpace));
   // Only iterate over significant tensor dimensions.
   auto decl = scop.promotedDecl(groupId);
   identityCopySchedule = dropDummyTensorDimensions(identityCopySchedule, decl);
-  auto readSchedule = isl::multi_union_pw_aff(
-      identityCopySchedule.set_tuple_id(isl::dim_type::in, readId));
-  auto writeSchedule = isl::multi_union_pw_aff(
-      identityCopySchedule.set_tuple_id(isl::dim_type::in, writeId));
+  auto readSpace = promotionSpace.wrap().set_set_tuple_id(readId);
+  auto writeSpace = promotionSpace.wrap().set_set_tuple_id(writeId);
+  auto readSchedule = isl::multi_union_pw_aff(identityCopySchedule.pullback(
+      isl::multi_aff::wrapped_range_map(readSpace)));
+  auto writeSchedule = isl::multi_union_pw_aff(identityCopySchedule.pullback(
+      isl::multi_aff::wrapped_range_map(writeSpace)));
 
   auto readBandNode = ScheduleTree::makeBand(readSchedule);
   auto writeBandNode = ScheduleTree::makeBand(writeSchedule);
@@ -508,18 +507,17 @@ ScheduleTree* insertCopiesUnder(
   auto promotedFootprint = group.promotedFootprint().set_tuple_id(groupId);
   auto scheduleUniverse =
       isl::set::universe(promotionSpace.domain().unwrap().domain());
-  auto arrayId =
-      promotionSpace.domain().unwrap().get_tuple_id(isl::dim_type::out);
+  auto arrayId = promotionSpace.domain().unwrap().get_map_range_tuple_id();
   auto approximatedRead =
       group.approximateScopedAccesses().intersect_range(tensorElements).wrap();
   approximatedRead = approximatedRead.product(promotedFootprint);
-  auto readExtension = extension.intersect_range(approximatedRead)
-                           .set_tuple_id(isl::dim_type::out, readId);
+  auto readExtension =
+      extension.intersect_range(approximatedRead).set_range_tuple_id(readId);
   auto writtenElements =
       group.scopedWrites().intersect_range(tensorElements).wrap();
   writtenElements = writtenElements.product(promotedFootprint);
-  auto writeExtension = extension.intersect_range(writtenElements)
-                            .set_tuple_id(isl::dim_type::out, writeId);
+  auto writeExtension =
+      extension.intersect_range(writtenElements).set_range_tuple_id(writeId);
 
   auto readFilterNode = ScheduleTree::makeFilter(
       isl::set::universe(readExtension.get_space().range()),