[LINALG] Fix: Incorrect linalg lowering for aten.convolution_transpose with negative effective padding #4369
+167
−30
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
The Bug
The
torch-to-linalglowering foraten.convolution(withtransposed=true) incorrectly handles cases where the effective padding is negative. The logic for this is contained increateTransposedInputPadding.The original implementation had two critical flaws:
Incorrect Math: The logic block for negative padding (if (anyDimensionPaddingIsNegative)) attempted to "pre-crop" the input tensor before un-striding. The math used to calculate these slice offsets and sizes was incorrect, resulting in
tensor.extract_sliceoperations with out-of-bounds offsets and negative sizes, causing the compiler to fail.Failed "Mixed-Mode" Logic: The code was built on an "all-or-nothing" assumption. It failed to handle "mixed-mode" padding, where one spatial dimension required padding (positive offset) while another required cropping (negative offset). It would enter the negative padding path and apply cropping logic to all dimensions, leading to out-of-bounds errors when it tried to crop a dimension that should have been padded.
The Fix
This patch refactors the logic into two clean, robust paths:
All-Padding Path (else block):
Trigger: All spatial dimensions have an effective padding offset >= 0.
Action: Retains the original, efficient "fast path." It uses a single
tensor.insert_sliceto perform both un-striding (with strides) and padding (with offsets) in one operation.Safe Path (if (anyDimensionPaddingIsNegative) block):
Trigger: At least one spatial dimension has a negative effective padding offset.
Action: This path is now a unified, robust 3-step process that correctly handles both all-crop and mixed-mode scenarios:
Create "Super-Tensor": It computes a maxSizes tensor, which is the "union" of the padded and un-strided sizes (i.e., max(innerSize, outerSize) for each dimension).
Pad & Un-stride: It performs a single
tensor.insert_sliceof the original input into this maxSizes tensor. This one operation correctly applies all positive padding (via insertSliceOffsets) and un-striding (via strideIndexValues).Crop: It performs a final
tensor.extract_sliceto crop the maxSizes tensor down to the final outerSizes. This correctly applies all negative padding (via extractSliceOffsets).This new logic resolves all known failure cases and is validated by the new TransposedConv{1,2,3}dNegativePadding test cases, which specifically target this functionality.