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

`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 llvm#131142.

Author: ftynse

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

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

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

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

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.

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

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

mlir/lib/Query/Query.cpp

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

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

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)

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