Keep Half and BFloat16 sums in float in portable matmul, var, cumsum and pooling kernels, and fix portable cat by shoumikhin · Pull Request #24401 · pytorch/executorch

shoumikhin · 2026-11-01T07:04:57Z

What is wrong today

Several portable CPU kernels add up their results in the tensor's own type. BFloat16 keeps only 8 significant bits, so once a running sum of ones reaches 256, adding 1 no longer changes it.

The most visible case is nn.Linear with a bias, which exports to aten.addmm and runs on the portable kernel when no backend takes it, as with BFloat16 and the XNNPACK partitioner's defaults. With 2560 inputs, its output was off by about 8% of the output range, where PyTorch is off by 0.25%.

What this change does

The affected kernels now sum Half and BFloat16 values in float and round once at the end. PyTorch and the portable sum, mean, softmax and avg_pool1d kernels already sum in a wider type. Other dtypes keep their exact results. For example, mixed_linear with a Float input and a Double output still multiplies in float and only sums in double, as before.

vec_ops.h: the dot product helpers behind addmm, mm, and the int8 mixed_mm and mixed_linear kernels. They also multiply in float, which for the int8 kernels stops a Half product like 600 * 127 from overflowing before it is scaled. The Cadence HiFi mm and bmm kernels use them too.
bmm: the same, in its own loop. ComplexHalf is unchanged; PyTorch has no ComplexHalf bmm on CPU.
addmm: every bias shape, including the 1-D bias of nn.Linear, is added in the same float pass, and alpha and beta stay in float, as in ATen. When beta is 0, the bias is now ignored, as PyTorch documents, so a NaN in it no longer turns the output into NaN.
var, var_mean, cumsum and avg_pool2d. With a Float input and a Half or BFloat16 output, var and var_mean still convert each input to the output dtype before summing, as ATen does. The 2D pooling helper gets the accumulator type parameter of its 1D counterpart. For Half and BFloat16, cumsum keeps its running sums in float: one sum along the last dim, or a block of column sums while it walks the rows; other dtypes keep the old loop.

This change also fixes the portable cat kernel. It skipped every shape check for empty inputs, but the output shape could come from an empty input, so cat of a 2x3 tensor and an empty 0x7 tensor read 14 values from an input that holds 6. check_cat_args and get_cat_out_target_size now follow PyTorch: only a 1-D empty tensor is skipped, and other empty inputs must match the other inputs outside the cat dim and count toward the output size. Eager PyTorch rejects every input that is now rejected. The portable kernel also copied raw bytes between dtypes of the same size, so an int32 input in a float32 cat gave bit patterns. It now converts. Its complex branch stopped at the first empty input, so cat([x, empty, y]) never copied y. It now keeps copying.

What was tested

A new test per kernel sums values with a known exact answer, such as 512 ones; the int8 ones also use a product that overflows Half. Each fails without this change. The cumsum dtype sweep now includes Half and BFloat16, and the mixed_mm test uses scales Half holds exactly (it assumed exact 0.2 and 0.4). The disabled mixed_linear dtype tests for Half input or output now run, and addmm with beta 0 and a NaN bias returns the product.
New cat tests, each failing without its fix: an empty input of another shape is rejected (without the fix, AddressSanitizer reports a heap buffer overflow), all-empty inputs get PyTorch's output shape, int32 converts into a float32 output, and complex cat copies past an empty input.
The portable, optimized, optimized-portable and quantized kernel test suites pass on Linux x86_64.
Largest error on random BFloat16 inputs, as a share of the largest output, against float64:

Kernel	Before	After	PyTorch
`addmm`, 1-D bias, 2560 inputs	8.1%	0.25%	0.25%
`addmm`, 2-D bias, 2560 inputs	7.3%	0.22%	0.22%
`mm`, 2560 inputs	9.8%	0.30%	0.30%
`var`, rows of 8192	51%	0.28%	0.28%
`cumsum`, rows of 4096	88%	0.39%	0.39%
`avg_pool2d`, 8x8 window	1.8%	0.36%	0.36%

The portable addmm, mm and bmm kernels, and the int8 mixed_mm and mixed_linear kernels, kept their running sums in the tensor's own type. BFloat16 keeps only 8 significant bits, so a running sum of ones stops growing at 256, and a long dot product loses most of its accuracy. A BFloat16 nn.Linear with a bias exports to addmm. With 2560 inputs its output was off by about 8% of the output range, where PyTorch is off by about 0.25%. These kernels now multiply and add Half and BFloat16 values in float and round the sum once, as ATen does. addmm with a 1-D bias, the nn.Linear case, still rounds that sum once more when it adds the bias, so it lands within one rounding step of PyTorch (0.37% against 0.25%). For the int8 kernels this also stops a Half product such as 600 * 127 from overflowing before it is scaled. Products use the input type and sums the output type, so every other dtype keeps its exact results, even when mixed_linear writes a different output dtype. The Cadence HiFi mm and bmm kernels call the same helper for Half and BFloat16 and change the same way. The existing mixed_mm test now uses scales that Half holds exactly. Its expected values assumed exact 0.2 and 0.4, and the Half result now rounds once instead of matching them by chance.

…avg_pool2d These kernels also kept their running sums in the tensor's own type. In BFloat16, var over rows of 8192 values was off by about 50% of the output range and cumsum over rows of 4096 values by about 90%, where PyTorch is off by about 0.3%. avg_pool2d with an 8x8 window was about five times worse than PyTorch. They now keep the running sums of Half and BFloat16 values in float and round once at the end, as the portable sum, mean, softmax and avg_pool1d kernels already do. PyTorch also sums these types in a wider type. The 2D pooling helper gets the same accumulator type parameter as its 1D counterpart. For Half and BFloat16, cumsum keeps the sums of a block of columns in float while it walks the rows. Other dtypes keep the old code paths and their exact results.

pytorch-bot · 2026-11-01T07:05:02Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24401

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit e09832c with merge base 05092c1 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

nn.Linear exports to addmm with a 1-D bias. That bias does not have the output's shape, so addmm took its broadcast path: it wrote the matmul result to the output in the tensor's own type, then added the bias in a second pass. For Half and BFloat16 that rounds every output twice, which left a BFloat16 nn.Linear with 2560 inputs at 0.37% of the output range, against PyTorch's 0.25%. A bias that is a single row, of shape [N] or [1, N], now goes through vec_addmm one output row at a time, so the sum and the bias meet in float and each output is rounded once. Column and scalar biases keep the broadcast path.

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

shoumikhin · 2026-11-01T09:15:43Z

@claude /pr-review

claude · 2026-11-01T09:16:27Z

Claude finished @shoumikhin's task in 23m 54s —— View job

Approve, no blockers: portable addmm still rounds more than once for column and scalar biases and for non-unit alpha/beta, and a few new sub-tests reuse an output tensor that already holds the expected value.

Detailed review

PR Review: #24401

Summary

The PR makes the portable addmm, mm, bmm, int8 mixed_mm/mixed_linear, var, var_mean, cumsum and avg_pool2d kernels sum Half and BFloat16 in float. Nothing blocks it. The issues below are all minor: two remaining rounding gaps in addmm, three spellings of the accumulator type, some test gaps, and an unmeasured access-pattern change in cumsum. Nothing was built or run for this review; the BFloat16 results below were worked out by hand from the code.

