viable/strict/1781292498: [DTensor] Preserve symbolic local layouts without DDE guards (#187026)

Compiled DTensor paths can see symbolic local layout metadata that is

semantically valid but not syntactically identical to the metadata saved during

forward propagation.

For to_local() backward, AOTAutograd can produce a local gradient stride like

(Max(1, u3), 1) while the saved DTensor metadata uses (u1, 1). The previous

backward path recomputed the global gradient stride before deciding whether it

could reuse the original DTensor spec. That forced

compute_global_tensor_info() to evaluate symbolic stride relations and could

raise a data-dependent guard for the default same-placement backward path.

Reuse the original DTensor spec only for the default same-placement backward

when the local gradient stride, saved forward local stride, and DTensor spec

stride are compatible. Exact/provable stride equality is accepted directly. For

contiguous symbolic stride forms such as Max(1, u*), use the existing

check_contiguous_sizes_strides(..., false_if_dde=True) helper so equivalent

contiguous layouts are recognized without requiring a brittle exact symbolic

match. If neither exact nor contiguous equivalence can be proven, emit

torch._check assertions for the required stride equalities before taking the

symbolic shortcut.

If the placement changes or the physical local stride cannot justify the

original spec layout, keep the existing recomputation path and build a fresh

spec from the observed gradient stride. This avoids the symbolic guard failure

without lying about memory layout. In particular, uneven channels-last shards

must keep using the recomputation path so autograd can repair the physical local

gradient layout correctly.

The same class of issue also appears in aten.t sharding propagation. The

single-dim strategy can enumerate candidate placements that move a symbolic

_StridedShard split factor onto the other tensor dimension. When that

candidate is not provably shardable, strategy expansion should reject it instead

of evaluating a Python bool on an unbacked expression such as 8 < 2*u0.

Register transpose with allow_unbacked_sharding=False so unproven candidates

are pruned while statically valid candidates, including _StridedShard(0, u0)

propagating to _StridedShard(1, u0), are still kept.

Fixes #187025

This PR was authored with assistance from an AI assistant.

Test Plan:

python -m pytest test/distributed/tensor/test_op_strategy.py::TestCostModel::test_t_prunes_unproven_unbacked_strided_shard_candidates test/distributed/tensor/test_op_strategy.py::TestCostModel::test_mm_strategies test/distributed/tensor/test_dtensor_compile.py::TestDTensorCompile::test_to_local_backward_unbacked_symbolic_stride test/distributed/tensor/test_tensor_ops.py::TestNewEmptyStridedUneven::test_backward_channels_last -q -s

lintrunner -a

Pull Request resolved: #187026

Approved by: https://github.com/pianpwk