[CGP] Make optimizeMemoryInst capable of handling multiple AddrModes

[john.brawn]

Summary

Currently optimizeMemoryInst requires that all of the AddrModes it sees are identical. This patch makes it capable of tracking multiple AddrModes, so long as they differ in at most one field.

This patch does nothing by itself, but later patches will make use of it to insert or reuse phi or select instructions for the differing fields.

[skatkov]

lib/CodeGen/CodeGenPrepare.cpp

4522

Can we move this to utility class with operation like diff? It will reduce the size of this method which is already big one.

4570

In reality you do not need iteration over found AddrModes. What you need here is just whether at least two AddrModes has different element.
So if Different* is true you do not need to compare this element more.

It actually means that you can always compare NewAddrMode against the first element of an array AddrMode.

Also I would put this functionality into utility class.

4619

I think it is not good that this bailout goes after debug message (three lines above) is printed.

[reames]

lib/CodeGen/CodeGenPrepare.cpp

4568–4569

If I understand your original patch correctly, you only support difference in one input across all AddrModes right?

If so, it might be better to track which fields is known to differ and reject eagerly any that differs in a different field. Using something like an Optional<FieldName> where FieldName is a enum might be a good representation.

Alternatively, this might make more sense as a post processing loop. (i.e. keep track of non-equal address modes in a set, then determine if all elements of the set are compatible.)

4569

From the comment, this really sounds like it's covering up a deficiency in the phi/select insertion logic. Would you mind leaving this case out for the moment or showing a test case where this is essential? In particular, Serguei's patch goes to lengths to try to reuse existing phi nodes.

When this is needed, this sounds like a function on AddressMode called isTrivial()

4570

We don't have the transform code yet to support this do we? If not, don't we need to remove this?

4573

For the moment, you only handle AddrMode.size() == 1. Just make that your exit condition here and remove the bailout below.

From this line, it looks like the trivial case might be a change in behaviour for the existing code. Is that so? If so, where's the test case for that?

4578

use empty()

[reames]

p.s. Thanks for doing the work to split this out and being so responsive to review comments. Much appreciated!

[skatkov]

lib/CodeGen/CodeGenPrepare.cpp

4569

I think it is about earlier bail out and compile time saving.

TrivialAddrMode actually means that if we found that all AddrModes are of the form only base reg or gv (no scale or offset) then later trying to find a common Phi node for, for example, base reg will end up with exactly the same Phi node we started from (addr). Yes my algorithm will detect it and will not create new Phi nodes but it will return the same Phi node addr. So all complex work in my patch is redundant.

This check just detect it in advance and bail out earlier.

I agree it should be moved to AddrMode class.

4570

I consider this is as a preparation to implementation. For now we allow processing later but will bail out after that.

Not sure that Addr->hasOneUse is a correct check here...

What if we have two instruction which can fold Addr after our modification?

[john.brawn]

Updated based on review comments:

• Move the field comparison into ExtAddrMode and use an enum to represent which field differs
• Move trivial AddrMode detection into ExtAddrMode
• Move everything else from optimizeMemoryInst into an AddressingModeCombiner class

[skatkov]

lib/CodeGen/CodeGenPrepare.cpp

2680

I cannot understand why you need a AddrModeUser. The only one use of this field I see is for checking that user of AddrMode is Phi or Select.

However according to usage of AddressingModeCombiner AddrModeUser cannot be anything else.

Do I miss anything?

3364

if AllAddrModesTrivial is false, you do not need to compute isTrivial. So you can re-write it as

AllAddrModesTrivial = AllAddrModesTrivial && NewAddrMode.isTrivial();

Or alternatively:

if (AllAddrModesTrivial)
  AllAddrModesTrivial = NewAddrMode.isTrivial();

4566

Redundant {}

[john.brawn]

lib/CodeGen/CodeGenPrepare.cpp

2680

It's there mostly for when we do actually want to do the transform to adjust the select / add a new one - I have an updated version of D38133 that builds on top of this (that I'm not intending to commit), to make sure that this patch actually does something useful, and it's used there. Depending on how D36073 goes it may or may not end up being useful though, so I'll remove it.

[john.brawn]

Update based on review comments.

[reames]

LGTM subject to the comment below being a clear bugfix addressed before commit.

lib/CodeGen/CodeGenPrepare.cpp

2721

Don't you also need to check that the offset and scale are zero? If not, then I'm still confused by what you mean by "trivial".