Code Quality

addmm still rounds two or three times for column and scalar biases (kernels/portable/cpu/op_addmm.cpp:95-121).
- Only a same-shape or [p]/[1, p] bias gets the new single-rounding pass. A [m, 1] bias, a [1] or [1, 1] bias with p > 1, or a 0-D bias still goes through vec_matmul, which rounds the float sum to CTYPE. Then val_a * alpha_val + val_b * beta_val runs with CTYPE_COMPUTE = CTYPE and rounds after each operation. So BFloat16 precision depends on the bias shape.
- Example: take the BFloat16RowBiasIsRoundedOnce inputs with self = ones({2, 1}). The result is 256: 257 rounds to 256, then 256 + 1 rounds back to 256. A row bias and ATen both give 258.
- The {2, 1} sub-case at kernels/test/op_addmm_test.cpp:120 uses a zero bias and K = 512, both exact, so it can't catch this.
- Suggested fix: tensor_is_broadcastable_to already limits in to 2-D. One vec_addmm-style pass that reads the bias with a row stride and a column stride (0 on broadcast dims) would cover every shape. It would also replace both the new per-row loop (lines 78-94) and the elementwise fallback.
Fix this →
alpha and beta are still rounded to 16 bits (op_addmm.cpp:60-61, passed at lines 68 and 84).
- vec_addmm now accumulates in float, but the scalars come from scalar_to<CTYPE>. ATen converts them to its opmath type (float) for the BFloat16 gemm.
- Example: with alpha = 0.1, the float sum is multiplied by 0.10009765625. That 0.1% bias moves many outputs one BFloat16 step away from ATen's result.
- nn.Linear uses alpha = beta = 1, so the PR's table holds. The match with PyTorch does not hold for other scalars.
- vec_addmm already takes a separate U, so passing the accumulator type there fixes it.
Fix this →
The accumulator type is spelled three ways.
- vec_ops.h:35-40 adds internal::opmath_t and a new c10/util/BFloat16-math.h include.
- matmul_ops_util.h:52, op_var.cpp:30, op_var_mean.cpp:30, op_cumsum.cpp:67 and op_avg_pool2d.cpp:75 use promote_types<T, T, /*half_to_float=*/true>::type.
- The kernels the description cites (op_sum.cpp:68, op_mean.cpp:65, op_softmax.cpp:51, op_avg_pool1d.cpp:60) use std::conditional_t<std::is_same_v<CTYPE, Half> || std::is_same_v<CTYPE, BFloat16>, float, CTYPE>.
- As a result, bmm_out_impl and vec_matmul (the same loop) use different spellings, avg_pool2d doesn't match avg_pool1d, and the names vary (ACC, Acc, AccT/AccU).
- Suggested fix: use one spelling and one name across the PR. The existing conditional_t idiom would make avg_pool2d match avg_pool1d.

Testing

Sub-cases reuse an output that already holds the answer.
- op_addmm_test.cpp:116-121, op_avg_pool2d_test.cpp:1132-1137 and op_var_mean_test.cpp:404-421 call the op several times on the same output tensor. After the first call it holds the expected value. The later calls, which the comments say take different paths, would still pass if their path wrote nothing.
- Suggested fix: use a fresh output for each call.
The blocked cumsum loop never runs with both offsets non-zero.
- No test has more than one leading index and more than one block at once. The new test is 2-D along dim 0. The dtype sweep's [2, 4] cases have either one leading index or one trailing position. So start_loc and first are never both non-zero.
- Suggested fix: add ones({2, 512, 65}) along dim 1.
Fix this →

Performance

The blocked Half/BFloat16 cumsum reads memory in a strided pattern (op_cumsum.cpp:89-108). This is low confidence and unmeasured.
- When trailing_dims > 64, the loop visits all dim_size rows once per 64-column block, and successive rows are trailing_dims elements apart.
- The doc at lines 35-37 says the kernel avoids this kind of jumpy access. The new note at lines 39-42 ("walk the rows in the same order") suggests the pattern didn't change.
- A wide input, such as BFloat16 [4096, 4096] along dim 0, would show whether it matters.
- If the blocked loop holds up, it could replace the old loop (lines 70-88) too. When ACC is CTYPE_OUT it computes the same values, so the if constexpr keeps two copies of the scan only for their access patterns.

Recommendation

Approve

Nothing blocks. The two addmm gaps (bias shape and alpha/beta) are the ones most worth fixing here, since the PR already rewrites that path. The rest can be follow-ups.

The portable cat kernel skipped every shape check for empty inputs. But get_cat_out_target_size can take the output shape from an empty input, and the copy loop sizes every input by the output shape. So cat of a 2x3 tensor and an empty 0x7 tensor gave a 2x7 output, and the copy read 14 values from an input that holds 6, past the end of its buffer. Follow PyTorch's rule in both functions: skip only 1-D empty tensors. Every other input, empty or not, must match the shape outside the cat dim, must have the cat dim in range, and adds its size along the cat dim to the output. Before, the output size added only non-empty inputs, so cat of empty 2x0 and 3x0 tensors gave 0x0 instead of 5x0. That wrong shape could not pass the stricter check in a later cat, so both change together. get_cat_out_target_size now also returns false for a dim that an input does not have, since the Cadence Fusion G3 kernel calls it before any other check. PyTorch rejects every input that this now rejects, so programs exported from PyTorch are not affected. Test: cat of a 2x3 tensor and an empty 0x7 tensor now fails. Before this change the kernel accepted it and read past the end of the 2x3 input. cat of a 1-D empty tensor with empty 2x0 and 3x0 tensors now gives 5x0, and an out-of-range dim with only an empty input now fails instead of aborting. The test that allowed an empty input of any shape now uses a 1-D empty tensor and an empty tensor of matching shape, and checks the values. Its old input, an empty tensor of another rank, now fails.

The portable cat kernel accepts inputs of different dtypes and converts each one to the output dtype. But when an input element had the same size as an output element, it copied the raw bytes instead. So int32 into a float32 output gave garbage, and so did int16 into half or int64 into double. Copy bytes only when the input and output types are the same, and convert otherwise, as the portable stack kernel already does. The choice is made at compile time, so each dtype pair keeps only the code it needs and the kernel does not grow much. Test: cat of a float32 and an int32 tensor into a float32 output now gives the converted values. Before this change, the int32 part held raw bits.

In the portable cat kernel, the complex branch returned from its loop at the first empty input. So cat([x, empty, y]) on complex tensors never copied y and left that part of the output untouched. Skip the empty input instead, as the real-dtype branch does. Test: the complex cat test now also concatenates x, an empty tensor and y, and expects the same result as x and y.

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

shoumikhin · 2026-11-01T19:11:09Z

@claude /pr-review

claude · 2026-11-01T19:12:01Z

Claude finished @shoumikhin's task in 23m 21s —— View job

Approve, no blockers: portable addmm still rounds column and scalar biases and non-unit alpha/beta more than once, and the Cadence cat kernels still have the dtype bug this PR fixes in portable.

