hlo-translate failure: segv:mlir::stablehlo::ReduceOp::build

[cxxhou]

hlo-translate fails to convert the attached StableHLO module to an HLO proto. run_hlo_module executes the same module successfully through its direct StableHLO frontend. Failure signature: segv:mlir::stablehlo::ReduceOp::build.

The same module runs successfully on:

• run_hlo_module on Interpreter (StableHLO input)
• run_hlo_module on CPU (StableHLO input)
• run_hlo_module on GPU (StableHLO input)

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 @KeepUsedSubcomputation attributes {mhlo.cross_program_prefetches = [], mhlo.input_output_alias = [], mhlo.is_dynamic = false, mhlo.use_auto_spmd_partitioning = false} {
  func.func private @reduction_subcomp(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
    %0 = stablehlo.add %arg0, %arg1 : tensor<f32>
    return %0 : tensor<f32>
  }
  func.func @main(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<f32> {
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f32>
    %0 = stablehlo.reduce(%arg0 init: %cst) across dimensions = [0] : (tensor<100xf32>, tensor<f32>) -> tensor<f32>
     reducer(%arg2: tensor<f32>, %arg3: tensor<f32>)  {
      %2 = stablehlo.reshape %arg3 : (tensor<f32>) -> tensor<1xf32>
      %3 = stablehlo.reshape %2 : (tensor<1xf32>) -> tensor<1x1xf32>
      %4 = stablehlo.broadcast_in_dim %3, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<108x1xf32>
      %5 = stablehlo.reshape %4 : (tensor<108x1xf32>) -> tensor<108xf32>
      %6 = stablehlo.reshape %5 : (tensor<108xf32>) -> tensor<3x3x3x4xf32>
      %7 = stablehlo.reshape %4 : (tensor<108x1xf32>) -> tensor<108xf32>
      %8 = stablehlo.slice %7 [0:12] : (tensor<108xf32>) -> tensor<12xf32>
      %9 = stablehlo.reshape %8 : (tensor<12xf32>) -> tensor<4x3xf32>
      %10 = stablehlo.reshape %3 : (tensor<1x1xf32>) -> tensor<1xf32>
      %11 = stablehlo.reshape %10 : (tensor<1xf32>) -> tensor<1x1xf32>
      %12 = stablehlo.broadcast_in_dim %11, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<18x1xf32>
      %13 = stablehlo.reshape %12 : (tensor<18x1xf32>) -> tensor<18xf32>
      %14 = stablehlo.reshape %13 : (tensor<18xf32>) -> tensor<3x6xf32>
      %15 = stablehlo.reshape %4 : (tensor<108x1xf32>) -> tensor<108xf32>
      %16 = stablehlo.slice %15 [0:48] : (tensor<108xf32>) -> tensor<48xf32>
      %17 = stablehlo.reshape %16 : (tensor<48xf32>) -> tensor<1x4x4x3xf32>
      %18:2 = func.call @wrapped_9451(%9, %14, %17, %6) : (tensor<4x3xf32>, tensor<3x6xf32>, tensor<1x4x4x3xf32>, tensor<3x3x3x4xf32>) -> (tensor<4x6xf32>, tensor<1x2x2x4xf32>)
      %19 = stablehlo.add %arg2, %arg3 : tensor<f32>
      %20 = stablehlo.reshape %18#1 : (tensor<1x2x2x4xf32>) -> tensor<16xf32>
      %21 = stablehlo.slice %20 [0:1] : (tensor<16xf32>) -> tensor<1xf32>
      %22 = stablehlo.reshape %21 : (tensor<1xf32>) -> tensor<f32>
      %23 = stablehlo.add %19, %22 : tensor<f32>
      stablehlo.return %23 : tensor<f32>
    }
    %cst_0 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
    %1 = stablehlo.reduce(%arg1 init: %cst_0) applies stablehlo.add across dimensions = [0] : (tensor<100xf32>, tensor<f32>) -> tensor<f32>
    return %1 : tensor<f32>
  }
  func.func private @wrapped_9451(%arg0: tensor<4x3xf32>, %arg1: tensor<3x6xf32>, %arg2: tensor<1x4x4x3xf32>, %arg3: tensor<3x3x3x4xf32>) -> (tensor<4x6xf32>, tensor<1x2x2x4xf32>) {
    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<4x3xf32>) -> tensor<3x4xf32>
    %1 = stablehlo.dot_general %0, %arg1, contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] : (tensor<3x4xf32>, tensor<3x6xf32>) -> tensor<4x6xf32>
    %2 = stablehlo.transpose %arg3, dims = [1, 0, 2, 3] : (tensor<3x3x3x4xf32>) -> tensor<3x3x3x4xf32>
    %3 = stablehlo.convolution(%arg2, %2) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[0, 0], [0, 0]], lhs_dilate = [1, 1], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]} : (tensor<1x4x4x3xf32>, tensor<3x3x3x4xf32>) -> tensor<1x2x2x4xf32>
    return %1, %3 : tensor<4x6xf32>, tensor<1x2x2x4xf32>
  }
}

