Enable TensorIndexer with the matmul scheduler tests

[naoyam]

No description provided.

[naoyam]

!test --diff

[github-actions]

Description

• Enable TensorIndexer (IdModel) for matmul scheduler tests by adding SetUp methods

• Fix circular buffer index variable allocation order in IdModel

• Add error checking for missing circular buffer index variables

• Enable IdModel option across multiple matmul test classes

Changes walkthrough

	Relevant files
Bug fix	id_model.cpp Fix circular buffer index variable allocation order            csrc/id_model/id_model.cpp Reorder circular buffer check to occur before loop_index check in allocateLoopIndexVariables Add error checking for missing circular buffer index variables in getLoopIndexVariable Improve robustness of circular buffer index variable handling +22/-16
Enhancement	test_matmul_scheduler.cpp Enable TensorIndexer in matmul test classes                            tests/cpp/test_matmul_scheduler.cpp Add SetUp method to MatmulSchedulerTest that enables IdModel option Add SetUp method to MatmulSchedulerPluginTest that enables IdModel option Add SetUp method to AllocationDomainTest that enables IdModel option Add EnableOptionsGuard line to HopperPlusMatmulSchedulerTest setup +17/-0

PR Reviewer Guide

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review

Code Movement Safety The circular buffer handling code has been moved earlier in the allocateLoopIndexVariables() function. While this appears to be a logical reordering (checking circular buffers before other loop index allocation), the reviewer should verify that this doesn't break any existing functionality or change the intended execution flow. if (GpuLower::current()->circularBufferInfo().isCircularBufferedIterDomain( loop_group->front()->as<IterDomain>())) { // Allocate index variable for each stage of the circular // buffered loop. auto indices = std::make_unique<CircularBufferIndices>(); for (auto i : arange(static_cast<int>(CircularBufferLoopStage::EndOfStages))) { indices->emplace( static_cast<CircularBufferLoopStage>(i), IrBuilder::create<Val>(DataType::Index)); } circular_buffered_loop_index_variable_map_[loop_group] = std::move(indices); continue; } Error Handling Enhancement A new NVF_ERROR check has been added to ensure circular buffer index variables are found. This is good defensive programming, but the reviewer should confirm this doesn't cause false positives in edge cases where circular buffer index variables might legitimately not exist. 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());

[greptile-apps]

Greptile Summary

This PR enables TensorIndexer (IdModel) for matmul scheduler tests and fixes a bug in circular buffer index allocation ordering.

Key changes:

• Bug fix in csrc/id_model/id_model.cpp: Moved circular buffer index allocation check earlier in allocateLoopIndexVariables() to ensure circular buffered loops always have their indices allocated before the function potentially returns early for other loop types. Added an NVF_ERROR check in getLoopIndexVariable() to catch cases where circular buffer indices are missing.
• Test enablement in tests/cpp/test_matmul_scheduler.cpp: Added SetUp() methods to four test classes (MatmulSchedulerTest, MatmulSchedulerPluginTest, AllocationDomainTest, HopperPlusMatmulSchedulerTest) to enable IdModel with {"all"} option.

The ordering fix prevents a potential bug where circular buffered loops could skip index allocation if they matched certain conditions (zero index, device dim parallelization, or thread parallelization) that would cause an early return.

Confidence Score: 4/5

• This PR is safe to merge with low risk - it fixes a potential indexing bug and enables tests
• The changes are well-structured: the circular buffer fix prevents a real bug by ensuring proper initialization order, and the test enablement is straightforward. Score is 4 rather than 5 due to the lack of explicit test coverage for the ordering bug fix itself
• No files require special attention

Important Files Changed

Filename	Overview
csrc/id_model/id_model.cpp	Moved circular buffer index allocation before other index handling to ensure indices are always created, added error check for missing indices
tests/cpp/test_matmul_scheduler.cpp	Enabled TensorIndexer (IdModel) for all matmul scheduler test classes via SetUp() methods

Sequence Diagram

sequenceDiagram
    participant Test as Matmul Scheduler Tests
    participant SetUp as Test SetUp()
    participant IdModel as IdModel
    participant Allocate as allocateLoopIndexVariables()
    participant Get as getLoopIndexVariable()

    Test->>SetUp: Initialize test fixture
    SetUp->>SetUp: Enable IdModel with {"all"}
    Note over SetUp: TensorIndexer enabled for all tests

    Test->>IdModel: Execute matmul scheduler test
    IdModel->>Allocate: Allocate loop index variables
    
    alt Circular Buffered IterDomain
        Allocate->>Allocate: Check isCircularBufferedIterDomain()
        Allocate->>Allocate: Create CircularBufferIndices for all stages
        Allocate->>Allocate: Store in circular_buffered_loop_index_variable_map_
        Note over Allocate: Early return - indices allocated
    else Non-circular Buffer Loop
        Allocate->>Allocate: Check parallel type & conditions
        Allocate->>Allocate: Allocate standard loop index
        Allocate->>Allocate: Store in loop_index_variable_map_
    end

    Test->>Get: Request loop index variable
    Get->>Get: Check if circular buffered
    alt Circular Buffered
        Get->>Get: Verify index exists (NVF_ERROR)
        Get->>Get: Return index for requested stage
    else Standard Loop
        Get->>Get: Return from loop_index_variable_map_
    end
    Get-->>Test: Return loop index

Loading

[greptile-apps]

Greptile's behavior is changing!

From now on, if a review finishes with no comments, we will not post an additional "statistics" comment to confirm that our review found nothing to comment on. However, you can confirm that we reviewed your changes in the status check section.

_{This feature can be toggled off in your Code Review Settings by deselecting "Create a status check for each PR".}