[NaryReassociate] Use new access type aware getGEPCost

[luke]

Summary

This was originally split out from D149889: Now that getGEPCost can more
accurately determine when an GEP will be folded, this patch updates
NaryReassociate to take advantage of it.
Unfortunately I wasn't able to create a test case for this where the GEP was
both foldable and reassociatable. Actually I couldn't find any instances of
this branch being taken since most targets can't fold more than a reg+reg add.
But posting this patch anyway for completeness sake.

[krzysz00]

This seems fine, but I don't know the code well enough to approve it

[ebrevnov]

llvm/include/llvm/Analysis/TargetTransformInfo.h

302

"for using getGEPCost" sounds a bit odd... maybe better say "// Helper function to get cost of already materialized GEP instruction"?

llvm/lib/Analysis/TargetTransformInfo.cpp

243

I don't think it's a good idea to have one more place where we define default value for AccessType. Especially taking into account it's incompatible with existing ones. The problem I see is that AccessType will be defaulted to different values depending on particular API used. It's unexpected to get different results depending if getInstructionCost or getGEPCost is used.

I believe we should leave a single place (inside TargetTransformInfoImplCRTPBase::getGEPCost) where default is selected.

llvm/lib/Transforms/Scalar/NaryReassociate.cpp

331

It seems to be possible to simplify all other existing calls to getGEPCost the same way. Are you going to fix them as well?

[luke]

llvm/lib/Analysis/TargetTransformInfo.cpp

243

Unfortunately we lose access to the GEP's users in TargetTransformInfoImplCRTPBase::getGEPCost, since it works on operands and not the GEP itself, so I'm not sure if we can calculate the default there.

It's unexpected to get different results depending if getInstructionCost or getGEPCost is used.

getInstructionCost actually duplicates this "check first user" logic, but I agree, it's annoying to have this code duplication. IIRC I ended up duplicating it because there was no easy way to call TargetTransformInfo::getGEPCost from inside TargetTransformInfoImplCRTPBase.

Maybe we could rejig this helper function to something like getGEPUserAccessType, which would then be used like getGEPCost(GEP->getSourceElementType(), GEP->getPointerOperand(), Ops, getGEPUserAccessType(GEP), CostKind);

llvm/lib/Transforms/Scalar/NaryReassociate.cpp

331

Ideally yes, there's at least one other place in NaryReassociate that I'm aware of. I was planning on leaving that to a separate patch

[ebrevnov]

llvm/lib/Analysis/TargetTransformInfo.cpp

243

Unfortunately we lose access to the GEP's users in TargetTransformInfoImplCRTPBase::getGEPCost, since it works on operands and not the GEP itself, so I'm not sure if we can calculate the default there.

Ok, I see.

getInstructionCost actually duplicates this "check first user" logic, but I agree, it's annoying to have this code duplication.

It's not quite true. getInstructionCost will use default if there is a single user while getGEPCost if there is at least one user.

Maybe we could rejig this helper function to something like getGEPUserAccessType, which would then be used like getGEPCost(GEP->getSourceElementType(), GEP->getPointerOperand(), Ops, getGEPUserAccessType(GEP), CostKind);

Thinking out it a bit more I find the idea of choosing AccessType implicitly inside TTI problematic. Even if the GEP itself is materialized it's really implementation specific if its users have been already updated or not. For that reason I think the best would be to take AccessType into account only if it is given and don't guess.

I understand having getGEPCost which takes GEP is convenient but extending TTI's API with new interfaces with complex interdependencies is error prune (due to type erased nature of TTI) . Probably the easiest solution would be to simply pass all args each time (as it's done now) or introduce say getFullyMaterializedGEPCost utility somewhere else (not inside TTI).

llvm/include/llvm/Analysis/TargetTransformInfo.h

 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_TARGETTRANSFORMINFO_H
 #define LLVM_ANALYSIS_TARGETTRANSFORMINFO_H
 
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/IR/FMF.h"
-#include "llvm/IR/InstrTypes.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/AtomicOrdering.h"
 #include "llvm/Support/BranchProbability.h"
 #include "llvm/Support/InstructionCost.h"
 #include <functional>
 #include <optional>
 #include <utility>
