Refactor layout constraint selection logic (#148104)

This PR:

- cleans up some existing comments that don't make sense anymore
- hooks up the "custom_op_default_layout_constraint" back (that seems to
have broken)
- cleans up the "lazy registration path" which seems to never get hit
anymore
- adds dislike_padding to nodes that require exact strides

Test Plan:
- tests + CI

disable padding

Pull Request resolved: https://github.com/pytorch/pytorch/pull/148104
Approved by: https://github.com/shunting314, https://github.com/eellison
This commit is contained in:
rzou
2025-05-02 08:57:47 -07:00
committed by PyTorch MergeBot
parent 0e59b594ee
commit 2b37a726e0
8 changed files with 79 additions and 69 deletions

View File

@ -3849,7 +3849,7 @@ Please use `add.register_fake` to add an fake impl.""",
# If multiple tags are provided, then we use the most constrained tag.
({flexible_layout, needs_fixed_stride_order}, needs_fixed_stride_order),
]
from torch._inductor.lowering import get_layout_constraint_tag
from torch._library.utils import get_layout_constraint_tag
for tags, expected in tests:
with torch.library._scoped_library("mylib", "FRAGMENT") as m: