[GlobalIsSel] Allow using PatFrags with multiple defs as the root of a combine rule

I had to tighten the restrictions on PatFrags a bit to make it consistent: instructions that
define the root of a PF can only have one def.

Reviewed By: arsenm

Differential Revision: https://reviews.llvm.org/D157700

Author: Pierre-vh

llvm/docs/GlobalISel/MIRPatterns.rst

@@ -110,6 +110,8 @@ This a non-exhaustive list of known issues with MIR patterns at this time.
 * Matching intrinsics is not yet possible.
 * Using ``GICombinePatFrag`` within another ``GICombinePatFrag`` is not
   supported.
+* ``GICombinePatFrag`` can only have a single root.
+* Instructions with multiple defs cannot be the root of a ``GICombinePatFrag``.
 * Using ``GICombinePatFrag`` in the ``apply`` pattern of a ``GICombineRule``
   is not supported.
 * Deleting the matched pattern in a ``GICombineRule`` needs to be done using

llvm/include/llvm/Target/GlobalISel/Combine.td

@@ -160,12 +160,15 @@ def copy_prop : GICombineRule<
 // Fold (freeze (freeze x)) -> (freeze x).
 // Fold (fabs (fabs x)) -> (fabs x).
 // Fold (fcanonicalize (fcanonicalize x)) -> (fcanonicalize x).
+def idempotent_prop_frags : GICombinePatFrag<
+  (outs root:$dst, $src), (ins),
+  !foreach(op, [G_FREEZE, G_FABS, G_FCANONICALIZE],
+           (pattern (op $dst, $src), (op $src, $x)))>;
+
 def idempotent_prop : GICombineRule<
-   (defs root:$mi),
-   (match (wip_match_opcode G_FREEZE, G_FABS, G_FCANONICALIZE):$mi,
-          [{ return MRI.getVRegDef(${mi}->getOperand(1).getReg())->getOpcode() ==
-                    ${mi}->getOpcode(); }]),
-   (apply [{ Helper.replaceSingleDefInstWithOperand(*${mi}, 1); }])>;
+   (defs root:$dst),
+   (match (idempotent_prop_frags $dst, $src)),
+   (apply (COPY $dst, $src))>;
 
 
 def extending_loads : GICombineRule<

llvm/test/TableGen/GlobalISelCombinerMatchTableEmitter/patfrag-errors.td

@@ -257,6 +257,20 @@ def inconsistent_arg_type : GICombineRule<
   (match (TypedParams $dst, i64:$k):$broken),
   (apply (COPY $dst, (i32 0)))>;
 
