[mlir] allow function type cloning to fail

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