Failure on CPU: common_pjrt_client.cc:1617] ExecutePrepareWithOomRetries failed: INVALID_ARGUMENT: Buffer…

[cxxhou]

Running the attached StableHLO module fails on CPU. Failure signature: common_pjrt_client.cc:1617] ExecutePrepareWithOomRetries failed: INVALID_ARGUMENT: Buffer…. The same module runs successfully on Interpreter.

The following backends fail without a crash signal:

• run_hlo_module on GPU (StableHLO input) — failed (rc=255)
• run_hlo_module on CPU (HLO proto input) — failed (rc=255)
• run_hlo_module on GPU (HLO proto input) — failed (rc=255)

The same module runs successfully on:

• run_hlo_module on Interpreter (StableHLO input)
• run_hlo_module on Interpreter (HLO proto input)

The module also compiles cleanly on:

• hlo_runner_main on CPU (HLO proto input, compile-only via --run=false)
• hlo_runner_main on GPU (HLO proto input, compile-only via --run=false)

The StableHLO→HLO-proto conversion:

• hlo-translate — succeeds

Command:

hlo-translate \
  --mlir-to-hlo \
  --emit-proto \
  b731ccd1e79c256a.mlir \
  -o \
  b731ccd1e79c256a.hlo
sed -i '1{/^goo\.gle\//d}' b731ccd1e79c256a.hlo

Environment

• CPU: Intel(R) Core(TM) i9-14900HX
• GPU: NVIDIA GeForce RTX 4060 Laptop GPU
• CUDA Driver: 580.126.09
• XLA commit: 6f1f2d916dcef0428bfa7ea66765be539126c4f2
• StableHLO commit: e6f81ebd06b3509f2c7fa6175430aadbd4d724ca

IR

