[RFC] Use `core_aten_decompositions()` as the FX-path decomposition table

Follows up on Discussion #4499 (future projects).

Problem

torch-mlir has two decomposition layers:

1. PT-side (pre-import): run_decompositions() in the FX exporter, configured by fx_decomp_util.py. Runs before any MLIR is produced.
2. MLIR-side (post-import): DecomposeComplexOps.cpp — 250 patterns that decompose Torch dialect ops into simpler Torch dialect ops.

Today the FX path uses a small hand-curated decomposition table (~50 ops in DEFAULT_DECOMPOSITIONS). Most ops arrive in the Torch dialect undecomposed, and DecomposeComplexOps handles them. This means:

• Redundant work for FX-imported programs: For ops like hardshrink, pixel_shuffle, layer_norm, etc., PyTorch already provides tested decompositions in core_aten_decompositions() (~900+ ops). The C++ patterns reimplement the same math. On the FX path, either layer could handle these ops — both produce equivalent results.
• Broader coverage available for free: core_aten_decompositions() covers ops that have no C++ pattern in torch-mlir at all. In the experiment below, 160 previously-xfailing tests now pass — these are ops where PyTorch provides a decomposition but no one had written the C++ equivalent.

What the C++ patterns do that PT-side decompositions cannot:

• The ONNX import path produces Torch dialect IR directly — no Python run_decompositions() step exists. DecomposeComplexOps is the only decomposition layer for ONNX. Any pattern needed by ONNX must remain in C++.
• Some C++ patterns produce backend-friendly IR that differs from PyTorch's decomposition. For example, roll decomposes to slice + cat in C++ (which all backends lower), but PyTorch decomposes it to arange + fmod + index_select (which StableHLO cannot lower). These patterns encode backend-aware choices that PT cannot make.
• Patterns that serve as transitive dependencies — other C++ decompositions produce ops that need further decomposition (e.g., ScaledDotProductAttention decomposes into softmax, CrossEntropyLoss into nll_loss_forward). These intermediate ops must still be decomposable in C++.
• Ops with no PyTorch decomposition available (~150 patterns) — these have no core_aten_decompositions() entry.

The proposal is narrow: for the FX path only, use core_aten_decompositions() to handle ops where the PT decomposition produces equivalent, backend-compatible IR — reducing the redundant overlap between the two layers without touching the cases where C++ patterns are load-bearing.

Proposed change

Switch get_decomposition_table() to return core_aten_decompositions() minus a small exclusion list:

def get_decomposition_table():
    return get_expanded_decomposition_table()

def get_expanded_decomposition_table():
    table = core_aten_decompositions()
    for op in _EXPANDED_DECOMP_EXCLUDE:
        table.pop(op, None)
    # Overlay decomps not in core (25 ops: addmm, native_layer_norm, grid_sampler_2d, etc.)
    table.update(get_decompositions(DEFAULT_DECOMPOSITIONS))
    return table

The C++ DecomposeComplexOps pass is unchanged — it still runs and still handles everything the ONNX path needs. For FX-imported programs, some of its patterns simply never fire because the ops were already decomposed PT-side.

Exclusion list

Ops whose PT decomposition produces IR that backends cannot lower:

Category	Count	Reason
FFT ops	22	Decompose to _fft_r2c/_fft_c2r which have no ODS definition
Norm ops	3	PT decomp introduces complex tensor ops that backends reject
Backend-incompatible	9	PT decomp produces ops specific backends can't lower (e.g., roll → arange+fmod+index_select fails on StableHLO)

Total: 34 ops excluded. Their C++ decompositions continue to handle them on all paths.

Impact on the ONNX path

None. The ONNX path does not call get_decomposition_table(). DecomposeComplexOps runs identically for ONNX-imported programs. No C++ patterns are removed.

Experimental results

All 4 e2e configurations at 0 unexpected failures:

Config	Passed	XFail	XPASS (newly passing)	Failed
onnx	1328	448	0	0
fx_importer	1681	98	22	0
fx_importer_stablehlo	1353	362	94	0
fx_importer_tosa	1390	384	44	0

160 tests that previously xfailed now pass. These are ops where core_aten_decompositions() provides a decomposition path that our C++ patterns didn't cover (or covered incompletely).

C++ patterns that become redundant on the FX path: 62 of 250

These 62 patterns never fire on FX-imported programs when core_aten_decompositions() is used, because the ops are decomposed before import:

Full list

scaled_dot_product_attention, hardshrink, softshrink, hstack, column_stack, nan_to_num, tanh_backward, mv, renorm, linalg_cross, pixel_shuffle, pixel_unshuffle, channel_shuffle, layer_norm, linspace, aminmax, std, zero, hardsigmoid, prelu, celu, fliplr, flipud, diag, trace, silu, heaviside, linear, bilinear, broadcast_tensors, clamp_max, cosine_similarity, fix, frac, baddbmm, select_scatter, count_nonzero, glu, mse_loss, selu, poisson_nll_loss, binary_cross_entropy_with_logits, kl_div, argsort, type_as, threshold, absolute, deg2rad, rad2deg, isneginf, isposinf, l1_loss, log_sigmoid, matmul, relu6, rot90, rrelu, special_expm1, t, var, var_mean, group_norm

These patterns remain registered and active — these are no-ops for ONNX-imported programs as well since these ops do not appear in the IR from an ONNX model.

Patterns that must stay active (all paths)

Reason	Examples
ONNX path produces the op	arange, pad, sgn, einsum, instance_norm, hardswish
Transitive dependency	softmax family (5), conv family (6), nll_loss_forward, embedding_bag, linalg_det
PT decomp is backend-incompatible	roll, one_hot, eye, empty_like, all, isfinite, logaddexp, logaddexp2, isclose
No PT decomp exists	~150 patterns

Milestones

M1: Switch FX decomposition table

Switch get_decomposition_table() to use core_aten_decompositions() with the exclusion list. Update xfail sets for the 160 newly-passing tests. No C++ code is modified. No patterns are removed. The 62 redundant C++ patterns remain registered — they simply don't fire on FX-imported programs because the ops are already decomposed before import.

M2: Remove redundant C++ patterns (M1 + 60 days)

After a 60-day window, remove the 62 C++ patterns that are fully covered by PT-side decomposition. During this window, downstream consumers can verify whether any PT decomposition produces worse code than its C++ equivalent for their backend. If a downstream consumer identifies such a case, they can add the op to _EXPANDED_DECOMP_EXCLUDE to suppress the PT decomposition — the C++ pattern will then fire as before. Ops where this happens will not be removed in M2; their C++ patterns stay.

The 60-day window provides:

• Time for downstream backends to run their own benchmarks/quality checks
• A simple escape hatch: add to the exclusion list → C++ decomp fires instead
• Confidence that removal is safe — any regression surfaces during this period

The 60-day period is negotiable — open to a longer or shorter timeframe if the community prefers.

Start with fx_decomp_util.py and the get_decomposition_table() entry point, then compare the proposed Python table with DecomposeComplexOps.cpp and the listed exclusions. Validate M1 across the onnx, fx_importer, fx_importer_stablehlo, and fx_importer_tosa configurations, confirming no unexpected failures and that the newly passing tests have updated xfail sets.