Runtime abort (OUT_OF_RANGE, length=2^39+N) invoking un-batched JAX Threefry on a runtime-input key; two or more lanes is fine

[maraxen]

What happened?

An exported StableHLO program that runs JAX's Threefry PRNG on a key supplied as
a runtime input compiles cleanly and then aborts at invocation in HAL
command-buffer validation:

iree/hal/command_buffer_validation.c:143: OUT_OF_RANGE; at least one command
attempted to access an address outside of the valid bound buffer range
(length=549755813984, end(inc)=549755813983, binding offset=0,
binding length=192, binding end(inc)=191); binding table slot 0;
while invoking native function hal.device.queue.execute.indirect

549755813984 is 2**39 + 224. A second spelling of the same program gives
2**39 + 200. The low bits track the shape while bit 39 is set spuriously,
which reads as a width or sign confusion in a buffer-size computation rather
than a genuine overflow. The binding itself is 192 bytes.

This was originally reported as "Additional context" inside #24927, where I said
I would split it out if it turned out to be a separate defect. Having now
characterised it, I believe it is — see Relationship to #24927 below. Nobody
asked me to open this; if you would rather it lived as a comment on #24927,
say so and I will close it.

The boundary is the lane count

The discriminator is how many lanes the draw has, not the rank of the key
input and not whether vmap is involved:

program	result
un-batched draw, (2,) key input	aborts
un-batched draw, (1,2) key input	aborts
un-batched draw, (4,2) key input	aborts
vmap, 1 lane	aborts
vmap, 2 lanes	runs, correct
vmap, 3 / 4 / 8 lanes	runs, correct

Other things that do not matter:

• Draw size. 2, 64 and 1024 elements all abort, so the bogus length does not
  track the requested output.
• Backend. vmvx, the reference backend, aborts identically — so this is
  above backend codegen.
• --iree-opt-const-eval. Disabling it changes nothing.

Things that do make it go away:

• A constant key. A key baked in as a literal is exact, and stays exact with
  --iree-opt-const-eval=false, so this is not simply "the draw got folded".
• JAX's non-partitionable Threefry lowering. Setting
  jax_threefry_partitionable=False produces a different module that runs
  correctly. This was verified paired against a red control in a single process,
  with differing MLIR hashes confirming the lowering actually changed.

Steps to reproduce

Both modules are attached below and need no JAX to reproduce. Both compile
successfully; only the first aborts at run time.

# the failing case: un-batched Threefry on a runtime key input
iree-compile repro_unbatched.mlir \
  --iree-hal-target-backends=llvm-cpu --iree-llvmcpu-target-cpu=host \
  -o repro_unbatched.vmfb
iree-run-module --module=repro_unbatched.vmfb --function=main --input=2xi32=0,0
#   -> OUT_OF_RANGE, length=549755813984 against a 192-byte binding

# the control: the same draw with two lanes
iree-compile control_two_lane.mlir \
  --iree-hal-target-backends=llvm-cpu --iree-llvmcpu-target-cpu=host \
  -o control_two_lane.vmfb
iree-run-module --module=control_two_lane.vmfb --function=main --input=2x2xi32=0,0,0,0
#   -> 2x8xui32=[4070199207 4202968722 1427181096 2012915765 2447653815 710830403 1332275837 2961296638][...]

The control's output matches eager JAX exactly, so the two-lane path is not just
"not crashing" — it is correct.

What component(s) does this issue relate to?

Runtime, and whatever computes the dispatch's buffer binding size. The failure
is in HAL command-buffer validation at invocation, not in compilation.

Version information

• iree-base-compiler 3.11.0
• iree-base-runtime 3.11.0
• jax / jaxlib 0.11.1 (used only to emit the StableHLO; not needed to reproduce)
• Linux x86-64 (WSL2), llvm-cpu and vmvx

Relationship to #24927

These look similar from the outside — both involve jax.random through IREE —
but they are separate defects and I do not think one fix covers both:

	#24927	this issue
symptom	silently wrong values	hard abort at invocation
trigger	a nested split feeding a draw, with a second draw live	un-batched threefry on a runtime-input key
jax_threefry_partitionable=False	still reproduces	fixes it
constant key	still reproduces	exact

The last row is the reason I think they are independent: this one is specific to
the partitionable lowering, and #24927 is not.

