[Mlir-commits] [mlir] [MLIR][TableGen] Use arg index in InferredResultType constructor. (PR #122717)

[Philipp Schilk]

https://github.com/schilkp created https://github.com/llvm/llvm-project/pull/122717

Trying to constrain two results to be of the same type using `AllTypesMatch` would cause `mlir-tablgen` to crash on this assertion[1].

Example:

```tblgen
def OpL5 : NS_Op<"op_with_same_but_unconstraint_results",
    [AllTypesMatch<["result_a", "result_b"]>]> {
  let results = (outs AnyType:$result_a, AnyType:$result_b);
}
```

This is because there was a small bug when constructing the `inferences` graph from these constraints: The sources should be specified by the combined arg/result index (in other words, with results negative) not with the result index.


[1] https://github.com/llvm/llvm-project/blob/99612a3a18e0c40aac9c52b68e67b106f97ed4fa/mlir/lib/TableGen/Operator.cpp#L526

>From e68bc77479803e58672aded105527ee4632b46fa Mon Sep 17 00:00:00 2001
From: schilkp <schilk.philipp at gmail.com>
Date: Mon, 13 Jan 2025 15:28:07 +0100
Subject: [PATCH] [MLIR][TableGen] Use arg index in InferredResultType
 constructor.

`InferredResultType` construct expects an "result and arg"-style index
(with results negative and args positive), and not a results index
(results nonegative).
---
 mlir/lib/TableGen/Operator.cpp     |  4 ++--
 mlir/test/mlir-tblgen/op-result.td | 21 +++++++++++++++++++++
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/TableGen/Operator.cpp b/mlir/lib/TableGen/Operator.cpp
index c360c61afd27bd..20a43ef15d09eb 100644
--- a/mlir/lib/TableGen/Operator.cpp
+++ b/mlir/lib/TableGen/Operator.cpp
@@ -503,8 +503,8 @@ void Operator::populateTypeInferenceInfo(
         for (int otherResultIndex : resultIndices) {
           if (resultIndex == otherResultIndex)
             continue;
-          inference[resultIndex].sources.emplace_back(otherResultIndex,
-                                                      "$_self");
+          inference[resultIndex].sources.emplace_back(
+              InferredResultType::unmapResultIndex(otherResultIndex), "$_self");
         }
       }
     }
diff --git a/mlir/test/mlir-tblgen/op-result.td b/mlir/test/mlir-tblgen/op-result.td
index 51f8b0671a328d..f668d9a5a6644a 100644
--- a/mlir/test/mlir-tblgen/op-result.td
+++ b/mlir/test/mlir-tblgen/op-result.td
@@ -180,6 +180,27 @@ def OpL4 : NS_Op<"two_inference_edges", [
 // CHECK: inferredReturnTypes[1] = odsInferredType1
 // CHECK: inferredReturnTypes[2] = odsInferredType2
 
+def OpL5 : NS_Op<"op_with_same_but_unconstraint_results",
+    [AllTypesMatch<["result_a", "result_b"]>]> {
+  let results = (outs AnyType:$result_a, AnyType:$result_b);
+}
+
+// CHECK-NOT: LogicalResult OpL5::inferReturnTypes
+
+def OpL6 : NS_Op<"op_with_same_and_constraint_results",
+    [AllTypesMatch<["result_a", "result_b", "result_c"]>]> {
+  let results = (outs AnyType:$result_a, AnyType:$result_b, I32:$result_c);
+}
+
+// CHECK-LABEL: LogicalResult OpL6::inferReturnTypes
+// CHECK-NOT: }
+// CHECK: odsInferredType0 = odsBuilder.getIntegerType(32);
+// CHECK: odsInferredType1 = odsBuilder.getIntegerType(32);
+// CHECK: odsInferredType2 = odsBuilder.getIntegerType(32);
+// CHECK: inferredReturnTypes[0] = odsInferredType0;
+// CHECK: inferredReturnTypes[1] = odsInferredType1;
+// CHECK: inferredReturnTypes[2] = odsInferredType2;
+
 def OpM : NS_Op<"mix_diff_size_variadic_and_normal_results_op", [AttrSizedResultSegments]> {
   let results = (outs Variadic<AnyTensor>:$output1, AnyTensor:$output2, Optional<AnyTensor>:$output3);
 }