[mlir] segfault in DialectConversion.cpp - undoing rewrites in EraseBlockRewrite::rollback #111751
Description
Hey everyone,
I'm hitting this segfault in my repository, but I'm sending a reproducer in XLS since the conversion pass they have is much simpler than mine.
Reproducer
// RUN: xls/contrib/mlir/xls_opt -index-type-conversion=index-bitwidth=64 %s 2>&1 | FileCheck %s
// CHECK-LABEL: @forloop
func.func @forloop(%arg0: i32, %arg1: i8, %arg2: i9) -> i32 attributes {xls = true} {
%0 = xls.for inits(%arg0) invariants(%arg1, %arg2) {
^bb0(%arg3: index, %arg4: i32, %arg5: i8, %arg6: i9):
%c0 = arith.constant 0 : index
%0 = tensor.empty() : tensor<1xindex>
tensor.insert %arg3 into %0[%c0] : tensor<1xindex>
%i = arith.index_cast %0 : tensor<1xindex> to tensor<1xi32>
%1 = tensor.extract %i[%c0] : tensor<1xi32>
xls.yield %1 : i32
} { trip_count = 6 : i64 } : (i32, i8, i9) -> i32
return %0 : i32
}
Trace
When I run the test, I get a segfault, running the command I see this stack trace:
/google/src/cloud/asraa/xls-crash/google3/third_party/xls/contrib/mlir/testdata/crash.mlir:9:5: error: failed to legalize operation 'tensor.insert'
tensor.insert %arg3 into %0[%c0] : tensor<1xindex>
^
/google/src/cloud/asraa/xls-crash/google3/third_party/xls/contrib/mlir/testdata/crash.mlir:9:5: note: see current operation: %6 = "tensor.insert"(<<UNKNOWN SSA VALUE>>, %5, %4) : (index, tensor<1xindex>, index) -> tensor<1xindex>
PLEASE submit a bug report to http://go/llvm-crash-bug and include the crash backtrace.
Stack dump:
0. Program arguments: /usr/local/google/_blaze_asraa/487a6a5ca35cd7eca3408d7acb73ab28/execroot/google3/blaze-out/k8-opt/bin/third_party/xls/contrib/mlir/xls_opt -index-type-conversion=index-bitwidth=64 /google/src/cloud/asraa/xls-crash/google3/third_party/xls/contrib/mlir/testdata/crash.mlir
#0 0x000056436c332a9a llvm::sys::RunSignalHandlers() (/usr/local/google/_blaze_asraa/487a6a5ca35cd7eca3408d7acb73ab28/execroot/google3/blaze-out/k8-opt/bin/third_party/xls/contrib/mlir/xls_opt+0x4932a9a)
#1 0x000056436c3353f1 SignalHandler(int) Signals.cpp:0:0
#2 0x00007f069470ae80 __restore_rt (/usr/grte/v5/lib64/libpthread.so.0+0x14e80)
#3 0x000056436bd2a53b (anonymous namespace)::EraseBlockRewrite::rollback() DialectConversion.cpp:0:0
#4 0x000056436bd19e1b mlir::detail::ConversionPatternRewriterImpl::undoRewrites(unsigned int) (/usr/local/google/_blaze_asraa/487a6a5ca35cd7eca3408d7acb73ab28/execroot/google3/blaze-out/k8-opt/bin/third_party/xls/contrib/mlir/xls_opt+0x4319e1b)
#5 0x000056436bd1e705 mlir::OperationConverter::convertOperations(llvm::ArrayRef<mlir::Operation*>) (/usr/local/google/_blaze_asraa/487a6a5ca35cd7eca3408d7acb73ab28/execroot/google3/blaze-out/k8-opt/bin/third_party/xls/contrib/mlir/xls_opt+0x431e705)
#6 0x000056436bd2405a mlir::applyFullConversion(mlir::Operation*, mlir::ConversionTarget const&, mlir::FrozenRewritePatternSet const&, mlir::ConversionConfig) (/usr/local/google/_blaze_asraa/487a6a5ca35cd7eca3408d7acb73ab28/execroot/google3/blaze-out/k8-opt/bin/third_party/xls/contrib/mlir/xls_opt+0x432405a)
#7 0x000056436865f5e6 mlir::WalkResult llvm::function_ref<mlir::WalkResult (mlir::Operation*)>::callback_fn<mlir::xls::(anonymous namespace)::IndexTypeConversionPass::runOnOperation()::'lambda'(mlir::Operation*)>(long, mlir::Operation*) index_type_conversion.cc:0:0
#8 0x000056436863dc2f mlir::WalkResult mlir::detail::walk<mlir::ForwardIterator>(mlir::Operation*, llvm::function_ref<mlir::WalkResult (mlir::Operation*)>, mlir::WalkOrder) (/usr/local/google/_blaze_asraa/487a6a5ca35cd7eca3408d7acb73ab28/execroot/google3/blaze-out/k8-opt/bin/third_party/xls/contrib/mlir/xls_opt+0xc3dc2f)
#9 0x000056436865dd32 mlir::xls::(anonymous namespace)::IndexTypeConversionPass::runOnOperation() index_type_conversion.cc:0:0
#10 0x000056436c19bfa6 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (/usr/local/google/_blaze_asraa/487a6a5ca35cd7eca3408d7acb73ab28/execroot/google3/blaze-out/k8-opt/bin/third_party/xls/contrib/mlir/xls_opt+0x479bfa6)
#11 0x000056436c19ec5c mlir::PassManager::run(mlir::Operation*) (/usr/local/google/_blaze_asraa/487a6a5ca35cd7eca3408d7acb73ab28/execroot/google3/blaze-out/k8-opt/bin/third_party/xls/contrib/mlir/xls_opt+0x479ec5c)
#12 0x000056436bd55b67 performActions(llvm::raw_ostream&, std::__u::shared_ptr<llvm::SourceMgr> const&, mlir::MLIRContext*, mlir::MlirOptMainConfig const&) MlirOptMain.cpp:0:0
#13 0x000056436bd55833 llvm::LogicalResult llvm::function_ref<llvm::LogicalResult (std::__u::unique_ptr<llvm::MemoryBuffer, std::__u::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>::callback_fn<mlir::MlirOptMain(llvm::raw_ostream&, std::__u::unique_ptr<llvm::MemoryBuffer, std::__u::default_delete<llvm::MemoryBuffer>>, mlir::DialectRegistry&, mlir::MlirOptMainConfig const&)::$_0>(long, std::__u::unique_ptr<llvm::MemoryBuffer, std::__u::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&) MlirOptMain.cpp:0:0
#14 0x000056436bd511a2 mlir::MlirOptMain(llvm::raw_ostream&, std::__u::unique_ptr<llvm::MemoryBuffer, std::__u::default_delete<llvm::MemoryBuffer>>, mlir::DialectRegistry&, mlir::MlirOptMainConfig const&) (/usr/local/google/_blaze_asraa/487a6a5ca35cd7eca3408d7acb73ab28/execroot/google3/blaze-out/k8-opt/bin/third_party/xls/contrib/mlir/xls_opt+0x43511a2)
#15 0x000056436bd51411 mlir::MlirOptMain(int, char**, llvm::StringRef, llvm::StringRef, mlir::DialectRegistry&) (/usr/local/google/_blaze_asraa/487a6a5ca35cd7eca3408d7acb73ab28/execroot/google3/blaze-out/k8-opt/bin/third_party/xls/contrib/mlir/xls_opt+0x4351411)
#16 0x000056436bd515f9 mlir::MlirOptMain(int, char**, llvm::StringRef, mlir::DialectRegistry&) (/usr/local/google/_blaze_asraa/487a6a5ca35cd7eca3408d7acb73ab28/execroot/google3/blaze-out/k8-opt/bin/third_party/xls/contrib/mlir/xls_opt+0x43515f9)
#17 0x00005643686353fc main (/usr/local/google/_blaze_asraa/487a6a5ca35cd7eca3408d7acb73ab28/execroot/google3/blaze-out/k8-opt/bin/third_party/xls/contrib/mlir/xls_opt+0xc353fc)
#18 0x00007f06945a63d4 __libc_start_main (/usr/grte/v5/lib64/libc.so.6+0x613d4)
#19 0x000056436863522a _start /usr/grte/v5/debug-src/src/csu/../sysdeps/x86_64/start.S:122:0
Segmentation fault
Root Cause
I know this failure is only occuring because there's an operation (arith.index_cast) that fails to legalize within an operation that was rewritten with LegalizeGeneralOps. I think rolling back the new blocks from this pattern is causing the segfault.
I couldn't really find any better way to rewrite the LegalizeGeneralOps pattern. <y repo also needs a similar pattern (or it could be a specific one for for loops, but specializing it doesn't help in this case). I also have no way of controlling the order that the dialect converter applies patterns, so I can't say to apply conversions to inner regions before applying the for loop pattern.