public:
   /// folded into the addressing mode of a load/store. If AccessType is null,
   /// then the resulting target type based off of PointeeType will be used as an
   /// approximation.
   InstructionCost
   getGEPCost(Type *PointeeType, const Value *Ptr,
              ArrayRef<const Value *> Operands, Type *AccessType = nullptr,
              TargetCostKind CostKind = TCK_SizeAndLatency) const;
 
+  // Helper function for using getGEPCost on an already materialized GEP
+  // instruction. Uses the first user as an approximation for AccessType.
+  InstructionCost
+  getGEPCost(const GetElementPtrInst *GEP,
+             TargetCostKind CostKind = TCK_SizeAndLatency) const;
+
   /// Describe known properties for a set of pointers.
   struct PointersChainInfo {
     /// All the GEPs in a set have same base address.
     unsigned IsSameBaseAddress : 1;
     /// These properties only valid if SameBaseAddress is set.
     /// True if all pointers are separated by a unit stride.
     unsigned IsUnitStride : 1;
     /// True if distance between any two neigbouring pointers is a known value.

llvm/lib/Analysis/TargetTransformInfo.cpp

 }
 
 InstructionCost TargetTransformInfo::getGEPCost(
     Type *PointeeType, const Value *Ptr, ArrayRef<const Value *> Operands,
     Type *AccessType, TTI::TargetCostKind CostKind) const {
   return TTIImpl->getGEPCost(PointeeType, Ptr, Operands, AccessType, CostKind);
 }
 
+InstructionCost
+TargetTransformInfo::getGEPCost(const GetElementPtrInst *GEP,
+                                TTI::TargetCostKind CostKind) const {
+  Type *AccessType = nullptr;
+  if (!GEP->user_empty()) {
+    // Only take into account the first user as a rough approximation to avoid
+    // O(N) complexity.
+    AccessType = GEP->user_back()->getAccessType();
+  }
+  SmallVector<const Value *> Ops(GEP->indices());
+  return getGEPCost(GEP->getSourceElementType(), GEP->getPointerOperand(), Ops,
+                    AccessType, CostKind);
+}
+
 InstructionCost TargetTransformInfo::getPointersChainCost(
     ArrayRef<const Value *> Ptrs, const Value *Base,
     const TTI::PointersChainInfo &Info, Type *AccessTy,
     TTI::TargetCostKind CostKind) const {
   assert((Base || !Info.isSameBase()) &&
          "If pointers have same base address it has to be provided.");
   return TTIImpl->getPointersChainCost(Ptrs, Base, Info, AccessTy, CostKind);
 }

llvm/lib/Transforms/Scalar/NaryReassociate.cpp

if ((ResI = matchAndReassociateMinOrMax<umin_pred_ty>(I, OrigSCEV)) ||
         (ResI = matchAndReassociateMinOrMax<smin_pred_ty>(I, OrigSCEV)) ||
         (ResI = matchAndReassociateMinOrMax<umax_pred_ty>(I, OrigSCEV)) ||
         (ResI = matchAndReassociateMinOrMax<smax_pred_ty>(I, OrigSCEV)))
       return ResI;
 
   return nullptr;
 }
 
-static bool isGEPFoldable(GetElementPtrInst *GEP,
-                          const TargetTransformInfo *TTI) {
-  SmallVector<const Value *, 4> Indices(GEP->indices());
-  return TTI->getGEPCost(GEP->getSourceElementType(), GEP->getPointerOperand(),
-                         Indices) == TargetTransformInfo::TCC_Free;
-}
-
 Instruction *NaryReassociatePass::tryReassociateGEP(GetElementPtrInst *GEP) {
   // Not worth reassociating GEP if it is foldable.
-  if (isGEPFoldable(GEP, TTI))
+  if (TTI->getGEPCost(GEP) == TTI::TCC_Free)
     return nullptr;
 
   gep_type_iterator GTI = gep_type_begin(*GEP);
   for (unsigned I = 1, E = GEP->getNumOperands(); I != E; ++I, ++GTI) {
     if (GTI.isSequential()) {
       if (auto *NewGEP = tryReassociateGEPAtIndex(GEP, I - 1,
                                                   GTI.getIndexedType())) {
         return NewGEP;