SinkAwaitToFirstConsumer hitting assertion while compiling Gemma

[mariecwhite]

What happened?

Failure when compiling Gemma for x86:

2024-06-30T22:33:18.2378706Z iree-compile: /work/third_party/llvm-project/llvm/include/llvm/Support/GenericDomTree.h:501: NodeT *llvm::DominatorTreeBase<mlir::Block, false>::findNearestCommonDominator(NodeT *, NodeT *) const [NodeT = mlir::Block, IsPostDom = false]: Assertion `NodeB && "B must be in the tree"' failed.
2024-06-30T22:33:18.2381771Z Please report issues to https://github.com/iree-org/iree/issues and include the crash backtrace.
2024-06-30T22:33:18.2383667Z  #0 0x00007f5bd717920b llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) /work/third_party/llvm-project/llvm/lib/Support/Unix/Signals.inc:723:13
2024-06-30T22:33:18.2385658Z  #1 0x00007f5bd7177260 llvm::sys::RunSignalHandlers() /work/third_party/llvm-project/llvm/lib/Support/Signals.cpp:106:18
2024-06-30T22:33:18.2387407Z  #2 0x00007f5bd71798ef SignalHandler(int) /work/third_party/llvm-project/llvm/lib/Support/Unix/Signals.inc:413:1
2024-06-30T22:33:18.2389152Z  #3 0x00007f5bdf4fd420 __restore_rt (/lib/x86_64-linux-gnu/libpthread.so.0+0x14420)
2024-06-30T22:33:18.2390287Z  #4 0x00007f5bd086f00b raise (/lib/x86_64-linux-gnu/libc.so.6+0x4300b)
2024-06-30T22:33:18.2391300Z  #5 0x00007f5bd084e859 abort (/lib/x86_64-linux-gnu/libc.so.6+0x22859)
2024-06-30T22:33:18.2392274Z  #6 0x00007f5bd084e729 (/lib/x86_64-linux-gnu/libc.so.6+0x22729)
2024-06-30T22:33:18.2393198Z  #7 0x00007f5bd085ffd6 (/lib/x86_64-linux-gnu/libc.so.6+0x33fd6)
2024-06-30T22:33:18.2395751Z  #8 0x00007f5bd7285dc6 llvm::DominatorTreeBase<mlir::Block, false>::findNearestCommonDominator(mlir::Block*, mlir::Block*) const /work/third_party/llvm-project/llvm/include/llvm/Support/GenericDomTree.h:501:5
2024-06-30T22:33:18.2399281Z  #9 0x00007f5bdad0171f mlir::iree_compiler::IREE::Stream::(anonymous namespace)::SinkAwaitToFirstConsumer::matchAndRewrite(mlir::iree_compiler::IREE::Stream::TimepointAwaitOp, mlir::PatternRewriter&) const /work/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp:3148:43
2024-06-30T22:33:18.2404122Z #10 0x00007f5bdb454be8 mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<mlir::LogicalResult (mlir::Pattern const&)>)::$_2::operator()() const /work/third_party/llvm-project/mlir/lib/Rewrite/PatternApplicator.cpp:212:13
2024-06-30T22:33:18.2409409Z #11 0x00007f5bdb45145d mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<mlir::LogicalResult (mlir::Pattern const&)>) /work/third_party/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:0:12
2024-06-30T22:33:18.2413288Z #12 0x00007f5bdb4342d9 (anonymous namespace)::GreedyPatternRewriteDriver::processWorklist() /work/third_party/llvm-project/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:0:0
2024-06-30T22:33:18.2416415Z #13 0x00007f5bdb43097b (anonymous namespace)::RegionPatternRewriteDriver::simplify(bool*) &&::$_8::operator()() const /work/third_party/llvm-project/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:874:11
2024-06-30T22:33:18.2419772Z #14 0x00007f5bdb43097b void llvm::function_ref<void ()>::callback_fn<(anonymous namespace)::RegionPatternRewriteDriver::simplify(bool*) &&::$_8>(long) /work/third_party/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:45:12
2024-06-30T22:33:18.2425208Z #15 0x00007f5bdb43097b llvm::function_ref<void ()>::operator()() const /work/third_party/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:68:12
2024-06-30T22:33:18.2428505Z #16 0x00007f5bdb43097b void mlir::MLIRContext::executeAction<(anonymous namespace)::GreedyPatternRewriteIteration, long&>(llvm::function_ref<void ()>, llvm::ArrayRef<mlir::IRUnit>, long&) /work/third_party/llvm-project/mlir/include/mlir/IR/MLIRContext.h:275:7
2024-06-30T22:33:18.2431937Z #17 0x00007f5bdb43097b (anonymous namespace)::RegionPatternRewriteDriver::simplify(bool*) && /work/third_party/llvm-project/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:872:10
2024-06-30T22:33:18.2435334Z #18 0x00007f5bdb43097b mlir::applyPatternsAndFoldGreedily(mlir::Region&, mlir::FrozenRewritePatternSet const&, mlir::GreedyRewriteConfig, bool*) /work/third_party/llvm-project/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:919:47
2024-06-30T22:33:18.2438203Z #19 0x00007f5bdb3b9b2b mlir::LogicalResult::failed() const /work/third_party/llvm-project/mlir/include/mlir/Support/LogicalResult.h:44:33
2024-06-30T22:33:18.2441268Z #20 0x00007f5bdb3b9b2b mlir::applyPatternsAndFoldGreedily(mlir::Operation*, mlir::FrozenRewritePatternSet const&, mlir::GreedyRewriteConfig, bool*) /work/third_party/llvm-project/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h:160:14
2024-06-30T22:33:18.2444286Z #21 0x00007f5bdb3b9b2b (anonymous namespace)::Canonicalizer::runOnOperation() /work/third_party/llvm-project/mlir/lib/Transforms/Canonicalizer.cpp:63:9
2024-06-30T22:33:18.2447048Z #22 0x00007f5bd732b246 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::$_7::operator()() const /work/third_party/llvm-project/mlir/lib/Pass/Pass.cpp:527:17
2024-06-30T22:33:18.2450678Z #23 0x00007f5bd732b246 void llvm::function_ref<void ()>::callback_fn<mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::$_7>(long) /work/third_party/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:45:12
2024-06-30T22:33:18.2453894Z #24 0x00007f5bd732b246 llvm::function_ref<void ()>::operator()() const /work/third_party/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:68:12
2024-06-30T22:33:18.2456793Z #25 0x00007f5bd732b246 void mlir::MLIRContext::executeAction<mlir::PassExecutionAction, mlir::Pass&>(llvm::function_ref<void ()>, llvm::ArrayRef<mlir::IRUnit>, mlir::Pass&) /work/third_party/llvm-project/mlir/include/mlir/IR/MLIRContext.h:275:7
2024-06-30T22:33:18.2460107Z #26 0x00007f5bd732b246 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) /work/third_party/llvm-project/mlir/lib/Pass/Pass.cpp:521:21
2024-06-30T22:33:18.2462534Z #27 0x00007f5bd732bb78 mlir::LogicalResult::failed() const /work/third_party/llvm-project/mlir/include/mlir/Support/LogicalResult.h:44:33
2024-06-30T22:33:18.2464606Z #28 0x00007f5bd732bb78 mlir::failed(mlir::LogicalResult) /work/third_party/llvm-project/mlir/include/mlir/Support/LogicalResult.h:72:58
2024-06-30T22:33:18.2468155Z #29 0x00007f5bd732bb78 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) /work/third_party/llvm-project/mlir/lib/Pass/Pass.cpp:593:9
2024-06-30T22:33:18.2472131Z #30 0x00007f5bd733181e mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::$_15::operator()(mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo&) const /work/third_party/llvm-project/mlir/lib/Pass/Pass.cpp:818:5
2024-06-30T22:33:18.2474911Z #31 0x00007f5bd73318cf mlir::LogicalResult::failed() const /work/third_party/llvm-project/mlir/include/mlir/Support/LogicalResult.h:44:33
2024-06-30T22:33:18.2477032Z #32 0x00007f5bd73318cf mlir::failed(mlir::LogicalResult) /work/third_party/llvm-project/mlir/include/mlir/Support/LogicalResult.h:72:58
2024-06-30T22:33:18.2487352Z #33 0x00007f5bd73318cf mlir::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> > >, mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::$_15&>(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)::$_15&)::'lambda'()::operator()() const /work/third_party/llvm-project/mlir/include/mlir/IR/Threading.h:62:11
2024-06-30T22:33:18.2505488Z #34 0x00007f5bd73318cf std::_Function_handler<void (), mlir::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> > >, mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::$_15&>(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)::$_15&)::'lambda'()>::_M_invoke(std::_Any_data const&) /usr/bin/../lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9/bits/std_function.h:300:2
2024-06-30T22:33:18.2516806Z #35 0x00007f5bd7284138 std::__future_base::_Task_setter<std::unique_ptr<std::__future_base::_Result<void>, std::__future_base::_Result_base::_Deleter>, std::thread::_Invoker<std::tuple<std::function<void ()> > >, void>::operator()() const /usr/bin/../lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9/future:1372:20
2024-06-30T22:33:18.2522061Z #36 0x00007f5bd7284138 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<std::function<void ()> > >, void> >::_M_invoke(std::_Any_data const&) /usr/bin/../lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9/bits/std_function.h:285:9
2024-06-30T22:33:18.2526846Z #37 0x00007f5bd7284097 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*) /usr/bin/../lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9/future:565:13
2024-06-30T22:33:18.2529226Z #38 0x00007f5bdf4fa4df __pthread_once_slow (/lib/x86_64-linux-gnu/libpthread.so.0+0x114df)
2024-06-30T22:33:18.2535138Z #39 0x00007f5bd7284448 void std::call_once<void (std::__future_base::_State_baseV2::*)(std::function<std::unique_ptr<std::__future_base::_Result_base, std::__future_base::_Result_base::_Deleter> ()>*, bool*), std::__future_base::_State_baseV2*, std::function<std::unique_ptr<std::__future_base::_Result_base, std::__future_base::_Result_base::_Deleter> ()>*, bool*>(std::once_flag&, void (std::__future_base::_State_baseV2::*&&)(std::function<std::unique_ptr<std::__future_base::_Result_base, std::__future_base::_Result_base::_Deleter> ()>*, bool*), std::__future_base::_State_baseV2*&&, std::function<std::unique_ptr<std::__future_base::_Result_base, std::__future_base::_Result_base::_Deleter> ()>*&&, bool*&&) /usr/bin/../lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9/mutex:690:11
2024-06-30T22:33:18.2541618Z #40 0x00007f5bd7284448 std::__future_base::_State_baseV2::_M_set_result(std::function<std::unique_ptr<std::__future_base::_Result_base, std::__future_base::_Result_base::_Deleter> ()>, bool) /usr/bin/../lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9/future:401:2
2024-06-30T22:33:18.2544884Z #41 0x00007f5bd7284448 std::__future_base::_Deferred_state<std::thread::_Invoker<std::tuple<std::function<void ()> > >, void>::_M_complete_async() /usr/bin/../lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9/future:1610:9
2024-06-30T22:33:18.2547806Z #42 0x00007f5bd72844c5 std::__atomic_base<unsigned int>::load(std::memory_order) const /usr/bin/../lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9/bits/atomic_base.h:419:9
2024-06-30T22:33:18.2550373Z #43 0x00007f5bd72844c5 std::__atomic_futex_unsigned<2147483648u>::_M_load(std::memory_order) /usr/bin/../lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9/bits/atomic_futex.h:80:22
2024-06-30T22:33:18.2553113Z #44 0x00007f5bd72844c5 std::__atomic_futex_unsigned<2147483648u>::_M_load_when_equal(unsigned int, std::memory_order) /usr/bin/../lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9/bits/atomic_futex.h:158:22
2024-06-30T22:33:18.2555534Z #45 0x00007f5bd72844c5 std::__future_base::_State_baseV2::wait() /usr/bin/../lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9/future:337:12
2024-06-30T22:33:18.2557518Z #46 0x00007f5bd72844c5 std::__basic_future<void>::wait() const /usr/bin/../lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9/future:692:19
2024-06-30T22:33:18.2560573Z #47 0x00007f5bd72844c5 std::shared_future<void> llvm::ThreadPoolInterface::asyncImpl<void>(std::function<void ()>, llvm::ThreadPoolTaskGroup*)::'lambda'()::operator()() const /work/third_party/llvm-project/llvm/include/llvm/Support/ThreadPool.h:114:38
2024-06-30T22:33:18.2564587Z #48 0x00007f5bd72844c5 std::_Function_handler<void (), std::shared_future<void> llvm::ThreadPoolInterface::asyncImpl<void>(std::function<void ()>, llvm::ThreadPoolTaskGroup*)::'lambda'()>::_M_invoke(std::_Any_data const&) /usr/bin/../lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9/bits/std_function.h:300:2
2024-06-30T22:33:18.2567836Z #49 0x00007f5bd7133e68 llvm::StdThreadPool::processTasks(llvm::ThreadPoolTaskGroup*) /work/third_party/llvm-project/llvm/lib/Support/ThreadPool.cpp:103:5
2024-06-30T22:33:18.2570900Z #50 0x00007f5bd713547c std::unique_ptr<std::tuple<llvm::StdThreadPool::grow(int)::$_0>, std::default_delete<std::tuple<llvm::StdThreadPool::grow(int)::$_0> > >::~unique_ptr() /usr/bin/../lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9/bits/unique_ptr.h:291:12
2024-06-30T22:33:18.2574186Z #51 0x00007f5bd713547c void llvm::thread::GenericThreadProxy<std::tuple<llvm::StdThreadPool::grow(int)::$_0> >(void*) /work/third_party/llvm-project/llvm/include/llvm/Support/thread.h:46:3
2024-06-30T22:33:18.2576852Z #52 0x00007f5bd713547c void* llvm::thread::ThreadProxy<std::tuple<llvm::StdThreadPool::grow(int)::$_0> >(void*) /work/third_party/llvm-project/llvm/include/llvm/Support/thread.h:55:5
2024-06-30T22:33:18.2578741Z #53 0x00007f5bdf4f1609 start_thread (/lib/x86_64-linux-gnu/libpthread.so.0+0x8609)
2024-06-30T22:33:18.2579889Z #54 0x00007f5bd094b353 clone (/lib/x86_64-linux-gnu/libc.so.6+0x11f353)
2024-06-30T22:33:18.2580632Z Aborted (core dumped)

Steps to reproduce your issue

1. Download mlir at https://storage.googleapis.com/iree-model-artifacts/jax/jax_models_0.4.25_1709787220/GEMMA2BIT_GREEDY_FP32_JAX_1X1024XI32_256XI32/stablehlo.mlirbc

2. Try compiling:

iree-compile \
  --iree-hal-target-backends=llvm-cpu \
  --iree-input-type=stablehlo \
  --iree-llvmcpu-target-triple=x86_64-unknown-linux-gnu \
  --iree-llvmcpu-target-cpu=cascadelake \
  --iree-opt-data-tiling=true \
  --iree-llvmcpu-enable-ukernels=all \
  /tmp/stablehlo.mlirbc -o /tmp/module.vmfb

What component(s) does this issue relate to?

Compiler

Version information

4ad00ef

Additional context

No response