Detailed review

PR Review: #24401

Summary

The PR makes the portable addmm, mm, bmm, var, var_mean, cumsum and avg_pool2d kernels and the int8 mixed_mm/mixed_linear kernels sum Half and BFloat16 in float, and fixes three bugs in portable cat. Nothing blocks it.

Two rounding gaps remain in addmm.
The accumulator type is written three different ways.
The Cadence cat kernels still have the bugs this PR fixes in portable. These bugs are older than the PR.
The blocked cumsum loop may be slower on wide inputs.

Nothing was built or run for this review. The BFloat16 results below were worked out by hand from the code.

Security

The Cadence cat kernels still have the bugs this PR fixes in portable. These bugs are older than the PR and outside the diff. check_cat_args accepts any input dtype that can be cast to the output dtype (copy_ops_util.cpp:103-104).
- HiFi NNLib path: it runs whenever out is Float or Int (backends/cadence/hifi/operators/op_cat.cpp:74-76) and reads every input as 32-bit words (:94). An int32 input in a float32 cat comes out as raw bit patterns, and reading a Half or int8 input runs past the end of its buffer.
- HiFi fallback: it copies inner * out.element_size() bytes from each input (:132-147). An int8 input into an int16 out therefore also reads past the end of its buffer.
- Fusion G3: its NNLib path (backends/cadence/fusion_g3/operators/op_cat.cpp:123-135) never calls check_cat_args, so the new shape check for empty inputs never runs there. {[2, 3], [0, 7]} along dim 0 still reaches xa_nn_cat with mismatched shapes. Low confidence: I did not check whether xa_nn_cat validates shapes itself.
- Suggested fix (here or as a follow-up): use the HiFi fast paths only when every input has the output dtype, and convert otherwise as portable now does. Run check_cat_args before the Fusion G3 fast path.
Fix this →

Code Quality

addmm still rounds column and scalar biases two to four times (kernels/portable/cpu/op_addmm.cpp:95-121).
- A bias is rounded once only if it has the output's shape or is a single row ([p], [1, p]).
- [m, 1], [1] and [1, 1] biases with p > 1, and 0-D biases, still go through vec_matmul, which rounds the float sum to CTYPE. Then val_a * alpha_val + val_b * beta_val is computed in CTYPE, with a rounding after each operation.
- The commit message says column and scalar biases keep this path on purpose. But the PR description says the result is rounded "once at the end", and BFloat16 accuracy now depends on the bias shape.
- Example: BFloat16 addmm(ones({2, 1}), ones({2, 257}), ones({257, 3})) gives 256, because 257 rounds to 256 and 256 + 1 rounds back to 256. A [3] bias gives 258, and so does PyTorch.
- Tests miss it:
  - The {2, 1} sub-case (kernels/test/op_addmm_test.cpp:120) adds a zero bias to 512, which BFloat16 holds exactly, so no rounding happens.
  - BFloat16RowBiasIsRoundedOnce (:124) only tests a [p] bias.
- Suggested fix: tensor_is_broadcastable_to already limits in to 2-D. One vec_addmm-style pass could read the bias through a row stride and a column stride, each 0 on a broadcast dim. That covers every bias shape and replaces both the new row branch (lines 78-94) and the elementwise fallback.
Fix this →
alpha and beta are still rounded to 16 bits (op_addmm.cpp:60-61, passed to vec_addmm with U = CTYPE at lines 68 and 84).
- vec_addmm now sums in float, but it multiplies by scalar_to<CTYPE>(alpha).
- ATen's CPU addmm converts both scalars to float for BFloat16. This is from memory of addmm_impl_cpu_; I did not re-check it.
- Example: BFloat16 addmm(zeros({1}), ones({1, 50}), ones({50, 1}), beta=1, alpha=0.3) multiplies by 0.30078125 and gives 15.0625. PyTorch gives 15.0.
- The PR's table uses nn.Linear, where alpha = beta = 1, so it cannot show this.
- Suggested fix: convert both scalars to the accumulator type and pass that type as the existing U parameter of vec_addmm.
Fix this →
The accumulator type is written three ways, under three names.
- vec_ops.h:35-40 adds internal::opmath_t, plus a new c10/util/BFloat16-math.h include. It is used as Acc and as AccT/AccU.
- matmul_ops_util.h:52, op_var.cpp:30, op_var_mean.cpp:30, op_cumsum.cpp:67 and op_avg_pool2d.cpp:75 use promote_types<T, T, /*half_to_float=*/true>::type, named ACC or Acc.
- The kernels the PR cites as precedent (op_sum.cpp:68, op_mean.cpp:65, op_softmax.cpp:51, op_avg_pool1d.cpp:60) use std::conditional_t<std::is_same_v<CTYPE, Half> || std::is_same_v<CTYPE, BFloat16>, float, CTYPE>, named ACC.
- As a result, bmm_out_impl and vec_matmul compute the type differently even though they run the same loop, and avg_pool2d no longer matches avg_pool1d.
- Also in vec_ops.h, line 249 multiplies by the scale (psum * s[...]) without casting it, while line 212 casts it.
- Suggested fix: use one spelling and one name across the PR. The existing conditional_t form would also make avg_pool2d match avg_pool1d.
Fix this →
The cat fix would be easier to cherry-pick or revert as its own PR.
- It fixes a heap over-read and changes which inputs cat accepts. Both changes go through helpers that HiFi and Fusion G3 share.
- The rest of the PR changes numerics in nine kernels.
- As one PR, reverting either part reverts both.

Performance

The blocked Half/BFloat16 cumsum loop may be slower on wide inputs (op_cumsum.cpp:89-108). Low confidence; not measured.
- Access pattern: with more than 64 trailing positions, each 64-wide block walks all dim_size rows, and consecutive rows are trailing_dims elements apart. For example, BFloat16 [4096, 4096] along dim 0 makes 64 passes, each over 4096 rows that are 8 KB apart.
- Comments: the doc at lines 35-37 gives sequential memory access as the reason for this loop order. The new note at lines 39-42 says the Half/BFloat16 loop walks "the rows in the same order", which is only true up to 64 trailing positions.
- Duplicate loops: when ACC is CTYPE_OUT, the blocked loop computes exactly what the old loop does, so the if constexpr keeps two versions of the scan only for their access patterns. A benchmark on a wide input would show whether the old loop can be removed, or whether the Half/BFloat16 path needs a wider block.

Recommendation

Approve

Nothing in the diff blocks this PR. The two addmm gaps are the most worth fixing here: the PR already rewrites that path and says the result is rounded once. The Cadence cat bugs are older than this PR and can be a follow-up.

This change named the accumulator type three ways: var, var_mean, cumsum and avg_pool2d used promote_types<T, T, /*half_to_float=*/true> as ACC, bmm used the same type as Acc, and vec_ops.h added its own opmath_t alias as Acc, AccT and AccU. Use the existing promote_types form and the name ACC everywhere, and drop opmath_t. In vec_ops.h the two input types are promoted together, so products and sums use one type. The grouped int8 matmul now also casts its scale to that type, as the other one does.

