intel/llvm
1
0
Fork 0
mirror of https://github.com/intel/llvm.git synced 2026-01-23 07:58:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

[mlir] allow function type cloning to fail (#137130)

Browse Source 
`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.
This commit is contained in:
Oleksandr "Alex" Zinenko
2025-04-30 09:26:42 +02:00
committed by GitHub
parent c91c3f930c
commit 7318074168
11 changed files with 111 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
flang/lib/Optimizer/Transforms/AbstractResult.cpp
View File

@@ -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);

3
mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.td
View File

@@ -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

81
mlir/include/mlir/Interfaces/FunctionInterfaces.td
View File

@@ -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<unsigned> 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<unsigned> 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<unsigned> 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<unsigned> 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<unsigned> argIndices, ::mlir::TypeRange argTypes,
::llvm::ArrayRef<unsigned> 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;

8
mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
View File

@@ -125,8 +125,12 @@ GPUFuncOpLowering::matchAndRewrite(gpu::GPUFuncOp gpuFuncOp, OpAdaptor adaptor,
// Perform signature modification
rewriter.modifyOpInPlace(
gpuFuncOp, [gpuFuncOp, &argIndices, &argTypes, &argAttrs, &argLocs]() {
static_cast<FunctionOpInterface>(gpuFuncOp).insertArguments(
argIndices, argTypes, argAttrs, argLocs);
LogicalResult inserted =
static_cast<FunctionOpInterface>(gpuFuncOp).insertArguments(
argIndices, argTypes, argAttrs, argLocs);
(void)inserted;
assert(succeeded(inserted) &&
"expected GPU funcs to support inserting any argument");
});
} else {
workgroupBuffers.reserve(gpuFuncOp.getNumWorkgroupAttributions());

3
mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
View File

@@ -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())

3
mlir/lib/Dialect/Bufferization/Transforms/DropEquivalentBufferResults.cpp
View File

@@ -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.

5
mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
View File

@@ -232,7 +232,10 @@ LLVMFunctionType::getChecked(function_ref<InFlightDiagnostic()> 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());
}

7
mlir/lib/Query/Query.cpp
View File

@@ -88,10 +88,11 @@ static Operation *extractFunction(std::vector<Operation *> &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;

7
mlir/lib/Transforms/RemoveDeadValues.cpp
View File

@@ -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

7
mlir/test/IR/test-func-erase-result.mlir
View File

@@ -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)

25
mlir/test/lib/IR/TestFunc.cpp
View File

@@ -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<int>(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<int>(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();
}
}
};