OpenXLA PR #49441: Why Cropping and Padding Could Not Be Added Together


OpenXLA PR #49441 fixes an algebraic rewrite that treated cropping as negative padding and then added it to a consumer’s window padding.

That arithmetic preserved output shapes. It did not preserve values.

The failing graph had this form:

input
→ pad with a negative edge amount
→ reduce-window or convolution with positive window padding

In XLA, a negative edge amount on pad removes elements from the tensor. It is a crop.

The simplifier merged the two operations by adding their padding numbers. If the pad cropped one element and the consumer added one padding element on the same side, the amounts canceled to zero.

The compiler then used the original input element where the source program required a padding value.

The program remained valid and executable. It silently computed a different function.

Positive padding can be merged by addition

For ordinary nonnegative padding, the rewrite is natural.

Suppose an input is padded with one zero on the left:

[x0, x1, x2]
        ↓
[0, x0, x1, x2]

and a following window operation adds one more zero on the same side.

The two operations can be represented as one window with two units of low padding:

[0, 0, x0, x1, x2]

The values added by the first operation and the values added by the second operation have the same role.

For reduce-window, the fold also checks that the pad value matches the reduction’s initial value when the window already has padding.

For convolution, the padding value is zero.

Under those conditions, adding nonnegative edge padding preserves the data seen by every window.

A negative pad removes a real element

A negative edge amount has different semantics.

Consider:

input:
[x0, x1, x2]

pad low = -1:
[x1, x2]

Now apply one unit of low window padding:

[0, x1, x2]

The old fold added the two amounts:

-1 + 1 = 0

and replaced the sequence with no padding on the original input:

[x0, x1, x2]

The lengths match.

The first value does not.

The correct graph discarded x0 and introduced zero. The folded graph kept x0.

Cropping followed by padding is not equivalent to the net arithmetic amount because the crop destroys information before the new padding value is inserted.

The reduce-window reproducer changed -763 into -94

The linked issue supplied a two-element example:

input:
[669, -763]

The pad crops the first element:

[-763]

The following reduce-window has width two and one unit of low padding, so it reduces:

[0, -763]

with addition.

The correct result is:

-763

The faulty simplifier combines the crop and the window padding:

-1 + 1 = 0

The pad disappears, and the optimized graph reduces the original input:

669 + (-763) = -94

No invalid shape is produced. Both paths return one output element.

That is why a verifier does not catch the problem. The rewrite is structurally legal and numerically wrong.

The convolution showed the exact element that survived

The convolution reproducer used:

input:
[5, 7, 11, 13]

kernel:
[1, 100]

After cropping the first input element:

[7, 11, 13]

the convolution adds one zero on the left.

Its first window should be:

[0, 7]

and the first output is:

0 × 1 + 7 × 100 = 700

The faulty fold removes both the crop and the positive padding.

The first window becomes:

[5, 7]

and the result is:

5 × 1 + 7 × 100 = 705

The remaining outputs are still:

1107
1311

Only the boundary that depended on the crop-and-pad ordering changes.

This pattern identifies the semantic error more clearly than a generic accuracy difference: an input value that should have been discarded appears exactly where zero padding belonged.

A phase-ordering detail exposed the invalid fold

XLA normally has another rewrite for negative padding.

It converts a cropping pad into a combination of padding and slicing, making the removal of input elements explicit before later consumers are simplified.

The issue included a no-op interior padding on a dimension of size one.

Interior padding inserts values between adjacent elements. A size-one dimension has no adjacent pair, so that interior amount changes nothing.

HandlePad clears the no-op interior padding and returns after that transformation. In the same simplifier sweep, the negative edge padding remains represented as a live pad.

The visitor then reaches the reduce-window or convolution and applies the invalid pad-fold before a later sweep can convert the crop into a slice.

The sequence is:

pad has:
  no-op interior padding
  negative edge padding
        ↓
HandlePad removes only the no-op interior component
        ↓
consumer still sees a negative pad
        ↓
consumer adds crop amount to window padding
        ↓
pad disappears before the next simplifier sweep

This explains why the bug was reachable even though XLA already had machinery for negative pads.

It also explains why fixing only that one early return would be incomplete. A negative pad can remain visible when the replacement option is disabled or when another rewrite creates one during a sweep.

The consumer fold itself needed a semantic precondition.

The patch rejects crops at both folding sites

The merged change adds one helper:

bool HasNegativePadding(const PaddingConfig& config)

It checks every dimension for:

edge_padding_low < 0

or

edge_padding_high < 0

HandleReduceWindow now returns without folding when the source pad crops any edge.

FoldConvInputPad returns false under the same condition.

The patch does not disable the positive-padding optimization.

Graphs with nonnegative edge padding can still merge the pad into the consumer’s window.

Graphs containing a crop remain explicit long enough for the normal negative-padding rewrite to turn the crop into a slice.

This is a conservative boundary: rather than trying to prove that one particular negative edge is harmless in one window configuration, the optimization declines the entire fold whenever cropping is present.

Losing one optimization opportunity is preferable to changing the source program.

The tests isolate the consumer-side rule

Both new tests configure the simplifier with:

enable_negative_padding_replacement = false

That choice is deliberate.

If the normal crop-to-slice rewrite ran first, the consumer would never receive a negative pad and the new guard would not be exercised directly.

The tests construct:

negative pad → convolution

and:

negative pad → reduce-window

then require the simplifier pass to report:

changed = false