Only a full or single-row bias went through the float pass. A column, 1x1 or 0-D bias took the broadcast path, which rounded the matmul result to the tensor's type, then rounded again after each multiply and add. In BFloat16, addmm of a [2, 1] bias of ones with ones [2, 257] and [257, 3] gave 256 instead of 258. alpha and beta were also rounded to the tensor's type before use, so alpha 0.3 became 0.30078125 in BFloat16. vec_addmm now reads the bias through a row and a column stride, with 0 for a broadcast dim, so every bias shape takes one pass and each output is rounded once. alpha and beta are passed in the accumulator type, as ATen does. The broadcast path and the row-bias special case go away. Test: the BFloat16 bias test now covers row, column, 1x1 and 0-D biases and a non-unit alpha. The column, 1x1, 0-D and alpha cases fail without this change.

The Half and BFloat16 cumsum loop keeps float sums for a block of trailing positions and walks every row once per block. With 64 positions per block, a wide input took many strided passes: on Linux x86_64, BFloat16 cumsum along dim 0 of [4096, 4096] took 127 ms against 75 ms before this change. With 1024 positions it takes 80 to 83 ms (Half: 177 to 191 ms against 175 ms), at the cost of 4 KB of stack for these two dtypes. Also correct the comment: the rows are walked once per block, not once. Test: the BFloat16 cumsum test now has 1025 columns, so it also covers a second, partial block.

shoumikhin · 2026-11-01T20:38:22Z

@claude /pr-review

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

claude · 2026-11-01T20:39:10Z

Claude finished @shoumikhin's task in 24m 13s —— View job

Approve, no blockers: the new Half/BFloat16 cumsum loop reserves 4 KB of stack, and the Cadence cat kernels still skip the fixes this PR makes to portable cat.

Detailed review

PR Review: #24401

Summary

The PR makes the portable addmm, mm, bmm, var, var_mean, cumsum and avg_pool2d kernels and the int8 mixed_mm and mixed_linear kernels sum Half and BFloat16 in float. It also fixes three bugs in portable cat. Nothing blocks it, and every finding below is minor. Nothing was built or run for this review; the numbers below were worked out by hand from the code.

Code Quality

vec_addmm no longer uses U for its math (kernels/portable/cpu/vec_ops.h:267-280).
- The accumulator is now promote_types<T, T, true>, and alpha and beta are cast to it. U only sets the parameter type, so a fractional scalar is truncated when T is an integer type.
- Example: vec_addmm<int64_t, double> with a dot product of 3, alpha = 0.5 and beta = 0.
  - Before: 3 * 0.5 ran in double and gave 1.5, stored as 1.
  - Now: alpha becomes 0, so the result is 0.
- No in-tree result changes: op_addmm passes U = ACC, and vec_ops_test.cpp:239 instantiates <int64_t, double> with whole-number scalars.
- vec_matmul and the int8 helpers in the same file promote T with U instead.
- Suggested fix: make U the accumulator type, with promote_types<T, T, true>::type as its default. That is the type op_addmm already passes.
Fix this →
The cat fixes ship together with the accumulator changes.
- Commits c8e8e02, ecb4cd1 and 6c63c1b fix a heap over-read and tighten checks that the Cadence kernels also use.
- The rest of the PR changes numerics in nine kernels.
- In one PR, reverting or cherry-picking either part to a release branch takes the other part with it. A separate cat PR avoids that.

Testing

No test uses a column bias whose rows differ (kernels/test/op_addmm_test.cpp:120, :133).
- [m, 1] biases used to go through the broadcast helper. They now depend on the new row stride.
- The only [m, 1] biases in the file are all zeros or all ones. A row stride of 0 reads self[0] for every row, and both tests would still pass.
- Suggested fix: add a Float case with bias [[10], [20]].
- The comment at :115 ("take different paths") has been out of date since 319b236: every bias shape now goes through the same vec_addmm call.
Fix this →
The cross-dtype mixed_linear paths changed, and no test covers them (kernels/portable/cpu/vec_ops.h:227-228, kernels/quantized/cpu/op_mixed_linear.cpp:93-94).
- ACC = promote_types<T, U, true> changes three input/output dtype pairs:
  - Float input, Half output: sums are now in float instead of Half.
  - Half input, Float output: products are now in float instead of Half.
  - Float input, Double output: products are now in double instead of float. A Double output is possible when dtype is unset (op_mixed_linear.cpp:37-42).
- The new test covers only Half input with Half output. FloatInputHalfOutput is still under #if 0 (op_mixed_linear_test.cpp:105).
- The Float-to-Double change also contradicts the description's "Other dtypes keep their exact results". For example, 1.1f × 3 gave 3.3000001907348633 before and gives 3.3000000715255737 now. The new value is more accurate, so only the description needs fixing.
Fix this →

Security

The Cadence cat kernels still bypass these fixes. The bugs are older than this PR and outside the diff, so a follow-up is fine.
- Fusion G3: the NNLib fast path (backends/cadence/fusion_g3/operators/op_cat.cpp:123-135) runs whenever every input has the output's dtype. check_cat_args only runs in the else branch (:137-141).
  - So the PR's own repro, float {[2, 3], [0, 7]} along dim 0, still reaches xa_nn_cat. With OP_ARG_CHECK defined, out is first resized to [2, 7].
  - The PR hardens get_cat_out_target_size for this caller (kernels/portable/cpu/util/copy_ops_util.cpp:148-152), but the shape check itself is in check_cat_args.
  - Low confidence: I did not check whether xa_nn_cat validates shapes itself.
- HiFi: check_cat_args accepts any input dtype that can be cast to the output's dtype (copy_ops_util.cpp:103-104), and neither HiFi path converts:
  - The NNLib path reads every input as 32-bit words when out is Float or Int (backends/cadence/hifi/operators/op_cat.cpp:74-94).
  - The fallback copies out.element_size() bytes per element (:129-150).
  - So an int32 input in a float32 cat gives raw bits, the bug this PR fixes in portable. A Half or int8 input is read past its end, since const_data_ptr<float>() does not check the dtype in portable mode.
- Suggested fix: call check_cat_args before the G3 fast path. Take the HiFi fast paths only when every input has the output's dtype.
Fix this →

Performance

The Half/BFloat16 cumsum loop reserves 4 KB of stack (kernels/portable/cpu/op_cumsum.cpp:90-91).
- ACC sums[1024] takes 4 KB on every Half or BFloat16 call. That includes cumsum along the last dim, which only uses sums[0].
- The other fixed-size stack arrays I found in kernels/portable/cpu hold kTensorDimensionLimit (16) elements. Portable kernels are the fallback on MCU targets, where thread stacks can be only a few KB.
- Low confidence: cumsum_tensors<CTYPE_OUT> is called once per dtype. If the compiler inlines it into cumsum_out, that frame may reserve the 4 KB for every dtype, not only the two that commit 40d26f4 names. Whether this happens depends on the compiler.
- Suggested fix: move the blocked loop into an ET_NOINLINE function, and handle trailing_dims == 1 with a scalar float accumulator outside it. Alternatively, try block sizes between 64 and 1024; only those two were measured.
Fix this →

Recommendation

Approve

Nothing blocks. The cumsum stack buffer is the one most worth settling here, since this PR adds it. The Cadence cat gaps are older than this PR and can be a follow-up.