module @MyModule attributes {mhlo.cross_program_prefetches = [], mhlo.input_output_alias = [], mhlo.is_dynamic = false, mhlo.use_auto_spmd_partitioning = false, mhlo.xla_entry_computation_parameter_layouts = [dense<[2, 1, 0]> : tensor<3xindex>, dense<[2, 1, 0]> : tensor<3xindex>], mhlo.xla_entry_computation_parameter_tiles = [[dense<[2, 128]> : tensor<2xindex>], [dense<[2, 128]> : tensor<2xindex>]], mhlo.xla_entry_computation_result_layout = [dense<[2, 1, 0]> : tensor<3xindex>], mhlo.xla_entry_computation_result_tiles = [[dense<[2, 128]> : tensor<2xindex>]]} {
  func.func @main(%arg0: tensor<32x1x1xf32>, %arg1: tensor<1x1x1xf32>) -> tensor<32x1x1xf32> {
    %c = stablehlo.constant dense<0> : tensor<i32>
    %0 = stablehlo.dynamic_update_slice %arg0, %arg1, %c, %c, %c {result_layout = dense<[2, 1, 0]> : tensor<3xindex>, xla_shape = "f32[32,1,1]{2,1,0:T(2,128)S(5)}"} : (tensor<32x1x1xf32>, tensor<1x1x1xf32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<32x1x1xf32>
    %1 = stablehlo.reshape %0 : (tensor<32x1x1xf32>) -> tensor<32xf32>
    %2 = stablehlo.reshape %1 : (tensor<32xf32>) -> tensor<1x32xf32>
    %3 = stablehlo.broadcast_in_dim %2, dims = [0, 1] : (tensor<1x32xf32>) -> tensor<497x32xf32>
    %4 = stablehlo.reshape %3 : (tensor<497x32xf32>) -> tensor<15904xf32>
    %5 = stablehlo.slice %4 [0:15875] : (tensor<15904xf32>) -> tensor<15875xf32>
    %6 = stablehlo.reshape %5 : (tensor<15875xf32>) -> tensor<127x125xf32>
    %7 = call @wrapped_1142(%6) : (tensor<127x125xf32>) -> tensor<127x125xcomplex<f32>>
    %8 = stablehlo.convert %7 : (tensor<127x125xcomplex<f32>>) -> tensor<127x125xf32>
    %9 = stablehlo.reshape %8 : (tensor<127x125xf32>) -> tensor<15875xf32>
    %10 = stablehlo.slice %9 [0:32] : (tensor<15875xf32>) -> tensor<32xf32>
    %11 = stablehlo.reshape %10 : (tensor<32xf32>) -> tensor<32x1x1xf32>
    return %11 : tensor<32x1x1xf32>
  }
  func.func private @wrapped_1142(%arg0: tensor<127x125xf32>) -> tensor<127x125xcomplex<f32>> {
    %cst = stablehlo.constant dense<0xFF800000> : tensor<f32>
    %0 = stablehlo.reduce(%arg0 init: %cst) applies stablehlo.maximum across dimensions = [1] : (tensor<127x125xf32>, tensor<f32>) -> tensor<127xf32>
    %1 = stablehlo.broadcast_in_dim %0, dims = [0] : (tensor<127xf32>) -> tensor<127x125xf32>
    %2 = stablehlo.subtract %arg0, %1 : tensor<127x125xf32>
    %3 = stablehlo.complex %2, %2 : tensor<127x125xcomplex<f32>>
    return %3 : tensor<127x125xcomplex<f32>>
  }
  func.func private @max_computation_1142(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
    %0 = stablehlo.maximum %arg0, %arg1 : tensor<f32>
    return %0 : tensor<f32>
  }
}

Reproduction (CPU (StableHLO input) — failed (rc=255))

Command:

export XLA_FLAGS='--xla_backend_extra_options=xla_cpu_use_multi_output_fusion --xla_enable_enzyme_comms_opt=true --xla_gpu_async_dot=true --xla_gpu_enable_cudnn_layer_norm=true --xla_gpu_enable_dus_accumulator_zero_init_elimination=true --xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_enable_scatter_determinism_expander=true --xla_gpu_experimental_stream_annotation=true --xla_gpu_experimental_use_ragged_dot_fusion --xla_gpu_gemm_rewrite_size_threshold=0 --xla_gpu_multi_streamed_windowed_einsum=true --xla_recognize_reduction_optimization_level=1 --xla_reduce_window_rewrite_base_length=1'
run_hlo_module \
  --platform=CPU \
  --reference_platform= \
  --input_format=stablehlo \
  b731ccd1e79c256a.mlir

Output:

 ** Running b731ccd1e79c256a.mlir**
Running HLO module with runner HloRunner...
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1786848678.520290 3843521 common_pjrt_client.cc:1617] ExecutePrepareWithOomRetries failed: INVALID_ARGUMENT: Buffer passed to Execute() as argument 0 to replica 0 has unexpected shape: f32[32,1,1]{2,1,0} (expected f32[32,1,1]{2,1,0:T(2,128)}).
... compiled and ran in 0.0261369s.
INVALID_ARGUMENT: Buffer passed to Execute() as argument 0 to replica 0 has unexpected shape: f32[32,1,1]{2,1,0} (expected f32[32,1,1]{2,1,0:T(2,128)}).
	Failed to execute on HloRunner

Reference run (Interpreter (StableHLO input) — succeeds)

Command:

export XLA_FLAGS='--xla_backend_extra_options=xla_cpu_use_multi_output_fusion --xla_enable_enzyme_comms_opt=true --xla_gpu_async_dot=true --xla_gpu_enable_cudnn_layer_norm=true --xla_gpu_enable_dus_accumulator_zero_init_elimination=true --xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_enable_scatter_determinism_expander=true --xla_gpu_experimental_stream_annotation=true --xla_gpu_experimental_use_ragged_dot_fusion --xla_gpu_gemm_rewrite_size_threshold=0 --xla_gpu_multi_streamed_windowed_einsum=true --xla_recognize_reduction_optimization_level=1 --xla_reduce_window_rewrite_base_length=1'
run_hlo_module \
  --platform=Interpreter \
  --reference_platform= \
  --input_format=stablehlo \
  b731ccd1e79c256a.mlir

Output:

 ** Running b731ccd1e79c256a.mlir**
Running HLO module with runner HloRunner...
... compiled and ran in 0.0150989s.
Skipping reference runner

Contact

• Email: ch395@njit.edu, zhihao.yao@njit.edu, benquike@gmail.com