From 230f633eb995d6690ba67c57117c3b94ea26f25b Mon Sep 17 00:00:00 2001
From: Naoya Maruyama <naoyam@users.noreply.github.com>
Date: Thu, 12 Dec 2024 21:58:54 -0800
Subject: [PATCH] Don't disable ExprSimplify in shouldUseZero (#3587)

Just a minor fix
---
 csrc/id_model/utils.h | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/csrc/id_model/utils.h b/csrc/id_model/utils.h
index 6926ea6e779..c0f73d0f8c3 100644
--- 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;