The previous push made vec_ops.h use promote_types, which broke every build of the ATen-mode quantized kernels: with ATen's ScalarType helpers also in scope, the calls inside promote_types are ambiguous. Give vec_ops.h a small alias that needs only c10 and gives the same type, float for Half and BFloat16 and T otherwise, and drop the scalar_type_util dependency again. The kernels keep promote_types.

shoumikhin · 2026-11-01T20:52:21Z

@claude /pr-review

claude · 2026-11-01T22:57:09Z

Claude finished @shoumikhin's task in 25m 30s —— View job

Approve, no blockers: two minor issues remain. Clang may put the new 4 KB Half/BFloat16 cumsum block in every dtype's stack frame, and the Fusion G3 cat fast path still skips the shape check this PR adds.

Detailed review

PR Review: #24401

Summary

The PR makes the portable addmm, mm, bmm, var, var_mean, cumsum and avg_pool2d kernels and the int8 mixed_mm and mixed_linear kernels sum Half and BFloat16 in float, and fixes three bugs in portable cat. Nothing blocks it, and both findings below are minor. Nothing was built or run for this review, and I could not run a compiler here to check the cumsum stack claim.

Security

The Fusion G3 cat fast path still skips the shape check this PR adds (kernels/portable/cpu/util/copy_ops_util.cpp:135-174). The gap is older than the PR, so a follow-up is fine.
- What G3 runs: before xa_nn_cat (backends/cadence/fusion_g3/operators/op_cat.cpp:123-135), G3 calls only get_cat_out_target_size, and only under OP_ARG_CHECK (:47-64). check_cat_args now rejects an empty input of another shape, but G3 runs it only in its fallback branch (:137-141).
- What get_cat_out_target_size misses: it takes the sizes outside the cat dim from the last input that isn't 1-D empty. It never compares the inputs with each other.
- Example: take the inputs of the new EmptyInputWithMismatchedShapeDies test: Int [2, 3] and [0, 7] along dim 0. With OP_ARG_CHECK defined, G3 still resizes out to [2, 7] and passes both shapes to xa_nn_cat. Low confidence in the impact: xa_nn_cat's source isn't in this checkout, so I don't know whether it rejects the mismatch itself.
- Untested check: commit c8e8e02 adds a dim-range check to this function for G3 (:148-152), but no test reaches it. Portable cat_out runs check_cat_args first (kernels/portable/cpu/op_cat.cpp:29, before :39), so the new DimBounds case is rejected before this check runs.
- Suggested fix: also return false when an input that isn't 1-D empty differs from the reference in rank or in a size outside dim. Test both checks by calling the function directly, as the overflow tests at kernels/test/op_cat_test.cpp:480 and :515 do.
Fix this →

Backward Compatibility

Clang builds may reserve the 4 KB cumsum block for every dtype (kernels/portable/cpu/op_cumsum.cpp:91-92). Low confidence: I could not compile it here.
- Stated cost: ACC sums[1024] is only instantiated for Half and BFloat16. The comment at lines 39-43 gives the 4 KB as the cost of those two dtypes.
- Why other dtypes may pay it too:
  - cumsum_tensors has internal linkage (the anonymous namespace at line 26), and each instantiation has one call site (line 150).
  - Clang usually inlines a function like that. When it does, it moves the function's fixed-size locals into the caller's frame.
  - cumsum_out would then reserve the 4 KB on every call, including Float and Int calls, which used no extra stack before this PR.
  - GCC limits inlines that grow a stack frame this much, so GCC builds are probably not affected.
- Example: a Float cumsum built with armclang or the LLVM embedded toolchain could overflow after this PR on a Cortex-M or Zephyr thread with a few KB of stack.
- Suggested fix: move the Half/BFloat16 loop into an ET_NOINLINE function, and check cumsum_out's frame with Clang's -fstack-usage.
Fix this →

Recommendation

Approve

Nothing blocks. The cumsum stack question is the one to settle in this PR, because this PR adds the block. The G3 gap is older than the PR.

The Half and BFloat16 cumsum loop keeps a 4 KB block of float sums on the stack. cumsum_tensors has one call site per dtype, so Clang inlined every instantiation into cumsum_out and reserved the block in its frame for every dtype. Measured with clang -O2 -fstack-usage on Linux x86_64: cumsum_out's frame grew from 216 bytes on main to 4296 bytes. Marking cumsum_tensors ET_NOINLINE keeps the block in the Half and BFloat16 instantiations only (4280 bytes); every other dtype uses 184 bytes or less.

BFloat16BiasIsRoundedOnce sums 257 ones, so it already fails without float accumulation, which was all BFloat16LargeDimAccumulatesInFloat checked. Add the full-shape bias it had to the merged test and drop it.

get_cat_out_target_size returns false for a dim that an input does not have, for callers that skip check_cat_args. Portable cat runs check_cat_args first, so no test reached the check. Call the function directly, as the overflow tests do.

mixed_linear accepts a Double output when no dtype is given. Since its products moved to the output's accumulator type, a Float input with a Double output multiplied in double, so its results changed although neither dtype is Half or BFloat16. Multiply in the input's accumulator type and sum in the output's, so only Half and BFloat16 results change.

var and var_mean now divide the float sum by the element count, and the var_mean fast path by the denominator, in float. The tests used 1024 elements, which BFloat16 holds exactly, so rounding the count to BFloat16 first still passed, and their tolerance hid small errors. Use 2058 elements, which BFloat16 rounds to 2064, with values whose mean and variance move when the count or the denominator is rounded, and compare exactly. var_mean uses a correction of 1, so its denominator 2057 is not exact in BFloat16 either.

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

shoumikhin · 2026-11-02T00:37:16Z

@claude /pr-review

claude · 2026-11-02T00:38:02Z

Claude finished @shoumikhin's task in 21m 41s —— View job

Approve, no blockers: three minor issues remain: Half/BFloat16 cumsum reserves 4 KB of stack even along the last dim, two vec_ops.h helpers handle mixed input and output types differently from the third, and the Fusion G3 cat fast path still gets no shape check.

Detailed review

PR Review: #24401

Summary

The PR makes the portable addmm, mm, bmm, var, var_mean, cumsum and avg_pool2d kernels and the int8 mixed_mm and mixed_linear kernels sum Half and BFloat16 in float, and fixes three bugs in portable cat. Nothing blocks it, and the three findings below are minor. Nothing was built or run for this review, because only read-only git commands were available.

Code Quality

Two vec_ops.h helpers now multiply in the output's type, and the third in the input's (kernels/portable/cpu/vec_ops.h:181-186, :202-208).
- vec_matmul and vec_quantized_matmul_int8 cast their U inputs to acc_t<T> before multiplying.
- vec_quantized_matmul_transb_int8 multiplies in acc_t<U> (:234). Its comment (:232-233) says that is so "a Double output does not change the products of a Float input".
- Before the PR, all three multiplied in U.
- Example: vec_matmul<double, float> now multiplies in double. vec_matmul<float, double> now cuts its double inputs to float before multiplying.
- Impact today: none. Every in-tree caller passes one type: op_mm.cpp:46, op_mixed_mm.cpp:86, and the Cadence HiFi op_mm.cpp:143 and op_bmm.cpp:159.
- Suggested fix: multiply in internal::acc_t<U> in both, as the transb helper does. Results don't change for current callers, where T == U.
Fix this →

