-
Notifications
You must be signed in to change notification settings - Fork 613
[LINALG] Fix: Incorrect linalg lowering for aten.convolution_transpose with negative effective padding #4369
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
|
@sahas3 and @ivangarcia44 Please help review and add additional reviewers. |
ivangarcia44
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
In general LGTM. Just minor comments.
|
The changes LGTM. No more comments on my side. Thank you for adding the description. It is easy to see now in the formula why a large stride is going to make a negative padding. |
zjgarvey
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the very clear PR. I don't have any comments, and the work looks correct to me.
sahas3
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LG! Thanks for fixing this.
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.