BroadcastIota: only fold when the broadcast keeps the iota's extent

[wsmoses]

BroadcastIota folds broadcast_in_dim(iota) into an iota of the result along the mapped dimension. That is only right when the broadcast leaves that dimension's size unchanged. If the iota dimension has size 1, the operand is all zeros and the broadcast repeats that 0. The fold instead produces 0, 1, 2, ….

%0 = stablehlo.iota dim = 0 : tensor<1x4xi64>
%1 = stablehlo.broadcast_in_dim %0, dims = [0, 1] : (tensor<1x4xi64>) -> tensor<3x4xi64>

On main, %1 becomes [[0,0,0,0],[1,1,1,1],[2,2,2,2]]; it should be all zeros. The pattern now requires the iota dimension to keep its extent. The new case in broadcastiota.mlir folds the zero away.

This came up as a miscompile in the MFEM quadrature-interpolator kernels (Derivatives2D with vdim = 1) under the exec-time transform pattern list:

1. broadcastindim_licm hoists a broadcast of the trip-1 vdim loop's induction variable out of the inner loop.
2. greedy_while_loop_batch_fission batches it as an iota of extent 1.
3. enzyme_hlo_unroll(1) unrolls the trip-1 loop.
4. The size-1 iota reaches the inner loop's broadcast to extent 2. broadcast_iota turns the repeated 0 into 0, 1, so the store index adds the inner index twice.

With this change the runtime's full exec pipeline gives the same values as the unoptimized module. The whole lit suite has no new failures.

🤖 Generated with Claude Code

https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 29.62%. Comparing base (9409410) to head (907293f).
⚠️ Report is 30 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #3410   +/-   ##
=======================================
  Coverage   29.61%   29.62%           
=======================================
  Files         240      240           
  Lines       48460    48474   +14     
=======================================
+ Hits        14351    14359    +8     
- Misses      34109    34115    +6     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.