Security

The Fusion G3 cat fast path still gets no shape check (kernels/portable/cpu/util/copy_ops_util.cpp:144-172). This gap is older than the PR.
- What G3 runs: its xa_nn_cat path (backends/cadence/fusion_g3/operators/op_cat.cpp:123-135) handles same-dtype Int, Short, Char, UInt32, UInt16, Byte and Float inputs.
  - It never calls check_cat_args, which is where this PR's new shape check lives (copy_ops_util.cpp:108-124).
  - Before that path, G3 calls only get_cat_out_target_size, and only when OP_ARG_CHECK is defined (op_cat.cpp:47-64).
- What get_cat_out_target_size misses: it takes the sizes outside dim from the last input that isn't 1-D empty, and never compares the inputs with each other. The new test comment at kernels/test/op_cat_test.cpp:493-494 names G3 as a caller, but the only check added for G3 is the dim-range one.
- Example: the PR's own case is Int [2, 3] and [0, 7] along dim 0. With OP_ARG_CHECK defined, G3 still resizes out to [2, 7] and passes both shapes to xa_nn_cat.
- Confidence: low. xa_nn_cat's source isn't in this checkout, so it may reject the mismatch itself.
- Suggested fix: return false when an input that isn't 1-D empty differs from the reference in rank or in a size outside dim. Test it directly, next to CatOutTargetSizeRejectsDimPastInputRank.
Fix this →

Backward Compatibility

Half and BFloat16 cumsum reserves 4 KB of stack even when it uses one float (kernels/portable/cpu/op_cumsum.cpp:90-91).
- ACC sums[kBlock] with kBlock = 1024 sits in the frame of every Half and BFloat16 call.
- The comment at lines 42-43 justifies the size with wide inputs. But a cumsum along the last dim has trailing_dims == 1 and uses only sums[0].
- Portable kernels are the fallback on Cortex-M and Zephyr targets, where thread stacks can be only a few KB. A Half or BFloat16 cumsum that fit there before can now overflow the stack.
- Suggested fix: handle trailing_dims == 1 with a single float accumulator. Move the blocked loop into its own ET_NOINLINE function so that path doesn't reserve the block.
Fix this →

Recommendation

Approve

Nothing blocks. The cumsum stack block is the one most worth settling here, because this PR adds it. The G3 gap is older than the PR, and the vec_ops.h mismatch affects no current caller.

Half and BFloat16 cumsum reserved its 4 KB block of float sums on every call, including a cumsum along the last dim, which needs a single sum. That is the common case, for example in top-p sampling, and portable kernels run on targets with stacks of a few KB. Keep one float sum along the last dim, and move the blocked loop into its own function that is not inlined, so only the calls that use the block reserve it. The BFloat16 test now also covers the last dim and a leading dim of 2.

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

shoumikhin · 2026-11-02T12:07:32Z

@claude /pr-review

claude · 2026-11-02T12:08:20Z

Claude finished @shoumikhin's task in 35m 29s —— View job

Approve, no blockers: the Cadence HiFi and Fusion G3 cat kernels miss what this PR fixes in portable cat (pre-existing, fine as a follow-up), and mixed-dtype var/var_mean now skip ATen's rounding of inputs to the output dtype.

Detailed review

PR Review: #24401

Summary

The PR keeps Half and BFloat16 sums in float in the portable matmul, var, cumsum and pooling kernels and fixes portable cat. No blockers; two minor issues below. Findings come from reading the code; nothing was built or run.

Security

The Cadence cat kernels miss what this PR fixes in portable cat (pre-existing and outside the diff, so a follow-up is fine).
- HiFi reads every input as the output dtype. backends/cadence/hifi/operators/op_cat.cpp:74-76 picks xa_nn_concat_32_32 from out's dtype alone, and :94 hands each input over as 32-bit data. The fallback at :129-150 memcpys with out.element_size(). check_cat_args only requires inputs to be castable to out (kernels/portable/cpu/util/copy_ops_util.cpp:102-104). So an Int input into a Float out has its raw bits copied, which is the bug fixed at kernels/portable/cpu/op_cat.cpp:110. A narrower input, such as Half into Float or Int into Long, is read past its end.
- The Fusion G3 xa_nn_cat fast path (backends/cadence/fusion_g3/operators/op_cat.cpp:123-135) is validated only by get_cat_out_target_size (:47-64, under OP_ARG_CHECK); check_cat_args runs only in the else branch (:137-141). The new comment at kernels/test/op_cat_test.cpp:493-494 names this caller, yet get_cat_out_target_size gained only a dim-range check (copy_ops_util.cpp:148-152). It still doesn't check ranks or the sizes outside dim, and it takes those sizes from the last input (:161-168): Int [2,7] and [1,3] along dim 0 give a 9-element out for 17 input elements. Whether xa_nn_cat then writes out of bounds depends on how it sizes its copies; its source isn't in this checkout, so this part is low confidence.
- Fix: run check_cat_args before the G3 fast path. In HiFi, use the nnlib path only when every input has out's dtype, and cast mismatched inputs in the fallback, as portable op_cat.cpp does.
Fix this →

Code Quality

Mixed-dtype var and var_mean no longer round inputs to the output dtype (kernels/portable/cpu/op_var.cpp:43, :50-51; kernels/portable/cpu/op_var_mean.cpp:80, :89-90). Main cast each input to CTYPE_OUT before reducing, as ATen does on CPU (make_reduction reduces over self.to(out_dtype); this is from ATen's ReduceOps.cpp, which isn't in this checkout). Casting straight to ACC skips that rounding for Half and BFloat16 outputs. Float [1000, 1001, 1002, 1003] into a BFloat16 out (unbiased) gives ≈1.664, while ATen rounds the inputs to [1000, 1000, 1000, 1004] and gives 4. The mixed-dtype path is tested (kernels/test/op_var_test.cpp:359-373), but only with values 0–23, which BFloat16 holds exactly. This PR's cumsum keeps the rounding (load_self converts to CTYPE_OUT), and static_cast<ACC>(static_cast<CTYPE_OUT>(v)) would do the same here. Sum and mean have the same pre-existing gap (op_sum.cpp:133, op_mean.cpp:107).

Fix this →

Recommendation

Approve

Neither issue blocks merging: the Cadence one predates this PR, and the var one is a small change to the map lambdas.

var and var_mean with a Float input and a Half or BFloat16 output used to convert each input to the output dtype and then sum, as ATen does. Summing in float converted the inputs straight to float, so Float [1000, 1001, 1002, 1003] gave a BFloat16 variance of about 1.66 instead of 4. Convert each input to the output dtype first, then to the float accumulator.

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

shoumikhin · 2026-11-02T12:59:44Z

@claude /pr-review

claude · 2026-11-02T13:00:34Z

Claude finished @shoumikhin's task in 20m 35s —— View job