Reproduction (hlo-translate — failed (rc=-11, SIGSEGV))

Command:

hlo-translate \
  --mlir-to-hlo \
  --emit-proto \
  008b68e90c0ff552.mlir

Output:

PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace and instructions to reproduce the bug.
 #0 0x00005b7558af244b llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (hlo-translate+0xae5a44b)
 #1 0x00005b7558aefc03 llvm::sys::RunSignalHandlers() (hlo-translate+0xae57c03)
 #2 0x00005b7558af318c SignalHandler(int, siginfo_t*, void*) Signals.cpp:0:0
 #3 0x000075b020245330 (/lib/x86_64-linux-gnu/libc.so.6+0x45330)
 #4 0x00005b7558820887 mlir::RankedTensorType::getShape() const (hlo-translate+0xab88887)
 #5 0x00005b755825bd71 mlir::hlo::verifyReduceOpInputsAndInferShape(std::optional<mlir::Location>, llvm::SmallVector<mlir::ShapedType, 3u>, llvm::ArrayRef<long>, llvm::SmallVector<long, 6u>&, mlir::Attribute&) (hlo-translate+0xa5c3d71)
 #6 0x00005b75580a22a7 mlir::stablehlo::ReduceOp::build(mlir::OpBuilder&, mlir::OperationState&, mlir::ValueRange, mlir::ValueRange, mlir::detail::DenseArrayAttrImpl<long>, mlir::TypeRange) (hlo-translate+0xa40a2a7)
 #7 0x00005b755810b344 mlir::stablehlo::ReduceOp::create(mlir::OpBuilder&, mlir::Location, mlir::ValueRange, mlir::ValueRange, mlir::detail::DenseArrayAttrImpl<long>, mlir::TypeRange) (hlo-translate+0xa473344)
 #8 0x00005b75580066f2 mlir::stablehlo::(anonymous namespace)::ReduceOpUnusedResultCanon::matchAndRewrite(mlir::stablehlo::ReduceOp, mlir::PatternRewriter&) const StablehloAggressiveSimplification.cpp:0:0
 #9 0x00005b75585aba20 void llvm::function_ref<void ()>::callback_fn<mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<llvm::LogicalResult (mlir::Pattern const&)>)::$_0>(long) PatternApplicator.cpp:0:0
