diff --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp index d9d69342e42a8..8655ed3005a93 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp @@ -95,12 +95,7 @@ getBufferizedFunctionArgType(FuncOp funcOp, int64_t index, /// Return the FuncOp called by `callOp`. static FuncOp getCalledFunction(CallOpInterface callOp, SymbolTableCollection &symbolTables) { - SymbolRefAttr sym = - llvm::dyn_cast_if_present(callOp.getCallableForCallee()); - if (!sym) - return nullptr; - return dyn_cast_or_null( - symbolTables.lookupNearestSymbolFrom(callOp, sym)); + return dyn_cast_or_null(callOp.resolveCallableInTable(&symbolTables)); } /// Return the FuncOp called by `callOp`. diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp index aa53f94fe839d..c233e24c2a151 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp @@ -285,12 +285,8 @@ static void removeBufferizationAttributes(BlockArgument bbArg) { static func::FuncOp getCalledFunction(func::CallOp callOp, mlir::SymbolTableCollection &symbolTable) { - SymbolRefAttr sym = - llvm::dyn_cast_if_present(callOp.getCallableForCallee()); - if (!sym) - return nullptr; return dyn_cast_or_null( - symbolTable.lookupNearestSymbolFrom(callOp, sym)); + callOp.resolveCallableInTable(&symbolTable)); } /// Return "true" if the given function signature has tensor semantics.