Approve, no blockers: three minor issues remain. Half and BFloat16 cumsum reserves 4 KB of stack even for narrow inputs. The Fusion G3 cat fast path still skips the new shape check. Two vec_ops.h helpers choose their product type differently from the third.

Detailed review

PR Review: #24401

Summary

The PR makes the portable addmm, mm, bmm, var, var_mean, cumsum and avg_pool2d kernels and the int8 mixed_mm and mixed_linear kernels sum Half and BFloat16 in float. It also fixes three bugs in portable cat. Nothing blocks it, and the three findings below are minor. Every changed file was read, but nothing was built or run.

Code Quality

Two vec_ops.h helpers multiply in the output's type, and the third in the input's (kernels/portable/cpu/vec_ops.h:181-186, :202-208).
- vec_matmul and vec_quantized_matmul_int8 cast their U operands to acc_t<T> before multiplying.
- vec_quantized_matmul_transb_int8 multiplies in acc_t<U> (:234). Its comment says this is so "a Double output does not change the products of a Float input".
- On main, all three multiplied in U.
- Example:
  - vec_matmul<double, float> now multiplies in double.
  - vec_matmul<float, double> now cuts its double inputs to float before multiplying.
  - Both change results, which contradicts the description's "Other dtypes keep their exact results".
- Impact today: none. Every in-tree caller of these two helpers passes a single type: op_mm.cpp:46, op_mixed_mm.cpp:86, and the Cadence HiFi op_mm.cpp:143 and op_bmm.cpp:159.
- Suggested fix: multiply in internal::acc_t<U> in both, as the transb helper does.
Fix this →

Security

The Fusion G3 cat fast path still skips the new shape check (kernels/portable/cpu/util/copy_ops_util.cpp:144-162). This gap is older than the PR.
- What G3 runs:
  - Before xa_nn_cat (backends/cadence/fusion_g3/operators/op_cat.cpp:123-135), G3 runs only get_cat_out_target_size, and only under OP_ARG_CHECK (:47-64).
  - The PR's shape check is in check_cat_args (copy_ops_util.cpp:108-124), which G3 runs only on its fallback path.
- What get_cat_out_target_size misses:
  - It takes the sizes outside dim from the last input that isn't 1-D empty, and never compares the inputs with each other.
  - The new test comment at kernels/test/op_cat_test.cpp:493-494 names G3 as a caller, but the only check added for G3 is the dim-range one.
- Example: take the inputs of the new EmptyInputWithMismatchedShapeDies test, Int [2, 3] and [0, 7] along dim 0. On G3 they still resize out to [2, 7] and reach xa_nn_cat with both shapes, as before this PR.
- Confidence: low on the impact. xa_nn_cat's source isn't in this checkout, so it may reject the mismatch itself.
- Suggested fix:
  - Return false from get_cat_out_target_size when an input that isn't 1-D empty differs from the first such input in rank or in a size outside dim.
  - Test it next to CatOutTargetSizeRejectsDimPastInputRank.
Fix this →

Backward Compatibility

Half and BFloat16 cumsum reserves 4 KB of stack even for narrow inputs (kernels/portable/cpu/op_cumsum.cpp:44-45).
- The cost:
  - cumsum_in_float_blocks declares float sums[1024].
  - Every Half or BFloat16 cumsum on a dim with more than one trailing position calls it.
  - A BFloat16 [N, 8] cumsum along dim 0 reserves 4 KB but uses only 8 floats.
- Why it matters: on main, these calls used no extra stack. A Half model whose cumsum fit a small Cortex-M or Zephyr thread stack can overflow that stack after this PR.
- What was measured: commit 40d26f4 tried only blocks of 64 and 1024 floats, on a [4096, 4096] input.
  - The 64 block was slow there because it walked the rows 64 times.
  - An input with 64 or fewer trailing positions takes a single pass with either block, so a smaller block costs it nothing.
- Suggested fix:
  - Make the block size a template parameter. Use a 64-float block (256 bytes) when trailing_dims fits in it, and keep 1024 for wider inputs.
  - Alternatively, measure a 256-float block, which takes 1 KB.
Fix this →

Recommendation

Approve

Nothing blocks. The cumsum stack cost is the one to settle here, because this PR adds it. The G3 gap is older than the PR, and the vec_ops.h mismatch affects no current caller.

shoumikhin · 2026-11-02T16:44:06Z

@pytorchbot cherry-pick --onto release/1.6 -c critical

pytorchbot · 2026-11-02T16:48:21Z

Cherry picking #24401

The cherry pick PR is at #24443 and it is recommended to link a critical cherry pick PR with an issue.

Details for Dev Infra team

Raised by workflow job

### What is wrong today `torch.cat` accepts tensors of different dtypes and promotes them: ```python torch.cat([torch.ones(2, dtype=torch.bfloat16), torch.ones(2)]).dtype # torch.float32 ``` `to_edge` rejects the same program: ``` SpecViolationError: These operators are taking Tensor inputs with mismatched dtypes: ... aten.cat.default ... {'tensors': torch.bfloat16, '__ret_0': torch.float32} ``` `torch.stack`, `torch.concat` and `torch.hstack` fail too, since they become `cat`. `cat([float32, int32])` is worse: it lowers without an error and gives wrong values. The portable `cat` kernel fix behind those values is in #24401. ### What this change does 1. `to_edge` casts the inputs of `cat` and `stack` to the dtype eager promoted to. The `RemoveMixedTypeOperators` pass already did this for binary ops like `add`. For `cat` and `stack` it takes that dtype from the op's output, so an empty input dropped by an earlier pass still counts. 2. The edge dtype check (`EdgeOpArgValidator`) now checks every tensor of an input list, not only the first. That gap let `cat([float32, int32])` pass. 3. `ConvertDtypePass` (used as `I64toI32`) left int64 topk indices in a `cat` it made int32, which the check in item 2 now rejects. It now casts the tensor-list inputs of ops it rewrites. Other inputs, like a `gather` index, get no new cast. Compatibility: - A pass after `to_edge` that changes only some `cat` or `stack` inputs to another dtype now fails the edge check. The Arm int64 passes, such as `ConvertInt64OutputOpsToInt32Pass`, can do this when passed as transform passes, as the Arm README shows. The Arm tester, scripts and recipes lower with `_check_ir_validity=False`, so they are not affected, and the README now says to do the same. An Arm test now expects this graph. - A dropped empty input with a wider dtype now widens the output, as in eager. - Mixed `cat` programs that already ran, like `cat([float32, int64])`, now carry a `_to_dim_order_copy` for each input of another dtype when re-exported. A selective build needs that operator to load them. ### What was tested - New tests, each failing without its fix: `exir/program/test/test_program.py` (float, integer, empty and repeated inputs), `exir/tests/test_arg_validator.py` and `backends/transforms/test/test_convert_dtype_pass.py`. - On the portable runtime, `cat` with mixed float, integer, bool and empty inputs, `stack`, `concat` and `hstack` match eager. - Apart from the new tests and the Arm test, `exir` and XNNPACK tests and the Arm, Cadence, Qualcomm, Vulkan and `backends/transforms` pass tests give the same results with and without this change. NXP and some Cadence pass tests could not run here.

@digantdesai

