From a9011b600bbe8346fc9f89b568aced552d92357a Mon Sep 17 00:00:00 2001 From: Naoya Maruyama Date: Fri, 19 Dec 2025 18:39:38 -0800 Subject: [PATCH 1/4] Enable TensorIndexer with all C++ tests --- tests/cpp/utils.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/cpp/utils.cpp b/tests/cpp/utils.cpp index 105e30513eb..b7d7093cfc6 100644 --- a/tests/cpp/utils.cpp +++ b/tests/cpp/utils.cpp @@ -56,6 +56,7 @@ void NVFuserTest::SetUp() { if (!deviceMajorMinorCheck(6)) { GTEST_SKIP() << "skipping tests on pre-PASCAL GPUs"; } + EnableOptionsGuard::getCurOptions().set(EnableOption::IdModel, {"all"}); } NVFuserTest::~NVFuserTest() { From 4af15b6cad8b20ffc4d92884feb518d11efce4fb Mon Sep 17 00:00:00 2001 From: Naoya Maruyama Date: Fri, 19 Dec 2025 21:30:10 -0800 Subject: [PATCH 2/4] Enable TensorIndexer --- tests/cpp/test_matmul_scheduler.cpp | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/cpp/test_matmul_scheduler.cpp b/tests/cpp/test_matmul_scheduler.cpp index 78f057329c2..1f5a7c83a58 100644 --- a/tests/cpp/test_matmul_scheduler.cpp +++ b/tests/cpp/test_matmul_scheduler.cpp @@ -30,6 +30,11 @@ class MatmulSchedulerTest : public NVFuserTest { protected: MatmulSchedulerTest() : optimization_guard_(false) {} + void SetUp() override { + NVFuserTest::SetUp(); + EnableOptionsGuard::getCurOptions().set(EnableOption::IdModel, {"all"}); + } + private: // Allocation order set by the pass breaks matmul tests // see issue https://github.com/NVIDIA/Fuser/issues/1810 @@ -2482,6 +2487,11 @@ class MatmulSchedulerPluginTest : public NVFuserTest { MatmulSchedulerPluginTest() : optimization_guard_(false), factory_guard_(testConfigFactory) {} + void SetUp() override { + NVFuserTest::SetUp(); + EnableOptionsGuard::getCurOptions().set(EnableOption::IdModel, {"all"}); + } + private: // Allocation order set by the pass breaks matmul tests // see issue https://github.com/NVIDIA/Fuser/issues/1810 @@ -2917,6 +2927,11 @@ class AllocationDomainTest mparams.circular_buffer_options.smem_circular_buffer_stage = 4; } + void SetUp() override { + NVFuserFixtureParamTest::SetUp(); + EnableOptionsGuard::getCurOptions().set(EnableOption::IdModel, {"all"}); + } + std::pair getInputTVs( int M, int N, @@ -3365,6 +3380,8 @@ class HopperPlusMatmulSchedulerTest mparams.circular_buffer_options.circular_buffer_smem_write = true; mparams.circular_buffer_options.circular_buffer_smem_read = true; mparams.circular_buffer_options.smem_circular_buffer_stage = 2; + + EnableOptionsGuard::getCurOptions().set(EnableOption::IdModel, {"all"}); } void TearDown() { From df46dac9c91b6ca12712a533c4d4c60fecb4cc18 Mon Sep 17 00:00:00 2001 From: Naoya Maruyama Date: Fri, 19 Dec 2025 21:30:16 -0800 Subject: [PATCH 3/4] fix --- csrc/id_model/id_model.cpp | 38 ++++++++++++++++++++++---------------- 1 file changed, 22 insertions(+), 16 deletions(-) diff --git a/csrc/id_model/id_model.cpp b/csrc/id_model/id_model.cpp index 7064b978ce6..9b48d485fbd 100644 --- a/csrc/id_model/id_model.cpp +++ b/csrc/id_model/id_model.cpp @@ -1345,6 +1345,22 @@ void IdModel::allocateLoopIndexVariables() { ParallelType ptype = getParallelType(loop_group); + if (GpuLower::current()->circularBufferInfo().isCircularBufferedIterDomain( + loop_group->front()->as())) { + // Allocate index variable for each stage of the circular + // buffered loop. + auto indices = std::make_unique(); + for (auto i : + arange(static_cast(CircularBufferLoopStage::EndOfStages))) { + indices->emplace( + static_cast(i), + IrBuilder::create(DataType::Index)); + } + circular_buffered_loop_index_variable_map_[loop_group] = + std::move(indices); + continue; + } + Val* loop_index = nullptr; // TODO: Cleanup needed. ir_utils::isMemoryPartitionedAcross @@ -1363,22 +1379,6 @@ void IdModel::allocateLoopIndexVariables() { continue; } - if (GpuLower::current()->circularBufferInfo().isCircularBufferedIterDomain( - loop_group->front()->as())) { - // Allocate index variable for each stage of the circular - // buffered loop. - auto indices = std::make_unique(); - for (auto i : - arange(static_cast(CircularBufferLoopStage::EndOfStages))) { - indices->emplace( - static_cast(i), - IrBuilder::create(DataType::Index)); - } - circular_buffered_loop_index_variable_map_[loop_group] = - std::move(indices); - continue; - } - // If enabled, allocate own indices. Otherwise, use the one // generated for ComputeAtMap for compatibility with the legacy // indexing @@ -1430,6 +1430,12 @@ Val* IdModel::getLoopIndexVariable( // stage defined, and we just default to using the main stage index. circular_buffer_loop_stage = CircularBufferLoopStage::Main; } + NVF_ERROR( + circular_buffered_loop_index_variable_map_.contains(loop_group), + "Failed to find circular buffer index var for: ", + nvfuser::toString(loop_group), + ", ", + loop_group->front()->toString()); return circular_buffered_loop_index_variable_map_.at(loop_group) ->at(circular_buffer_loop_stage); } else { From 8e9bd5f93cddb6fafb3fa4dcf85fb172fe112d87 Mon Sep 17 00:00:00 2001 From: Naoya Maruyama Date: Fri, 19 Dec 2025 21:53:00 -0800 Subject: [PATCH 4/4] Enable TensorIndexer with the stream tests Loop index fix --- csrc/id_model/id_model.cpp | 2 +- tests/cpp/test_stream.cpp | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/csrc/id_model/id_model.cpp b/csrc/id_model/id_model.cpp index 7064b978ce6..c75ff672ffa 100644 --- a/csrc/id_model/id_model.cpp +++ b/csrc/id_model/id_model.cpp @@ -1354,7 +1354,7 @@ void IdModel::allocateLoopIndexVariables() { if (shouldUseZeroIndex(loop_group, *this) || isParallelTypeDeviceDim(ptype)) { loop_index = fusion_->zeroVal(); - } else if (isParallelTypeThread(ptype)) { + } else if (isParallelTypeThread(ptype) || ptype == ParallelType::Stream) { loop_index = NamedScalar::getParallelIndex(ptype); } diff --git a/tests/cpp/test_stream.cpp b/tests/cpp/test_stream.cpp index 49a4653c69c..ce28bf77294 100644 --- a/tests/cpp/test_stream.cpp +++ b/tests/cpp/test_stream.cpp @@ -28,6 +28,7 @@ class StreamTest : public NVFuserTest { public: StreamTest() { EnableOptionsGuard::getCurOptions().set(EnableOption::HostIrLowering); + EnableOptionsGuard::getCurOptions().set(EnableOption::IdModel, {"all"}); } };