Minimal reproducer — repro_unbatched.mlir

#loc = loc(unknown)
#loc1 = loc("k")
#loc2 = loc("/home/marielle/.claude/jobs/a73ad07c/tmp/probe5108/dump_mlir.py":44:14 to :49)
#loc7 = loc("main.<locals>.<lambda>"(#loc2))
module @jit__lambda attributes {jax.uses_shape_polymorphism = false, mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<2xui32> loc("k")) -> (tensor<8xui32> {jax.result_info = "result"}) {
    %c = stablehlo.constant dense<32> : tensor<ui64> loc(#loc33)
    %c_0 = stablehlo.constant dense<1> : tensor<ui64> loc(#loc33)
    %0 = stablehlo.slice %arg0 [0:1] : (tensor<2xui32>) -> tensor<1xui32> loc(#loc34)
    %1 = stablehlo.reshape %0 : (tensor<1xui32>) -> tensor<ui32> loc(#loc34)
    %2 = stablehlo.slice %arg0 [1:2] : (tensor<2xui32>) -> tensor<1xui32> loc(#loc34)
    %3 = stablehlo.reshape %2 : (tensor<1xui32>) -> tensor<ui32> loc(#loc34)
    %4 = stablehlo.iota dim = 0 : tensor<8xui64> loc(#loc33)
    %5 = stablehlo.broadcast_in_dim %c_0, dims = [] : (tensor<ui64>) -> tensor<8xui64> loc(#loc33)
    %6 = stablehlo.multiply %5, %4 : tensor<8xui64> loc(#loc33)
    %7 = stablehlo.broadcast_in_dim %c, dims = [] : (tensor<ui64>) -> tensor<8xui64> loc(#loc33)
    %8 = stablehlo.shift_right_logical %6, %7 : tensor<8xui64> loc(#loc33)
    %9 = stablehlo.convert %6 : (tensor<8xui64>) -> tensor<8xui32> loc(#loc33)
    %10 = stablehlo.convert %8 : (tensor<8xui64>) -> tensor<8xui32> loc(#loc33)
    %11:2 = call @threefry2x32(%1, %3, %10, %9) : (tensor<ui32>, tensor<ui32>, tensor<8xui32>, tensor<8xui32>) -> (tensor<8xui32>, tensor<8xui32>) loc(#loc35)
    %12 = stablehlo.xor %11#0, %11#1 : tensor<8xui32> loc(#loc36)
    return %12 : tensor<8xui32> loc(#loc28)
  } loc(#loc)
  func.func private @threefry2x32(%arg0: tensor<ui32> loc("main.<locals>.<lambda>"(#loc2)), %arg1: tensor<ui32> loc("main.<locals>.<lambda>"(#loc2)), %arg2: tensor<8xui32> loc("main.<locals>.<lambda>"(#loc2)), %arg3: tensor<8xui32> loc("main.<locals>.<lambda>"(#loc2))) -> (tensor<8xui32>, tensor<8xui32>) {
    %c = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %c_0 = stablehlo.constant dense<5> : tensor<i32> loc(#loc)
    %c_1 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %c_2 = stablehlo.constant dense<466688986> : tensor<ui32> loc(#loc11)
    %c_3 = stablehlo.constant dense<[13, 15, 26, 6]> : tensor<4xui32> loc(#loc11)
    %c_4 = stablehlo.constant dense<[17, 29, 16, 24]> : tensor<4xui32> loc(#loc11)
    %0 = stablehlo.xor %arg0, %arg1 : tensor<ui32> loc(#loc13)
    %1 = stablehlo.xor %0, %c_2 : tensor<ui32> loc(#loc13)
    %2 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<ui32>) -> tensor<8xui32> loc(#loc14)
    %3 = stablehlo.add %arg2, %2 : tensor<8xui32> loc(#loc14)
    %4 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<ui32>) -> tensor<8xui32> loc(#loc14)
    %5 = stablehlo.add %arg3, %4 : tensor<8xui32> loc(#loc14)
    %6:9 = stablehlo.while(%iterArg = %c_1, %iterArg_5 = %c_1, %iterArg_6 = %3, %iterArg_7 = %5, %iterArg_8 = %arg1, %iterArg_9 = %1, %iterArg_10 = %arg0, %iterArg_11 = %c_3, %iterArg_12 = %c_4) : tensor<i32>, tensor<i32>, tensor<8xui32>, tensor<8xui32>, tensor<ui32>, tensor<ui32>, tensor<ui32>, tensor<4xui32>, tensor<4xui32>
    cond {
      %7 = stablehlo.compare LT, %iterArg, %c_0, SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc29)
      stablehlo.return %7 : tensor<i1> loc(#loc24)
    } do {
      %7:8 = func.call @eval_jaxpr(%iterArg_5, %iterArg_6, %iterArg_7, %iterArg_8, %iterArg_9, %iterArg_10, %iterArg_11, %iterArg_12) : (tensor<i32>, tensor<8xui32>, tensor<8xui32>, tensor<ui32>, tensor<ui32>, tensor<ui32>, tensor<4xui32>, tensor<4xui32>) -> (tensor<i32>, tensor<8xui32>, tensor<8xui32>, tensor<ui32>, tensor<ui32>, tensor<ui32>, tensor<4xui32>, tensor<4xui32>) loc(#loc30)
      %8 = stablehlo.add %iterArg, %c : tensor<i32> loc(#loc31)
      stablehlo.return %8, %7#0, %7#1, %7#2, %7#3, %7#4, %7#5, %7#6, %7#7 : tensor<i32>, tensor<i32>, tensor<8xui32>, tensor<8xui32>, tensor<ui32>, tensor<ui32>, tensor<ui32>, tensor<4xui32>, tensor<4xui32> loc(#loc24)
    } loc(#loc24)
    return %6#2, %6#3 : tensor<8xui32>, tensor<8xui32> loc(#loc7)
  } loc(#loc7)
  func.func private @eval_jaxpr(%arg0: tensor<i32> loc(unknown), %arg1: tensor<8xui32> loc(unknown), %arg2: tensor<8xui32> loc(unknown), %arg3: tensor<ui32> loc(unknown), %arg4: tensor<ui32> loc(unknown), %arg5: tensor<ui32> loc(unknown), %arg6: tensor<4xui32> loc(unknown), %arg7: tensor<4xui32> loc(unknown)) -> (tensor<i32>, tensor<8xui32>, tensor<8xui32>, tensor<ui32>, tensor<ui32>, tensor<ui32>, tensor<4xui32>, tensor<4xui32>) {
    %c = stablehlo.constant dense<32> : tensor<ui32> loc(#loc)
    %c_0 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %0 = stablehlo.add %arg0, %c_0 : tensor<i32> loc(#loc14)
    %1 = stablehlo.slice %arg6 [0:1] : (tensor<4xui32>) -> tensor<1xui32> loc(#loc17)
    %2 = stablehlo.reshape %1 : (tensor<1xui32>) -> tensor<ui32> loc(#loc17)
    %3 = stablehlo.slice %arg6 [1:2] : (tensor<4xui32>) -> tensor<1xui32> loc(#loc17)
    %4 = stablehlo.reshape %3 : (tensor<1xui32>) -> tensor<ui32> loc(#loc17)
    %5 = stablehlo.slice %arg6 [2:3] : (tensor<4xui32>) -> tensor<1xui32> loc(#loc17)
    %6 = stablehlo.reshape %5 : (tensor<1xui32>) -> tensor<ui32> loc(#loc17)
    %7 = stablehlo.slice %arg6 [3:4] : (tensor<4xui32>) -> tensor<1xui32> loc(#loc17)
    %8 = stablehlo.reshape %7 : (tensor<1xui32>) -> tensor<ui32> loc(#loc17)
    %9 = stablehlo.add %arg1, %arg2 : tensor<8xui32> loc(#loc14)
    %10 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<ui32>) -> tensor<8xui32> loc(#loc18)
    %11 = stablehlo.shift_left %arg2, %10 : tensor<8xui32> loc(#loc18)
    %12 = stablehlo.subtract %c, %2 : tensor<ui32> loc(#loc19)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<ui32>) -> tensor<8xui32> loc(#loc20)
    %14 = stablehlo.shift_right_logical %arg2, %13 : tensor<8xui32> loc(#loc20)
    %15 = stablehlo.or %11, %14 : tensor<8xui32> loc(#loc21)
    %16 = stablehlo.xor %9, %15 : tensor<8xui32> loc(#loc13)
    %17 = stablehlo.add %9, %16 : tensor<8xui32> loc(#loc14)
    %18 = stablehlo.broadcast_in_dim %4, dims = [] : (tensor<ui32>) -> tensor<8xui32> loc(#loc18)
    %19 = stablehlo.shift_left %16, %18 : tensor<8xui32> loc(#loc18)
    %20 = stablehlo.subtract %c, %4 : tensor<ui32> loc(#loc19)
    %21 = stablehlo.broadcast_in_dim %20, dims = [] : (tensor<ui32>) -> tensor<8xui32> loc(#loc20)
    %22 = stablehlo.shift_right_logical %16, %21 : tensor<8xui32> loc(#loc20)
    %23 = stablehlo.or %19, %22 : tensor<8xui32> loc(#loc21)
    %24 = stablehlo.xor %17, %23 : tensor<8xui32> loc(#loc13)
    %25 = stablehlo.add %17, %24 : tensor<8xui32> loc(#loc14)
    %26 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor<ui32>) -> tensor<8xui32> loc(#loc18)
    %27 = stablehlo.shift_left %24, %26 : tensor<8xui32> loc(#loc18)
    %28 = stablehlo.subtract %c, %6 : tensor<ui32> loc(#loc19)
    %29 = stablehlo.broadcast_in_dim %28, dims = [] : (tensor<ui32>) -> tensor<8xui32> loc(#loc20)
    %30 = stablehlo.shift_right_logical %24, %29 : tensor<8xui32> loc(#loc20)
    %31 = stablehlo.or %27, %30 : tensor<8xui32> loc(#loc21)
    %32 = stablehlo.xor %25, %31 : tensor<8xui32> loc(#loc13)
    %33 = stablehlo.add %25, %32 : tensor<8xui32> loc(#loc14)
    %34 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<ui32>) -> tensor<8xui32> loc(#loc18)
    %35 = stablehlo.shift_left %32, %34 : tensor<8xui32> loc(#loc18)
    %36 = stablehlo.subtract %c, %8 : tensor<ui32> loc(#loc19)
    %37 = stablehlo.broadcast_in_dim %36, dims = [] : (tensor<ui32>) -> tensor<8xui32> loc(#loc20)
    %38 = stablehlo.shift_right_logical %32, %37 : tensor<8xui32> loc(#loc20)
    %39 = stablehlo.or %35, %38 : tensor<8xui32> loc(#loc21)
    %40 = stablehlo.xor %33, %39 : tensor<8xui32> loc(#loc13)
    %41 = stablehlo.broadcast_in_dim %arg3, dims = [] : (tensor<ui32>) -> tensor<8xui32> loc(#loc14)
    %42 = stablehlo.add %33, %41 : tensor<8xui32> loc(#loc14)
    %43 = stablehlo.broadcast_in_dim %arg4, dims = [] : (tensor<ui32>) -> tensor<8xui32> loc(#loc14)
    %44 = stablehlo.add %40, %43 : tensor<8xui32> loc(#loc14)
    %45 = stablehlo.add %arg0, %c_0 : tensor<i32> loc(#loc14)
    %46 = stablehlo.convert %45 : (tensor<i32>) -> tensor<ui32> loc(#loc22)
    %47 = stablehlo.broadcast_in_dim %46, dims = [] : (tensor<ui32>) -> tensor<8xui32> loc(#loc14)
    %48 = stablehlo.add %44, %47 : tensor<8xui32> loc(#loc14)
    return %0, %42, %48, %arg4, %arg5, %arg3, %arg7, %arg6 : tensor<i32>, tensor<8xui32>, tensor<8xui32>, tensor<ui32>, tensor<ui32>, tensor<ui32>, tensor<4xui32>, tensor<4xui32> loc(#loc)
  } loc(#loc26)
} loc(#loc)
#loc3 = loc("/home/marielle/.claude/jobs/a73ad07c/tmp/probe5108/dump_mlir.py":23:13 to :49)
#loc4 = loc("/home/marielle/.claude/jobs/a73ad07c/tmp/probe5108/dump_mlir.py":42:2 to 46:3)
#loc5 = loc("/home/marielle/.claude/jobs/a73ad07c/tmp/probe5108/dump_mlir.py":55:2 to :8)
#loc6 = loc("threefry2x32")
#loc8 = loc("emit"(#loc3))
#loc9 = loc("main"(#loc4))
#loc10 = loc("<module>"(#loc5))
#loc11 = loc("threefry2x32:"(#loc6))
#loc12 = loc(callsite(#loc9 at #loc10))
#loc13 = loc("xor"(#loc7))
#loc14 = loc("add"(#loc7))
#loc15 = loc(callsite(#loc7 at #loc7))
#loc16 = loc("while/body/eval_jaxpr"(#loc7))
#loc17 = loc("unstack"(#loc7))
#loc18 = loc("shift_left"(#loc7))
#loc19 = loc("sub"(#loc7))
#loc20 = loc("shift_right_logical"(#loc7))
#loc21 = loc("or"(#loc7))
#loc22 = loc("convert_element_type"(#loc7))
#loc23 = loc(callsite(#loc8 at #loc12))
#loc24 = loc("while"(#loc15))
#loc25 = loc(callsite(#loc15 at #loc7))
#loc26 = loc("eval_jaxpr:"(#loc16))
#loc27 = loc(callsite(#loc7 at #loc23))
#loc28 = loc("jit(<lambda>)"(#loc23))
#loc29 = loc("while/cond/lt"(#loc25))
#loc30 = loc("while/body/eval_jaxpr"(#loc25))
#loc31 = loc("while/body/add"(#loc25))
#loc32 = loc(callsite(#loc7 at #loc27))
#loc33 = loc("jit(<lambda>)/iota_2x32_shape"(#loc32))
#loc34 = loc("jit(<lambda>)/unstack"(#loc32))
#loc35 = loc("jit(<lambda>)"(#loc32))
#loc36 = loc("jit(<lambda>)/xor"(#loc32))

The two-lane control (control_two_lane.mlir) is the same program under
jax.vmap over a (2,2) key input. I have left it out of the body for length;
it is the identical module with the loop's tensors gaining a leading 2 and the
entry point taking tensor<2x2xi32>. Happy to attach it, or a --dump-compilation-phases
trace, if that would help.

What I have not done

I have not tried to localize this inside IREE — I do not know whether the bad
length originates in the ABI wrapper, in the dispatch's workgroup-count
computation, or in the binding-table construction. The 2**39 bit and the
192-byte binding are the only two concrete facts I have about the mechanism.

---

AI tool use disclosure

Assisted-by: Claude Opus 5 (Claude Code)

Per the project's
AI tool use policy.
The investigation, the minimization, and the text of this report were
substantially AI-assisted. Every number quoted above is machine-measured rather
than asserted — the commands shown are the ones that produced them, and the
whole thing is re-derivable from the attached MLIR with iree-compile +
iree-run-module alone. Happy to trim or restructure this if a shorter report
would be more useful.

[github-actions]

Hello @maraxen 👋

Thank you for submitting an issue to IREE! It looks like it's your first one.

The volume of issues is huge, and triage times may vary. If you need quicker assistance on the issue, consider posting it on IREE Discord and providing short context there.
Feel free to identify the most frequent contributors to the corresponding part of the codebase via Git history/CODEOWNERS and (gently!) tag them below.

---

Action required: acknowledge IREE project policies

IREE is a Linux Foundation project. All participants are expected to follow the LF Projects Code of Conduct.

Please note that all contributions to IREE, including GitHub issues, must follow our IREE AI Tool Use Policy. This means:

• Issue reporters must fully understand the primary issue and the impact of the problem.
  For bug reports, specifics of the use-case, workflow and/or crash scenario usually suffice.
  For improvement suggestions, this usually implies deeper understanding of suboptimal code, IR design choices, assumed performance tradeoffs, etc.
• Contributors must write the summary of the issue themselves, so that the exact impact and scenario is clear to other participants.
• Issue authors are expected to review and understand all LLM/AI-assisted segments of the report, if any. E.g.: CLI commands and IR reproducers, model/framework details, lowering pipeline details, assumptions about problematic behaviors of IREE components.
• Substantial use of LLM/generative AI tools must be disclaimed in the header or footer of the issue description.

Before the issue can be triaged, we kindly ask you to reply to this message and confirm that you understand and accept the proposed policy. As required, please adjust your initial issue description.

---

If you have any questions, feel free to leave a comment here, or ask away on IREE Discord.

Thank you,
The IREE Community