Enable TensorIndexer with all python_direct tests #5781

naoyam · 2026-01-09T00:40:19Z

No description provided.

github-actions · 2026-01-09T00:41:11Z

Review updated until commit 6fb3cc6

Description

Enable IdModel by default for all python_direct tests
Replace compute-at-map with IdModel for concrete ID mapping
Add helper function getConcreteMappedId for consistent ID resolution
Update test expectations to match new indexing variable naming
Remove id_model disable option and simplify configuration logic

Changes walkthrough

Relevant files

Enhancement

4 files

lower2device.cpp `Always build IdModel without conditional check`	+6/-8
predicate_compute.cpp `Replace compute-at-map with IdModel for ID mapping`	+23/-16
__init__.py `Add id_model as default enable option`	+6/-1
utils.h `Simplify IdModel option detection logic`	+8/-35

Configuration changes

3 files

options.cpp `Remove id_model from disable options`	+0/-1
id_model_options.h `Remove build_id_model configuration option`	+4/-27
options.h `Remove IdModel from disable option enum`	+0/-1

Tests

2 files

test_indexing.cpp `Update test expectations for new variable naming`	+3/-2
test_python_direct.py `Update expected CUDA kernel for new indexing`	+4/-6

PR Reviewer Guide

Here are some key observations to aid the review process:

🧪 PR contains tests
⚡ Recommended focus areas for review
IdModel Always Enabled The IdModel is now unconditionally created instead of being conditional. This is a significant architectural change that should be validated for performance impact and correctness across all use cases. info().set(std::make_unique<IdModel>( fusion_, /build_graphs=/true, /allow_self_mapping=/false, /validate=/false)); info().idModel().validateAndPropagatePType(); API Migration from caMap to idModel Multiple functions have been migrated from using compute-at map (caMap) to using the IdModel API. This is a substantial change that should be thoroughly tested to ensure predicate computation remains correct. IterDomain* getConcreteMappedId(IterDomain* id) { NVF_ERROR(GpuLower::hasCurrent()); return GpuLower::current() ->info() .idModel() .idGraph(IdMappingMode::EXACT) .toGroup(id) ->front() ->as<IterDomain>(); } Default Enable Options The Python API now defaults to enabling "id_model" option. This changes the default behavior for all Python users and should be documented and validated for compatibility. # Add "id_model" as a default enable option default_enable_options = ["id_model"] merged_enable_options = default_enable_options + _enable_options return self.fec.execute( inputs, device=self._get_device_index(device), _enable_options=merged_enable_options,

Test failures (partial, pipeline still running)

(Medium, 4) NVFuser internal assert in ComputeAtMap for stride_order_with_explicit_broadcast (python/direct frontend tests)

Test Name	A100	GB200	Source
tests.python.direct.test_python_frontend.test_stride_order_with_explicit_broadcast[nvfuser_direct_test=eager]	❌	❌
tests.python.direct.test_python_frontend.test_stride_order_with_explicit_broadcast[nvfuser_direct_test=lru_cache]	❌	❌

(Medium, 2) NVFuser CUDA kernel codegen text mismatch in tests.python.direct.test_python_direct.test_fusion_execution_cache

Test Name A100 GB200 Source

tests.python.direct.test_python_direct.test_fusion_execution_cache ❌ ❌

greptile-apps

Greptile Overview

Greptile Summary

Enables the TensorIndexer (IdModel) by default for all Python direct tests by adding "id_model" as a default enable option in FusionDefinition.execute(). Simplifies the C++ option parsing logic in getIdModelEnabledOptions() by switching from an explicit opt-in model (checking for specific arguments like "consumer_index", "producer_index", etc.) to an opt-out model using "predicate_only" and "index_only" flags, removing redundant validation checks.

Confidence Score: 4/5

