One shot bufferization on call op segfaults when bufferizing function boundaries for a callee with tensor arguments and forcing copy before write

[paul0403]

The llvm version I am using is commit f8cb7987c64dcffb72414a40560055cb717dbf74 f8cb798

The following fails

  // RUN: mlir-opt --pass-pipeline=
  //                          "builtin.module(one-shot-bufferize{
  //                              bufferize-function-boundaries  
  //                              function-boundary-type-conversion=identity-layout-map  
  //                              copy-before-write})" %s
  func.func public @main(%arg0: tensor<2xf64>) -> f64 {
    %0 = call @my_four(%arg0) : (tensor<2xf64>) -> f64
    return %0 : f64
  }
  func.func public @my_four(%arg0: tensor<2xf64>) -> f64 {
    %four = arith.constant 4.2 : f64
    return %four : f64
  }

PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
Stack dump:
0.	Program arguments: /home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt "--pass-pipeline=builtin.module(one-shot-bufferize{bufferize-function-boundaries  function-boundary-type-conversion=identity-layout-map  copy-before-write})" ./multiple_qnodes/3_QuantumCompilationPass.mlir
 #0 0x00005e5cef4d58a7 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x37188a7)
 #1 0x00005e5cef4d34f5 llvm::sys::RunSignalHandlers() (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x37164f5)
 #2 0x00005e5cef4d60d5 SignalHandler(int, siginfo_t*, void*) Signals.cpp:0:0
 #3 0x0000753708642520 (/lib/x86_64-linux-gnu/libc.so.6+0x42520)
 #4 0x00005e5cefaadfe0 mlir::bufferization::func_ext::FuncAnalysisState* mlir::bufferization::OneShotAnalysisState::getExtension<mlir::bufferization::func_ext::FuncAnalysisState>() FuncBufferizableOpInterfaceImpl.cpp:0:0
 #5 0x00005e5cefaadc83 mlir::bufferization::func_ext::getCalledFunction(mlir::CallOpInterface, mlir::bufferization::AnalysisState const&) FuncBufferizableOpInterfaceImpl.cpp:0:0
 #6 0x00005e5cefaae11a mlir::bufferization::func_ext::CallOpInterface::bufferizesToMemoryWrite(mlir::Operation*, mlir::OpOperand&, mlir::bufferization::AnalysisState const&) const FuncBufferizableOpInterfaceImpl.cpp:0:0
 #7 0x00005e5cefa3d531 mlir::bufferization::AnalysisState::isInPlace(mlir::OpOperand&) const (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x3c80531)
 #8 0x00005e5cefa3ac82 mlir::bufferization::BufferizableOpInterface::resolveTensorOpOperandConflicts(mlir::RewriterBase&, mlir::bufferization::AnalysisState const&, mlir::bufferization::BufferizationState const&) (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x3c7dc82)
 #9 0x00005e5cefaaef01 mlir::bufferization::detail::BufferizableOpInterfaceInterfaceTraits::ExternalModel<mlir::bufferization::func_ext::CallOpInterface, mlir::func::CallOp>::resolveConflicts(mlir::Operation*, mlir::RewriterBase&, mlir::bufferization::AnalysisState const&, mlir::bufferization::BufferizationState const&) const (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x3cf1f01)
