Skip to content

Failure on CPU: copy_insertion.cc:427 param != root #47392

Description

@cxxhou

Running the attached StableHLO module aborts run_hlo_module on CPU through an internal XLA check. Failure signature: copy_insertion.cc:427 param != root. 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)

The following backends do not terminate within the replay timeout of 30s:

  • run_hlo_module on Interpreter (StableHLO input)

The same module runs successfully on:

  • run_hlo_module on Interpreter (HLO proto input)
  • run_hlo_module on CPU (HLO proto input)
  • run_hlo_module on GPU (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 \
  dedc9236f4bf09bd.mlir \
  -o \
  dedc9236f4bf09bd.hlo
sed -i '1{/^goo\.gle\//d}' dedc9236f4bf09bd.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 @IrToHlo.185 attributes {mhlo.cross_program_prefetches = [], mhlo.input_output_alias = [], mhlo.is_dynamic = false, mhlo.use_auto_spmd_partitioning = false} {
  func.func @main(%arg0: tensor<2x1x1x1xf32>, %arg1: tensor<2x1x1x1xf32>, %arg2: tensor<2x1x1x1x1xf32>, %arg3: tensor<2x1x1x1xf32>, %arg4: tensor<4xf32>, %arg5: tensor<4xf32>, %arg6: tensor<4xf32>, %arg7: tensor<4xf32>) -> (tensor<2x1x1x1xf32>, tensor<2x1x1x1xf32>, tensor<2x1x1x1xf32>, tensor<2x4x1x1xf32>, tensor<2x4x1x1xf32>, tensor<2x1x1xi64>) {
    %cst = stablehlo.constant dense<1.000000e+00> : tensor<2x1x1x1xf32>
    %cst_0 = stablehlo.constant dense<0xFF800000> : tensor<f32>
    %cst_1 = stablehlo.constant dense<3.40282347E+38> : tensor<f32>
    %c = stablehlo.constant dense<0> : tensor<i32>
    %cst_2 = stablehlo.constant dense<0x7F800000> : tensor<f32>
    %0 = stablehlo.reduce(%arg2 init: %cst_1) applies stablehlo.minimum across dimensions = [4] : (tensor<2x1x1x1x1xf32>, tensor<f32>) -> tensor<2x1x1x1xf32>
    %1 = "stablehlo.reduce_window"(%0, %cst_0) <{base_dilations = array<i64: 1, 1, 1, 1>, padding = dense<0> : tensor<4x2xi64>, window_dilations = array<i64: 1, 1, 1, 1>, window_dimensions = array<i64: 1, 1, 1, 1>, window_strides = array<i64: 1, 1, 1, 1>}> ({
    ^bb0(%arg8: tensor<f32>, %arg9: tensor<f32>):
      %12 = stablehlo.maximum %arg8, %arg9 : tensor<f32>
      stablehlo.return %12 : tensor<f32>
    }) : (tensor<2x1x1x1xf32>, tensor<f32>) -> tensor<2x1x1x1xf32>
    %2 = stablehlo.concatenate %0, %arg1, %arg0, %arg3, dim = 1 : (tensor<2x1x1x1xf32>, tensor<2x1x1x1xf32>, tensor<2x1x1x1xf32>, tensor<2x1x1x1xf32>) -> tensor<2x4x1x1xf32>
    %3 = "stablehlo.reduce_window"(%2, %cst_0) <{base_dilations = array<i64: 1, 1, 1, 1>, padding = dense<0> : tensor<4x2xi64>, window_dilations = array<i64: 1, 1, 1, 1>, window_dimensions = array<i64: 1, 1, 1, 1>, window_strides = array<i64: 1, 1, 1, 1>}> ({
    ^bb0(%arg8: tensor<f32>, %arg9: tensor<f32>):
      %12 = stablehlo.maximum %arg8, %arg9 : tensor<f32>
      stablehlo.return %12 : tensor<f32>
    }) : (tensor<2x4x1x1xf32>, tensor<f32>) -> tensor<2x4x1x1xf32>
    %4 = stablehlo.tan %arg0 : tensor<2x1x1x1xf32>
    %5 = stablehlo.atan2 %arg1, %cst : tensor<2x1x1x1xf32>
    %6 = stablehlo.round_nearest_even %2 : tensor<2x4x1x1xf32>
    %7 = "stablehlo.batch_norm_inference"(%6, %arg7, %arg6, %arg5, %arg4) <{epsilon = 9.99999974E-6 : f32, feature_index = 1 : i64}> : (tensor<2x4x1x1xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>) -> tensor<2x4x1x1xf32>
    %8 = stablehlo.iota dim = 0 : tensor<4xi32>
    %9 = stablehlo.broadcast_in_dim %8, dims = [1] : (tensor<4xi32>) -> tensor<2x4x1x1xi32>
    %10:2 = stablehlo.reduce(%6 init: %cst_2), (%9 init: %c) across dimensions = [1] : (tensor<2x4x1x1xf32>, tensor<2x4x1x1xi32>, tensor<f32>, tensor<i32>) -> (tensor<2x1x1xf32>, tensor<2x1x1xi32>)
     reducer(%arg8: tensor<f32>, %arg10: tensor<f32>) (%arg9: tensor<i32>, %arg11: tensor<i32>)  {
      %12 = stablehlo.compare  LE, %arg8, %arg10 : (tensor<f32>, tensor<f32>) -> tensor<i1>
      %13 = stablehlo.select %12, %arg8, %arg10 : tensor<i1>, tensor<f32>
      %14 = stablehlo.compare  EQ, %arg8, %arg10 : (tensor<f32>, tensor<f32>) -> tensor<i1>
      %15 = stablehlo.minimum %arg9, %arg11 : tensor<i32>
      %16 = stablehlo.reshape %13 : (tensor<f32>) -> tensor<1xf32>
      %17 = stablehlo.reshape %16 : (tensor<1xf32>) -> tensor<1x1xf32>
      %18 = stablehlo.broadcast_in_dim %17, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<8x1xf32>
      %19 = stablehlo.reshape %18 : (tensor<8x1xf32>) -> tensor<8xf32>
      %20 = func.call @wrapped_5149(%arg9, %19) : (tensor<i32>, tensor<8xf32>) -> tensor<i32>
      %21 = stablehlo.select %12, %arg9, %arg11 : tensor<i1>, tensor<i32>
      %22 = stablehlo.add %15, %20 : tensor<i32>
      %23 = stablehlo.select %14, %22, %21 : tensor<i1>, tensor<i32>
      stablehlo.return %13, %23 : tensor<f32>, tensor<i32>
    }
    %11 = stablehlo.convert %10#1 : (tensor<2x1x1xi32>) -> tensor<2x1x1xi64>
    return %4, %5, %1, %3, %7, %11 : tensor<2x1x1x1xf32>, tensor<2x1x1x1xf32>, tensor<2x1x1x1xf32>, tensor<2x4x1x1xf32>, tensor<2x4x1x1xf32>, tensor<2x1x1xi64>
  }
  func.func private @wrapped_5149(%arg0: tensor<i32>, %arg1: tensor<8xf32>) -> tensor<i32> {
    %0:2 = stablehlo.while(%iterArg = %arg0, %iterArg_0 = %arg1) : tensor<i32>, tensor<8xf32>
    cond {
      %c = stablehlo.constant dense<10> : tensor<i32>
      %3 = stablehlo.compare  LT, %iterArg, %c : (tensor<i32>, tensor<i32>) -> tensor<i1>
      stablehlo.return %3 : tensor<i1>
    } do {
      %3 = stablehlo.convert %iterArg : (tensor<i32>) -> tensor<f32>
      %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<8xf32>
      %5 = stablehlo.add %iterArg_0, %4 : tensor<8xf32>
      stablehlo.return %iterArg, %5 : tensor<i32>, tensor<8xf32>
    }
    %1:2 = stablehlo.while(%iterArg = %arg0, %iterArg_0 = %arg1) : tensor<i32>, tensor<8xf32>
    cond {
      %c = stablehlo.constant dense<10> : tensor<i32>
      %3 = stablehlo.compare  LT, %iterArg, %c : (tensor<i32>, tensor<i32>) -> tensor<i1>
      stablehlo.return %3 : tensor<i1>
    } do {
      %3 = stablehlo.convert %iterArg : (tensor<i32>) -> tensor<f32>
      %4 = stablehlo.broadcast_in_dim %3, dims = [] : (tensor<f32>) -> tensor<8xf32>
      %5 = stablehlo.add %iterArg_0, %4 : tensor<8xf32>
      stablehlo.return %iterArg, %5 : tensor<i32>, tensor<8xf32>
    }
    %2 = stablehlo.add %0#0, %1#0 : tensor<i32>
    return %2 : tensor<i32>
  }
  func.func private @DependentTupleElements_OneReadOnly_TwoLoops_EntryParams.Body_5149(%arg0: tensor<i32>, %arg1: tensor<8xf32>) -> (tensor<i32>, tensor<8xf32>) {
    %0 = stablehlo.convert %arg0 : (tensor<i32>) -> tensor<f32>
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<8xf32>
    %2 = stablehlo.add %arg1, %1 : tensor<8xf32>
    return %arg0, %2 : tensor<i32>, tensor<8xf32>
  }
  func.func private @DependentTupleElements_OneReadOnly_TwoLoops_EntryParams.Condition_5149(%arg0: tensor<i32>, %arg1: tensor<8xf32>) -> tensor<i1> {
    %c = stablehlo.constant dense<10> : tensor<i32>
    %0 = stablehlo.compare  LT, %arg0, %c : (tensor<i32>, tensor<i32>) -> tensor<i1>
    return %0 : tensor<i1>
  }
  func.func private @DependentTupleElements_OneReadOnly_TwoLoops_EntryParams.Body.1_5149(%arg0: tensor<i32>, %arg1: tensor<8xf32>) -> (tensor<i32>, tensor<8xf32>) {
    %0 = stablehlo.convert %arg0 : (tensor<i32>) -> tensor<f32>
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<8xf32>
    %2 = stablehlo.add %arg1, %1 : tensor<8xf32>
    return %arg0, %2 : tensor<i32>, tensor<8xf32>
  }
  func.func private @DependentTupleElements_OneReadOnly_TwoLoops_EntryParams.Condition.1_5149(%arg0: tensor<i32>, %arg1: tensor<8xf32>) -> tensor<i1> {
    %c = stablehlo.constant dense<10> : tensor<i32>
    %0 = stablehlo.compare  LT, %arg0, %c : (tensor<i32>, tensor<i32>) -> tensor<i1>
    return %0 : tensor<i1>
  }
}

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 \
  dedc9236f4bf09bd.mlir

Output:


 ** Running dedc9236f4bf09bd.mlir**
Running HLO module with runner HloRunner...
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1786849928.663291 3895706 status_macros.cc:58] INTERNAL: RET_CHECK failure (xla/service/copy_insertion.cc:427) param != root
*** Begin stack trace ***

























	__libc_start_main

*** End stack trace ***

... compiled and ran in 0.00360393s.
INTERNAL: RET_CHECK failure (xla/service/copy_insertion.cc:427) param != root
	Failed to execute on HloRunner

Reference run (Interpreter (HLO proto 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=pbtxt \
  dedc9236f4bf09bd.hlo

Output:

WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1786849929.501610 3895954 text_format.cc:378] Error parsing text-format xla.HloSnapshot: 1:12: Message type "xla.HloSnapshot" has no field named "hlo_module".
I0000 00:00:1786849929.502542 3895954 run_hlo_module.cc:578] Using input data from the user-provided snapshot.
E0000 00:00:1786849929.502563 3895954 text_format.cc:378] Error parsing text-format xla.HloSnapshot: 1:12: Message type "xla.HloSnapshot" has no field named "hlo_module".

 ** Running dedc9236f4bf09bd.hlo**
Running HLO module with runner HloRunner...
... compiled and ran in 0.0046433s.
Skipping reference runner

Contact

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    CPURelated to XLA on CPUbugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions