[Autocast] Optimize _add_cast runtime

[aboubezari]

What does this PR do?

Type of change: Bug fix

Overview: Optimize the runtime of the _add_cast function. This function will greedily re-compute the producers & consumers of each node, leading to a O(N^2) runtime. This new implementation will efficiently pre-compute all the producers and consumers of the graph and look them up at constant time while looping through the tensors. There is no need to greedily compute producers and consumers, since we process each tensor only once, and only affect the consumers of that tensor during each iteration.

Before:

After:

Testing

Precision converter unittest

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
• Did you update Changelog?: No

[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.

[galagam]

This is indeed more efficient, but it's a bit risky. It relies on the assumption that the graph is static, meaning the list of consumers and producers for each tensor is constant. However, when we inject cast nodes, we render this data invalid.

In the limited scope where this is applied, it's probably OK, because we call self._remove_preexisting_casts() which prevents "chains" of cast nodes (cast->cast->cast). But we cannot replace all calls to utils.get_consumer_nodes and utils.get_consumer_nodes, nor can we modify all _add_cast instances (which you probably know, because you avoided that in this PR).
I think this at least warrants a comment to warn unsuspecting developers.

[aboubezari]

You're right, we can only do this optimization since there are no cast nodes in the graph and the producers/consumers and guaranteed not to be affected from iteration to iteration, and it's why I made sure to keep the params optional for this specific case.

It's up to you how you want to proceed, we can:

1. Keep as is, I added a warning to other devs, I feel that devs using this function can just leave those optional params empty to keep the safe behavior.
2. We write a separate function or move this logic out of _add_cast to make it super explicit for this use case.

[galagam]

Thanks for adding the comments. Approved.

[codecov]

Codecov Report

❌ Patch coverage is 84.61538% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.39%. Comparing base (41de55f) to head (8b8c735).
⚠️ Report is 6 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/onnx/autocast/precisionconverter.py	84.61%	2 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #469   +/-   ##
=======================================
  Coverage   73.38%   73.39%           
=======================================
  Files         180      180           
  Lines       18110    18138   +28     
=======================================
+ Hits        13290    13312   +22     
- Misses       4820     4826    +6     

☔ 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.

[galagam]

/ok to test 8b8c735