### What is wrong today `torch.cat` accepts tensors of different dtypes and promotes them: ```python torch.cat([torch.ones(2, dtype=torch.bfloat16), torch.ones(2)]).dtype # torch.float32 ``` `to_edge` rejects the same program: ``` SpecViolationError: These operators are taking Tensor inputs with mismatched dtypes: ... aten.cat.default ... {'tensors': torch.bfloat16, '__ret_0': torch.float32} ``` `torch.stack`, `torch.concat` and `torch.hstack` fail too, since they become `cat`. `cat([float32, int32])` is worse: it lowers without an error and gives wrong values. The portable `cat` kernel fix behind those values is in #24401. ### What this change does 1. `to_edge` casts the inputs of `cat` and `stack` to the dtype eager promoted to. The `RemoveMixedTypeOperators` pass already did this for binary ops like `add`. For `cat` and `stack` it takes that dtype from the op's output, so an empty input dropped by an earlier pass still counts. 2. The edge dtype check (`EdgeOpArgValidator`) now checks every tensor of an input list, not only the first. That gap let `cat([float32, int32])` pass. 3. `ConvertDtypePass` (used as `I64toI32`) left int64 topk indices in a `cat` it made int32, which the check in item 2 now rejects. It now casts the tensor-list inputs of ops it rewrites. Other inputs, like a `gather` index, get no new cast. Compatibility: - A pass after `to_edge` that changes only some `cat` or `stack` inputs to another dtype now fails the edge check. The Arm int64 passes, such as `ConvertInt64OutputOpsToInt32Pass`, can do this when passed as transform passes, as the Arm README shows. The Arm tester, scripts and recipes lower with `_check_ir_validity=False`, so they are not affected, and the README now says to do the same. An Arm test now expects this graph. - A dropped empty input with a wider dtype now widens the output, as in eager. - Mixed `cat` programs that already ran, like `cat([float32, int64])`, now carry a `_to_dim_order_copy` for each input of another dtype when re-exported. A selective build needs that operator to load them. ### What was tested - New tests, each failing without its fix: `exir/program/test/test_program.py` (float, integer, empty and repeated inputs), `exir/tests/test_arg_validator.py` and `backends/transforms/test/test_convert_dtype_pass.py`. - On the portable runtime, `cat` with mixed float, integer, bool and empty inputs, `stack`, `concat` and `hstack` match eager. - Apart from the new tests and the Arm test, `exir` and XNNPACK tests and the Arm, Cadence, Qualcomm, Vulkan and `backends/transforms` pass tests give the same results with and without this change. NXP and some Cadence pass tests could not run here. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Co-authored-by: Anthony Shoumikhin <anthony@shoumikh.in>

shoumikhin added 2 commits October 31, 2026 23:57

shoumikhin added the release notes: ops & kernels Changes to the opset and any new / changed kernel implementations label Nov 1, 2026

Copilot AI balanced review requested due to automatic review settings November 1, 2026 07:04

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 1, 2026

Copilot AI reviewed Nov 1, 2026

Copilot AI balanced review requested due to automatic review settings November 1, 2026 07:28

Copilot AI reviewed Nov 1, 2026

shoumikhin mentioned this pull request Nov 1, 2026

Fix mixed-dtype cat and stack in to_edge and the edge check #24381

Merged

shoumikhin added 3 commits November 1, 2026 10:59

Copilot AI balanced review requested due to automatic review settings November 1, 2026 19:10

shoumikhin changed the title ~~Keep Half and BFloat16 sums in float in portable matmul, var, cumsum and pooling kernels~~ Keep Half and BFloat16 sums in float in portable matmul, var, cumsum and pooling kernels, and fix portable cat Nov 1, 2026

Copilot AI reviewed Nov 1, 2026

shoumikhin added 3 commits November 1, 2026 11:48

Copilot AI balanced review requested due to automatic review settings November 1, 2026 20:38

Copilot AI reviewed Nov 1, 2026

shoumikhin added 5 commits November 1, 2026 16:21

Merge the two portable addmm BFloat16 tests

24ec32c

BFloat16BiasIsRoundedOnce sums 257 ones, so it already fails without float accumulation, which was all BFloat16LargeDimAccumulatesInFloat checked. Add the full-shape bias it had to the merged test and drop it.

Test the dim check in get_cat_out_target_size

ac757d2

get_cat_out_target_size returns false for a dim that an input does not have, for callers that skip check_cat_args. Portable cat runs check_cat_args first, so no test reached the check. Call the function directly, as the overflow tests do.

Copilot AI balanced review requested due to automatic review settings November 2, 2026 00:36

Copilot AI reviewed Nov 2, 2026

Copilot AI balanced review requested due to automatic review settings November 2, 2026 12:07

Copilot AI reviewed Nov 2, 2026

Copilot AI balanced review requested due to automatic review settings November 2, 2026 12:59

Copilot AI reviewed Nov 2, 2026

shoumikhin mentioned this pull request Nov 2, 2026

TensorRT: LeRobot policies that do not run yet, and what blocks them #24353

Open

SS-JIA approved these changes Nov 2, 2026

View reviewed changes

shoumikhin merged commit db46084 into pytorch:main Nov 2, 2026
226 checks passed

shoumikhin mentioned this pull request Nov 2, 2026

[v1.6.0] Release Schedule and Tracker #23608

Open

pytorchbot mentioned this pull request Nov 2, 2026

Fix mixed-dtype cat and stack in to_edge and the edge check #24450

Merged

shoumikhin mentioned this pull request Nov 2, 2026

OpenVINO: LeRobot policies that do not run yet, and what blocks them #24325

Open

Conversation

shoumikhin commented Nov 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What is wrong today

What this change does

What was tested

Uh oh!

pytorch-bot Bot commented Nov 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24401

✅ No Failures

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

shoumikhin commented Nov 1, 2026

Uh oh!

claude Bot commented Nov 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Review: #24401

Summary

Code Quality

Testing

Performance

Recommendation

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

shoumikhin commented Nov 1, 2026

Uh oh!

claude Bot commented Nov 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Review: #24401

Summary

Security

Code Quality

Performance

Recommendation

Uh oh!

shoumikhin commented Nov 1, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

claude Bot commented Nov 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Review: #24401

Summary

Code Quality

Testing

Security

Performance

Recommendation

Uh oh!

shoumikhin commented Nov 1, 2026

Uh oh!

claude Bot commented Nov 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Review: #24401

Summary

Security

Backward Compatibility

Recommendation

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

shoumikhin commented Nov 2, 2026

Uh oh!

claude Bot commented Nov 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Review: #24401

Summary

Code Quality

Security

shoumikhin commented Nov 1, 2026 •

edited

Loading

pytorch-bot Bot commented Nov 1, 2026 •

edited

Loading

claude Bot commented Nov 1, 2026 •

edited

Loading

claude Bot commented Nov 1, 2026 •

edited

Loading

claude Bot commented Nov 1, 2026 •

edited

Loading

claude Bot commented Nov 1, 2026 •

edited

Loading

claude Bot commented Nov 2, 2026 •

edited

Loading

claude Bot commented Nov 2, 2026 •

edited

Loading

claude Bot commented Nov 2, 2026 •

edited

Loading