[gicombiner] Add support for arbitrary match data being passed from match to apply

[dsanders]

Summary

This is used by the extending_loads combine to tell the apply step which
use is the preferred one to fold and the other uses should be re-written
to consume.

Depends on D69117

[volkan]

LGTM.

llvm/test/CodeGen/AArch64/GlobalISel/prelegalizercombiner-concat-vectors.mir

22 ↗

(On Diff #225539)

These changes doesn't look relevant, could you double check?

[dsanders]

llvm/test/CodeGen/AArch64/GlobalISel/prelegalizercombiner-concat-vectors.mir

22 ↗

(On Diff #225539)

Looks like it was caused by a bad rebase that lost the G_CONCAT_VECTORS case. I'm testing a fix for that and if all goes well I'll push this tomorrow.

[dsanders]

Update before commit (and because arcanist keeps pulling down the wrong commit message)

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

 class GIDefKindWithArgs;
 /// Declare a root node. There must be at least one of these in every combine
 /// rule.
 /// TODO: The plan is to elide `root` definitions and determine it from the DAG
 ///       itself with an overide for situations where the usual determination
 ///       is incorrect.
 def root : GIDefKind;
 
+/// Declares data that is passed from the match stage to the apply stage.
+class GIDefMatchData<string type> : GIDefKind {
+  /// A C++ type name indicating the storage type.
+  string Type = type;
+}
+
+def extending_load_matchdata : GIDefMatchData<"PreferredTuple">;
+
 /// The operator at the root of a GICombineRule.Match dag.
 def match;
 /// All arguments of the match operator must be either:
 /// * A subclass of GIMatchKind
 /// * A subclass of GIMatchKindWithArgs
+/// * A subclass of Instruction
 /// * A MIR code block (deprecated)
 /// The GIMatchKind and GIMatchKindWithArgs cases are described in more detail
 /// in their definitions below.
 /// For the Instruction case, these are collected into a DAG where operand names
 /// that occur multiple times introduce edges.
 class GIMatchKind;
 class GIMatchKindWithArgs;
 
 /// The operator at the root of a GICombineRule.Apply dag.
 def apply;
 /// All arguments of the apply operator must be subclasses of GIApplyKind, or
 /// sub-dags whose operator is GIApplyKindWithArgs, or an MIR block
 /// (deprecated).
 class GIApplyKind;
 class GIApplyKindWithArgs;
 
 def copy_prop : GICombineRule<
   (defs root:$d),
   (match [{ return Helper.matchCombineCopy(${d}); }]),
   (apply [{ Helper.applyCombineCopy(${d}); }])>;
 def trivial_combines : GICombineGroup<[copy_prop]>;
 
+def extending_loads : GICombineRule<
+  (defs root:$root, extending_load_matchdata:$matchinfo),
+  (match [{ return Helper.matchCombineExtendingLoads(${root}, ${matchinfo}); }]),
+  (apply [{ Helper.applyCombineExtendingLoads(${root}, ${matchinfo}); }])>;
+
 // FIXME: Is there a reason this wasn't in tryCombine? I've left it out of
 //        all_combines because it wasn't there.
 def elide_br_by_inverting_cond : GICombineRule<
   (defs root:$d),
   (match [{ return Helper.matchElideBrByInvertingCond(${d}); }]),
   (apply [{ Helper.applyElideBrByInvertingCond(${d}); }])>;
 
-def all_combines : GICombineGroup<[trivial_combines]>;
+def combines_for_extload: GICombineGroup<[extending_loads]>;
+
+def all_combines : GICombineGroup<[trivial_combines, combines_for_extload]>;

llvm/lib/Target/AArch64/AArch64PreLegalizerCombiner.cpp

 };
 
 bool AArch64PreLegalizerCombinerInfo::combine(GISelChangeObserver &Observer,
                                               MachineInstr &MI,
                                               MachineIRBuilder &B) const {
   CombinerHelper Helper(Observer, B, KB, MDT);
 
   switch (MI.getOpcode()) {
-  case TargetOpcode::G_CONCAT_VECTORS:
-    return Helper.tryCombineConcatVectors(MI);
-  case TargetOpcode::G_SHUFFLE_VECTOR:
-    return Helper.tryCombineShuffleVector(MI);
-  case TargetOpcode::G_LOAD:
-  case TargetOpcode::G_SEXTLOAD:
-  case TargetOpcode::G_ZEXTLOAD: {
-    bool Changed = false;
-    Changed |= Helper.tryCombineExtendingLoads(MI);
-    Changed |= Helper.tryCombineIndexedLoadStore(MI);
-    return Changed;
-  }
-  case TargetOpcode::G_STORE:
-    return Helper.tryCombineIndexedLoadStore(MI);
   case TargetOpcode::G_INTRINSIC_W_SIDE_EFFECTS:
     switch (MI.getIntrinsicID()) {
     case Intrinsic::memcpy:
     case Intrinsic::memmove:
     case Intrinsic::memset: {
       // If we're at -O0 set a maxlen of 32 to inline, otherwise let the other
       // heuristics decide.
       unsigned MaxLen = EnableOpt ? 0 : 32;
       // Try to inline memcpy type calls if optimizations are enabled.
       return (!EnableMinSize) ? Helper.tryCombineMemCpyFamily(MI, MaxLen)
                               : false;
     }
     default:
       break;
     }
   }
 
   if (Generated.tryCombineAll(Observer, MI, B))
     return true;
 
+  switch (MI.getOpcode()) {
+  case TargetOpcode::G_CONCAT_VECTORS:
+    return Helper.tryCombineConcatVectors(MI);
+  case TargetOpcode::G_SHUFFLE_VECTOR:
+    return Helper.tryCombineShuffleVector(MI);
+  case TargetOpcode::G_LOAD:
+  case TargetOpcode::G_SEXTLOAD:
+  case TargetOpcode::G_ZEXTLOAD:
+  case TargetOpcode::G_STORE:
+    return Helper.tryCombineIndexedLoadStore(MI);
+  }
+
   return false;
 }
 
 #define AARCH64PRELEGALIZERCOMBINERHELPER_GENCOMBINERHELPER_CPP
 #include "AArch64GenGICombiner.inc"
 #undef AARCH64PRELEGALIZERCOMBINERHELPER_GENCOMBINERHELPER_CPP
 
 // Pass boilerplate

llvm/utils/TableGen/GICombinerEmitter.cpp

 StringSet<> StrTab;
 
 StringRef insertStrTab(StringRef S) {
   if (S.empty())
     return S;
   return StrTab.insert(S).first->first();
 }
 
+/// Declares data that is passed from the match stage to the apply stage.
+class MatchDataInfo {
+  /// The symbol used in the tablegen patterns
+  StringRef PatternSymbol;
+  /// The data type for the variable
+  StringRef Type;
+  /// The name of the variable as declared in the generated matcher.
+  std::string VariableName;
+
+public:
+  MatchDataInfo(StringRef PatternSymbol, StringRef Type, StringRef VariableName)
+      : PatternSymbol(PatternSymbol), Type(Type), VariableName(VariableName) {}
+
+  StringRef getPatternSymbol() const { return PatternSymbol; };
+  StringRef getType() const { return Type; };
+  StringRef getVariableName() const { return VariableName; };
+};
+
 class RootInfo {
   StringRef PatternSymbol;
 
 public:
   RootInfo(StringRef PatternSymbol) : PatternSymbol(PatternSymbol) {}
 
   StringRef getPatternSymbol() const { return PatternSymbol; }
 };
 
 class CombineRule {
+public:
+
+  using const_matchdata_iterator = std::vector<MatchDataInfo>::const_iterator;
+
   struct VarInfo {
     const GIMatchDagInstr *N;
     const GIMatchDagOperand *Op;
     const DagInit *Matcher;
 
   public:
     VarInfo(const GIMatchDagInstr *N, const GIMatchDagOperand *Op,
             const DagInit *Matcher)
protected:
   std::vector<RootInfo> Roots;
 
   GIMatchDag MatchDag;
 
   /// A block of arbitrary C++ to finish testing the match.
   /// FIXME: This is a temporary measure until we have actual pattern matching
   const CodeInit *MatchingFixupCode = nullptr;
 
+  /// The MatchData defined by the match stage and required by the apply stage.
+  /// This allows the plumbing of arbitrary data from C++ predicates between the
+  /// stages.
+  ///
+  /// For example, suppose you have:
+  ///   %A = <some-constant-expr>
+  ///   %0 = G_ADD %1, %A
+  /// you could define a GIMatchPredicate that walks %A, constant folds as much
+  /// as possible and returns an APInt containing the discovered constant. You
+  /// could then declare:
+  ///   def apint : GIDefMatchData<"APInt">;
+  /// add it to the rule with:
+  ///   (defs root:$root, apint:$constant)
+  /// evaluate it in the pattern with a C++ function that takes a
+  /// MachineOperand& and an APInt& with:
+  ///   (match [{MIR %root = G_ADD %0, %A }],
+  ///             (constantfold operand:$A, apint:$constant))
+  /// and finally use it in the apply stage with:
+  ///   (apply (create_operand
+  ///                [{ MachineOperand::CreateImm(${constant}.getZExtValue());
+  ///                ]}, apint:$constant),
+  ///             [{MIR %root = FOO %0, %constant }])
+  std::vector<MatchDataInfo> MatchDataDecls;
+
+  void declareMatchData(StringRef PatternSymbol, StringRef Type,
+                        StringRef VarName);
+
   bool parseInstructionMatcher(const CodeGenTarget &Target, StringInit *ArgName,
                                const Init &Arg,
                                StringMap<std::vector<VarInfo>> &NamedEdgeDefs,
                                StringMap<std::vector<VarInfo>> &NamedEdgeUses);
 
 public:
   CombineRule(const CodeGenTarget &Target, GIMatchDagContext &Ctx, RuleID ID,
               const Record &R)
public:
 
   using const_root_iterator = std::vector<RootInfo>::const_iterator;
   const_root_iterator roots_begin() const { return Roots.begin(); }
   const_root_iterator roots_end() const { return Roots.end(); }
   iterator_range<const_root_iterator> roots() const {
     return llvm::make_range(Roots.begin(), Roots.end());
   }
 
+  iterator_range<const_matchdata_iterator> matchdata_decls() const {
+    return make_range(MatchDataDecls.begin(), MatchDataDecls.end());
+  }
+
+  /// Export expansions for this rule
+  void declareExpansions(CodeExpansions &Expansions) const {
+    for (const auto &I : matchdata_decls())
+      Expansions.declare(I.getPatternSymbol(), I.getVariableName());
+  }
+
   /// The matcher will begin from the roots and will perform the match by
   /// traversing the edges to cover the whole DAG. This function reverses DAG
   /// edges such that everything is reachable from a root. This is part of the
   /// preparation work for flattening the DAG into a tree.
   void reorientToRoots() {
     SmallSet<const GIMatchDagInstr *, 5> Roots;
     SmallSet<const GIMatchDagInstr *, 5> Visited;
     SmallSet<GIMatchDagEdge *, 20> EdgesRemaining;
StringRef makeDebugName(CombineRule &Rule, StringRef Name) {
   return insertStrTab(Name.empty() ? makeNameForAnonInstr(Rule) : StringRef(Name));
 }
 
 StringRef makeNameForAnonPredicate(CombineRule &Rule) {
   return insertStrTab(
       to_string(format("__anonpred%d_%d", Rule.getID(), Rule.allocUID())));
 }
 
+void CombineRule::declareMatchData(StringRef PatternSymbol, StringRef Type,
+                                   StringRef VarName) {
+  MatchDataDecls.emplace_back(PatternSymbol, Type, VarName);
+}
+
 bool CombineRule::parseDefs() {
   NamedRegionTimer T("parseDefs", "Time spent parsing the defs", "Rule Parsing",
                      "Time spent on rule parsing", TimeRegions);
   DagInit *Defs = TheDef.getValueAsDag("Defs");
 
   if (Defs->getOperatorAsDef(TheDef.getLoc())->getName() != "defs") {
     PrintError(TheDef.getLoc(), "Expected defs operator");
     return false;
   }
 
   for (unsigned I = 0, E = Defs->getNumArgs(); I < E; ++I) {
     // Roots should be collected into Roots
     if (isSpecificDef(*Defs->getArg(I), "root")) {
       Roots.emplace_back(Defs->getArgNameStr(I));
       continue;
     }
 
+    // Subclasses of GIDefMatchData should declare that this rule needs to pass
+    // data from the match stage to the apply stage, and ensure that the
+    // generated matcher has a suitable variable for it to do so.
+    if (Record *MatchDataRec =
+            getDefOfSubClass(*Defs->getArg(I), "GIDefMatchData")) {
+      declareMatchData(Defs->getArgNameStr(I),
+                       MatchDataRec->getValueAsString("Type"),
+                       llvm::to_string(llvm::format("MatchData%d", ID)));
+      continue;
+    }
+
     // Otherwise emit an appropriate error message.
     if (getDefOfSubClass(*Defs->getArg(I), "GIDefKind"))
       PrintError(TheDef.getLoc(),
                  "This GIDefKind not implemented in tablegen");
     else if (getDefOfSubClass(*Defs->getArg(I), "GIDefKindWithArgs"))
       PrintError(TheDef.getLoc(),
                  "This GIDefKindWithArgs not implemented in tablegen");
     else
void GICombinerEmitter::generateCodeForRule(raw_ostream &OS,
 
     OS << Indent << "// Rule: " << RuleDef.getName() << "\n"
        << Indent << "if (!isRuleDisabled(" << Rule->getID() << ")) {\n";
 
     CodeExpansions Expansions;
     for (const RootInfo &Root : Rule->roots()) {
       Expansions.declare(Root.getPatternSymbol(), "MI");
     }
+    Rule->declareExpansions(Expansions);
+
     DagInit *Applyer = RuleDef.getValueAsDag("Apply");
     if (Applyer->getOperatorAsDef(RuleDef.getLoc())->getName() !=
         "apply") {
       PrintError(RuleDef.getLoc(), "Expected apply operator");
       return;
     }
 
     OS << Indent << "  if (1\n";
OS << "bool " << getClassName() << "::tryCombineAll(\n"
      << "    MachineInstr &MI,\n"
      << "    MachineIRBuilder &B) const {\n"
      << "  CombinerHelper Helper(Observer, B);\n"
      << "  MachineBasicBlock *MBB = MI.getParent();\n"
      << "  MachineFunction *MF = MBB->getParent();\n"
      << "  MachineRegisterInfo &MRI = MF->getRegInfo();\n"
      << "  (void)MBB; (void)MF; (void)MRI;\n\n";
 
+  OS << "  // Match data\n";
+  for (const auto &Rule : Rules)
+    for (const auto &I : Rule->matchdata_decls())
+      OS << "  " << I.getType() << " " << I.getVariableName() << ";\n";
+  OS << "\n";
+
   for (const auto &Rule : Rules)
     generateCodeForRule(OS, Rule.get(), "  ");
   OS << "\n  return false;\n"
      << "}\n"
      << "#endif // ifdef " << Name.upper() << "_GENCOMBINERHELPER_CPP\n";
 }
 
 } // end anonymous namespace