Integrate Automated QDQ placement tool - part 2.2

[willg-nv]

What does this PR do?

This PR implements RegionSearch class. RegionSearch could help partition big ONNX model into small region. QDQ autouning will be performed on the regions.

Overview: ?

Usage

# Add a code snippet demonstrating how to use this

Testing

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes
• Did you write any new necessary tests?: Yes
• Did you add or update any necessary documentation?: No, document updates is in Part 4.
• Did you update Changelog?: CHANGELOG will be updated when all changes are ready.

Additional Information

Summary by CodeRabbit

Release Notes

Refactor

• Improved ONNX quantization backend with new optimization framework and extensive test coverage to enhance internal graph processing capabilities.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

A new ONNX graph region search module is introduced, implementing a two-phase hierarchical partitioning framework. The implementation includes bottom-up partitioning via RegionPartitioner, top-down refinement via TopDownRegionBuilder, orchestrated by CombinedRegionSearch, with supporting graph analysis in RegionSearchBase. Comprehensive unit tests validate all components.

Changes

Cohort / File(s)	Summary
Region Search Framework modelopt/onnx/quantization/autotune/region_search.py	Introduces 4 new classes: RegionSearchBase (graph analysis utilities including root region construction, divergence checks, forward-reachability, and convergence detection), RegionPartitioner (bottom-up graph traversal creating leaf regions and handling divergent nodes), TopDownRegionBuilder (hierarchical refinement splitting and merging regions), and CombinedRegionSearch (orchestrates two-phase process). Includes constants DEFAULT_MAX_STEPS, DEFAULT_MAX_NODES_TO_SHOW, MAX_PROBE_STEPS_AFTER_CONVERGE, and region boundary/tree printing utilities.
Region Search Tests tests/unit/onnx/quantization/autotune/test_region_search.py	Comprehensive test suite with helpers for constructing ONNX graphs, tests for RegionPartitioner, TopDownRegionBuilder, and CombinedRegionSearch covering simple/divergent graph patterns, edge cases, region hierarchy validation, and tree printing output validation.

Sequence Diagram(s)

sequenceDiagram
    actor User
    participant CombinedRegionSearch
    participant RegionPartitioner
    participant RegionSearchBase as RegionSearchBase<br/>(Analysis)
    participant TopDownRegionBuilder
    participant Graph

    User->>CombinedRegionSearch: search_regions()
    CombinedRegionSearch->>RegionPartitioner: __init__() + partition_graph()
    RegionPartitioner->>RegionSearchBase: _build_forward_reachable_nodes_map()<br/>_is_tensor_divergent()
    RegionPartitioner->>Graph: traverse nodes, detect sequences
    Graph-->>RegionPartitioner: nodes, connectivity
    RegionPartitioner->>RegionPartitioner: _build_sequence_from_node()<br/>_build_small_converged_region()
    RegionPartitioner-->>CombinedRegionSearch: list[Region] (leaf regions)

    CombinedRegionSearch->>TopDownRegionBuilder: __init__(root_region)<br/>build_composite_region()
    TopDownRegionBuilder->>RegionSearchBase: compute_region_boundaries()
    TopDownRegionBuilder->>TopDownRegionBuilder: _split_sequence_regions()<br/>_merge_converged_regions()
    TopDownRegionBuilder-->>CombinedRegionSearch: Region (hierarchical composite)

    CombinedRegionSearch-->>User: final hierarchical regions

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title references the main change (integrating an automated QDQ placement tool as part 2.2), and the changeset introduces core region search and partitioning components essential for region-local QDQ autotuning, which aligns with the stated objective.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 2

🤖 Fix all issues with AI agents

In `@modelopt/onnx/quantization/autotune/region_search.py`:
- Around line 36-56: In __init__, _build_root_region() can call
compute_region_boundaries which expects self.tensor_users_map to exist, so
initialize tensor_users_map before calling _build_root_region: compute and
assign self.tensor_users_map (using get_tensor_consumer_node_indices(self.graph)
when tensor_users_map is None) prior to calling self._build_root_region(); then
proceed to construct root and finally set/compute forward_reachable_nodes_map as
currently done.
- Around line 425-455: In _build_small_converged_region the divergent node
(start_node_idx) is removed from visited_nodes and never appended to the region;
fix this by explicitly adding the divergent node to the region and marking it
visited: after visited_nodes.remove(start_node_idx) call
self._append_node_to_region(start_node_idx) and
self.visited_nodes.add(start_node_idx) (so it isn't skipped by the subsequent
loop), keeping the rest of the logic for appending branch nodes and the converge
node as-is.

🧹 Nitpick comments (2)

tests/unit/onnx/quantization/autotune/test_region_search.py (2)

126-133: Use self.assert* methods instead of bare assert in unittest tests.

Bare assert statements are compiled away when Python runs with the -O or -OO optimization flags, causing these checks to be skipped entirely in optimized mode. self.assert* methods (e.g., self.assertEqual(), self.assertTrue()) are regular method calls that execute regardless of optimization flags and provide better diagnostic context on failure.

Replace assert partitioner is not None with self.assertIsNotNone(partitioner) and assert partitioner.graph == graph with self.assertEqual(partitioner.graph, graph).

---

23-30: Remove the sys.path.insert() call; it's unnecessary given the pytest configuration in pyproject.toml.

The project already configures pythonpath = ["tests/"] in pyproject.toml (line 137), which handles import resolution properly. The sys.path.insert() on line 29 is redundant since modelopt is at the repository root and will be discoverable by pytest's standard import resolution. Removing this keeps the test file cleaner and avoids manual path manipulation.

[gcunhase]

@willg-nv can you please recover important information from docstrings here? Thanks!

[gcunhase]

LGTM, deferring final approval to @ajrasane, thanks.

[ajrasane]

/ok to test 00e8202

[codecov]

Codecov Report

❌ Patch coverage is 85.19362% with 65 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.70%. Comparing base (e53ca61) to head (bdedacf).
⚠️ Report is 8 commits behind head on main.

Files with missing lines	Patch %	Lines
...delopt/onnx/quantization/autotune/region_search.py	85.19%	65 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #845      +/-   ##
==========================================
+ Coverage   73.43%   73.70%   +0.26%     
==========================================
  Files         197      198       +1     
  Lines       20651    21096     +445     
==========================================
+ Hits        15166    15549     +383     
- Misses       5485     5547      +62     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[cjluo-nv]

/ok to test f67d957

[ajrasane]

/ok to test bdedacf