[inductor] respect layout tags for ops with registered lowerings #159134
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/159134
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 50f10e8 with merge base 7ac70ac ( 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. |
torch/_inductor/lowering.py
Outdated
| @@ -2629,6 +2629,7 @@ def is_aligned(x): | |||
| make_fallback(aten._adaptive_avg_pool3d) # @isuruf | |||
| make_fallback(aten.adaptive_max_pool3d) # @isuruf | |||
| make_fallback(aten._scaled_dot_product_attention_math_for_mps) # @malfet | |||
| make_fallback(aten._scaled_grouped_mm, constrain_to_fx_strides) | |||
There was a problem hiding this comment.
constrain_to_fx_strides isn't good enough. You want constrain_to_fake_tensors (aka needs_fixed_strides).
constrain_to_fx_strides just makes sure the tensors have the same stride order. The tensor might still not be row major/column major.
There was a problem hiding this comment.
If you want to merge this as a workaround, feel free, but there is going to need to be work to be done to audit and migrate usages of constrain_fx_strides to constrain_to_fake_tensors. cc @eellison
|
@ellison that's what i added at first, but it looks like we only check the tags when the ops have no lowering, we should always respect the op tags right? |
|
@xmfan - yea, if there is no lowering, we should respect the tag. |
test/inductor/test_torchinductor.py
Outdated
| use_fast_accum=True, | ||
| ) | ||
|
|
||
| fn() |
There was a problem hiding this comment.
@xmfan can you add a test like in https://github.com/pytorch/pytorch/pull/159174/files ?
The idea is to run a pattern matching pass to swizzle an operation so that we are sure that we're testing the right thing.
I have done e2e tests like what is currently in this PR and they don't work out sometimes or inductor changes and they end up not testing the right thing.
There was a problem hiding this comment.
i don't think i understand that test. it still passes even when i comment out the impl2
def impl2(x):
# return x.clone(memory_format=torch.channels_last)
return x|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Merge failedReason: 1 jobs have failed, first few of them are: trunk / linux-jammy-rocm-py3.10 / test (default, 2, 2, linux.rocm.gpu.2) Details for Dev Infra teamRaised by workflow job |
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
…9134) scaled_grouped_mm's kernel only supports column-major on the second operand. I -think- this is just for efficiency reasons. But inductor treats that buffer as flexible and may tweak the strides to be row-major instead, as seen in the issue. ~Tagging the op as "needs_fixed_stride_order"/"needs_exact_strides" does not work. Inductor only considers those tags for ops that don't have registered lowering (not sure if this is intended). scaled_grouped_mm does have a lowering, so we never check its tags.~ From discussion below, the op tags are expected to work. FIXES #159097 Pull Request resolved: #159134 Approved by: https://github.com/eellison
Useful helper function for stage 1 export -> manual partitioner -> stage 2 compile users Pull Request resolved: #159705 Approved by: https://github.com/zou3519 ghstack dependencies: #159134
Stack from ghstack (oldest at bottom):
scaled_grouped_mm's kernel only supports column-major on the second operand. I -think- this is just for efficiency reasons. But inductor treats that buffer as flexible and may tweak the strides to be row-major instead, as seen in the issue.
Tagging the op as "needs_fixed_stride_order"/"needs_exact_strides" does not work. Inductor only considers those tags for ops that don't have registered lowering (not sure if this is intended). scaled_grouped_mm does have a lowering, so we never check its tags.From discussion below, the op tags are expected to work.FIXES #159097
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @mlazos