[CodeGenPrep]Sink sext when foldable in user GEPs

[junbuml]

Summary

Sink sign extensions to user blocks if foldable in all users when placed in the same block. This change was initially proposed in D26524 and now I do this in CGP based on D27853 and D28680.

include/llvm/Analysis/TargetTransformInfo.h

public:
   unsigned getRegisterBitWidth(bool Vector) const;
 
   /// \return True if it should be considered for address type promotion.
   /// \p AllowPromotionWithoutCommonHeader Set true if promoting \p I is
   /// profitable without finding other extensions fed by the same input.
   bool shouldConsiderAddressTypePromotion(
       const Instruction &I, bool &AllowPromotionWithoutCommonHeader) const;
 
+  /// \return True if the extension instrution \p I is foldable in all users.
+  bool isExtFoldableInAllUsers(const Instruction &I) const;
+
   /// \return The size of a cache line in bytes.
   unsigned getCacheLineSize() const;
 
   /// \return How much before a load we should place the prefetch instruction.
   /// This is currently measured in number of instructions.
   unsigned getPrefetchDistance() const;
 
   /// \return Some HW prefetchers can handle accesses up to a certain constant
public:
   virtual int getIntImmCost(unsigned Opc, unsigned Idx, const APInt &Imm,
                             Type *Ty) = 0;
   virtual int getIntImmCost(Intrinsic::ID IID, unsigned Idx, const APInt &Imm,
                             Type *Ty) = 0;
   virtual unsigned getNumberOfRegisters(bool Vector) = 0;
   virtual unsigned getRegisterBitWidth(bool Vector) = 0;
   virtual bool shouldConsiderAddressTypePromotion(
       const Instruction &I, bool &AllowPromotionWithoutCommonHeader) = 0;
