Open
Description
I have the following MLIR program:
test.mlir:
mesh.mesh @mesh_2d(shape = 2x4)
func.func @element_wise_on_def(%arg0: tensor<8x16xf32>, %arg1: tensor<4096x14336xf32>) -> tensor<8x16xf32> {
%0 = tosa.sigmoid %arg0 : (tensor<8x16xf32>) -> tensor<8x16xf32>
%1 = mesh.shard %0 to <@mesh_2d, [[0], [1]]> : tensor<8x16xf32>
return %1 : tensor<8x16xf32>
}
The program does not use %arg1,but when running the --pass-pipeline="builtin.module(func.func(sharding-propagation,cse))"
The correct result is as follows
module {
mesh.mesh @mesh_2d(shape = 2x4)
func.func @element_wise_on_def(%arg0: tensor<8x16xf32>, %arg1: tensor<4096x14336xf32>) -> tensor<8x16xf32> {
%sharding_annotated = mesh.shard %arg0 to <@mesh_2d, [[0], [1]]> annotate_for_users : tensor<8x16xf32>
%0 = tosa.sigmoid %sharding_annotated : (tensor<8x16xf32>) -> tensor<8x16xf32>
%sharding_annotated_0 = mesh.shard %0 to <@mesh_2d, [[0], [1]]> : tensor<8x16xf32>
return %sharding_annotated_0 : tensor<8x16xf32>
}
}
The above MLIR program will cause a crash when using the following command:
mlir-opt --pass-pipeline="builtin.module(func.func(mesh-spmdization,test-constant-fold))" test.mlir
And the crash backtrace is:
`mlir-opt: mlir/lib/Dialect/Mesh/Transforms/Spmdization.cpp:536: mlir::mesh::shardedBlockArgumentTypes(mlir::Block&, mlir::SymbolTableCollection&)::<lambda(mlir::BlockArgument)>: Assertion `rankedTensorArg.hasOneUse()' failed.
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
Stack dump:
0. Program arguments: /home/home/BiShengIR/build/bin/mlir-opt --pass-pipeline=builtin.module(func.func(mesh-spmdization,test-constant-fold)) test_oneuse_spmd.mlir
#0 0x0000000001fc58c8 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) /home/home/BiShengIR/llvm/lib/Support/Unix/Signals.inc:723:22
#1 0x0000000001fc5c5c PrintStackTraceSignalHandler(void*) /home/home/BiShengIR/llvm/lib/Support/Unix/Signals.inc:798:1
#2 0x0000000001fc388c llvm::sys::RunSignalHandlers() /home/home/BiShengIR/llvm/lib/Support/Signals.cpp:105:20
#3 0x0000000001fc524c SignalHandler(int) /BiShengIR/llvm/lib/Support/Unix/Signals.inc:413:1
#4 0x0000ffff805f98dc (linux-vdso.so.1+0x8dc)
#5 0x0000ffff80168bd0 __pthread_kill_implementation /usr/src/debug/glibc-2.34-70.oe2203.aarch64/nptl/pthread_kill.c:44:76
#6 0x0000ffff80124bfc gsignal /usr/src/debug/glibc-2.34-70.oe2203.aarch64/signal/../sysdeps/posix/raise.c:27:6
#7 0x0000ffff80112c70 abort /usr/src/debug/glibc-2.34-70.oe2203.aarch64/stdlib/abort.c:81:7
#8 0x0000ffff8011e624 __assert_fail_base /usr/src/debug/glibc-2.34-70.oe2203.aarch64/assert/assert.c:89:7
#9 0x0000ffff8011e68c (/usr/lib64/libc.so.6+0x3668c)
#10 0x0000000003da91b0 mlir::mesh::shardedBlockArgumentTypes(mlir::Block&, mlir::SymbolTableCollection&)::'lambda'(mlir::BlockArgument)::operator()(mlir::BlockArgument) const /BiShengIR/mlir/lib/Dialect/Mesh/Transforms/Spmdization.cpp:537:54
#11 0x0000000003daa3d4 std::back_insert_iterator<llvm::SmallVector<mlir::Type, 6u>> std::transform<mlir::BlockArgument*, std::back_insert_iterator<llvm::SmallVector<mlir::Type, 6u>>, mlir::mesh::shardedBlockArgumentTypes(mlir::Block&, mlir::SymbolTableCollection&)::'lambda'(mlir::BlockArgument)>(mlir::BlockArgument*, mlir::BlockArgument*, std::back_insert_iterator<llvm::SmallVector<mlir::Type, 6u>>, mlir::mesh::shardedBlockArgumentTypes(mlir::Block&, mlir::SymbolTableCollection&)::'lambda'(mlir::BlockArgument)) /usr/include/c++/10.3.1/bits/stl_algo.h:4320:24
#12 0x0000000003daa00c std::back_insert_iterator<llvm::SmallVector<mlir::Type, 6u>> llvm::transform<llvm::MutableArrayRef<mlir::BlockArgument>, std::back_insert_iterator<llvm::SmallVector<mlir::Type, 6u>>, mlir::mesh::shardedBlockArgumentTypes(mlir::Block&, mlir::SymbolTableCollection&)::'lambda'(mlir::BlockArgument)>(llvm::MutableArrayRef<mlir::BlockArgument>&&, std::back_insert_iterator<llvm::SmallVector<mlir::Type, 6u>>, mlir::mesh::shardedBlockArgumentTypes(mlir::Block&, mlir::SymbolTableCollection&)::'lambda'(mlir::BlockArgument)) /BiShengIR/llvm/include/llvm/ADT/STLExtras.h:1939:1
#13 0x0000000003da92c8 mlir::mesh::shardedBlockArgumentTypes(mlir::Block&, mlir::SymbolTableCollection&) /BiShengIR/mlir/lib/Dialect/Mesh/Transforms/Spmdization.cpp:544:10
#14 0x0000000003da99bc mlir::mesh::spmdizeBlock(mlir::Block&, mlir::IRMapping&, mlir::SymbolTableCollection&, mlir::OpBuilder&) /BiShengIR/mlir/lib/Dialect/Mesh/Transforms/Spmdization.cpp:667:40
#15 0x0000000003da9c44 mlir::mesh::spmdizeFuncOp(mlir::FunctionOpInterface, mlir::IRMapping&, mlir::SymbolTableCollection&) BiShengIR/mlir/lib/Dialect/Mesh/Transforms/Spmdization.cpp:699:15
#16 0x0000000003da9e6c mlir::mesh::(anonymous namespace)::Spmdization::runOnOperation() /home/home/BiShengIR/mlir/lib/Dialect/Mesh/Transforms/Spmdization.cpp:736:15
#17 0x00000000061ace1c mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::'lambda0'()::operator()() const /home/home/BiShengIR/mlir/lib/Pass/Pass.cpp:516:57
#18 0x00000000061afdcc void llvm::function_ref<void ()>::callback_fn<mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::'lambda0'()>(long) /home/home/BiShengIR/llvm/include/llvm/ADT/STLFunctionalExtras.h:46:40
#19 0x0000000001fe3158 llvm::function_ref<void ()>::operator()() const /home/home/BiShengIR/llvm/include/llvm/ADT/STLFunctionalExtras.h:68:62
#20 0x00000000061b6f18 void mlir::MLIRContext::executeAction<mlir::PassExecutionAction, mlir::Pass&>(llvm::function_ref<void ()>, llvm::ArrayRef<mlir::IRUnit>, mlir::Pass&) /BiShengIR/mlir/include/mlir/IR/MLIRContext.h:276:3
#21 0x00000000061ad108 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) /home/home/BiShengIR/mlir/lib/Pass/Pass.cpp:525:23
#22 0x00000000061ad33c mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) /BiShengIR/mlir/lib/Pass/Pass.cpp:585:15
#23 0x00000000061ae104 mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::'lambda'(mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo&)::operator()(mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo&) const /home/home/BiShengIR/mlir/lib/Pass/Pass.cpp:807:69
#24 0x00000000061b0740 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)::'lambda'(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)::'lambda'(mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo&)&) /home/home/BiShengIR/mlir/include/mlir/IR/Threading.h:46:17
#25 0x00000000061afbac mlir::LogicalResult mlir::failableParallelForEach<std::vector<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo, std::allocator<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo>>&, mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::'lambda'(mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo&)&>(mlir::MLIRContext*, std::vector<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo, std::allocator<mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo>>&, mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::'lambda'(mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo&)&) /home/home/BiShengIR/mlir/include/mlir/IR/Threading.h:94:1
#26 0x00000000061ae490 mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool) /home/home/BiShengIR/mlir/lib/Pass/Pass.cpp:815:13
#27 0x00000000061adc20 mlir::detail::OpToOpPassAdaptor::runOnOperation(bool) /home/home/BiShengIR/mlir/lib/Pass/Pass.cpp:709:1
#28 0x00000000061acdf0 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::'lambda0'()::operator()() const /home/home/BiShengIR/mlir/lib/Pass/Pass.cpp:517:34
#29 0x00000000061afdcc void llvm::function_ref<void ()>::callback_fn<mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int)::'lambda0'()>(long) /home/home/BiShengIR/llvm/include/llvm/ADT/STLFunctionalExtras.h:46:40
#30 0x0000000001fe3158 llvm::function_ref<void ()>::operator()() const /home/home/BiShengIR/llvm/include/llvm/ADT/STLFunctionalExtras.h:68:62
#31 0x00000000061b6f18 void mlir::MLIRContext::executeAction<mlir::PassExecutionAction, mlir::Pass&>(llvm::function_ref<void ()>, llvm::ArrayRef<mlir::IRUnit>, mlir::Pass&) /home/home/BiShengIR/mlir/include/mlir/IR/MLIRContext.h:276:3
#32 0x00000000061ad108 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) /home/home/BiShengIR/mlir/lib/Pass/Pass.cpp:525:23
#33 0x00000000061ad33c mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) /home/home/BiShengIR/mlir/lib/Pass/Pass.cpp:585:15
#34 0x00000000061aeaf8 mlir::PassManager::runPasses(mlir::Operation*, mlir::AnalysisManager) /home/home/BiShengIR/mlir/lib/Pass/Pass.cpp:898:1
#35 0x00000000061ae9a8 mlir::PassManager::run(mlir::Operation*) /home/home/BiShengIR/mlir/lib/Pass/Pass.cpp:876:76
#36 0x00000000061a2e4c performActions(llvm::raw_ostream&, std::shared_ptr<llvm::SourceMgr> const&, mlir::MLIRContext*, mlir::MlirOptMainConfig const&) /home/home/BiShengIR/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:394:13
#37 0x00000000061a3310 processBuffer(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, mlir::MlirOptMainConfig const&, mlir::DialectRegistry&, llvm::ThreadPool*) /home/home/BiShengIR/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:459:58
#38 0x00000000061a36c4 mlir::MlirOptMain(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, mlir::DialectRegistry&, mlir::MlirOptMainConfig const&)::'lambda'(std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)::operator()(std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&) const /home/home/BiShengIR/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:531:36
#39 0x00000000061a4434 mlir::LogicalResult llvm::function_ref<mlir::LogicalResult (std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>::callback_fn<mlir::MlirOptMain(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, mlir::DialectRegistry&, mlir::MlirOptMainConfig const&)::'lambda'(std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>(long, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&) /home/home/BiShengIR/llvm/include/llvm/ADT/STLFunctionalExtras.h:45:52
#40 0x00000000062ff168 llvm::function_ref<mlir::LogicalResult (std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>::operator()(std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&) const /home/home/BiShengIR/llvm/include/llvm/ADT/STLFunctionalExtras.h:68:12
#41 0x00000000062febdc mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<mlir::LogicalResult (std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>, llvm::raw_ostream&, bool, bool) /home/home/BiShengIR/mlir/lib/Support/ToolUtilities.cpp:28:60
#42 0x00000000061a3814 mlir::MlirOptMain(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, mlir::DialectRegistry&, mlir::MlirOptMainConfig const&) /home/home/BiShengIR/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:535:61
#43 0x00000000061a3a20 mlir::MlirOptMain(int, char**, llvm::StringRef, llvm::StringRef, mlir::DialectRegistry&) /home/home/BiShengIR/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:568:13
#44 0x00000000061a3b3c mlir::MlirOptMain(int, char**, llvm::StringRef, mlir::DialectRegistry&) /home/home/BiShengIR/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp:584:73
#45 0x0000000001f3d41c main /home/home/BiShengIR/mlir/tools/mlir-opt/mlir-opt.cpp:297:0
#46 0x0000ffff80112f3c __libc_start_call_main /usr/src/debug/glibc-2.34-70.oe2203.aarch64/csu/../sysdeps/nptl/libc_start_call_main.h:74:3
#47 0x0000ffff80113008 __libc_start_main /usr/src/debug/glibc-2.34-70.oe2203.aarch64/csu/../csu/libc-start.c:389:24
#48 0x0000000001f3d0b0 _start (/home/home/BiShengIR/build/bin/mlir-opt+0x1f3d0b0)`
Because I didn't use %arg1, an assertion failure was triggered. Why is this one-use restriction designed?