In XLA, a pass that reports false must leave the module unmodified.

The first review requested removal of additional structural assertions because they duplicated that pass contract. The final tests retain the focused EXPECT_FALSE checks and received approval.

The committed tests are compiler-structure tests. They do not execute the JAX reproducer on CPU or GPU.

The issue and pull-request description provide the runtime results separately.

I reproduced the wrong results on an older JAX build

I ran the issue’s minimized programs on:

Python 3.13.5
JAX 0.9.0.1
jaxlib 0.9.0.1
CPU backend

The reduce-window result was:

eager:
[[-763]]

jit:
[[-94]]

The optimized HLO contained a plain reduction over the original input and no remaining pad or slice in that path.

The convolution results were:

int32 eager:
[700, 1107, 1311]

int32 jit:
[705, 1107, 1311]

and:

float32 eager:
[700.0, 1107.0, 1311.0]

float32 jit:
[705.0, 1107.0, 1311.0]

The optimized convolution had no window padding and consumed the original input.

This local run reproduces the pre-fix failure mechanism.

It does not validate the merged OpenXLA commit. I did not build the fixed source locally.

The pull request reports that, with the fix, the negative pad becomes a slice and the consumer retains its positive window padding, producing the expected results.

The observed failure was backend-independent

The issue was minimized on the JAX CPU path, and the pull request states that the wrong result was checked on CPU and GPU.

The defect lived in the shared HLO algebraic simplifier, before backend-specific code generation.

That makes backend-independent impact plausible from the source location.

It does not prove that every deployment, pass pipeline, or input graph reaches the same rewrite.

The trigger requires a negative pad to remain visible to one of the two affected consumer folds.

The compiler contract restored by the change

Algebraic simplification can combine operations only when their effects compose in the same semantic domain.

The corrected decision is:

inspect the source pad
        ↓
if every edge amount is nonnegative:
    consider merging it with window padding
        ↓
if any edge amount is negative:
    keep the crop explicit
        ↓
allow the crop-to-slice rewrite to preserve data removal

Numeric padding amounts may be added only when both operations add values rather than remove them.

OpenXLA PR #49441 restores that distinction for reduce-window and convolution and prevents a shape-correct optimization from silently retaining data the source program discarded.

Sources

https://github.com/openxla/xla/pull/49441

https://github.com/openxla/xla/pull/49351

https://github.com/openxla/xla/issues/48715

https://github.com/openxla/xla/commit/603c0a7c8889873acb43ab323c9575eef6ab22ad

https://github.com/openxla/xla/blob/603c0a7c8889873acb43ab323c9575eef6ab22ad/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc

https://github.com/openxla/xla/blob/603c0a7c8889873acb43ab323c9575eef6ab22ad/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc

Popular posts from this blog

AMD Is Trying to Turn Kernel Diversity Into One Software Asset

The first broad integration was rolled back, but the direction of the software stack is difficult to ignore The GPU software race is often described as a contest over who can write the fastest kernel. That is only part of the problem. A company may have several teams capable of producing high-performance matrix-multiplication kernels. Some may work through template libraries. Others may use compiler-generated code, handwritten assembly, or model-specific optimization stacks. The existence of those kernels does not automatically create one coherent platform. They still need to be: discovered validated matched to the correct problem given the correct arguments launched through a stable runtime and delivered through a reliable software package Recent public code changes inside AMD’s software ecosystem showed an attempt to connect these layers more directly. The central idea was to allow hipBLASLt’s GEMM selection system to choose kernels created by multiple internal kernel technologi...
Read more

PyTorch #188031 — Why Did Forward Use 64-Bit Indexing While Backward Still Used uint32_t?

PyTorch Spatial Softmax, Part 1 of 3 — An asymmetric indexing contract and the patch that aligned it A recent PyTorch change fixed a 64-bit indexing problem in the spatial softmax backward path. The bug was not in the softmax formula itself. It was in the width of the integer used to calculate memory locations. For the same large tensor, the forward path already selected either 32-bit or 64-bit indexing based on the input size. One backward path, however, continued to calculate sizes, strides, and offsets using uint32_t . That created a possible split in behavior: Forward → 64-bit address calculation → accesses the intended elements Backward → 32-bit address calculation → offset wraparound → gradients read from or written to different locations A tensor could therefore complete the forward pass correctly while the corresponding backward kernel could no longer represent every address in that tensor. The changes described in PyTorch #188031 were later landed in PyTorch’s main branch. S...
Read more

NVIDIA CUTLASS #3017 — Why load() and store() Mapped the Same Tile Coordinate to Different Addresses

CUTLASS #3017, Part 1 of 3 — An asymmetry in address calculation and the patch that fixed it In February 2026, while reading through NVIDIA CUTLASS’s RegularTileIterator<layout::PitchLinear> implementation, I noticed an odd asymmetry. load() and store() accepted the same tile coordinate, yet they calculated the corresponding memory address differently. The difference was a single line. load() divided the contiguous coordinate by kElementsPerAccess . store() did not. When each memory access handled only one element, the difference remained invisible. But in vectorized paths, where a single access handles multiple elements, the same logical tile coordinate could resolve to two different physical addresses. I reported the issue as CUTLASS Issue #3017. The missing division was later added to store() in PR #3049. CUTLASS and tile iterators CUTLASS is NVIDIA’s open-source library for high-performance CUDA linear algebra. It provides C++ template abstractions for composing GPU ope...
Read more