[GlobalISel] Combine (logic_op (op x...), (op y...)) -> (op (logic_op x, y))

[paquette]

Summary

This implements

(logic_op (op x...), (op y...)) -> (op (logic_op x, y))

when op is an extend, a shift, or an and.

This is similar to DAGCombiner::hoistLogicOpWithSameOpcodeHands (with a bunch of missing cases, e.g. G_TRUNC, G_BITCAST, etc.)

This is implemented so it works both pre and post-legalization.

I'm not entirely happy with the idea of creating (but not inserting) instructions in the match step, but it seems cleaner than passing around a bunch of registers etc. No strong opinions either way.

[aemerson]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

1834

I wonder if we can do this in a more functional style, and instead of storing the actual MachineInstr in the match info, store a function/lambda that can build and insert the instructions using a given MachineIRBuilder argument. You would have to modify the legality check here to use the other API call instead of passing an MachineInstr.

[paquette]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

1834

Yeah, I was thinking of yanking the code from ComplexRendererFns and doing it that way. Wasn't sure about the best way to do it, but I can give it another shot.

[arsenm]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

1795

Shouldn't construct a single use MachineIRBuilder. There should be one in CombinerHelper already

1824

Why not use the form with all operands in the initial construction?

1846

I don't think vectors are a problem?

[paquette]

• Use a more functional style + generalize a bit more so this can be used for other combines
• Allow vectors + update test
• Add a helper function for legality checks

[aemerson]

Tests missing?

llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h

60

= default?

69

same.

[paquette]

• Add back tests
• Use = default for default constructors
• Add some asserts to applyBuildInstructionSteps

[aemerson]

LGTM, thanks.

llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h

 #include "llvm/CodeGen/LowLevelType.h"
 #include "llvm/CodeGen/Register.h"
 #include "llvm/Support/Alignment.h"
 
 namespace llvm {
 
 class GISelChangeObserver;
 class MachineIRBuilder;
+class MachineInstrBuilder;
 class MachineRegisterInfo;
 class MachineInstr;
 class MachineOperand;
 class GISelKnownBits;
 class MachineDominatorTree;
 class LegalizerInfo;
+struct LegalityQuery;
 
 struct PreferredTuple {
   LLT Ty;                // The result type of the extend.
   unsigned ExtendOpcode; // G_ANYEXT/G_SEXT/G_ZEXT
   MachineInstr *MI;
 };
 
 struct IndexedLoadStoreMatchInfo {
   Register Addr;
   Register Base;
   Register Offset;
   bool IsPre;
 };
 
 struct PtrAddChain {
   int64_t Imm;
   Register Base;
 };
 
+using OperandBuildSteps =
+    SmallVector<std::function<void(MachineInstrBuilder &)>, 4>;
+struct InstructionBuildSteps {
+  unsigned Opcode;           /// The opcode for the produced instruction.
+  OperandBuildSteps StepFns; /// Operands to be added to the instruction.
+  InstructionBuildSteps() {}
+  InstructionBuildSteps(unsigned Opcode, const OperandBuildSteps &StepFns)
+      : Opcode(Opcode), StepFns(StepFns) {}
+};
+
+struct InstructionStepsMatchInfo {
+  /// Describes instructions to be built during a combine.
+  /// These are built in order from first to last.
+  SmallVector<InstructionBuildSteps, 2> InstrsToBuild;
+  InstructionStepsMatchInfo() {}
+  InstructionStepsMatchInfo(
+      std::initializer_list<InstructionBuildSteps> InstrsToBuild)
+      : InstrsToBuild(InstrsToBuild) {}
+};
+
 class CombinerHelper {
 protected:
   MachineIRBuilder &Builder;
   MachineRegisterInfo &MRI;
   GISelChangeObserver &Observer;
   GISelKnownBits *KB;
   MachineDominatorTree *MDT;
   const LegalizerInfo *LI;
 
 public:
   CombinerHelper(GISelChangeObserver &Observer, MachineIRBuilder &B,
                  GISelKnownBits *KB = nullptr,
                  MachineDominatorTree *MDT = nullptr,
                  const LegalizerInfo *LI = nullptr);
 
   GISelKnownBits *getKnownBits() const {
     return KB;
   }
 
+  /// \return true if the combine is running prior to legalization, or if \p
+  /// Query is legal on the target.
+  bool isLegalOrBeforeLegalizer(const LegalityQuery &Query) const;
+
   /// MachineRegisterInfo::replaceRegWith() and inform the observer of the changes
   void replaceRegWith(MachineRegisterInfo &MRI, Register FromReg, Register ToReg) const;
 
   /// Replace a single register operand with a new register and inform the
   /// observer of the changes.
   void replaceRegOpWith(MachineRegisterInfo &MRI, MachineOperand &FromRegOp,
                         Register ToReg) const;
 
public:
   bool eraseInst(MachineInstr &MI);
 
   /// Return true if MI is a G_ADD which can be simplified to a G_SUB.
   bool matchSimplifyAddToSub(MachineInstr &MI,
                              std::tuple<Register, Register> &MatchInfo);
   bool applySimplifyAddToSub(MachineInstr &MI,
                              std::tuple<Register, Register> &MatchInfo);
 
+  /// Match (logic_op (op x...), (op y...)) -> (op (logic_op x, y))
+  bool
+  matchHoistLogicOpWithSameOpcodeHands(MachineInstr &MI,
+                                       InstructionStepsMatchInfo &MatchInfo);
+
+  /// Replace \p MI with a series of instructions described in \p MatchInfo.
+  bool applyBuildInstructionSteps(MachineInstr &MI,
+                                  InstructionStepsMatchInfo &MatchInfo);
+
   /// Try to transform \p MI by using all of the above
   /// combine functions. Returns true if changed.
   bool tryCombine(MachineInstr &MI);
 
 private:
   // Memcpy family optimization helpers.
   bool optimizeMemcpy(MachineInstr &MI, Register Dst, Register Src,
                       unsigned KnownLen, Align DstAlign, Align SrcAlign,

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

 /// 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">;
 def indexed_load_store_matchdata : GIDefMatchData<"IndexedLoadStoreMatchInfo">;
+def instruction_steps_matchdata: GIDefMatchData<"InstructionStepsMatchInfo">;
 
 /// 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)

 def i2p_to_p2i_matchinfo: GIDefMatchData<"Register">;
 def i2p_to_p2i: GICombineRule<
   (defs root:$root, i2p_to_p2i_matchinfo:$info),
   (match (wip_match_opcode G_PTRTOINT):$root,
     [{ return Helper.matchCombineP2IToI2P(*${root}, ${info}); }]),
   (apply [{ return Helper.applyCombineP2IToI2P(*${root}, ${info}); }])
 >;
 
+// Simplify: (logic_op (op x...), (op y...)) -> (op (logic_op x, y))
+def hoist_logic_op_with_same_opcode_hands: GICombineRule <
+  (defs root:$root, instruction_steps_matchdata:$info),
+  (match (wip_match_opcode G_AND, G_OR, G_XOR):$root,
+    [{ return Helper.matchHoistLogicOpWithSameOpcodeHands(*${root}, ${info}); }]),
+  (apply [{ return Helper.applyBuildInstructionSteps(*${root}, ${info});}])
+>;
+
 // FIXME: These should use the custom predicate feature once it lands.
 def undef_combines : GICombineGroup<[undef_to_fp_zero, undef_to_int_zero,
                                      undef_to_negative_one,
                                      propagate_undef_any_op,
                                      propagate_undef_all_ops,
                                      propagate_undef_shuffle_mask,
                                      erase_undef_store]>;
 
 def identity_combines : GICombineGroup<[select_same_val, right_identity_zero,
                                         binop_same_val, binop_left_to_zero,
-                                        binop_right_to_zero, p2i_to_i2p, 
+                                        binop_right_to_zero, p2i_to_i2p,
                                         i2p_to_p2i]>;
 
 def trivial_combines : GICombineGroup<[copy_prop, mul_to_shl]>;
 def all_combines : GICombineGroup<[trivial_combines, ptr_add_immed_chain,
     combines_for_extload, combine_indexed_load_store, undef_combines,
-    identity_combines, simplify_add_to_sub]>;
+    identity_combines, simplify_add_to_sub,
+    hoist_logic_op_with_same_opcode_hands]>;

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

 
 // Option to allow testing of the combiner while no targets know about indexed
 // addressing.
 static cl::opt<bool>
     ForceLegalIndexing("force-legal-indexing", cl::Hidden, cl::init(false),
                        cl::desc("Force all indexed operations to be "
                                 "legal for the GlobalISel combiner"));
 
-
 CombinerHelper::CombinerHelper(GISelChangeObserver &Observer,
                                MachineIRBuilder &B, GISelKnownBits *KB,
                                MachineDominatorTree *MDT,
                                const LegalizerInfo *LI)
     : Builder(B), MRI(Builder.getMF().getRegInfo()), Observer(Observer),
       KB(KB), MDT(MDT), LI(LI) {
   (void)this->KB;
 }
 
+bool CombinerHelper::isLegalOrBeforeLegalizer(
+    const LegalityQuery &Query) const {
+  return !LI || LI->getAction(Query).Action == LegalizeActions::Legal;
+}
+
 void CombinerHelper::replaceRegWith(MachineRegisterInfo &MRI, Register FromReg,
                                     Register ToReg) const {
   Observer.changingAllUsesOfReg(MRI, FromReg);
 
   if (MRI.constrainRegAttrs(ToReg, FromReg))
     MRI.replaceRegWith(FromReg, ToReg);
   else
     Builder.buildCopy(ToReg, FromReg);
bool CombinerHelper::applySimplifyAddToSub(
   Builder.setInstr(MI);
   Register SubLHS, SubRHS;
   std::tie(SubLHS, SubRHS) = MatchInfo;
   Builder.buildSub(MI.getOperand(0).getReg(), SubLHS, SubRHS);
   MI.eraseFromParent();
   return true;
 }
 
+bool CombinerHelper::matchHoistLogicOpWithSameOpcodeHands(
+    MachineInstr &MI, InstructionStepsMatchInfo &MatchInfo) {
+  // Matches: logic (hand x, ...), (hand y, ...) -> hand (logic x, y), ...
+  //
+  // Creates the new hand + logic instruction (but does not insert them.)
+  //
+  // On success, MatchInfo is populated with the new instructions. These are
+  // inserted in applyHoistLogicOpWithSameOpcodeHands.
+  unsigned LogicOpcode = MI.getOpcode();
+  assert(LogicOpcode == TargetOpcode::G_AND ||
+         LogicOpcode == TargetOpcode::G_OR ||
+         LogicOpcode == TargetOpcode::G_XOR);
+  MachineIRBuilder MIB(MI);
+  Register Dst = MI.getOperand(0).getReg();
+  Register LHSReg = MI.getOperand(1).getReg();
+  Register RHSReg = MI.getOperand(2).getReg();
+
+  // Don't recompute anything.
+  if (!MRI.hasOneNonDBGUse(LHSReg) || !MRI.hasOneNonDBGUse(RHSReg))
+    return false;
+
+  // Make sure we have (hand x, ...), (hand y, ...)
+  MachineInstr *LeftHandInst = getDefIgnoringCopies(LHSReg, MRI);
+  MachineInstr *RightHandInst = getDefIgnoringCopies(RHSReg, MRI);
+  if (!LeftHandInst || !RightHandInst)
+    return false;
+  unsigned HandOpcode = LeftHandInst->getOpcode();
+  if (HandOpcode != RightHandInst->getOpcode())
+    return false;
+  if (!LeftHandInst->getOperand(1).isReg() ||
+      !RightHandInst->getOperand(1).isReg())
+    return false;
+
+  // Make sure the types match up, and if we're doing this post-legalization,
+  // we end up with legal types.
+  Register X = LeftHandInst->getOperand(1).getReg();
+  Register Y = RightHandInst->getOperand(1).getReg();
+  LLT XTy = MRI.getType(X);
+  LLT YTy = MRI.getType(Y);
+  if (XTy != YTy)
+    return false;
+  if (!isLegalOrBeforeLegalizer({LogicOpcode, {XTy, YTy}}))
+    return false;
+
+  // Optional extra source register.
+  Register ExtraHandOpSrcReg;
+  switch (HandOpcode) {
+  default:
+    return false;
+  case TargetOpcode::G_ANYEXT:
+  case TargetOpcode::G_SEXT:
+  case TargetOpcode::G_ZEXT: {
+    // Match: logic (ext X), (ext Y) --> ext (logic X, Y)
+    break;
+  }
+  case TargetOpcode::G_AND:
+  case TargetOpcode::G_ASHR:
+  case TargetOpcode::G_LSHR:
+  case TargetOpcode::G_SHL: {
+    // Match: logic (binop x, z), (binop y, z) -> binop (logic x, y), z
+    MachineOperand &ZOp = LeftHandInst->getOperand(2);
+    if (!matchEqualDefs(ZOp, RightHandInst->getOperand(2)))
+      return false;
+    ExtraHandOpSrcReg = ZOp.getReg();
+    break;
+  }
+  }
+
+  // Record the steps to build the new instructions.
+  //
+  // Steps to build (logic x, y)
+  auto NewLogicDst = MRI.createGenericVirtualRegister(XTy);
+  OperandBuildSteps LogicBuildSteps = {
+      [=](MachineInstrBuilder &MIB) { MIB.addDef(NewLogicDst); },
+      [=](MachineInstrBuilder &MIB) { MIB.addReg(X); },
+      [=](MachineInstrBuilder &MIB) { MIB.addReg(Y); }};
+  InstructionBuildSteps LogicSteps(LogicOpcode, LogicBuildSteps);
+
+  // Steps to build hand (logic x, y), ...z
+  OperandBuildSteps HandBuildSteps = {
+      [=](MachineInstrBuilder &MIB) { MIB.addDef(Dst); },
+      [=](MachineInstrBuilder &MIB) { MIB.addReg(NewLogicDst); }};
+  if (ExtraHandOpSrcReg.isValid())
+    HandBuildSteps.push_back(
+        [=](MachineInstrBuilder &MIB) { MIB.addReg(ExtraHandOpSrcReg); });
+  InstructionBuildSteps HandSteps(HandOpcode, HandBuildSteps);
+
+  MatchInfo = InstructionStepsMatchInfo({LogicSteps, HandSteps});
+  return true;
+}
+
+bool CombinerHelper::applyBuildInstructionSteps(
+    MachineInstr &MI, InstructionStepsMatchInfo &MatchInfo) {
+  Builder.setInstr(MI);
+  for (auto &InstrToBuild : MatchInfo.InstrsToBuild) {
+    MachineInstrBuilder Instr = Builder.buildInstr(InstrToBuild.Opcode);
+    for (auto &StepFn : InstrToBuild.StepFns)
+      StepFn(Instr);
+  }
+  MI.eraseFromParent();
+  return true;
+}
+
 bool CombinerHelper::tryCombine(MachineInstr &MI) {
   if (tryCombineCopy(MI))
     return true;
   if (tryCombineExtendingLoads(MI))
     return true;
   if (tryCombineIndexedLoadStore(MI))
     return true;
   return false;
 }

llvm/lib/Target/AArch64/AArch64Combine.td

 
 // Combines which replace a G_SHUFFLE_VECTOR with a target-specific pseudo
 // instruction.
 def shuffle_vector_pseudos : GICombineGroup<[dup, rev, ext, zip, uzp, trn]>;
 
 def AArch64PostLegalizerCombinerHelper
     : GICombinerHelper<"AArch64GenPostLegalizerCombinerHelper",
                        [erase_undef_store, combines_for_extload,
-                        sext_trunc_sextload, shuffle_vector_pseudos]> {
+                        sext_trunc_sextload, shuffle_vector_pseudos,
+                        hoist_logic_op_with_same_opcode_hands]> {
   let DisableRuleOption = "aarch64postlegalizercombiner-disable-rule";
 }