From 7318074168e324097b7f6ec93e34915463bf70f1 Mon Sep 17 00:00:00 2001 From: "Oleksandr \"Alex\" Zinenko" Date: Wed, 30 Apr 2025 09:26:42 +0200 Subject: [PATCH] [mlir] allow function type cloning to fail (#137130) `FunctionOpInterface` assumed the fact that the function type (attribute of the operation) can be cloned with arbirary lists of function arguments and results to support argument and result list mutation. This is not always correct, in particular, LLVM dialect functions require exactly one result making it impossible to erase the result. Allow function type cloning to fail and propagate this failure through various APIs that use it. The common assumption is that existing IR has not been modified. Fixes #131142. Reland a8c7ecdcbc3e89b493b495c6831cc93671c3b844 / #136300. --- .../Optimizer/Transforms/AbstractResult.cpp | 8 +- mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.td | 3 +- .../mlir/Interfaces/FunctionInterfaces.td | 81 ++++++++++++------- .../Conversion/GPUCommon/GPUOpsLowering.cpp | 8 +- .../Transforms/BufferResultsToOutParams.cpp | 3 +- .../DropEquivalentBufferResults.cpp | 3 +- mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp | 5 +- mlir/lib/Query/Query.cpp | 7 +- mlir/lib/Transforms/RemoveDeadValues.cpp | 7 +- mlir/test/IR/test-func-erase-result.mlir | 7 +- mlir/test/lib/IR/TestFunc.cpp | 25 ++++-- 11 files changed, 111 insertions(+), 46 deletions(-) diff --git a/flang/lib/Optimizer/Transforms/AbstractResult.cpp b/flang/lib/Optimizer/Transforms/AbstractResult.cpp index f8badfa639f9..59e2eeb76c71 100644 --- a/flang/lib/Optimizer/Transforms/AbstractResult.cpp +++ b/flang/lib/Optimizer/Transforms/AbstractResult.cpp @@ -387,8 +387,12 @@ public: mlir::OpBuilder rewriter(context); auto resultType = funcTy.getResult(0); auto argTy = getResultArgumentType(resultType, shouldBoxResult); - func.insertArgument(0u, argTy, {}, loc); - func.eraseResult(0u); + llvm::LogicalResult res = func.insertArgument(0u, argTy, {}, loc); + (void)res; + assert(llvm::succeeded(res) && "failed to insert function argument"); + res = func.eraseResult(0u); + (void)res; + assert(llvm::succeeded(res) && "failed to erase function result"); mlir::Value newArg = func.getArgument(0u); if (mustEmboxResult(resultType, shouldBoxResult)) { auto bufferType = fir::ReferenceType::get(resultType); diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.td index e9f01ba5cd4e..072fc7be07dd 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.td +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.td @@ -104,7 +104,8 @@ def LLVMFunctionType : LLVMType<"LLVMFunction", "func"> { bool isVarArg() const { return getVarArg(); } /// Returns a clone of this function type with the given argument - /// and result types. + /// and result types. Returns null if the resulting function type would + /// not verify. LLVMFunctionType clone(TypeRange inputs, TypeRange results) const; /// Returns the result type of the function as an ArrayRef, enabling better diff --git a/mlir/include/mlir/Interfaces/FunctionInterfaces.td b/mlir/include/mlir/Interfaces/FunctionInterfaces.td index 1e2ca9a8b449..cad0e790fdd9 100644 --- a/mlir/include/mlir/Interfaces/FunctionInterfaces.td +++ b/mlir/include/mlir/Interfaces/FunctionInterfaces.td @@ -255,79 +255,105 @@ def FunctionOpInterface : OpInterface<"FunctionOpInterface", [ BlockArgListType getArguments() { return getFunctionBody().getArguments(); } /// Insert a single argument of type `argType` with attributes `argAttrs` and - /// location `argLoc` at `argIndex`. - void insertArgument(unsigned argIndex, ::mlir::Type argType, ::mlir::DictionaryAttr argAttrs, - ::mlir::Location argLoc) { - insertArguments({argIndex}, {argType}, {argAttrs}, {argLoc}); + /// location `argLoc` at `argIndex`. Returns failure if the function cannot be + /// updated to have the new signature. + ::llvm::LogicalResult insertArgument( + unsigned argIndex, ::mlir::Type argType, ::mlir::DictionaryAttr argAttrs, + ::mlir::Location argLoc) { + return insertArguments({argIndex}, {argType}, {argAttrs}, {argLoc}); } /// Inserts arguments with the listed types, attributes, and locations at the /// listed indices. `argIndices` must be sorted. Arguments are inserted in the /// order they are listed, such that arguments with identical index will - /// appear in the same order that they were listed here. - void insertArguments(::llvm::ArrayRef argIndices, ::mlir::TypeRange argTypes, - ::llvm::ArrayRef<::mlir::DictionaryAttr> argAttrs, - ::llvm::ArrayRef<::mlir::Location> argLocs) { + /// appear in the same order that they were listed here. Returns failure if + /// the function cannot be updated to have the new signature. + ::llvm::LogicalResult insertArguments( + ::llvm::ArrayRef argIndices, ::mlir::TypeRange argTypes, + ::llvm::ArrayRef<::mlir::DictionaryAttr> argAttrs, + ::llvm::ArrayRef<::mlir::Location> argLocs) { unsigned originalNumArgs = $_op.getNumArguments(); ::mlir::Type newType = $_op.getTypeWithArgsAndResults( argIndices, argTypes, /*resultIndices=*/{}, /*resultTypes=*/{}); + if (!newType) + return ::llvm::failure(); ::mlir::function_interface_impl::insertFunctionArguments( $_op, argIndices, argTypes, argAttrs, argLocs, originalNumArgs, newType); + return ::llvm::success(); } - /// Insert a single result of type `resultType` at `resultIndex`. - void insertResult(unsigned resultIndex, ::mlir::Type resultType, - ::mlir::DictionaryAttr resultAttrs) { - insertResults({resultIndex}, {resultType}, {resultAttrs}); + /// Insert a single result of type `resultType` at `resultIndex`.Returns + /// failure if the function cannot be updated to have the new signature. + ::llvm::LogicalResult insertResult( + unsigned resultIndex, ::mlir::Type resultType, + ::mlir::DictionaryAttr resultAttrs) { + return insertResults({resultIndex}, {resultType}, {resultAttrs}); } /// Inserts results with the listed types at the listed indices. /// `resultIndices` must be sorted. Results are inserted in the order they are /// listed, such that results with identical index will appear in the same - /// order that they were listed here. - void insertResults(::llvm::ArrayRef resultIndices, ::mlir::TypeRange resultTypes, - ::llvm::ArrayRef<::mlir::DictionaryAttr> resultAttrs) { + /// order that they were listed here. Returns failure if the function + /// cannot be updated to have the new signature. + ::llvm::LogicalResult insertResults( + ::llvm::ArrayRef resultIndices, + ::mlir::TypeRange resultTypes, + ::llvm::ArrayRef<::mlir::DictionaryAttr> resultAttrs) { unsigned originalNumResults = $_op.getNumResults(); ::mlir::Type newType = $_op.getTypeWithArgsAndResults( /*argIndices=*/{}, /*argTypes=*/{}, resultIndices, resultTypes); + if (!newType) + return ::llvm::failure(); ::mlir::function_interface_impl::insertFunctionResults( $_op, resultIndices, resultTypes, resultAttrs, originalNumResults, newType); + return ::llvm::success(); } - /// Erase a single argument at `argIndex`. - void eraseArgument(unsigned argIndex) { + /// Erase a single argument at `argIndex`. Returns failure if the function + /// cannot be updated to have the new signature. + ::llvm::LogicalResult eraseArgument(unsigned argIndex) { ::llvm::BitVector argsToErase($_op.getNumArguments()); argsToErase.set(argIndex); - eraseArguments(argsToErase); + return eraseArguments(argsToErase); } - /// Erases the arguments listed in `argIndices`. - void eraseArguments(const ::llvm::BitVector &argIndices) { + /// Erases the arguments listed in `argIndices`. Returns failure if the + /// function cannot be updated to have the new signature. + ::llvm::LogicalResult eraseArguments(const ::llvm::BitVector &argIndices) { ::mlir::Type newType = $_op.getTypeWithoutArgs(argIndices); + if (!newType) + return ::llvm::failure(); ::mlir::function_interface_impl::eraseFunctionArguments( $_op, argIndices, newType); + return ::llvm::success(); } - /// Erase a single result at `resultIndex`. - void eraseResult(unsigned resultIndex) { + /// Erase a single result at `resultIndex`. Returns failure if the function + /// cannot be updated to have the new signature. + LogicalResult eraseResult(unsigned resultIndex) { ::llvm::BitVector resultsToErase($_op.getNumResults()); resultsToErase.set(resultIndex); - eraseResults(resultsToErase); + return eraseResults(resultsToErase); } - /// Erases the results listed in `resultIndices`. - void eraseResults(const ::llvm::BitVector &resultIndices) { + /// Erases the results listed in `resultIndices`. Returns failure if the + /// function cannot be updated to have the new signature. + ::llvm::LogicalResult eraseResults(const ::llvm::BitVector &resultIndices) { ::mlir::Type newType = $_op.getTypeWithoutResults(resultIndices); + if (!newType) + return ::llvm::failure(); ::mlir::function_interface_impl::eraseFunctionResults( $_op, resultIndices, newType); + return ::llvm::success(); } /// Return the type of this function with the specified arguments and /// results inserted. This is used to update the function's signature in /// the `insertArguments` and `insertResults` methods. The arrays must be - /// sorted by increasing index. + /// sorted by increasing index. Return nullptr if the updated type would + /// not be valid. ::mlir::Type getTypeWithArgsAndResults( ::llvm::ArrayRef argIndices, ::mlir::TypeRange argTypes, ::llvm::ArrayRef resultIndices, ::mlir::TypeRange resultTypes) { @@ -341,7 +367,8 @@ def FunctionOpInterface : OpInterface<"FunctionOpInterface", [ /// Return the type of this function without the specified arguments and /// results. This is used to update the function's signature in the - /// `eraseArguments` and `eraseResults` methods. + /// `eraseArguments` and `eraseResults` methods. Return nullptr if the + /// updated type would not be valid. ::mlir::Type getTypeWithoutArgsAndResults( const ::llvm::BitVector &argIndices, const ::llvm::BitVector &resultIndices) { ::llvm::SmallVector<::mlir::Type> argStorage, resultStorage; diff --git a/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp b/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp index f22ad1fd70db..dea3d99704d9 100644 --- a/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp +++ b/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp @@ -125,8 +125,12 @@ GPUFuncOpLowering::matchAndRewrite(gpu::GPUFuncOp gpuFuncOp, OpAdaptor adaptor, // Perform signature modification rewriter.modifyOpInPlace( gpuFuncOp, [gpuFuncOp, &argIndices, &argTypes, &argAttrs, &argLocs]() { - static_cast(gpuFuncOp).insertArguments( - argIndices, argTypes, argAttrs, argLocs); + LogicalResult inserted = + static_cast(gpuFuncOp).insertArguments( + argIndices, argTypes, argAttrs, argLocs); + (void)inserted; + assert(succeeded(inserted) && + "expected GPU funcs to support inserting any argument"); }); } else { workgroupBuffers.reserve(gpuFuncOp.getNumWorkgroupAttributions()); diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp index 1c95ab77b9f3..acf5f7767d12 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp @@ -92,7 +92,8 @@ updateFuncOp(func::FuncOp func, } // Erase the results. - func.eraseResults(erasedResultIndices); + if (failed(func.eraseResults(erasedResultIndices))) + return failure(); // Add the new arguments to the entry block if the function is not external. if (func.isExternal()) diff --git a/mlir/lib/Dialect/Bufferization/Transforms/DropEquivalentBufferResults.cpp b/mlir/lib/Dialect/Bufferization/Transforms/DropEquivalentBufferResults.cpp index 9bc75267e70e..ae011904cb97 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/DropEquivalentBufferResults.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/DropEquivalentBufferResults.cpp @@ -113,7 +113,8 @@ mlir::bufferization::dropEquivalentBufferResults(ModuleOp module) { } // Update function. - funcOp.eraseResults(erasedResultIndices); + if (failed(funcOp.eraseResults(erasedResultIndices))) + return failure(); returnOp.getOperandsMutable().assign(newReturnValues); // Update function calls. diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp index 29cf38c1fefe..e2bccf7f6493 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp @@ -232,7 +232,10 @@ LLVMFunctionType::getChecked(function_ref emitError, LLVMFunctionType LLVMFunctionType::clone(TypeRange inputs, TypeRange results) const { - assert(results.size() == 1 && "expected a single result type"); + if (results.size() != 1 || !isValidResultType(results[0])) + return {}; + if (!llvm::all_of(inputs, isValidArgumentType)) + return {}; return get(results[0], llvm::to_vector(inputs), isVarArg()); } diff --git a/mlir/lib/Query/Query.cpp b/mlir/lib/Query/Query.cpp index 869ee8f2ae1d..73f313cd37fd 100644 --- a/mlir/lib/Query/Query.cpp +++ b/mlir/lib/Query/Query.cpp @@ -88,10 +88,11 @@ static Operation *extractFunction(std::vector &ops, // Remove unused function arguments size_t currentIndex = 0; while (currentIndex < funcOp.getNumArguments()) { + // Erase if possible. if (funcOp.getArgument(currentIndex).use_empty()) - funcOp.eraseArgument(currentIndex); - else - ++currentIndex; + if (succeeded(funcOp.eraseArgument(currentIndex))) + continue; + ++currentIndex; } return funcOp; diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp index 92c05d87a002..4939dda4a20f 100644 --- a/mlir/lib/Transforms/RemoveDeadValues.cpp +++ b/mlir/lib/Transforms/RemoveDeadValues.cpp @@ -698,8 +698,11 @@ static void cleanUpDeadVals(RDVFinalCleanupList &list) { // 3. Functions for (auto &f : list.functions) { - f.funcOp.eraseArguments(f.nonLiveArgs); - f.funcOp.eraseResults(f.nonLiveRets); + // Some functions may not allow erasing arguments or results. These calls + // return failure in such cases without modifying the function, so it's okay + // to proceed. + (void)f.funcOp.eraseArguments(f.nonLiveArgs); + (void)f.funcOp.eraseResults(f.nonLiveRets); } // 4. Operands diff --git a/mlir/test/IR/test-func-erase-result.mlir b/mlir/test/IR/test-func-erase-result.mlir index 8fe40b2fdfcb..a32866227547 100644 --- a/mlir/test/IR/test-func-erase-result.mlir +++ b/mlir/test/IR/test-func-erase-result.mlir @@ -1,4 +1,4 @@ -// RUN: mlir-opt %s -test-func-erase-result -split-input-file | FileCheck %s +// RUN: mlir-opt %s -test-func-erase-result -split-input-file -verify-diagnostics | FileCheck %s // CHECK: func private @f(){{$}} // CHECK-NOT: attributes{{.*}}result @@ -66,3 +66,8 @@ func.func private @f() -> ( f32 {test.erase_this_result}, tensor<3xf32> ) + +// ----- + +// expected-error @below {{failed to erase results}} +llvm.func @llvm_func(!llvm.ptr, i64) diff --git a/mlir/test/lib/IR/TestFunc.cpp b/mlir/test/lib/IR/TestFunc.cpp index 2ade47249863..94a461036586 100644 --- a/mlir/test/lib/IR/TestFunc.cpp +++ b/mlir/test/lib/IR/TestFunc.cpp @@ -45,8 +45,12 @@ struct TestFuncInsertArg : unknownLoc); } func->removeAttr("test.insert_args"); - func.insertArguments(indicesToInsert, typesToInsert, attrsToInsert, - locsToInsert); + if (succeeded(func.insertArguments(indicesToInsert, typesToInsert, + attrsToInsert, locsToInsert))) + continue; + + emitError(func->getLoc()) << "failed to insert arguments"; + return signalPassFailure(); } } }; @@ -79,7 +83,12 @@ struct TestFuncInsertResult : DictionaryAttr::get(&getContext())); } func->removeAttr("test.insert_results"); - func.insertResults(indicesToInsert, typesToInsert, attrsToInsert); + if (succeeded(func.insertResults(indicesToInsert, typesToInsert, + attrsToInsert))) + continue; + + emitError(func->getLoc()) << "failed to insert results"; + return signalPassFailure(); } } }; @@ -100,7 +109,10 @@ struct TestFuncEraseArg for (auto argIndex : llvm::seq(0, func.getNumArguments())) if (func.getArgAttr(argIndex, "test.erase_this_arg")) indicesToErase.set(argIndex); - func.eraseArguments(indicesToErase); + if (succeeded(func.eraseArguments(indicesToErase))) + continue; + emitError(func->getLoc()) << "failed to erase arguments"; + return signalPassFailure(); } } }; @@ -122,7 +134,10 @@ struct TestFuncEraseResult for (auto resultIndex : llvm::seq(0, func.getNumResults())) if (func.getResultAttr(resultIndex, "test.erase_this_result")) indicesToErase.set(resultIndex); - func.eraseResults(indicesToErase); + if (succeeded(func.eraseResults(indicesToErase))) + continue; + emitError(func->getLoc()) << "failed to erase results"; + return signalPassFailure(); } } };