+// CHECK: :[[@LINE+2]]:{{[0-9]+}}: error: all instructions that define root 'foo' in 'RootDefHasMultiDefs' can only have a single output operand
+// CHECK: :[[@LINE+1]]:{{[0-9]+}}: error: Could not parse GICombinePatFrag 'RootDefHasMultiDefs'
+def RootDefHasMultiDefs: GICombinePatFrag<
+    (outs root:$foo),
+    (ins gi_imm:$cst),
+    [
+      (pattern (G_UNMERGE_VALUES $foo, $z, $y))
+    ]>;
+// CHECK: :[[@LINE+1]]:{{[0-9]+}}: error: Failed to parse pattern: '(RootDefHasMultiDefs ?:$root, (i32 10))'
+def root_def_has_multi_defs : GICombineRule<
+  (defs root:$root),
+  (match (RootDefHasMultiDefs $root, (i32 10))),
+  (apply (COPY $root, (i32 0)))>;
+
 // CHECK: error: Failed to parse one or more rules
 
 def MyCombiner: GICombinerHelper<"GenMyCombiner", [
@@ -278,5 +292,6 @@ def MyCombiner: GICombinerHelper<"GenMyCombiner", [
   expected_mo_namedimm,
   patfrag_in_apply,
   patfrag_cannot_be_root,
-  inconsistent_arg_type
+  inconsistent_arg_type,
+  root_def_has_multi_defs
 ]>;

llvm/test/TableGen/GlobalISelCombinerMatchTableEmitter/pattern-errors.td

@@ -119,14 +119,13 @@ def undef_in_apply : GICombineRule<
   (match (COPY $a, $b):$d),
   (apply (COPY $a, $x))>;
 
-// CHECK: :[[@LINE+2]]:{{[0-9]+}}: error: def of pattern root 'a' is not redefined in the apply pattern!
-// CHECK: :[[@LINE+1]]:{{[0-9]+}}: note: match pattern root is 'foo'
+// CHECK: :[[@LINE+1]]:{{[0-9]+}}: error: 'a' must be redefined in the 'apply' pattern
 def no_redef_in_apply : GICombineRule<
   (defs root:$a),
   (match (COPY $a, $b):$foo),
   (apply (COPY $x, $b))>;
 
-// CHECK: :[[@LINE+1]]:{{[0-9]+}}: error:  def of pattern root 'b' is not redefined in the apply pattern!
+// CHECK: :[[@LINE+1]]:{{[0-9]+}}: error: 'b' must be redefined in the 'apply' pattern
 def no_redef_in_apply_multidefroot : GICombineRule<
   (defs root:$a),
   (match (G_UNMERGE_VALUES $a, $b, $c)),

llvm/utils/TableGen/GlobalISelCombinerMatchTableEmitter.cpp

@@ -581,6 +581,16 @@ class InstructionPattern : public Pattern {
   InstructionOperand &getOperand(unsigned K) { return Operands[K]; }
   const InstructionOperand &getOperand(unsigned K) const { return Operands[K]; }
 
+  /// When this InstructionPattern is used as the match root, returns the
+  /// operands that must be redefined in the 'apply' pattern for the rule to be
+  /// valid.
+  ///
+  /// For CodeGenInstructionPatterns, this just returns the defs of the CGI.
+  /// For PatFrag this only returns the root of the PF.
+  ///
+  /// Returns an empty array on error.
+  virtual ArrayRef<InstructionOperand> getApplyDefsNeeded() const = 0;
+
   auto named_operands() {
     return make_filter_range(Operands,
                              [&](auto &O) { return O.isNamedOperand(); });
@@ -760,6 +770,8 @@ class CodeGenInstructionPattern : public InstructionPattern {
   unsigned getNumInstDefs() const override;
   unsigned getNumInstOperands() const override;
 
+  ArrayRef<InstructionOperand> getApplyDefsNeeded() const override;
+
   const CodeGenInstruction &getInst() const { return I; }
   StringRef getInstName() const override { return I.TheDef->getName(); }
 
@@ -798,6 +810,11 @@ unsigned CodeGenInstructionPattern::getNumInstOperands() const {
                       : NumCGIOps;
 }
 
+ArrayRef<InstructionOperand>
+CodeGenInstructionPattern::getApplyDefsNeeded() const {
+  return {operands().begin(), getNumInstDefs()};
+}
+
 //===- OperandTypeChecker -------------------------------------------------===//
 
 /// This is a trivial type checker for all operands in a set of
@@ -1078,12 +1095,24 @@ bool PatFrag::checkSemantics() {
 
     // Check this operand is defined in all alternative's patterns.
     for (const auto &Alt : Alts) {
-      if (!Alt.OpTable.getDef(Op.Name)) {
+      const auto *OpDef = Alt.OpTable.getDef(Op.Name);
+      if (!OpDef) {
         PrintError("output parameter '" + Op.Name +
                    "' must be defined by all alternative patterns in '" +
                    Def.getName() + "'");
         return false;
       }
+
+      if (Op.Kind == PK_Root && OpDef->getNumInstDefs() != 1) {
+        // The instruction that defines the root must have a single def.
+        // Otherwise we'd need to support multiple roots and it gets messy.
+        //
+        // e.g. this is not supported:
+        //   (pattern (G_UNMERGE_VALUES $x, $root, $vec))
+        PrintError("all instructions that define root '" + Op.Name + "' in '" +
+                   Def.getName() + "' can only have a single output operand");
+        return false;
+      }
     }
 
     SeenOps.insert(Op.Name);
@@ -1232,6 +1261,8 @@ class PatFragPattern : public InstructionPattern {
   unsigned getNumInstDefs() const override { return PF.num_out_params(); }
   unsigned getNumInstOperands() const override { return PF.num_params(); }
 
+  ArrayRef<InstructionOperand> getApplyDefsNeeded() const override;
+
   bool checkSemantics(ArrayRef<SMLoc> DiagLoc) override;
 
   /// Before emitting the patterns inside the PatFrag, add all necessary code
@@ -1258,6 +1289,16 @@ class PatFragPattern : public InstructionPattern {
   const PatFrag &PF;
 };
 
+ArrayRef<InstructionOperand> PatFragPattern::getApplyDefsNeeded() const {
+  assert(PF.num_roots() == 1);
+  // Only roots need to be redef.
+  for (auto [Idx, Param] : enumerate(PF.out_params())) {
+    if (Param.Kind == PatFrag::PK_Root)
+      return getOperand(Idx);
+  }
+  llvm_unreachable("root not found!");
+}
+
 bool PatFragPattern::checkSemantics(ArrayRef<SMLoc> DiagLoc) {
   if (!InstructionPattern::checkSemantics(DiagLoc))
     return false;
@@ -1990,16 +2031,14 @@ bool CombineRuleBuilder::findRoots() {
 
     // Collect all redefinitions of the MatchRoot's defs and put them in
     // ApplyRoots.
-    for (unsigned K = 0; K < IPRoot->getNumInstDefs(); ++K) {
-      auto &O = IPRoot->getOperand(K);
-      assert(O.isDef() && O.isNamedOperand());
-      StringRef Name = O.getOperandName();
+    const auto DefsNeeded = IPRoot->getApplyDefsNeeded();
+    for (auto &Op : DefsNeeded) {
+      assert(Op.isDef() && Op.isNamedOperand());
+      StringRef Name = Op.getOperandName();
 
       auto *ApplyRedef = ApplyOpTable.getDef(Name);
       if (!ApplyRedef) {
-        PrintError("def of pattern root '" + Name +
-                   "' is not redefined in the apply pattern!");
-        PrintNote("match pattern root is '" + MatchRoot->getName() + "'");
+        PrintError("'" + Name + "' must be redefined in the 'apply' pattern");
         return false;
       }