Add option to assert if kernel is not fully fused in foreach_map #159213
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/159213
Note: Links to docs will display an error until the docs builds have been completed. ❌ 11 New Failures, 1 Unrelated FailureAs of commit 93485fc with merge base 3ced107 ( NEW FAILURES - The following jobs have failed:
UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
Going to defer to @mlazos |
| @@ -9,15 +9,15 @@ class ForeachMap(BaseHOP): | |||
| def __init__(self): | |||
| super().__init__("foreach_map") | |||
|
|
|||
| def __call__(self, fn, *operands, **kwargs): # type: ignore[override] | |||
| def __call__(self, fn, *operands, assert_fused=False, **kwargs): # type: ignore[override] | |||
There was a problem hiding this comment.
We've discussed things that "check if something is fused" before and the conclusion was always that "Inductor makes no promises". If so, this should be named something like a debug API. Thoughts @eellison ?
There was a problem hiding this comment.
Yea, especially internally, we do not want this to get used, something changed, and then we get a huge number of failures.
There was a problem hiding this comment.
Marking this as debug is fine. Will that work for you guys? I think that's the main ask from users anyway, they just want some way to test that it's working as intended.
There was a problem hiding this comment.
Mark as debug and then document that we will break this would be good
| @@ -1104,6 +1104,23 @@ def ref_fn(xs): | |||
|
|
|||
| self.assertEqual(torch._inductor.metrics.generated_kernel_count, 5) | |||
|
|
|||
| def test_foreach_map_assert_fused_passes(self): | |||
There was a problem hiding this comment.
Can you do one more test where you have a slightly more complex sequence of pointwise functions (maybe multiply the result as well? I want to make sure vertical fusion is tested as well.
test/inductor/test_foreach.py
Outdated
| return a + b | ||
| x = torch.randn(10) | ||
| y = torch.randn(10) | ||
| torch.compile(foreach_map, fullgraph=True)(fn, (x,), (y,), assert_fused=True) |
There was a problem hiding this comment.
Can we also count the number of kernels and verify the accuracy like in the other tests? Other than that this looks good!
|
@haveheartt looks good, last thing can you preface |
|
Looks good, please prefix the arg with debug as I mentioned above before merging, and mention that we don't have BC guarantees on fusion behvaior. Thanks! |
Fixes #158970
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben