[Mlir-commits] [mlir] [mlir] Validate function insertions (PR #206009)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: mygitljf

<details>
<summary>Changes</summary>

I added validation before function argument/result insertion updates the signature. Invalid insertion positions now fail cleanly with diagnostics instead of reaching lower-level type list handling and crashing.
Added focused tests for both argument and result insertion, covering out-of-range and unsorted insertion positions.
Fixes #<!-- -->205743

---
Full diff: https://github.com/llvm/llvm-project/pull/206009.diff


4 Files Affected:

- (modified) mlir/include/mlir/Interfaces/FunctionInterfaces.h (+13) 
- (modified) mlir/include/mlir/Interfaces/FunctionInterfaces.td (+11) 
- (added) mlir/test/IR/test-func-insert-arg-invalid.mlir (+17) 
- (added) mlir/test/IR/test-func-insert-result-invalid.mlir (+13) 


``````````diff
diff --git a/mlir/include/mlir/Interfaces/FunctionInterfaces.h b/mlir/include/mlir/Interfaces/FunctionInterfaces.h
index e10e9bd342702..8d10157af049b 100644
--- a/mlir/include/mlir/Interfaces/FunctionInterfaces.h
+++ b/mlir/include/mlir/Interfaces/FunctionInterfaces.h
@@ -28,6 +28,19 @@ class FunctionOpInterface;
 
 namespace function_interface_impl {
 
+inline bool areValidInsertionIndices(ArrayRef<unsigned> indices,
+                                     unsigned originalSize) {
+  unsigned previousIndex = 0;
+  bool isFirstIndex = true;
+  for (unsigned index : indices) {
+    if (index > originalSize || (!isFirstIndex && index < previousIndex))
+      return false;
+    previousIndex = index;
+    isFirstIndex = false;
+  }
+  return true;
+}
+
 /// Returns the dictionary attribute corresponding to the argument at 'index'.
 /// If there are no argument attributes at 'index', a null attribute is
 /// returned.
diff --git a/mlir/include/mlir/Interfaces/FunctionInterfaces.td b/mlir/include/mlir/Interfaces/FunctionInterfaces.td
index f701e828ed641..009b1fc08afa9 100644
--- a/mlir/include/mlir/Interfaces/FunctionInterfaces.td
+++ b/mlir/include/mlir/Interfaces/FunctionInterfaces.td
@@ -273,6 +273,12 @@ def FunctionOpInterface : OpInterface<"FunctionOpInterface", [
         ::llvm::ArrayRef<::mlir::DictionaryAttr> argAttrs,
         ::llvm::ArrayRef<::mlir::Location> argLocs) {
       unsigned originalNumArgs = $_op.getNumArguments();
+      if (argIndices.size() != argTypes.size() ||
+          argIndices.size() != argLocs.size() ||
+          (!argAttrs.empty() && argIndices.size() != argAttrs.size()) ||
+          !::mlir::function_interface_impl::areValidInsertionIndices(
+              argIndices, originalNumArgs))
+        return ::llvm::failure();
       ::mlir::Type newType = $_op.getTypeWithArgsAndResults(
           argIndices, argTypes, /*resultIndices=*/{}, /*resultTypes=*/{});
       if (!newType)
@@ -301,6 +307,11 @@ def FunctionOpInterface : OpInterface<"FunctionOpInterface", [
         ::mlir::TypeRange resultTypes,
         ::llvm::ArrayRef<::mlir::DictionaryAttr> resultAttrs) {
       unsigned originalNumResults = $_op.getNumResults();
+      if (resultIndices.size() != resultTypes.size() ||
+          (!resultAttrs.empty() && resultIndices.size() != resultAttrs.size()) ||
+          !::mlir::function_interface_impl::areValidInsertionIndices(
+              resultIndices, originalNumResults))
+        return ::llvm::failure();
       ::mlir::Type newType = $_op.getTypeWithArgsAndResults(
         /*argIndices=*/{}, /*argTypes=*/{}, resultIndices, resultTypes);
       if (!newType)
diff --git a/mlir/test/IR/test-func-insert-arg-invalid.mlir b/mlir/test/IR/test-func-insert-arg-invalid.mlir
new file mode 100644
index 0000000000000..d4c2769035552
--- /dev/null
+++ b/mlir/test/IR/test-func-insert-arg-invalid.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt %s -test-func-insert-arg -split-input-file -verify-diagnostics
+
+// expected-error @below {{failed to insert arguments}}
+func.func @f() attributes {test.insert_args = [
+  [0, i1, {test.A}],
+  [1, i2, {test.B}]]} {
+  return
+}
+
+// -----
+
+// expected-error @below {{failed to insert arguments}}
+func.func @f(%arg0: i1 {test.A}) attributes {test.insert_args = [
+  [1, i2, {test.B}],
+  [0, i3, {test.C}]]} {
+  return
+}
diff --git a/mlir/test/IR/test-func-insert-result-invalid.mlir b/mlir/test/IR/test-func-insert-result-invalid.mlir
new file mode 100644
index 0000000000000..176bd94aa1575
--- /dev/null
+++ b/mlir/test/IR/test-func-insert-result-invalid.mlir
@@ -0,0 +1,13 @@
+// RUN: mlir-opt %s -test-func-insert-result -split-input-file -verify-diagnostics
+
+// expected-error @below {{failed to insert results}}
+func.func private @f() attributes {test.insert_results = [
+  [0, f32, {test.A}],
+  [1, f32, {test.B}]]}
+
+// -----
+
+// expected-error @below {{failed to insert results}}
+func.func private @f() -> (f32 {test.A}) attributes {test.insert_results = [
+  [1, f32, {test.B}],
+  [0, f32, {test.C}]]}

``````````

</details>


https://github.com/llvm/llvm-project/pull/206009