[MLIR]Crashed on converting tosa to linalg with assertion failure `idx < size()'

[BealZephyr]

Test on commit: 3b4cb1e
The same crash occurred in LLVM 16

steps to reproduce:

mlir-opt -pass-pipeline=builtin.module(func.func(tosa-to-linalg)) test.mlir

test case:

func.func @main(%arg0: tensor<?xf32>) -> tensor<?xi32> {
  %0 = "tosa.argmax"(%arg0) {axis = 1 : i64} : (tensor<?xf32>) -> tensor<?xi32>
  return %0 : tensor<?xi32>
}

crash trace:

mlir-opt: /home/verify/llvm-project-3b4cb1e96c645bb833fe710856479c31383859bb/llvm/include/llvm/ADT/SmallVector.h:294: T& llvm::SmallVectorTemplateCommon<T, <template-parameter-1-2> >::operator[](llvm::SmallVectorTemplateCommon<T, <template-parameter-1-2> >::size_type) [with T = mlir::utils::IteratorType; <template-parameter-1-2> = void; llvm::SmallVectorTemplateCommon<T, <template-parameter-1-2> >::reference = mlir::utils::IteratorType&; llvm::SmallVectorTemplateCommon<T, <template-parameter-1-2> >::size_type = long unsigned int]: Assertion `idx < size()' 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/verify/llvm-project-3b4cb1e96c645bb833fe710856479c31383859bb/build/bin/mlir-opt -pass-pipeline=builtin.module(func.func(tosa-to-linalg)) test.mlir
 #0 0x0000561fa2394f1a llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (/home/verify/llvm-project-3b4cb1e96c645bb833fe710856479c31383859bb/build/bin/mlir-opt+0x221f1a)
 #1 0x0000561fa239241c SignalHandler(int) Signals.cpp:0:0
 #2 0x00007f10d2ed5980 __restore_rt (/lib/x86_64-linux-gnu/libpthread.so.0+0x12980)
 #3 0x00007f10d1fcde87 raise /build/glibc-CVJwZb/glibc-2.27/signal/../sysdeps/unix/sysv/linux/raise.c:51:0
 #4 0x00007f10d1fcf7f1 abort /build/glibc-CVJwZb/glibc-2.27/stdlib/abort.c:81:0
 #5 0x00007f10d1fbf3fa __assert_fail_base /build/glibc-CVJwZb/glibc-2.27/assert/assert.c:89:0
 #6 0x00007f10d1fbf472 (/lib/x86_64-linux-gnu/libc.so.6+0x30472)
 #7 0x0000561fa380fee3 (anonymous namespace)::ArgMaxConverter::matchAndRewrite(mlir::tosa::ArgMaxOp, mlir::PatternRewriter&) const TosaToLinalg.cpp:0:0
 #8 0x0000561fa37e7a78 mlir::detail::OpOrInterfaceRewritePatternBase<mlir::tosa::ArgMaxOp>::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&) const (/home/verify/llvm-project-3b4cb1e96c645bb833fe710856479c31383859bb/build/bin/mlir-opt+0x1674a78)
 #9 0x0000561fa4029389 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/verify/llvm-project-3b4cb1e96c645bb833fe710856479c31383859bb/build/bin/mlir-opt+0x1eb6389)
#10 0x0000561fa3cc7c67 (anonymous namespace)::OperationLegalizer::legalize(mlir::Operation*, mlir::ConversionPatternRewriter&) DialectConversion.cpp:0:0
#11 0x0000561fa3cc86b7 (anonymous namespace)::OperationConverter::convertOperations(llvm::ArrayRef<mlir::Operation*>, llvm::function_ref<void (mlir::Diagnostic&)>) DialectConversion.cpp:0:0
#12 0x0000561fa3ccc37f mlir::applyFullConversion(mlir::Operation*, mlir::ConversionTarget&, mlir::FrozenRewritePatternSet const&) (/home/verify/llvm-project-3b4cb1e96c645bb833fe710856479c31383859bb/build/bin/mlir-opt+0x1b5937f)
#13 0x0000561fa37dfcbb (anonymous namespace)::TosaToLinalg::runOnOperation() TosaToLinalgPass.cpp:0:0
#14 0x0000561fa3c4c759 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (/home/verify/llvm-project-3b4cb1e96c645bb833fe710856479c31383859bb/build/bin/mlir-opt+0x1ad9759)
#15 0x0000561fa3c4ca7a mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) (/home/verify/llvm-project-3b4cb1e96c645bb833fe710856479c31383859bb/build/bin/mlir-opt+0x1ad9a7a)
#16 0x0000561fa3c4cdf8 mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::'lambda'(mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo&)::operator()(mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool)::OpPMInfo&) const Pass.cpp:0:0
#17 0x0000561fa3c4bf35 mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool) (/home/verify/llvm-project-3b4cb1e96c645bb833fe710856479c31383859bb/build/bin/mlir-opt+0x1ad8f35)
#18 0x0000561fa3c4c6ca mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (/home/verify/llvm-project-3b4cb1e96c645bb833fe710856479c31383859bb/build/bin/mlir-opt+0x1ad96ca)
#19 0x0000561fa3c4d5b3 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) (.constprop.495) Pass.cpp:0:0
#20 0x0000561fa3c4dc9d mlir::PassManager::run(mlir::Operation*) (/home/verify/llvm-project-3b4cb1e96c645bb833fe710856479c31383859bb/build/bin/mlir-opt+0x1adac9d)
#21 0x0000561fa3c40ca3 performActions(llvm::raw_ostream&, std::shared_ptr<llvm::SourceMgr> const&, mlir::MLIRContext*, mlir::MlirOptMainConfig const&) MlirOptMain.cpp:0:0
#22 0x0000561fa3c413d9 processBuffer(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, mlir::MlirOptMainConfig const&, mlir::DialectRegistry&, llvm::ThreadPool*) MlirOptMain.cpp:0:0
#23 0x0000561fa3c41701 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&) MlirOptMain.cpp:0:0
#24 0x0000561fa3d06380 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/verify/llvm-project-3b4cb1e96c645bb833fe710856479c31383859bb/build/bin/mlir-opt+0x1b93380)
#25 0x0000561fa3c3e3e3 mlir::MlirOptMain(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, mlir::DialectRegistry&, mlir::MlirOptMainConfig const&) (/home/verify/llvm-project-3b4cb1e96c645bb833fe710856479c31383859bb/build/bin/mlir-opt+0x1acb3e3)
#26 0x0000561fa3c41ae8 mlir::MlirOptMain(int, char**, llvm::StringRef, mlir::DialectRegistry&, bool) (/home/verify/llvm-project-3b4cb1e96c645bb833fe710856479c31383859bb/build/bin/mlir-opt+0x1aceae8)
#27 0x0000561fa2361330 main (/home/verify/llvm-project-3b4cb1e96c645bb833fe710856479c31383859bb/build/bin/mlir-opt+0x1ee330)
#28 0x00007f10d1fb0c87 __libc_start_main /build/glibc-CVJwZb/glibc-2.27/csu/../csu/libc-start.c:344:0
#29 0x0000561fa2373fca _start (/home/verify/llvm-project-3b4cb1e96c645bb833fe710856479c31383859bb/build/bin/mlir-opt+0x200fca)

[llvmbot]

@llvm/issue-subscribers-mlir-linalg

[BealZephyr]

still reproducible at b92839c!

[GeorgeARM]

Axis seems to be wrong in the example above as we have a 1d tensor and probably we need to perform reduction on axis 0 not 1.
Think this #68410 resolves the issue by adding a verifier