[skatkov]

lib/CodeGen/CodeGenPrepare.cpp

2721

OriginalValue = gv + base + scale * Index + offset.

If OriginalValue == base => scale == 0, offset == 0, gv == 0

the same if OriginalValue == gv

[skatkov]

lib/CodeGen/CodeGenPrepare.cpp

2697–2722

Don't you want to check type compatibility also here?
Say if type of base regs are different we can return Incompatible types here.

The similar for others...

Or we can do it as a separate patch...

[john.brawn]

lib/CodeGen/CodeGenPrepare.cpp

2697–2722

Will do.

2721

Yes, scale and offset are expected to be zero in a trivial addrmode. I'll add a couple of asserts to make sure that's the case.

lib/CodeGen/CodeGenPrepare.cpp

 
 namespace {
 
 /// This is an extended version of TargetLowering::AddrMode
 /// which holds actual Value*'s for register values.
 struct ExtAddrMode : public TargetLowering::AddrMode {
   Value *BaseReg = nullptr;
   Value *ScaledReg = nullptr;
+  Value *OriginalValue = nullptr;
+  User *AddrModeUser = nullptr;
+
+  enum FieldName {
+    NoField        = 0x00,
+    BaseRegField   = 0x01,
+    BaseGVField    = 0x02,
+    BaseOffsField  = 0x04,
+    ScaledRegField = 0x08,
+    ScaleField     = 0x10,
+    MultipleFields = 0xff
+  };
 
   ExtAddrMode() = default;
 
   void print(raw_ostream &OS) const;
   void dump() const;
 
-  bool operator==(const ExtAddrMode& O) const {
-    return (BaseReg == O.BaseReg) && (ScaledReg == O.ScaledReg) &&
-           (BaseGV == O.BaseGV) && (BaseOffs == O.BaseOffs) &&
-           (HasBaseReg == O.HasBaseReg) && (Scale == O.Scale);
+  FieldName compare(const ExtAddrMode &other) {
+    unsigned Result = NoField;
+    if (BaseReg != other.BaseReg)
+      Result |= BaseRegField;
+    if (BaseGV != other.BaseGV)
+      Result |= BaseGVField;
+    if (BaseOffs != other.BaseOffs)
+      Result |= BaseOffsField;
+    if (ScaledReg != other.ScaledReg)
+      Result |= ScaledRegField;
+    // Don't count 0 as being a different scale, because that actually means
+    // unscaled (which will already be counted by having no ScaledReg).
+    if (Scale && other.Scale && Scale != other.Scale)
+      Result |= ScaleField;
+
+    if (countPopulation(Result) > 1)
+      return MultipleFields;
+    else
+      return static_cast<FieldName>(Result);
+  }
+
+  // AddrModes with a base reg or gv where the reg/gv is just the original
+  // value are trivial.
+  bool isTrivial() {
+    return (BaseGV && BaseGV == OriginalValue) ||
+      (BaseReg && BaseReg == OriginalValue);
   }
 };
 
 } // end anonymous namespace
 
 #ifndef NDEBUG
 static inline raw_ostream &operator<<(raw_ostream &OS, const ExtAddrMode &AM) {
   AM.print(OS);
private:
   bool isProfitableToFoldIntoAddressingMode(Instruction *I,
                                             ExtAddrMode &AMBefore,
                                             ExtAddrMode &AMAfter);
   bool valueAlreadyLiveAtInst(Value *Val, Value *KnownLive1, Value *KnownLive2);
   bool isPromotionProfitable(unsigned NewCost, unsigned OldCost,
                              Value *PromotedOperand) const;
 };
 
+/// \brief A helper class for combining addressing modes.
+class AddressingModeCombiner {
+private:
+  /// The addressing modes we've collected.
+  SmallVector<ExtAddrMode, 16> AddrModes;
+
+  /// The field in which the AddrModes differ, when we have more than one.
+  ExtAddrMode::FieldName DifferentField = ExtAddrMode::NoField;
+
+  /// Are the AddrModes that we have all just equal to their original values?
+  bool AllAddrModesTrivial = true;
+
+public:
+  /// \brief Get the combined AddrMode
+  const ExtAddrMode &getAddrMode() const {
+    return AddrModes[0];
+  }
+
+  /// \brief Add a new AddrMode if it's compatible with the AddrModes we already
+  /// have.
+  /// \return True iff we succeeded in doing so.
+  bool addNewAddrMode(ExtAddrMode &NewAddrMode) {
+    // Take note of if we have any non-trivial AddrModes, as we need to detect
+    // when all AddrModes are trivial as then we would introduce a phi or select
+    // which just duplicates what's already there.
+    if (!NewAddrMode.isTrivial())
+      AllAddrModesTrivial = false;
+
+    // If this is the first addrmode then everything is fine.
+    if (AddrModes.empty()) {
+      AddrModes.emplace_back(NewAddrMode);
+      return true;
+    }
+
+    // Figure out how different this is from the other address modes, which we
+    // can do just by comparing against the first one given that we only care
+    // about the cumulative difference.
+    ExtAddrMode::FieldName ThisDifferentField =
+      AddrModes[0].compare(NewAddrMode);
+    if (DifferentField == ExtAddrMode::NoField)
+      DifferentField = ThisDifferentField;
+    else if (DifferentField != ThisDifferentField)
+      DifferentField = ExtAddrMode::MultipleFields;
+
+    // If this AddrMode is the same as all the others then everything is fine
+    // (which should only happen when there is actually only one AddrMode).
+    if (DifferentField == ExtAddrMode::NoField) {
+      assert(AddrModes.size() == 1);
+      return true;
+    }
+
+    // If NewAddrMode differs in only one dimension and comes from a phi or
+    // select then we can handle it by inserting a phi/select later on.
+    if ((isa<PHINode>(NewAddrMode.AddrModeUser) ||
+         isa<SelectInst>(NewAddrMode.AddrModeUser)) &&
+        DifferentField != ExtAddrMode::MultipleFields) {
+      AddrModes.emplace_back(NewAddrMode);
+      return true;
+    }
+
+    // We couldn't combine NewAddrMode with the rest, so return failure.
+    AddrModes.clear();
+    return false;
+  }
+
+  /// \brief Combine the addressing modes we've collected into a single
+  /// addressing mode.
+  /// \return True iff we successfully combined them or we only had one so
+  /// didn't need to combine them anyway.
+  bool combineAddrModes() {
+    // If we have no AddrModes then they can't be combined.
+    if (AddrModes.size() == 0)
+      return false;
+
+    // A single AddrMode can trivially be combined.
+    if (AddrModes.size() == 1)
+      return true;
+
+    // If the AddrModes we collected are all just equal to the value they are
+    // derived from then combining them wouldn't do anything useful.
+    if (AllAddrModesTrivial)
+      return false;
+
+    // TODO: Combine multiple AddrModes by inserting a select or phi for the
+    // field in which the AddrModes differ.
+    return false;
+  }
+};
+
 } // end anonymous namespace
 
 /// Try adding ScaleReg*Scale to the current addressing mode.
 /// Return true and update AddrMode if this addr mode is legal for the target,
 /// false if not.
 bool AddressingModeMatcher::matchScaledValue(Value *ScaleReg, int64_t Scale,
                                              unsigned Depth) {
   // If Scale is 1, then this is the same as adding ScaleReg to the addressing

 /// constraint above), and reduce register pressure globally (by removing the
 /// addressing mode computation from the fast path entirely.).
 bool CodeGenPrepare::optimizeMemoryInst(Instruction *MemoryInst, Value *Addr,
                                         Type *AccessTy, unsigned AddrSpace) {
   Value *Repl = Addr;
 
   // Try to collapse single-value PHI nodes.  This is necessary to undo
   // unprofitable PRE transformations.
-  SmallVector<Value*, 8> worklist;
+  SmallVector<std::pair<Value*, User*>, 8> worklist;
   SmallPtrSet<Value*, 16> Visited;
-  worklist.push_back(Addr);
+  worklist.push_back({Addr,MemoryInst});
 
   // Use a worklist to iteratively look through PHI and select nodes, and
   // ensure that the addressing mode obtained from the non-PHI/select roots of
-  // the graph are equivalent.
-  bool AddrModeFound = false;
+  // the graph are compatible.
   bool PhiOrSelectSeen = false;
   SmallVector<Instruction*, 16> AddrModeInsts;
-  ExtAddrMode AddrMode;
+  AddressingModeCombiner AddrModes;
   TypePromotionTransaction TPT(RemovedInsts);
   TypePromotionTransaction::ConstRestorationPt LastKnownGood =
       TPT.getRestorationPoint();
   while (!worklist.empty()) {
-    Value *V = worklist.back();
+    Value *V = worklist.back().first;
+    User *U = worklist.back().second;
     worklist.pop_back();
 
     // We allow traversing cyclic Phi nodes.
     // In case of success after this loop we ensure that traversing through
     // Phi nodes ends up with all cases to compute address of the form
     //    BaseGV + Base + Scale * Index + Offset
     // where Scale and Offset are constans and BaseGV, Base and Index
     // are exactly the same Values in all cases.
     // It means that BaseGV, Scale and Offset dominate our memory instruction
     // and have the same value as they had in address computation represented
     // as Phi. So we can safely sink address computation to memory instruction.
     if (!Visited.insert(V).second)
       continue;
 
     // For a PHI node, push all of its incoming values.
     if (PHINode *P = dyn_cast<PHINode>(V)) {
       for (Value *IncValue : P->incoming_values())
-        worklist.push_back(IncValue);
+        worklist.push_back({IncValue,P});
       PhiOrSelectSeen = true;
       continue;
     }
     // Similar for select.
     if (SelectInst *SI = dyn_cast<SelectInst>(V)) {
-      worklist.push_back(SI->getFalseValue());
-      worklist.push_back(SI->getTrueValue());
+      worklist.push_back({SI->getFalseValue(),SI});
+      worklist.push_back({SI->getTrueValue(),SI});
       PhiOrSelectSeen = true;
       continue;
     }
 
     // For non-PHIs, determine the addressing mode being computed.  Note that
     // the result may differ depending on what other uses our candidate
     // addressing instructions might have.
     AddrModeInsts.clear();
     ExtAddrMode NewAddrMode = AddressingModeMatcher::Match(
         V, AccessTy, AddrSpace, MemoryInst, AddrModeInsts, *TLI, *TRI,
         InsertedInsts, PromotedInsts, TPT);
+    NewAddrMode.OriginalValue = V;
+    NewAddrMode.AddrModeUser = U;
 
-    if (!AddrModeFound) {
-      AddrModeFound = true;
-      AddrMode = NewAddrMode;
-      continue;
-    }
-    if (NewAddrMode == AddrMode)
-      continue;
-
-    AddrModeFound = false;
+    if (!AddrModes.addNewAddrMode(NewAddrMode)) {
       break;
     }
+  }
 
-  // If the addressing mode couldn't be determined, or if multiple different
-  // ones were determined, bail out now.
-  if (!AddrModeFound) {
+  // Try to combine the AddrModes we've collected. If we couldn't collect any,
+  // or we have multiple but either couldn't combine them or combining them
+  // wouldn't do anything useful, bail out now.
+  if (!AddrModes.combineAddrModes()) {
     TPT.rollback(LastKnownGood);
     return false;
   }
   TPT.commit();
 
+  // Get the combined AddrMode (or the only AddrMode, if we only had one).
+  ExtAddrMode AddrMode = AddrModes.getAddrMode();
+
   // If all the instructions matched are already in this BB, don't do anything.
   // If we saw a Phi node then it is not local definitely, and if we saw a select
   // then we want to push the address calculation past it even if it's already
   // in this BB.
   if (!PhiOrSelectSeen && none_of(AddrModeInsts, [&](Value *V) {
         return IsNonLocalValue(V, MemoryInst->getParent());
                   })) {
     DEBUG(dbgs() << "CGP: Found      local addrmode: " << AddrMode << "\n");
if (SunkAddr->getType() != Addr->getType())
               SubtargetInfo->useAA())) {
     // By default, we use the GEP-based method when AA is used later. This
     // prevents new inttoptr/ptrtoint pairs from degrading AA capabilities.
     DEBUG(dbgs() << "CGP: SINKING nonlocal addrmode: " << AddrMode << " for "
                  << *MemoryInst << "\n");
     Type *IntPtrTy = DL->getIntPtrType(Addr->getType());
     Value *ResultPtr = nullptr, *ResultIndex = nullptr;
 
     // First, find the pointer.
     if (AddrMode.BaseReg && AddrMode.BaseReg->getType()->isPointerTy()) {
       ResultPtr = AddrMode.BaseReg;
       AddrMode.BaseReg = nullptr;
     }
 
     if (AddrMode.Scale && AddrMode.ScaledReg->getType()->isPointerTy()) {
       // We can't add more than one pointer together, nor can we scale a
       // pointer (both of which seem meaningless).