+  virtual bool isExtFoldableInAllUsers(const Instruction &I) = 0;
   virtual unsigned getCacheLineSize() = 0;
   virtual unsigned getPrefetchDistance() = 0;
   virtual unsigned getMinPrefetchStride() = 0;
   virtual unsigned getMaxPrefetchIterationsAhead() = 0;
   virtual unsigned getMaxInterleaveFactor(unsigned VF) = 0;
   virtual unsigned
   getArithmeticInstrCost(unsigned Opcode, Type *Ty, OperandValueKind Opd1Info,
                          OperandValueKind Opd2Info,
public:
   unsigned getRegisterBitWidth(bool Vector) override {
     return Impl.getRegisterBitWidth(Vector);
   }
   bool shouldConsiderAddressTypePromotion(
       const Instruction &I, bool &AllowPromotionWithoutCommonHeader) override {
     return Impl.shouldConsiderAddressTypePromotion(
         I, AllowPromotionWithoutCommonHeader);
   }
+  bool isExtFoldableInAllUsers(const Instruction &I) override {
+    return Impl.isExtFoldableInAllUsers(I);
+  }
   unsigned getCacheLineSize() override {
     return Impl.getCacheLineSize();
   }
   unsigned getPrefetchDistance() override { return Impl.getPrefetchDistance(); }
   unsigned getMinPrefetchStride() override {
     return Impl.getMinPrefetchStride();
   }
   unsigned getMaxPrefetchIterationsAhead() override {

include/llvm/Analysis/TargetTransformInfoImpl.h

public:
 
   bool
   shouldConsiderAddressTypePromotion(const Instruction &I,
                                      bool &AllowPromotionWithoutCommonHeader) {
     AllowPromotionWithoutCommonHeader = false;
     return false;
   }
 
+  bool isExtFoldableInAllUsers(const Instruction &I) {
+    return false;
+  }
+
   unsigned getCacheLineSize() { return 0; }
 
   unsigned getPrefetchDistance() { return 0; }
 
   unsigned getMinPrefetchStride() { return 1; }
 
   unsigned getMaxPrefetchIterationsAhead() { return UINT_MAX; }
 

lib/Analysis/TargetTransformInfo.cpp

 }
 
 bool TargetTransformInfo::shouldConsiderAddressTypePromotion(
     const Instruction &I, bool &AllowPromotionWithoutCommonHeader) const {
   return TTIImpl->shouldConsiderAddressTypePromotion(
       I, AllowPromotionWithoutCommonHeader);
 }
 
+bool TargetTransformInfo::isExtFoldableInAllUsers(const Instruction &I) const {
+  return TTIImpl->isExtFoldableInAllUsers(I);
+}
+
 unsigned TargetTransformInfo::getCacheLineSize() const {
   return TTIImpl->getCacheLineSize();
 }
 
 unsigned TargetTransformInfo::getPrefetchDistance() const {
   return TTIImpl->getPrefetchDistance();
 }
 

lib/CodeGen/CodeGenPrepare.cpp

bool CodeGenPrepare::simplifyOffsetableRelocate(Instruction &I) {
   for (auto &Item : RelocateInstMap)
     // Item.first is the RelocatedBase to offset against
     // Item.second is the vector of Targets to replace
     MadeChange = simplifyRelocatesOffABase(Item.first, Item.second);
   return MadeChange;
 }
 
 /// SinkCast - Sink the specified cast instruction into its user blocks
-static bool SinkCast(CastInst *CI) {
+static bool SinkCast(CastInst *&CI, SetOfInstrs &InsertedInsts) {
   BasicBlock *DefBB = CI->getParent();
 
   /// InsertedCasts - Only insert a cast in each block once.
   DenseMap<BasicBlock*, CastInst*> InsertedCasts;
 
   bool MadeChange = false;
   for (Value::user_iterator UI = CI->user_begin(), E = CI->user_end();
        UI != E; ) {
for (Value::user_iterator UI = CI->user_begin(), E = CI->user_end();
     // If we have already inserted a cast into this block, use it.
     CastInst *&InsertedCast = InsertedCasts[UserBB];
 
     if (!InsertedCast) {
       BasicBlock::iterator InsertPt = UserBB->getFirstInsertionPt();
       assert(InsertPt != UserBB->end());
       InsertedCast = CastInst::Create(CI->getOpcode(), CI->getOperand(0),
                                       CI->getType(), "", &*InsertPt);
+      InsertedInsts.insert(InsertedCast);
     }
 
     // Replace a use of the cast with a use of the new cast.
     TheUse = InsertedCast;
     MadeChange = true;
     ++NumCastUses;
   }
 
   // If we removed all uses, nuke the cast.
   if (CI->use_empty()) {
     CI->eraseFromParent();
+    CI = nullptr;
     MadeChange = true;
   }
 
   return MadeChange;
 }
 
 /// If the specified cast instruction is a noop copy (e.g. it's casting from
 /// one pointer type to another, i32->i8 on PPC), sink it into user blocks to
 /// reduce the number of virtual registers that must be created and coalesced.
 ///
 /// Return true if any changes are made.
 ///
 static bool OptimizeNoopCopyExpression(CastInst *CI, const TargetLowering &TLI,
-                                       const DataLayout &DL) {
+                                       const DataLayout &DL, SetOfInstrs &InsertedInsts) {
   // Sink only "cheap" (or nop) address-space casts.  This is a weaker condition
   // than sinking only nop casts, but is helpful on some platforms.
   if (auto *ASC = dyn_cast<AddrSpaceCastInst>(CI)) {
     if (!TLI.isCheapAddrSpaceCast(ASC->getSrcAddressSpace(),
                                   ASC->getDestAddressSpace()))
       return false;
   }
 
static bool OptimizeNoopCopyExpression(CastInst *CI, const TargetLowering &TLI,
   if (TLI.getTypeAction(CI->getContext(), DstVT) ==
       TargetLowering::TypePromoteInteger)
     DstVT = TLI.getTypeToTransformTo(CI->getContext(), DstVT);
 
   // If, after promotion, these are the same types, this is a noop copy.
   if (SrcVT != DstVT)
     return false;
 
-  return SinkCast(CI);
+  return SinkCast(CI, InsertedInsts);
 }
 
 /// Try to combine CI into a call to the llvm.uadd.with.overflow intrinsic if
 /// possible.
 ///
 /// Return true if any changes were made.
 static bool CombineUAddWithOverflow(CmpInst *CI) {
   Value *A, *B;

 /// f = sext i32 b to i64
 /// a = add nsw i64 f, 3
 /// e = getelementptr ..., i64 a
 /// \endcode
 ///
 /// \p Inst[in/out] the extension may be modified during the process if some
 /// promotions apply.
 bool CodeGenPrepare::optimizeExt(Instruction *&Inst) {
+   bool Changed = false;
   // ExtLoad formation and address type promotion infrastructure requires TLI to
   // be effective.
   if (!TLI)
     return false;
 
+  if (Inst->getOpcode() == Instruction::SExt) { 
+    if (TTI->isExtFoldableInAllUsers(*Inst)) {
+      // Sink the current extension into user blocks if foldable with all users.
+      CastInst *SI = cast<CastInst>(Inst);
+      Changed = SinkCast(SI, InsertedInsts);
+      // Check if the instruction is erased.
+      if (!SI) {
+        Inst = nullptr;
+        return true;
+      }
+    }
+  }
+
   bool AllowPromotionWithoutCommonHeader = false;
   /// See if it is an interesting sext operations for the address type
   /// promotion before trying to promote it, e.g., the ones with the right
   /// type and used in memory accesses.
   bool ATPConsiderable = TTI->shouldConsiderAddressTypePromotion(
       *Inst, AllowPromotionWithoutCommonHeader);
   TypePromotionTransaction TPT(RemovedInsts);
   TypePromotionTransaction::ConstRestorationPt LastKnownGood =
bool CodeGenPrepare::optimizeExt(Instruction *&Inst) {
 
   // Continue promoting SExts if known as considerable depending on targets.
   if (ATPConsiderable &&
       performAddressTypePromotion(Inst, AllowPromotionWithoutCommonHeader,
                                   HasPromoted, TPT, SpeculativelyMovedExts))
     return true;
 
   TPT.rollback(LastKnownGood);
-  return false;
+  return Changed;
 }
 
 // Perform address type promotion if doing so is profitable.
 // If AllowPromotionWithoutCommonHeader == false, we should find other sext
 // instructions that sign extended the same initial value. However, if
 // AllowPromotionWithoutCommonHeader == true, we expect promoting the
 // extension is just profitable.
 bool CodeGenPrepare::performAddressTypePromotion(
for (auto VisitedSExt : UnhandledExts) {
         SeenChainsForSExt[HeadOfChain] = nullptr;
         ValToSExtendedUses[HeadOfChain].push_back(I);
       }
     }
   return Promoted;
 }
 
 bool CodeGenPrepare::optimizeExtUses(Instruction *I) {
+  if (!I)
+    return false;
+
   BasicBlock *DefBB = I->getParent();
 
   // If the result of a {s|z}ext and its source are both live out, rewrite all
   // other uses of the source with result of extension.
   Value *Src = I->getOperand(0);
   if (Src->hasOneUse())
     return false;
 
if (CastInst *CI = dyn_cast<CastInst>(I)) {
     // already been constant folded.  The only reason NOT to constant fold
     // it is if something (e.g. LSR) was careful to place the constant
     // evaluation in a block other than then one that uses it (e.g. to hoist
     // the address of globals out of a loop).  If this is the case, we don't
     // want to forward-subst the cast.
     if (isa<Constant>(CI->getOperand(0)))
       return false;
 
-    if (TLI && OptimizeNoopCopyExpression(CI, *TLI, *DL))
+    if (TLI && OptimizeNoopCopyExpression(CI, *TLI, *DL, InsertedInsts))
       return true;
 
     if (isa<ZExtInst>(I) || isa<SExtInst>(I)) {
       /// Sink a zext or sext into its user blocks if the target type doesn't
       /// fit in one register
       if (TLI &&
           TLI->getTypeAction(CI->getContext(),
                              TLI->getValueType(*DL, CI->getType())) ==
               TargetLowering::TypeExpandInteger) {
-        return SinkCast(CI);
+        return SinkCast(CI, InsertedInsts);
       } else {
         bool MadeChange = optimizeExt(I);
         return MadeChange | optimizeExtUses(I);
       }
     }
     return false;
   }
 

lib/Target/AArch64/AArch64TargetTransformInfo.h

public:
   int getInterleavedMemoryOpCost(unsigned Opcode, Type *VecTy, unsigned Factor,
                                  ArrayRef<unsigned> Indices, unsigned Alignment,
                                  unsigned AddressSpace);
 
   bool
   shouldConsiderAddressTypePromotion(const Instruction &I,
                                      bool &AllowPromotionWithoutCommonHeader);
 
+  bool isExtFoldableInAllUsers(const Instruction &I);
+
   unsigned getCacheLineSize();
 
   unsigned getPrefetchDistance();
 
   unsigned getMinPrefetchStride();
 
   unsigned getMaxPrefetchIterationsAhead();
   /// @}
 };
 
 } // end namespace llvm
 
 #endif

lib/Target/AArch64/AArch64TargetTransformInfo.cpp

if (const GetElementPtrInst *GEPInst = dyn_cast<GetElementPtrInst>(U)) {
         AllowPromotionWithoutCommonHeader = true;
         break;
       }
     }
   }
   return Considerable;
 }
 
+bool AArch64TTIImpl::isExtFoldableInAllUsers(const Instruction &I) {
+  // FIXME: we handle only SExt for now.
+  if (I.getOpcode() != Instruction::SExt)
+    return false;
+  // SExt will be free when it is foldable in all users.
+  return TLI->isExtFree(&I);
+}
+
 unsigned AArch64TTIImpl::getCacheLineSize() {
   return ST->getCacheLineSize();
 }
 
 unsigned AArch64TTIImpl::getPrefetchDistance() {
   return ST->getPrefetchDistance();
 }
 
 unsigned AArch64TTIImpl::getMinPrefetchStride() {
   return ST->getMinPrefetchStride();
 }
 
 unsigned AArch64TTIImpl::getMaxPrefetchIterationsAhead() {
   return ST->getMaxPrefetchIterationsAhead();
 }

test/CodeGen/AArch64/aarch64-address-type-promotion-sink.ll

+; RUN: llc < %s -o - | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64--linux-gnu"
+
+%struct.16B = type { i16, i16 , i16, i16, i16, i16, i16, i16}
+define i32 @func_16B(i16 %c, i16 %c2, i16* %base, i32 %i, i16 %v16, %struct.16B* %P) {
+; CHECK-LABEL: @func_16B
+
+entry:
+  %s_ext = sext i32 %i to i64
+
+; CHECK-LABEL: %entry
+; CHECK: ldrh   w{{[0-9]+}}, [x[[ADDR0:[0-9]+]]]
+
+  %addr0 = getelementptr inbounds %struct.16B, %struct.16B* %P, i64 %s_ext, i32 0
+  %cc = load i16, i16* %addr0
+  %cmp = icmp eq i16 %cc, %c
+  br i1 %cmp, label %if.then, label %out
+
+if.then:
+; CHECK-LABEL: %if.then
+; CHECK-NOT: sxtw x{{[0-9]+}}, w{{[0-9]+}}
+; CHECK: ldrh   w{{[0-9]+}}, [x{{[0-9]+}}, w{{[0-9]+}}, sxtw #1]
+  %addr1 = getelementptr inbounds i16, i16* %base, i64 %s_ext
+  %v = load i16, i16* %addr1
+  %cmp2 = icmp eq i16 %v, %c2
+  br i1 %cmp2, label %if.then2, label %out
+
+if.then2:
+; CHECK-LABEL: %if.then2
+; CHECK-NOT: add x{{[0-9]+}}, x{{[0-9]+}}, x{{[0-9]+}}
+; CHECK: strh   w{{[0-9]+}}, [x[[ADDR0]]]
+; CHECK: strh   w{{[0-9]+}}, [x[[ADDR0]], #2]
+
+  %addr2 = getelementptr inbounds %struct.16B, %struct.16B* %P, i64 %s_ext, i32 1
+  store i16 %v16, i16* %addr0
+  store i16 %v16, i16* %addr2
+  ret i32  0
+
+out:
+  ret i32 0
+}
+
+%struct.8B = type { i16, i16 , i16, i16}
+define i32 @func_8B(i16 %c, i16 %c2, i16* %base, i32 %i, i16 %v16, %struct.8B* %P) {
+; CHECK-LABEL: @func_8B
+
+entry:
+  %s_ext = sext i32 %i to i64
+
+; CHECK-LABEL: %entry
+; CHECK: ldrh   w{{[0-9]+}}, [x[[ADDR0:[0-9]+]]]
+
+  %addr0 = getelementptr inbounds %struct.8B, %struct.8B* %P, i64 %s_ext, i32 0
+  %cc = load i16, i16* %addr0
+  %cmp = icmp eq i16 %cc, %c
+  br i1 %cmp, label %if.then, label %out
+
+if.then:
+; CHECK-LABEL: %if.then
+; CHECK-NOT: sxtw x{{[0-9]+}}, w{{[0-9]+}}
+; CHECK: ldrh   w{{[0-9]+}}, [x{{[0-9]+}}, w{{[0-9]+}}, sxtw #1]
+  %addr1 = getelementptr inbounds i16, i16* %base, i64 %s_ext
+  %v = load i16, i16* %addr1
+  %cmp2 = icmp eq i16 %v, %c2
+  br i1 %cmp2, label %if.then2, label %out
+
+if.then2:
+; CHECK-LABEL: %if.then2
+; CHECK-NOT: add x{{[0-9]+}}, x{{[0-9]+}}, x{{[0-9]+}}
+; CHECK: strh   w{{[0-9]+}}, [x[[ADDR0]]]
+; CHECK: strh   w{{[0-9]+}}, [x[[ADDR0]], #2]
+
+  %addr2 = getelementptr inbounds %struct.8B, %struct.8B* %P, i64 %s_ext, i32 1
+  store i16 %v16, i16* %addr0
+  store i16 %v16, i16* %addr2
+  ret i32  0
+
+out:
+  ret i32 0
+}
+
+%struct.4B = type { i16, i16 }
+define i32 @func_4B(i16 %c, i16 %c2, i16* %base, i32 %i, i16 %v16, %struct.4B* %P) {
+; CHECK-LABEL: @func_4B
+
+entry:
+  %s_ext = sext i32 %i to i64
+
+; CHECK-LABEL: %entry
+; CHECK: ldrh   w{{[0-9]+}}, [x[[ADDR0:[0-9]+]]]
+
+  %addr0 = getelementptr inbounds %struct.4B, %struct.4B* %P, i64 %s_ext, i32 0
+  %cc = load i16, i16* %addr0
+  %cmp = icmp eq i16 %cc, %c
+  br i1 %cmp, label %if.then, label %out
+
+if.then:
+; CHECK-LABEL: %if.then
+; CHECK-NOT: sxtw x{{[0-9]+}}, w{{[0-9]+}}
+; CHECK: ldrh   w{{[0-9]+}}, [x{{[0-9]+}}, w{{[0-9]+}}, sxtw #1]
+  %addr1 = getelementptr inbounds i16, i16* %base, i64 %s_ext
+  %v = load i16, i16* %addr1
+  %cmp2 = icmp eq i16 %v, %c2
+  br i1 %cmp2, label %if.then2, label %out
+
+if.then2:
+; CHECK-LABEL: %if.then2
+; CHECK-NOT: add x{{[0-9]+}}, x{{[0-9]+}}, x{{[0-9]+}}
+; CHECK: strh   w{{[0-9]+}}, [x[[ADDR0]]]
+; CHECK: strh   w{{[0-9]+}}, [x[[ADDR0]], #2]
+
+  %addr2 = getelementptr inbounds %struct.4B, %struct.4B* %P, i64 %s_ext, i32 1
+  store i16 %v16, i16* %addr0
+  store i16 %v16, i16* %addr2
+  ret i32  0
+
+out:
+  ret i32 0
+}
+