Receiving "Check failed: LayoutUtil::IsMonotonicWithDim0Major" error after upgrading from 0.4.29 to >=0.4.30

[Chulabhaya]

Hi there! My code that was functioning at Jax versions <=0.4.29 is now failing after upgrading to any version beyond 0.4.30. In particular, my code now crashes after a short period with the following:

F0922 22:52:09.660401 1439192 reduction.cc:147] Check failed: LayoutUtil::IsMonotonicWithDim0Major( hero.instruction().operand(i)->shape().layout()) reduction-layout-normalizer must run before code generation
*** Check failure stack trace: ***
    @     0x7f669b0001d4  absl::lts_20230802::log_internal::LogMessage::SendToLog()
    @     0x7f669b0000d4  absl::lts_20230802::log_internal::LogMessage::Flush()
    @     0x7f669b000579  absl::lts_20230802::log_internal::LogMessageFatal::~LogMessageFatal()
    @     0x7f6696e8e13e  xla::gpu::(anonymous namespace)::ReductionEmitter::ReductionEmitter()
    @     0x7f6696e8d307  xla::gpu::ReductionFusion::EmitInitializers()
    @     0x7f6696fb1dfa  xla::gpu::KernelFusionEmitterBase::Emit()
    @     0x7f6693a704a4  xla::gpu::IrEmitterUnnested::EmitFusion()
    @     0x7f6693a7b135  xla::gpu::IrEmitterUnnested::EmitHloInstruction()
    @     0x7f6693a5dcbe  xla::gpu::IrEmitterUnnested::EmitHloComputation()
    @     0x7f6693a613f4  xla::gpu::IrEmitterUnnested::EmitCommandBufferThunk()
    @     0x7f6693a7b42c  xla::gpu::IrEmitterUnnested::EmitHloInstruction()
    @     0x7f6693a5dcbe  xla::gpu::IrEmitterUnnested::EmitHloComputation()
    @     0x7f6693a72653  xla::gpu::IrEmitterUnnested::BuildWhileThunk()
    @     0x7f6693a722b6  xla::gpu::IrEmitterUnnested::EmitWhile()
    @     0x7f6693a7b0b6  xla::gpu::IrEmitterUnnested::EmitHloInstruction()
    @     0x7f6693a5dcbe  xla::gpu::IrEmitterUnnested::EmitHloComputation()
    @     0x7f6693a5d566  xla::gpu::IrEmitterUnnested::EmitConditional()
    @     0x7f6693a7afbf  xla::gpu::IrEmitterUnnested::EmitHloInstruction()
    @     0x7f6693a5dcbe  xla::gpu::IrEmitterUnnested::EmitHloComputation()
    @     0x7f6693a72653  xla::gpu::IrEmitterUnnested::BuildWhileThunk()
    @     0x7f6693a722b6  xla::gpu::IrEmitterUnnested::EmitWhile()
    @     0x7f6693a7b0b6  xla::gpu::IrEmitterUnnested::EmitHloInstruction()
    @     0x7f6693a5dcbe  xla::gpu::IrEmitterUnnested::EmitHloComputation()
    @     0x7f6693838d2a  xla::gpu::CompileModuleToLlvmIr()
    @     0x7f6693819920  xla::gpu::GpuCompiler::CompileToBackendResult()
    @     0x7f669381c4fe  xla::gpu::GpuCompiler::RunBackend()
    @     0x7f66937dd982  xla::Service::BuildExecutable()
    @     0x7f66937a5f55  xla::LocalService::CompileExecutables()
    @     0x7f6693799a14  xla::LocalClient::Compile()
    @     0x7f669373edcb  xla::PjRtStreamExecutorClient::CompileInternal()
    @     0x7f669373fe7e  xla::PjRtStreamExecutorClient::Compile()
    @     0x7f669369766a  std::__detail::__variant::__gen_vtable_impl<>::__visit_invoke()
    @     0x7f66936872d4  pjrt::PJRT_Client_Compile()
    @     0x7f66a74fd4dd  xla::InitializeArgsAndCompile()
    @     0x7f66a74fd9be  xla::PjRtCApiClient::Compile()
    @     0x7f66ad0a016c  xla::ifrt::PjRtLoadedExecutable::Create()
    @     0x7f66ad09b5a1  xla::ifrt::PjRtCompiler::Compile()
    @     0x7f66ac8439fe  xla::PyClient::CompileIfrtProgram()
    @     0x7f66ac8448de  xla::PyClient::Compile()
    @     0x7f66ac84cbef  nanobind::detail::func_create<>()::{lambda()#1}::__invoke()
    @     0x7f66ad079631  nanobind::detail::nb_func_vectorcall_complex()
    @     0x7f6776f7d478  nanobind::detail::nb_bound_method_vectorcall()
    @     0x5605f5266fbf  PyObject_Vectorcall
Aborted

Would anyone happen to know what would cause this error when upgrading from 0.4.29 to a newer version of Jax?

[yashk2810]

What was your code? Do you have a minimal repro? Is this on CPU/GPU/TPU/something else? Maybe try with jax 0.4.33?

[Chulabhaya]

Hi @nouiz , thank you for the reply! I've attached the HLO dump as a tar.gz file to this message, and this was run on a single GPU, it is consistently reproducible.

hlo_dump.tar.gz

[Chulabhaya]

Please let me know if that HLO data was generated correctly or if you would like me to re-run it or attach something different!

[akuegel]

From the stack trace, it looks like this is using the legacy reduction emitter. Since Jax version 0.4.32, we have replaced it with a new reduction emitter that should be able to handle non-default reduction layouts. I have run your HLO dump with XLA at head, and it seems to be working fine. However I also cannot reproduce it with the legacy emitter with the provided dump. Can you please try out with the newest Jax version, and provide a reproducer with that version if it still doesn't work? You can try the flag --xla_dump_hlo_as_proto instead of --xla_dump_hlo_as_text in case this is tooling related and the text dump doesn't contain some important piece.

[jreiffers]

This crash is in legacy emitters, which is strange, since they should be disabled by default now. Do you set any special XLA_FLAGS? Can you try setting --xla_gpu_mlir_emitter_level=4? But normally, this should be the default with recent versions of JAX, and the code from your stack trace shouldn't be used. Just to confirm: you are using an Nvidia GPU, right? Judging from your installed plugins that should be the case, but just to double check.

See https://jax.readthedocs.io/en/latest/xla_flags.html#configuring-xla-in-jax for instructions on how to configure XLA_FLAGS.

[Chulabhaya]

Hi @akuegel and @jreiffers , thank you guys for helping me investigate this. First off, here is the HLO dump as proto:
hlo_dump_proto.tar.gz

As far as I know I'm not setting any special XLA_FLAGS, unless some package somewhere is doing it behind-the-scenes.

Setting --xla_gpu_mlir_emitter_level=4 actually did the trick and my code is running! So this appears to be the issue. I'm definitely using an Nvidia GPU as well. I'm running 0.4.33 currently, so I'm not sure why I need to enable this flag manually? Do I just need to set this variable manually every time now going forwards?

[akuegel]

0.4.33 should theoretically use a XLA revision that has the flag set to 4 by default. We are not sure why it does not work. Currently, this is guarded with a compilation flag of if_cuda_is_configured (as it doesn't work on ROCM yet), maybe for some reason in the jaxlib build this configuration for cuda was not done? Unfortunately I don't know the JAX side.

[Chulabhaya]

I see, thank you for the clarification. Are there people who are familiar with the JAX side who I should be reporting this issue to for further investigation? While setting the --xla_gpu_mlir_emitter_level=4 does work, it would be nice to not have to do that in my code going forwards with newer versions of JAX.

[akuegel]

We have a plan to eventually remove the legacy emitters (and the flag). When exactly I cannot tell yet, but at least in some newer version of JAX, setting the flag will not be needed anymore.