From 1b121d636e7340a0171e770c3fa8b34091e8a243 Mon Sep 17 00:00:00 2001
From: "Yu, Guangye" <guangye.yu@intel.com>
Date: Fri, 17 Oct 2025 17:16:48 +0000
Subject: [PATCH] Fix AllocatorConfig parse roundup division bug (#165304)

* #165288
Pull Request resolved: https://github.com/pytorch/pytorch/pull/165304
Approved by: https://github.com/albanD
ghstack dependencies: #165288, #165289, #165291, #165298
---
 c10/core/AllocatorConfig.cpp           | 3 ++-
 c10/test/core/AllocatorConfig_test.cpp | 8 ++++----
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/c10/core/AllocatorConfig.cpp b/c10/core/AllocatorConfig.cpp
index 26058cf389a5..4d8bc9b8746e 100644
--- a/c10/core/AllocatorConfig.cpp
+++ b/c10/core/AllocatorConfig.cpp
@@ -128,7 +128,8 @@ size_t AcceleratorAllocatorConfig::parseRoundUpPower2Divisions(
         std::fill(
             std::next(
                 roundup_power2_divisions_.begin(),
-                static_cast<std::vector<size_t>::difference_type>(last_index)),
+                static_cast<std::vector<size_t>::difference_type>(
+                    last_index + 1)),
             roundup_power2_divisions_.end(),
             value);
       } else {
diff --git a/c10/test/core/AllocatorConfig_test.cpp b/c10/test/core/AllocatorConfig_test.cpp
index 5f6804639067..049d9921cd5e 100644
--- a/c10/test/core/AllocatorConfig_test.cpp
+++ b/c10/test/core/AllocatorConfig_test.cpp
@@ -67,8 +67,8 @@ TEST(AllocatorConfigTest, allocator_config_test) {
   EXPECT_EQ(AcceleratorAllocatorConfig::roundup_power2_divisions(128 * kMB), 2);
   EXPECT_EQ(AcceleratorAllocatorConfig::roundup_power2_divisions(256 * kMB), 4);
   EXPECT_EQ(AcceleratorAllocatorConfig::roundup_power2_divisions(512 * kMB), 2);
-  // EXPECT_EQ(
-  //     AcceleratorAllocatorConfig::roundup_power2_divisions(1024 * kMB), 4);
+  EXPECT_EQ(
+      AcceleratorAllocatorConfig::roundup_power2_divisions(1024 * kMB), 4);
   EXPECT_EQ(
       AcceleratorAllocatorConfig::roundup_power2_divisions(2048 * kMB), 1);
   EXPECT_EQ(
@@ -101,8 +101,8 @@ TEST(AllocatorConfigTest, allocator_config_test) {
   EXPECT_EQ(AcceleratorAllocatorConfig::roundup_power2_divisions(512 * kMB), 1);
   EXPECT_EQ(
       AcceleratorAllocatorConfig::roundup_power2_divisions(1024 * kMB), 0);
-  // EXPECT_EQ(
-  //     AcceleratorAllocatorConfig::roundup_power2_divisions(2048 * kMB), 8);
+  EXPECT_EQ(
+      AcceleratorAllocatorConfig::roundup_power2_divisions(2048 * kMB), 8);
   EXPECT_EQ(
       AcceleratorAllocatorConfig::roundup_power2_divisions(4096 * kMB), 2);