Don't disable ExprSimplify in shouldUseZero (#3587)

Just a minor fix
NVIDIA · Dec 13, 2024 · 230f633 · 230f633
1 parent 79a088a
commit 230f633
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/csrc/id_model/utils.h b/csrc/id_model/utils.h
@@ -159,6 +159,12 @@ inline bool shouldUseZeroIndex(
   // Trivial loop
   auto promotion_id =
       getLoopPromotion(loop_group->front()->as<IterDomain>(), id_model);
+
+  // ExprSimplify should be disabled here as it would fail to
+  // recognize size-one IterDomain.
+  DisableOptionsGuard options_guard;
+  DisableOptionsGuard::getCurOptions().unset(DisableOption::ExprSimplify);
+
   if (promotion_id->isBroadcast() ||
       simplifyExpr(promotion_id->extent())->isOneInt()) {
     return true;