#10 0x00005e5cefad72c8 mlir::WalkResult llvm::function_ref<mlir::WalkResult (mlir::Operation*)>::callback_fn<mlir::bufferization::insertTensorCopies(mlir::Operation*, mlir::bufferization::AnalysisState const&, mlir::bufferization::BufferizationState const&)::$_1>(long, mlir::Operation*) TensorCopyInsertion.cpp:0:0
#11 0x00005e5cef580f87 mlir::WalkResult mlir::detail::walk<mlir::ForwardIterator>(mlir::Operation*, llvm::function_ref<mlir::WalkResult (mlir::Operation*)>, mlir::WalkOrder) (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x37c3f87)
#12 0x00005e5cefad717a mlir::bufferization::insertTensorCopies(mlir::Operation*, mlir::bufferization::AnalysisState const&, mlir::bufferization::BufferizationState const&) (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x3d1a17a)
#13 0x00005e5cefa8e8db mlir::bufferization::bufferizeOp(mlir::Operation*, mlir::bufferization::BufferizationOptions const&, mlir::bufferization::BufferizationState&, mlir::bufferization::BufferizationStatistics*) (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x3cd18db)
#14 0x00005e5cefac88b8 mlir::bufferization::bufferizeModuleOp(mlir::ModuleOp, mlir::bufferization::OneShotBufferizationOptions const&, mlir::bufferization::BufferizationState&, mlir::bufferization::BufferizationStatistics*) (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x3d0b8b8)
#15 0x00005e5cefac9720 mlir::bufferization::runOneShotModuleBufferize(mlir::ModuleOp, mlir::bufferization::OneShotBufferizationOptions const&, mlir::bufferization::BufferizationState&, mlir::bufferization::BufferizationStatistics*) (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x3d0c720)
#16 0x00005e5cefa92bc4 (anonymous namespace)::OneShotBufferizePass::runOnOperation() Bufferize.cpp:0:0
#17 0x00005e5cf2f1f755 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x7162755)
#18 0x00005e5cf2f202a1 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x71632a1)
#19 0x00005e5cf2f22a6b mlir::PassManager::run(mlir::Operation*) (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x7165a6b)
#20 0x00005e5cf2f19dd1 performActions(llvm::raw_ostream&, std::shared_ptr<llvm::SourceMgr> const&, mlir::MLIRContext*, mlir::MlirOptMainConfig const&) MlirOptMain.cpp:0:0
#21 0x00005e5cf2f19a24 llvm::LogicalResult llvm::function_ref<llvm::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&)::$_3>(long, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&) MlirOptMain.cpp:0:0
#22 0x00005e5cf2fd6bb5 mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<llvm::LogicalResult (std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>, llvm::raw_ostream&, llvm::StringRef, llvm::StringRef) (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x7219bb5)
#23 0x00005e5cf2f12b01 mlir::MlirOptMain(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, mlir::DialectRegistry&, mlir::MlirOptMainConfig const&) (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x7155b01)
#24 0x00005e5cf2f12db8 mlir::MlirOptMain(int, char**, llvm::StringRef, llvm::StringRef, mlir::DialectRegistry&) (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x7155db8)
#25 0x00005e5cf2f12fd2 mlir::MlirOptMain(int, char**, llvm::StringRef, mlir::DialectRegistry&) (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x7155fd2)
#26 0x00005e5cef421dc3 main (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x3664dc3)
#27 0x0000753708629d90 __libc_start_call_main ./csu/../sysdeps/nptl/libc_start_call_main.h:58:16
#28 0x0000753708629e40 call_init ./csu/../csu/libc-start.c:128:20
#29 0x0000753708629e40 __libc_start_main ./csu/../csu/libc-start.c:379:5
#30 0x00005e5cef421865 _start (/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../mlir/llvm-project/build/bin/mlir-opt+0x3664865)
Segmentation fault (core dumped)

The logs led me to #138125, which I believe updated mlir::bufferization::func_ext::getCalledFunction(mlir::CallOpInterface, mlir::bufferization::AnalysisState const&) responsible for the segfault.

---

Any one of the following will make the segfault go away:

1. Making the callee argument not a tensor, or removing the argument altogether. (Of course, without tensors there's nothing to bufferize)

2. Turning off copy-before-write, i.e.

mlir-opt --pass-pipeline="builtin.module(one-shot-bufferize{
   bufferize-function-boundaries  
   function-boundary-type-conversion=identity-layout-map  })" test.mlir

This produces the correct results:

module {
  func.func public @main(%arg0: memref<2xf64>) -> f64 {
    %0 = call @my_four(%arg0) : (memref<2xf64>) -> f64
    return %0 : f64
  }
  func.func public @my_four(%arg0: memref<2xf64>) -> f64 {
    %cst = arith.constant 4.200000e+00 : f64
    return %cst : f64
  }
}

Unfortunately for my purposes this is not an option, since I am doing something with asynchronous execution so I have to make conservative buffer decisions.

3. Do not bufferize function boundaries, i.e.

mlir-opt --pass-pipeline="builtin.module(one-shot-bufferize{
   copy-before-write  })" test.mlir

module {
  func.func public @main(%arg0: tensor<2xf64>) -> f64 {
    %0 = call @my_four(%arg0) : (tensor<2xf64>) -> f64
    return %0 : f64
  }
  func.func public @my_four(%arg0: tensor<2xf64>) -> f64 {
    %cst = arith.constant 4.200000e+00 : f64
    return %cst : f64
  }
}

But this essentially amounts to giving up, instead of solving the problem.

*. As a special note, if we don't set identity layout map and still force copy before write, the pass hangs indefinitely:

mlir-opt --pass-pipeline="builtin.module(one-shot-bufferize{
   bufferize-function-boundaries  
   copy-before-write  })" test.mlir

// Nothing happens, hangs indefinitely

[paul0403]

Just an update that I confirmed #138125 was the cause, reverting back to the old getCalledFunction solved the issue for me.

--- a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
@@ -89,16 +89,12 @@ static FuncOp getCalledFunction(CallOpInterface callOp,
 /// Return the FuncOp called by `callOp`.
 static FuncOp getCalledFunction(CallOpInterface callOp,
                                 const AnalysisState &state) {
-  auto &oneShotAnalysisState = static_cast<const OneShotAnalysisState &>(state);
-
-  if (auto *funcAnalysisState =
-          oneShotAnalysisState.getExtension<FuncAnalysisState>()) {
-    // Use the cached symbol tables.
-    return getCalledFunction(callOp, funcAnalysisState->symbolTables);
-  }
-
-  SymbolTableCollection symbolTables;
-  return getCalledFunction(callOp, symbolTables);
+  SymbolRefAttr sym =
+      llvm::dyn_cast_if_present<SymbolRefAttr>(callOp.getCallableForCallee());
+  if (!sym)
+    return nullptr;
+  return dyn_cast_or_null<FuncOp>(
+      SymbolTable::lookupNearestSymbolFrom(callOp, sym));
 }
 
 /// Get FuncAnalysisState.