#10 0x00005b75585a9211 mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<llvm::LogicalResult (mlir::Pattern const&)>) (hlo-translate+0xa911211)
#11 0x00005b755857fd23 (anonymous namespace)::GreedyPatternRewriteDriver::processWorklist() GreedyPatternRewriteDriver.cpp:0:0
#12 0x00005b755857f4d9 void llvm::function_ref<void ()>::callback_fn<(anonymous namespace)::RegionPatternRewriteDriver::simplify(bool*) &&::$_2>(long) GreedyPatternRewriteDriver.cpp:0:0
#13 0x00005b755857dbf5 mlir::applyPatternsGreedily(mlir::Region&, mlir::FrozenRewritePatternSet const&, mlir::GreedyRewriteConfig, bool*) (hlo-translate+0xa8e5bf5)
#14 0x00005b755800c38d mlir::stablehlo::StablehloTargetIndependentOptimizationPass::runOnOperation() (hlo-translate+0xa37438d)
#15 0x00005b755875d235 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (hlo-translate+0xaac5235)
#16 0x00005b755875d7d6 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) (hlo-translate+0xaac57d6)
#17 0x00005b7558762b99 void llvm::detail::UniqueFunctionBase<void>::CallImpl<llvm::LogicalResult mlir::failableParallelForEach<__gnu_cxx::__normal_iterator<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo*, std::vector<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo, std::allocator<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo>>>, void mlir::parallelForEach<__gnu_cxx::__normal_iterator<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo*, std::vector<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo, std::allocator<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo>>>, mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::$_0>(mlir::MLIRContext*, __gnu_cxx::__normal_iterator<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo*, std::vector<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo, std::allocator<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo>>>, __gnu_cxx::__normal_iterator<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo*, std::vector<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo, std::allocator<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo>>>, mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::$_0&&)::'lambda'(__gnu_cxx::__normal_iterator<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo*, std::vector<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo, std::allocator<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo>>>&&)>(mlir::MLIRContext*, __gnu_cxx::__normal_iterator<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo*, std::vector<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo, std::allocator<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo>>>, __gnu_cxx::__normal_iterator<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo*, std::vector<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo, std::allocator<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo>>>, mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::$_0&&)::'lambda'()>(void*) Pass.cpp:0:0
#18 0x00005b755858749c std::__future_base::_Task_setter<std::unique_ptr<std::__future_base::_Result<void>, std::__future_base::_Result_base::_Deleter>, std::thread::_Invoker<std::tuple<llvm::unique_function<void ()>>>, void>::operator()() const (hlo-translate+0xa8ef49c)
#19 0x00005b7558587422 std::_Function_handler<std::unique_ptr<std::__future_base::_Result_base, std::__future_base::_Result_base::_Deleter> (), std::__future_base::_Task_setter<std::unique_ptr<std::__future_base::_Result<void>, std::__future_base::_Result_base::_Deleter>, std::thread::_Invoker<std::tuple<llvm::unique_function<void ()>>>, void>>::_M_invoke(std::_Any_data const&) (hlo-translate+0xa8ef422)
#20 0x00005b7552919c4f std::__future_base::_State_baseV2::_M_do_set(std::function<std::unique_ptr<std::__future_base::_Result_base, std::__future_base::_Result_base::_Deleter> ()>*, bool*) (hlo-translate+0x4c81c4f)
#21 0x000075b0202a1ed3 __pthread_once_slow ./nptl/pthread_once.c:118:7
#22 0x00005b75585878e8 std::__future_base::_Deferred_state<std::thread::_Invoker<std::tuple<llvm::unique_function<void ()>>>, void>::_M_complete_async() (hlo-translate+0xa8ef8e8)
#23 0x00005b7558587a37 void llvm::detail::UniqueFunctionBase<void>::CallImpl<std::shared_future<void> llvm::ThreadPoolInterface::asyncImpl<void>(llvm::unique_function<void ()>, llvm::ThreadPoolTaskGroup*)::'lambda'()>(void*) (hlo-translate+0xa8efa37)
#24 0x00005b7558b04f22 llvm::StdThreadPool::processTasks(llvm::ThreadPoolTaskGroup*) (hlo-translate+0xae6cf22)
#25 0x00005b7558b069d9 void* llvm::thread::ThreadProxy<std::tuple<llvm::StdThreadPool::grow(int)::$_0>>(void*) ThreadPool.cpp:0:0
#26 0x000075b02029caa4 start_thread ./nptl/pthread_create.c:447:8
#27 0x000075b020329c6c clone3 ./misc/../sysdeps/unix/sysv/linux/x86_64/clone3.S:80:0

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 \
  008b68e90c0ff552.mlir

Output:

 ** Running 008b68e90c0ff552.mlir**
Running HLO module with runner HloRunner...
... compiled and ran in 0.0148678s.
Skipping reference runner

Contact

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