[llvm] [GISel] Infer the type of an immediate when there is one element in TEC (PR #77399)

[Sergei Barannikov via llvm-commits]

https://github.com/s-barannikov created https://github.com/llvm/llvm-project/pull/77399

When there is just one element in the type equivalence class (TEC), `inferNamedOperandType` fails because it does not consider the passed operand as a suitable one. This is incorrect when inferring the type of an (unnamed) immediate operand.

>From 8785396045132dbb3aae6b5fd749c7939fdae071 Mon Sep 17 00:00:00 2001
From: Sergei Barannikov <s.barannikov at module.ru>
Date: Tue, 9 Jan 2024 03:56:24 +0300
Subject: [PATCH] [GISel] Infer the type of an immediate when there is one
 element in TEC

When there is just one element in the type equivalence class (TEC),
`inferNamedOperandType` fails because it does not consider the passed
operand as a suitable one. This is incorrect when inferring the type of
an (unnamed) immediate operand.
---
 .../type-inference.td                         | 27 ++++++++++++++++---
 .../TableGen/GlobalISelCombinerEmitter.cpp    | 15 ++++++-----
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/llvm/test/TableGen/GlobalISelCombinerEmitter/type-inference.td b/llvm/test/TableGen/GlobalISelCombinerEmitter/type-inference.td
index c9ffe4e7adb3db..ed4e0e411c7afd 100644
--- a/llvm/test/TableGen/GlobalISelCombinerEmitter/type-inference.td
+++ b/llvm/test/TableGen/GlobalISelCombinerEmitter/type-inference.td
@@ -1,5 +1,5 @@
 // RUN: llvm-tblgen -I %p/../../../include -gen-global-isel-combiner \
-// RUN:     -gicombiner-debug-typeinfer -combiners=MyCombiner %s 2>&1 | \
+// RUN:     -gicombiner-debug-typeinfer -combiners=MyCombiner %s 2>&1 >/dev/null | \
 // RUN: FileCheck %s
 
 // Checks reasoning of the inference rules.
@@ -10,13 +10,15 @@ include "llvm/Target/GlobalISel/Combine.td"
 def MyTargetISA : InstrInfo;
 def MyTarget : Target { let InstructionSet = MyTargetISA; }
 
+// This also checks that the type of a def is preferred when inferring the type
+// of an immediate.
 // CHECK:      Rule Operand Type Equivalence Classes for inference_mul_by_neg_one:
 // CHECK-NEXT:         Groups for __inference_mul_by_neg_one_match_0:             [dst, x]
 // CHECK-NEXT:         Groups for __inference_mul_by_neg_one_apply_0:             [dst, x]
 // CHECK-NEXT: Final Type Equivalence Classes: [dst, x]
-// CHECK-NEXT: INFER: imm 0 -> GITypeOf<$x>
+// CHECK-NEXT: INFER: imm 0 -> GITypeOf<$dst>
 // CHECK-NEXT: Apply patterns for rule inference_mul_by_neg_one after inference:
-// CHECK-NEXT:   (CodeGenInstructionPattern name:__inference_mul_by_neg_one_apply_0 G_SUB operands:[<def>$dst, (GITypeOf<$x> 0), $x])
+// CHECK-NEXT:   (CodeGenInstructionPattern name:__inference_mul_by_neg_one_apply_0 G_SUB operands:[<def>$dst, (GITypeOf<$dst> 0), $x])
 def inference_mul_by_neg_one: GICombineRule <
   (defs root:$dst),
   (match (G_MUL $dst, $x, -1)),
@@ -61,8 +63,25 @@ def infer_variadic_outs: GICombineRule <
          (COPY $dst, $tmp))
 >;
 
+// Check that the type of an immediate is inferred when there is just one
+// element in the corresponding equivalence class.
+// CHECK:      Rule Operand Type Equivalence Classes for infer_imm_0:
+// CHECK-NEXT:   Groups for __infer_imm_0_match_0: [dst]
+// CHECK-NEXT:   Groups for __infer_imm_0_apply_0: [dst]
+// CHECK-NEXT: Final Type Equivalence Classes: [dst]
+// CHECK-NEXT: INFER: imm 1 -> GITypeOf<$dst>
+// CHECK-NEXT: INFER: imm 2 -> GITypeOf<$dst>
+// CHECK-NEXT: Apply patterns for rule infer_imm_0 after inference:
+// CHECK-NEXT:   (CodeGenInstructionPattern name:__infer_imm_0_apply_0 G_ADD operands:[<def>$dst, (GITypeOf<$dst> 1), (GITypeOf<$dst> 2)])
+def infer_imm_0 : GICombineRule<
+  (defs root:$dst),
+  (match (G_ADD $dst, 0, 3)),
+  (apply (G_ADD $dst, 1, 2))
+>;
+
 def MyCombiner: GICombiner<"GenMyCombiner", [
   inference_mul_by_neg_one,
   infer_complex_tempreg,
-  infer_variadic_outs
+  infer_variadic_outs,
+  infer_imm_0,
 ]>;
diff --git a/llvm/utils/TableGen/GlobalISelCombinerEmitter.cpp b/llvm/utils/TableGen/GlobalISelCombinerEmitter.cpp
index 89aca87a28ec0d..9121139be93540 100644
--- a/llvm/utils/TableGen/GlobalISelCombinerEmitter.cpp
+++ b/llvm/utils/TableGen/GlobalISelCombinerEmitter.cpp
@@ -284,7 +284,8 @@ class CombineRuleOperandTypeChecker : private OperandTypeChecker {
   /// succeed.
   PatternType inferNamedOperandType(const InstructionPattern &IP,
                                     StringRef OpName,
-                                    const TypeEquivalenceClasses &TECs) const;
+                                    const TypeEquivalenceClasses &TECs,
+                                    bool AllowSelf = false) const;
 
   const Record &RuleDef;
   SmallVector<InstructionPattern *, 8> MatchPats;
@@ -428,7 +429,8 @@ PatternType CombineRuleOperandTypeChecker::inferImmediateType(
 
       // Named operand with the same name, try to infer that.
       if (PatternType InferTy =
-              inferNamedOperandType(IP, Op.getOperandName(), TECs))
+              inferNamedOperandType(IP, Op.getOperandName(), TECs,
+                                    /*IncludeSelf=*/true))
         return InferTy;
     }
   }
@@ -438,16 +440,17 @@ PatternType CombineRuleOperandTypeChecker::inferImmediateType(
 
 PatternType CombineRuleOperandTypeChecker::inferNamedOperandType(
     const InstructionPattern &IP, StringRef OpName,
-    const TypeEquivalenceClasses &TECs) const {
+    const TypeEquivalenceClasses &TECs, bool AllowSelf) const {
   // This is the simplest possible case, we just need to find a TEC that
-  // contains OpName. Look at all other operands in equivalence class and try to
-  // find a suitable one.
+  // contains OpName. Look at all operands in equivalence class and try to
+  // find a suitable one. If `AllowSelf` is true, the operand itself is also
+  // considered suitable.
 
   // Check for a def of a matched pattern. This is guaranteed to always
   // be a register so we can blindly use that.
   StringRef GoodOpName;
   for (auto It = TECs.findLeader(OpName); It != TECs.member_end(); ++It) {
-    if (*It == OpName)
+    if (!AllowSelf && *It == OpName)
       continue;
 
     const auto LookupRes = MatchOpTable.lookup(*It);