partialScheduleMupa: do not assume the given node is a band

The original implementation assumed that the given node is a band and
used a CHECK macro to enforce that.  In practice, it may be useful to
call partialScheduleMupa to obtain the partial schedule including the
schedule of the current node if it is a band, and the prefix schedule
otherwise.  Replace the CHECK by an early return of the result of
prefixScheduleMupa.

Author: ftynse

tc/core/polyhedral/schedule_transforms.cc

@@ -486,9 +486,9 @@ isl::multi_union_pw_aff prefixScheduleMupa(
 isl::multi_union_pw_aff partialScheduleMupa(
     const detail::ScheduleTree* root,
     const detail::ScheduleTree* tree) {
+  auto prefix = prefixScheduleMupa(root, tree);
   auto band = tree->elemAs<ScheduleTreeElemBand>();
-  TC_CHECK(band);
-  return prefixScheduleMupa(root, tree).flat_range_product(band->mupa_);
+  return band ? prefix.flat_range_product(band->mupa_) : prefix;
 }
 
 void updateTopLevelContext(detail::ScheduleTree* root, isl::set context) {