Aggressive fusion causes fused op to have wrong number of results

What happened?

When compiling a model I get:

iree-compile: repo/compiler/src/iree/compiler/Dialect/Flow/Transforms/FusionOfTensorOps.cpp:133: virtual mlir::LogicalResult mlir::iree_compiler::IREE::Flow::(anonymous namespace)::FuseElementwiseOpsWithMultipleUses::matchAndRewrite(linalg::GenericOp, mlir::PatternRewriter &) const: Assertion `fusedOperation.value()->getNumResults() == producerOp->getNumResults() + consumerOp->getNumResults()' failed.
Please report issues to https://github.com/iree-org/iree/issues and include the crash backtrace.
 #0 0x00007f006ea96537 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) /home/petkantchin/ws/iree/repo/third_party/llvm-project/llvm/lib/Support/Unix/Signals.inc:567:13
 #1 0x00007f006ea94910 llvm::sys::RunSignalHandlers() /home/petkantchin/ws/iree/repo/third_party/llvm-project/llvm/lib/Support/Signals.cpp:105:18
 #2 0x00007f006ea96bbf SignalHandler(int) /home/petkantchin/ws/iree/repo/third_party/llvm-project/llvm/lib/Support/Unix/Signals.inc:412:1
 #3 0x00007f006f121420 __restore_rt (/lib/x86_64-linux-gnu/libpthread.so.0+0x14420)
 #4 0x00007f006608900b raise /build/glibc-SzIz7B/glibc-2.31/signal/../sysdeps/unix/sysv/linux/raise.c:51:1
 #5 0x00007f0066068859 abort /build/glibc-SzIz7B/glibc-2.31/stdlib/abort.c:81:7
 #6 0x00007f0066068729 get_sysdep_segment_value /build/glibc-SzIz7B/glibc-2.31/intl/loadmsgcat.c:509:8
 #7 0x00007f0066068729 _nl_load_domain /build/glibc-SzIz7B/glibc-2.31/intl/loadmsgcat.c:970:34
 #8 0x00007f0066079fd6 (/lib/x86_64-linux-gnu/libc.so.6+0x33fd6)
 #9 0x00007f006b6b41d6 (/home/petkantchin/ws/iree/build/ninja/RelWithDebInfo/lib/libIREECompiler.so.0+0x527f1d6)
#10 0x00007f006e720048 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&)>) /home/petkantchin/ws/iree/repo/third_party/llvm-project/mlir/lib/Rewrite/PatternApplicator.cpp:200:25
#11 0x00007f006e6f2a36 (anonymous namespace)::GreedyPatternRewriteDriver::simplify(llvm::MutableArrayRef<mlir::Region>) /home/petkantchin/ws/iree/repo/third_party/llvm-project/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:292:19
#12 0x00007f006e6f2a36 mlir::applyPatternsAndFoldGreedily(llvm::MutableArrayRef<mlir::Region>, mlir::FrozenRewritePatternSet const&, mlir::GreedyRewriteConfig) /home/petkantchin/ws/iree/repo/third_party/llvm-project/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:452:27
#13 0x00007f006b6b2cfd mlir::iree_compiler::IREE::Flow::(anonymous namespace)::fuseMultiUseProducers(mlir::Operation*, mlir::MLIRContext*, mlir::DominanceInfo&) /home/petkantchin/ws/iree/repo/compiler/src/iree/compiler/Dialect/Flow/Transforms/FusionOfTensorOps.cpp:193:14
#14 0x00007f006b6b2cfd mlir::iree_compiler::IREE::Flow::(anonymous namespace)::FusionOfTensorOpsPass::runOnOperation() /home/petkantchin/ws/iree/repo/compiler/src/iree/compiler/Dialect/Flow/Transforms/FusionOfTensorOps.cpp:361:13
#15 0x00007f006e7486ce mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) /home/petkantchin/ws/iree/repo/third_party/llvm-project/mlir/lib/Pass/Pass.cpp:0:11
#16 0x00007f006e748e28 mlir::LogicalResult::failed() const /home/petkantchin/ws/iree/repo/third_party/llvm-project/mlir/include/mlir/Support/LogicalResult.h:44:33
#17 0x00007f006e748e28 mlir::failed(mlir::LogicalResult) /home/petkantchin/ws/iree/repo/third_party/llvm-project/mlir/include/mlir/Support/LogicalResult.h:72:58
#18 0x00007f006e748e28 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) /home/petkantchin/ws/iree/repo/third_party/llvm-project/mlir/lib/Pass/Pass.cpp:534:9
#19 0x00007f006e74dc35 mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::$_14::operator()(mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo&) const /home/petkantchin/ws/iree/repo/third_party/llvm-project/mlir/lib/Pass/Pass.cpp:759:5
#20 0x00007f006e74dc35 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)::$_14&>(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)::$_14&)::'lambda'()::operator()() const /home/petkantchin/ws/iree/repo/third_party/llvm-project/mlir/include/mlir/IR/Threading.h:62:18
#21 0x00007f006e74dc35 __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> > > std::__invoke_impl<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)::$_14&>(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)::$_14&)::'lambda'()&>(std::__invoke_other, mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::$_14&, 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)::$_14&>(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)::$_14&)::'lambda'()&...) /usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/invoke.h:60:14
#22 0x00007f006e74dc35 std::enable_if<is_invocable_r_v<__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)::$_14&, 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)::$_14&>(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)::$_14&)::'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> > > >::type std::__invoke_r<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)::$_14&>(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)::$_14&)::'lambda'()&>(mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::$_14&, 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)::$_14&>(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)::$_14&)::'lambda'()&...) /usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/invoke.h:110:2
#23 0x00007f006e74dc35 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)::$_14&>(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)::$_14&)::'lambda'()>::_M_invoke(std::_Any_data const&) /usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/std_function.h:291:9
#24 0x00007f0069b9f486 std::__shared_ptr<std::promise<void>, (__gnu_cxx::_Lock_policy)2>::get() const /usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/shared_ptr_base.h:1325:16
#25 0x00007f0069b9f486 std::__shared_ptr_access<std::promise<void>, (__gnu_cxx::_Lock_policy)2, false, false>::_M_get() const /usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/shared_ptr_base.h:1024:66
#26 0x00007f0069b9f486 std::__shared_ptr_access<std::promise<void>, (__gnu_cxx::_Lock_policy)2, false, false>::operator->() const /usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/shared_ptr_base.h:1018:9
#27 0x00007f0069b9f486 llvm::ThreadPool::createTaskAndFuture(std::function<void ()>)::'lambda'()::operator()() const /home/petkantchin/ws/iree/repo/third_party/llvm-project/llvm/include/llvm/Support/ThreadPool.h:136:15
#28 0x00007f0069b9f486 void std::__invoke_impl<void, llvm::ThreadPool::createTaskAndFuture(std::function<void ()>)::'lambda'()&>(std::__invoke_other, llvm::ThreadPool::createTaskAndFuture(std::function<void ()>)::'lambda'()&) /usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/invoke.h:60:14
#29 0x00007f0069b9f486 std::enable_if<is_invocable_r_v<void, llvm::ThreadPool::createTaskAndFuture(std::function<void ()>)::'lambda'()&>, void>::type std::__invoke_r<void, llvm::ThreadPool::createTaskAndFuture(std::function<void ()>)::'lambda'()&>(llvm::ThreadPool::createTaskAndFuture(std::function<void ()>)::'lambda'()&) /usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/invoke.h:110:2
#30 0x00007f0069b9f486 std::_Function_handler<void (), llvm::ThreadPool::createTaskAndFuture(std::function<void ()>)::'lambda'()>::_M_invoke(std::_Any_data const&) /usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/std_function.h:291:9
#31 0x00007f006ea48175 llvm::ThreadPool::processTasks(llvm::ThreadPoolTaskGroup*) /home/petkantchin/ws/iree/repo/third_party/llvm-project/llvm/lib/Support/ThreadPool.cpp:102:5
#32 0x00007f006ea48ee3 std::default_delete<std::tuple<llvm::ThreadPool::grow(int)::$_0> >::operator()(std::tuple<llvm::ThreadPool::grow(int)::$_0>*) const /usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/unique_ptr.h:85:2
#33 0x00007f006ea48ee3 std::unique_ptr<std::tuple<llvm::ThreadPool::grow(int)::$_0>, std::default_delete<std::tuple<llvm::ThreadPool::grow(int)::$_0> > >::~unique_ptr() /usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/unique_ptr.h:361:4
#34 0x00007f006ea48ee3 void llvm::thread::GenericThreadProxy<std::tuple<llvm::ThreadPool::grow(int)::$_0> >(void*) /home/petkantchin/ws/iree/repo/third_party/llvm-project/llvm/include/llvm/Support/thread.h:46:3
#35 0x00007f006ea48ee3 void* llvm::thread::ThreadProxy<std::tuple<llvm::ThreadPool::grow(int)::$_0> >(void*) /home/petkantchin/ws/iree/repo/third_party/llvm-project/llvm/include/llvm/Support/thread.h:55:5
#36 0x00007f006f115609 start_thread /build/glibc-SzIz7B/glibc-2.31/nptl/pthread_create.c:478:7
#37 0x00007f0066165133 clone /build/glibc-SzIz7B/glibc-2.31/misc/../sysdeps/unix/sysv/linux/x86_64/clone.S:97:0
Aborted (core dumped)

Without --iree-flow-enable-aggressive-fusion I don't get this assert failure.

Steps to reproduce your issue

1. Build IREE with assertions enabled. E.g. cmake -DIREE_ENABLE_ASSERTIONS=ON
2. Get model from here. Note that it is 4.5 GiB.
3. compile with

iree-compile \
  stable_diffusion_training.mlir \
  -o=stable_diffusion_training_llvm-cpu.vmfb \
  --iree-input-type=mhlo \
  --iree-hal-target-backends=llvm-cpu \
  --iree-flow-enable-aggressive-fusion

What component(s) does this issue relate to?

MLIR, Compiler

Version information

fd2a62e1452b1371eb559f6ecb157fb3853dea97

Additional context

No response