Safe to merge - clean refactoring that enables IdModel by default for Python tests
The changes are straightforward and follow the pattern established in the previous PR (#5724) that enabled TensorIndexer for C++ tests. The C++ logic simplification removes complex conditional checks and validation code, making it easier to maintain. The Python change simply adds "id_model" to the default enable options list, which is non-breaking since users can still override via _enable_options/_disable_options parameters. No logical errors or edge cases detected.
No files require special attention

Important Files Changed

File Analysis

Filename	Score	Overview
csrc/id_model/utils.h	4/5	Simplifies IdModel option logic from explicit opt-in to opt-out with predicate_only/index_only flags
python/nvfuser_direct/init.py	4/5	Adds "id_model" as default enable option for all FusionDefinition.execute() calls

Sequence Diagram

sequenceDiagram
    participant Python as Python FusionDefinition.execute()
    participant FEC as FusionExecutorCache
    participant CPP as getIdModelEnabledOptions()
    
    Python->>Python: "Add 'id_model' to default_enable_options"
    Python->>FEC: "execute(merged_enable_options)"
    FEC->>CPP: "Check IdModel enable option"
    CPP->>CPP: "Check !predicate_only → Enable Index opts"
    CPP->>CPP: "Check !index_only → Enable Predicate opts"
    CPP->>CPP: "Check both unset → Enable Loop opt"
    CPP-->>FEC: "Return enabled options set"
    FEC-->>Python: "Execute with TensorIndexer enabled"

naoyam · 2026-01-09T00:43:34Z

!test --diff

naoyam · 2026-01-09T06:11:40Z

!test

greptile-apps

Greptile Overview

Greptile Summary

Enables TensorIndexer by default for all python_direct tests by making IdModel always build and removing the disable option. The PR removes the build_id_model_ flag from IdModelOptions and the DisableOption::IdModel enum, unconditionally building IdModel during lowering. Python's FusionDefinition.execute() now adds "id_model" as a default enable option.

Confidence Score: 2/5

Breaking change to IdModel option semantics requires test updates
The PR makes significant changes to how IdModel options are interpreted. The new getIdModelEnabledOptions() function changes from an opt-in model (explicit arguments like "all", "index") to a negation-based model ("predicate_only", "index_only"). This breaks existing usage in tests/python/direct/test_with_id_model_indexer.py which uses NVFUSER_ENABLE=id_model(all), but the "all" argument is no longer recognized. The old code would enable specific options based on explicit arguments, while the new code enables everything by default unless restricted. This semantic change warrants careful testing before merge.
csrc/id_model/utils.h - breaks backward compatibility with option arguments; tests/python/direct/test_with_id_model_indexer.py - needs update for new semantics

Important Files Changed

File Analysis

Filename	Score	Overview
csrc/device_lower/id_model_options.h	4/5	Removes `build_id_model_` flag and related methods, simplifying the options structure
csrc/device_lower/lower2device.cpp	4/5	Unconditionally builds IdModel by removing the conditional check for `buildIdModel()`
csrc/id_model/utils.h	4/5	Simplifies option logic by checking `EnableOption::IdModel` first and using simpler argument checks
python/nvfuser_direct/init.py	4/5	Adds "id_model" as default enable option for all fusion executions

Sequence Diagram

sequenceDiagram
    participant Python as Python FusionDefinition
    participant FEC as FusionExecutorCache
    participant Lower as GpuLower
    participant IdModel as IdModel
    participant Options as IdModelOptions
    
    Python->>Python: "Add 'id_model' to enable options"
    Python->>FEC: "execute(inputs, enable=['id_model'])"
    FEC->>Lower: "analysis(fusion)"
    Lower->>IdModel: "Unconditionally construct IdModel"
    IdModel->>Options: "Initialize IdModelOptions"
    Options->>Options: "Check EnableOption::IdModel"
    alt IdModel not enabled
        Options->>Options: "Return empty options"
    else IdModel enabled
        alt predicate_only arg
            Options->>Options: "Enable predicates only"
        else index_only arg
            Options->>Options: "Enable indexing only"
        else No restrictive args
            Options->>Options: "Enable all options"
        end
    end

greptile-apps · 2026-01-09T06:14:27Z

csrc/id_model/utils.h

 inline std::unordered_set<IdModelEnableOption> getIdModelEnabledOptions() {
+  if (!isOptionEnabled(EnableOption::IdModel)) {
+    return {};
+  }
+
  std::unordered_set<IdModelEnableOption> opts;

-  if (hasEnableOptionArgument(EnableOption::IdModel, "consumer_index") ||
-      hasEnableOptionArgument(EnableOption::IdModel, "index") ||
-      hasEnableOptionArgument(EnableOption::IdModel, "all")) {
+  if (!hasEnableOptionArgument(EnableOption::IdModel, "predicate_only")) {
    opts.insert(IdModelEnableOption::ConsumerIndex);
-  }
-
-  if (hasEnableOptionArgument(EnableOption::IdModel, "producer_index") ||
-      hasEnableOptionArgument(EnableOption::IdModel, "index") ||
-      hasEnableOptionArgument(EnableOption::IdModel, "all")) {
    opts.insert(IdModelEnableOption::ProducerIndex);
  }

-  if (hasEnableOptionArgument(EnableOption::IdModel, "inline_predicate") ||
-      hasEnableOptionArgument(EnableOption::IdModel, "predicate") ||
-      hasEnableOptionArgument(EnableOption::IdModel, "all")) {
+  if (!hasEnableOptionArgument(EnableOption::IdModel, "index_only")) {
    opts.insert(IdModelEnableOption::InlinePredicate);
-  }
-
-  if (hasEnableOptionArgument(EnableOption::IdModel, "unswitch_predicate") ||
-      hasEnableOptionArgument(EnableOption::IdModel, "predicate") ||
-      hasEnableOptionArgument(EnableOption::IdModel, "all")) {
    opts.insert(IdModelEnableOption::UnswitchPredicate);
  }

-  if (hasEnableOptionArgument(EnableOption::IdModel, "loop") ||
-      hasEnableOptionArgument(EnableOption::IdModel, "all")) {
+  if (!hasEnableOptionArgument(EnableOption::IdModel, "predicate_only") &&
+      !hasEnableOptionArgument(EnableOption::IdModel, "index_only")) {
    opts.insert(IdModelEnableOption::Loop);
  }

-  // Loop requires ConsumerIndex, ProducerIndex, InlinePredicate and
-  // UnswitchPredicate
-  if (opts.find(IdModelEnableOption::Loop) != opts.end()) {
-    NVF_ERROR(
-        opts.find(IdModelEnableOption::ConsumerIndex) != opts.end(),
-        "ConsumerIndex required for Loop");
-    NVF_ERROR(
-        opts.find(IdModelEnableOption::ProducerIndex) != opts.end(),
-        "ProducerIndex required for Loop");
-    NVF_ERROR(
-        opts.find(IdModelEnableOption::InlinePredicate) != opts.end(),
-        "InlinePredicate required for Loop");
-    NVF_ERROR(
-        opts.find(IdModelEnableOption::UnswitchPredicate) != opts.end(),
-        "UnswitchPredicate required for Loop");
-  }
-
  return opts;
 }


The new implementation changes the semantics of EnableOption::IdModel arguments. The old code supported explicit arguments like "all", "index", "predicate", "consumer_index", "producer_index", etc. The new code only recognizes "predicate_only" and "index_only" as restrictive flags.

This breaks existing usage like NVFUSER_ENABLE=id_model(all) found in tests/python/direct/test_with_id_model_indexer.py:183. Consider either:

Updating the test to use NVFUSER_ENABLE=id_model (no arguments, which enables everything by default)

Adding backward compatibility for the "all" argument

naoyam · 2026-01-09T08:38:28Z

!test

greptile-apps

Greptile Overview

Greptile Summary

This PR enables TensorIndexer with all Python direct tests by making IdModel always-on and simplifying its configuration.

Key Changes:

Removes DisableOption::IdModel - IdModel is now always built during GPU lowering, cannot be disabled
Simplifies IdModelOptions - Removes build_id_model_ flag and related methods, streamlines consistency checking
Updates enable option parsing - Replaces granular options (consumer_index, producer_index, etc.) with simpler predicate_only/index_only flags
Migrates from ComputeAtMap to IdModel APIs - Changes predicate_compute.cpp to use idModel().idGraph() instead of caMap() for getting concrete IDs and checking mappings
Enables id_model by default in Python API - Adds "id_model" to default enable options for FusionDefinition.execute()
Updates test expectations - Adjusts expected variable names and generated CUDA code to match new indexing behavior

The changes are part of a larger effort to transition from ComputeAtMap to IdModel/TensorIndexer as the primary indexing mechanism.

Confidence Score: 5/5

Safe to merge - well-tested refactoring that enables TensorIndexer consistently across Python tests
This PR removes the ability to disable IdModel and simplifies its configuration. All changes are internally consistent: DisableOption::IdModel is removed from enum, options map, and all conditional checks. The migration from caMap to idModel APIs in predicate_compute.cpp follows established patterns seen elsewhere in the codebase. Test updates reflect expected behavior changes from the new indexing approach. The Python API change safely adds id_model to default options (duplicates are handled gracefully by the options map).
No files require special attention

Important Files Changed

File Analysis

Filename	Score	Overview
csrc/device_lower/id_model_options.h	5/5	Removes build_id_model option and related code, simplifying IdModelOptions to always build IdModel
csrc/device_lower/lower2device.cpp	5/5	Removes conditional IdModel building, now always builds IdModel during lowering
csrc/id_model/utils.h	5/5	Simplifies getIdModelEnabledOptions to use predicate_only/index_only flags instead of granular options
csrc/options.cpp	5/5	Removes id_model from DisableOptions map
csrc/options.h	5/5	Removes IdModel enum value from DisableOption
csrc/predicate_compute.cpp	5/5	Replaces caMap with idModel for getting concrete IDs and checking ID mappings
python/nvfuser_direct/init.py	5/5	Adds id_model to default enable options for Python API execution
tests/cpp/test_indexing.cpp	5/5	Updates expected index variable names in test (i98-i100 to i114-i116)
tests/python/direct/test_python_direct.py	5/5	Updates expected CUDA kernel code to match new simplified indexing patterns

Sequence Diagram

sequenceDiagram
    participant User
    participant PythonAPI as Python FusionDefinition
    participant Lower as GpuLower
    participant IdModelOpts as IdModelOptions
    participant IdModel
    participant PredicateCompute
    participant TensorIndexer

    User->>PythonAPI: execute(inputs)
    PythonAPI->>PythonAPI: Add "id_model" to enable_options
    PythonAPI->>Lower: FusionExecutorCache.execute()
    
    Lower->>IdModelOpts: Construct with options
    Note over IdModelOpts: No longer checks DisableOption::IdModel<br/>Always initializes for TensorIndexer
    
    Lower->>IdModel: Build IdModel (unconditional)
    Note over IdModel: Previously conditional on buildIdModel()
    IdModel-->>Lower: IdModel graphs built
    
    Lower->>PredicateCompute: Generate predicates
    PredicateCompute->>IdModel: getConcreteMappedId(id)
    Note over PredicateCompute: Uses idModel().idGraph().toGroup()->front()<br/>Instead of caMap().getConcreteMappedID()
    IdModel-->>PredicateCompute: concrete ID
    
    PredicateCompute->>IdModel: Check ID mapping
    Note over PredicateCompute: Uses idGraph().disjointValSets().strictAreMapped()<br/>Instead of caMap().areMapped()
    IdModel-->>PredicateCompute: mapping result
    
    PredicateCompute-->>Lower: Predicates with TensorIndexer
    Lower-->>User: Compiled kernel

Enable TensorIndexer with all python_direct tests

562d1f1

greptile-apps bot reviewed Jan 9, 2026

View reviewed changes

naoyam added 2 commits January 8, 2026 18:49

fix test

f0c8bbe

fix

af1b4b7

greptile-apps bot reviewed Jan 9, 2026

View reviewed changes

naoyam added 2 commits January 9, 2026 00:31

cleanup

e4b6f22

test fix

6fb3cc6

greptile-apps bot reviewed Jan 9, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Enable TensorIndexer with all python_direct tests #5781

Enable TensorIndexer with all python_direct tests #5781

naoyam commented Jan 9, 2026

Uh oh!

github-actions bot commented Jan 9, 2026 •

edited by xwang233

Loading

Changes walkthrough

PR Reviewer Guide

Test failures (partial, pipeline still running)

Uh oh!

greptile-apps bot left a comment

Uh oh!

naoyam commented Jan 9, 2026

Uh oh!

naoyam commented Jan 9, 2026

Uh oh!

greptile-apps bot left a comment

Uh oh!

greptile-apps bot Jan 9, 2026

Uh oh!

naoyam commented Jan 9, 2026

Uh oh!

greptile-apps bot left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Enable TensorIndexer with all python_direct tests #5781

Are you sure you want to change the base?

Enable TensorIndexer with all python_direct tests #5781

Conversation

naoyam commented Jan 9, 2026

Uh oh!

github-actions bot commented Jan 9, 2026 • edited by xwang233 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Changes walkthrough

PR Reviewer Guide

Test failures (partial, pipeline still running)

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Greptile Overview

Greptile Summary

Confidence Score: 4/5

Important Files Changed

Sequence Diagram

Uh oh!

naoyam commented Jan 9, 2026

Uh oh!

naoyam commented Jan 9, 2026

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Greptile Overview

Greptile Summary

Confidence Score: 2/5

Important Files Changed

Sequence Diagram

Uh oh!

greptile-apps bot Jan 9, 2026

Choose a reason for hiding this comment

Uh oh!

naoyam commented Jan 9, 2026

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Greptile Overview

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Sequence Diagram

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

github-actions bot commented Jan 9, 2026 •

edited by xwang233

Loading