Merge pull request #574 from facebookresearch/pr/clean-up

clean up  updateTopLevelContext

Author: skimo-openhub

tc/core/polyhedral/schedule_transforms.cc

@@ -286,11 +286,10 @@ ScheduleTree* insertTopLevelEmptyBand(ScheduleTree* root) {
 
 void updateTopLevelContext(detail::ScheduleTree* root, isl::set context) {
   if (!matchOne(tc::polyhedral::domain(tc::polyhedral::context(any())), root)) {
-    root->appendChild(ScheduleTree::makeContext(
-        isl::set::universe(context.get_space()), root->detachChildren()));
+    root->appendChild(
+        ScheduleTree::makeContext(context, root->detachChildren()));
   }
-  auto contextElem = const_cast<detail::ScheduleTreeContext*>(
-      root->child({0})->as<detail::ScheduleTreeContext>());
+  auto contextElem = root->child({0})->as<detail::ScheduleTreeContext>();
   TC_CHECK(contextElem) << "Expected domain(context(any()))";
   contextElem->context_ = contextElem->context_ & context;
 }

tc/core/polyhedral/schedule_transforms.h

@@ -109,7 +109,7 @@ detail::ScheduleTree* insertTopLevelEmptyBand(detail::ScheduleTree* root);
 
 // Update the top-level context node by intersecting it with "context".  The
 // top-level context node must be located directly under the root of the tree.
-// If there is no such node, insert one with universe context first.
+// If there is no such node, insert one first.
 void updateTopLevelContext(detail::ScheduleTree* root, isl::set context);
 
 // In a tree starting at "root", insert a sequence node with