Skip to content

Revert "[cortex_m] Fix linear weight layout: transpose in AOT pass, a… #16910

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.

Sign up for GitHub

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

Merged
psiddh merged 1 commit into from
Jan 27, 2026
Merged

Revert "[cortex_m] Fix linear weight layout: transpose in AOT pass, a… #16910

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions backends/cortex_m/ops/operators.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -352,7 +352,7 @@ def quantized_linear_meta(
activation_min,
) -> torch.Tensor:

shape = (*input.shape[:-1], weights.shape[1])
shape = (*input.shape[:-1], weights.shape[0])
Copy link

Copilot AI Jan 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The meta function now uses weights.shape[0] which assumes weights are in [out_features, in_features] format (PyTorch convention). However, this is inconsistent with the C++ runtime expectation. If weights were properly transposed in the AOT pass (as PR #16782 attempted), this should use weights.shape[1] instead. The current implementation will produce incorrect output shapes when used with the C++ runtime.

Suggested change
shape = (*input.shape[:-1], weights.shape[0])
shape = (*input.shape[:-1], weights.shape[1])

Copilot uses AI. Check for mistakes.
return torch.empty(shape, dtype=input.dtype, device=input.device)


Expand Down Expand Up @@ -386,7 +386,7 @@ def quantized_linear_impl(
input_reshaped = input_int32.reshape(new_shape)

lhs_sum = torch.sum(input_reshaped, dim=-1, keepdim=True) * filter_offset
output = torch.mm(input_reshaped, weights_int32) + lhs_sum + kernel_sum
output = torch.mm(input_reshaped, weights_int32.T) + lhs_sum + kernel_sum
output_shape = (*input.shape[:-1], output.shape[-1])
output_reshaped = output.reshape(output_shape)
else:
Expand All @@ -396,7 +396,7 @@ def quantized_linear_impl(
new_shape = (prod(input.shape[:-1]), input.shape[-1])
input_reshaped = input_int32.reshape(new_shape)

output = torch.mm(input_reshaped, weights_int32)
output = torch.mm(input_reshaped, weights_int32.T)
if bias is not None:
output = output + bias
output_shape = (*input.shape[:-1], output.shape[-1])
Expand Down
25 changes: 4 additions & 21 deletions backends/cortex_m/passes/convert_to_cortex_m_pass.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,19 +33,14 @@ class ConvertToCortexMPass(XNNPACKPass):
by call_operator.
"""

def _compute_kernel_sum(
self, weights_transposed, bias, input_offset, weight_offset
):
def _compute_kernel_sum(self, weights, bias, input_offset, weight_offset):
"""
Computes the precomputed kernel sum term (bias optional)
a * sum_j(wij + b) + ci

for i = (1, ..., n), where j indexes the input activations.

Args:
weights_transposed: Weights already in [in_features, out_features] format
"""
# No transpose needed - weights already transposed by caller
weights_transposed = weights.T
weights_int32 = weights_transposed.to(torch.int32)
offset_weights = weights_int32 + weight_offset
kernel_sum = torch.sum(offset_weights, dim=0, keepdim=True, dtype=torch.int32)
Expand Down Expand Up @@ -115,12 +110,8 @@ def _get_linear_replacement(self, node):
if len(node.args) > 2
else None
)
# Transpose weights once from PyTorch format [out_features, in_features]
# to CMSIS-NN format [in_features, out_features]
weights_transposed = weights_tensor.T.contiguous()
# Pass already-transposed weights to kernel_sum computation
kernel_sum_tensor = self._compute_kernel_sum(
weights_transposed, bias_tensor, -input_zp, -weight_zp
weights_tensor, bias_tensor, -input_zp, -weight_zp
)
with node.graph.inserting_after(weights):
kernel_sum = create_constant_placeholder(
Expand All @@ -131,17 +122,9 @@ def _get_linear_replacement(self, node):
kernel_sum_tensor,
)

weights_transposed_node = create_constant_placeholder(
self.exported_program,
node.graph,
node.name + "_weights_transposed",
InputKind.PARAMETER,
weights_transposed,
)

args = (
node.args[0],
weights_transposed_node,
weights,
Copy link

Copilot AI Jan 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This revert creates an inconsistency between the Python reference implementation and the C++ CMSIS-NN runtime. After this revert, weights are passed to the C++ runtime in PyTorch's [out_features, in_features] format without transposition. However, the C++ runtime at op_quantized_linear.cpp:80 sets filter_dims to {in_feat, 1, 1, out_feat}, indicating CMSIS-NN expects weights in [in_features, out_features] format. The reference implementation compensates by adding .T at runtime (operators.py:389, 399), but the C++ runtime does not perform this transpose, which will cause incorrect results when running on actual hardware.

Copilot uses AI. Check for mistakes.
None,
kernel_sum,
-input_zp,
Expand Down
Loading