[AArch64] Sink sext when foldable in user GEPs

[junbuml]

Summary

Sink sign extensions to user blocks if foldable as a part of address calculation by the local ISel when placed in the same block. With this change, 3% performance gain was observed in spec2006/astar. The test case was reduced from spec2006/astar wayobj::makebound2() function.

[jmolloy]

lib/Target/AArch64/AArch64AddressTypePromotion.cpp

140

"performSink" is so generic - can we call it something more specific like "sinkSExt" ?

377

Same here - "isAllFoldable" - can we make that clear we're referring to SExts?

386

You shouldn't need these clear() operations, because this array will not be used if we return false.

395

Shouldn't this function only succeed if the SExt was the *last* operand in the GEP? We can't represent a GEP such as:

%0 = sext i16 %idx to i32
%1 = gep inbounds i32* %base, i32 %0, i32 6

So I think this function is not conservative enough.

[junbuml]

Addressed James' comments. Thanks James for the review.

[junbuml]

lib/Target/AArch64/AArch64AddressTypePromotion.cpp

395

If we are accessing an element from a struct pointer, such GEP could be used where a constant is placed after the sext to index the element in struct. The test case shows such case. For example, if you see func_16B(), %addr2 is to access the 2nd element in %struct.16B :

%addr2 = getelementptr inbounds %struct.16B, %struct.16B* %P, i64 %s_ext, i32 1

Since all other indexes are constants, it will end up with a simple add in the form of %base + TypeSize * sext + constant, where the sext is used to index a struct object from the struct array.

[mcrosier]

lib/Target/AArch64/AArch64AddressTypePromotion.cpp

131

It would probably be a good idea to rename this method and update the comment accordingly now that this logic both hoists common sexts and sinks foldable sexts. Unfortunately, I can't come up with a good name off the top of my head.

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

246

I don't think you're using the CHECK-LABEL directive the way it is intended here (i.e., the %if.then label isn't unique to this file). I suggest using a regular CHECK.

246

You can ignore this comment. I tried to remove it, but Phab wasn't taking the change. I think your usage is okay..

[qcolombet]

Hi,

When I first implemented this pass the intent was to get insights on how to do that in a generic way. Then, we added some logic to do that transformation as part of CGP. The goal was to kill that pass once and for all.

The bottom line is instead of improving this pass, could teach the generic framework in CGP?

Cheers,
Quentin

[junbuml]

When I first implemented this pass the intent was to get insights on how to do that in a generic way. Then, we added some logic to do that transformation as part of CGP. The goal was to kill that pass once and for all. The bottom line is instead of improving this pass, could teach the generic framework in CGP?

That makes sense. Let me try to make this more general in CGP. Thanks Quentin.

[junbuml]

Before we perform this transformation in CGP, I think we need to move this pass into CGP first. So, my plan is to extend moveExtToFormExtLoad() to cover the transformation done in this pass. Instead of just looking for a mergeable load in extLdPromotion() , I will look up heads of chains like analyzeSExtension() in aarch64-type-promotion using existing code (TypePromotionTransaction and TypePromotionHelper) and then promote SExts which derived from the same header. The new extension in CGP will be enabled in AArch64 at first. Please let me know if there is any suggestion or concern regarding the migration of this pass into CGP.

[junbuml]

My plan is to move aarch64-type-promotion pass into CGP first, and then perform the transformation proposed in this change in CGP. As a part of that, I first restructured the current implementation in CGP in D27853.

lib/Target/AArch64/AArch64AddressTypePromotion.cpp

 // FIXME: This pass may be useful for other targets too.
 // ===---------------------------------------------------------------------===//
 
 #include "AArch64.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/GetElementPtrTypeIterator.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
public:
   bool runOnFunction(Function &F) override;
 
 private:
   /// The current function.
   Function *Func;
   /// Filter out all sexts that does not have this type.
   /// Currently initialized with Int64Ty.
   Type *ConsideredSExtType;
+  const DataLayout *DL;
 
   // This transformation requires dominator info.
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
     AU.addRequired<DominatorTreeWrapperPass>();
     AU.addPreserved<DominatorTreeWrapperPass>();
     FunctionPass::getAnalysisUsage(AU);
   }
private:
   /// type and used in memory accesses.
   /// More precisely, a sext instruction is considered as interesting if it
   /// is used in a "complex" getelementptr or it exits at least another
   /// sext instruction that sign extended the same initial value.
   /// A getelementptr is considered as "complex" if it has more than 2
   // operands.
   void analyzeSExtension(Instructions &SExtInsts);
 
   /// Merge redundant sign extension operations in common dominator.
   void mergeSExts(ValueToInsts &ValToSExtendedUses,
                   SetOfInstructions &ToRemove);
+
+  /// Check if the sign extension is foldable as a part of address calculation
+  /// in all GEPs which take the sign extension as an index.
+  bool isSExtFoldableInAllUsers(Instruction *SExtInst,
+                                Instructions &GEPsInOtherBlock);
+
+  /// Sink a foldable sign extension to user blocks.
+  void sinkFoldableSExt(Instruction *SExtInst, Instructions &GEPsInOtherBlock);
 };
 } // end anonymous namespace.
 
 char AArch64AddressTypePromotion::ID = 0;
 
 INITIALIZE_PASS_BEGIN(AArch64AddressTypePromotion, "aarch64-type-promotion",
                       AARCH64_TYPE_PROMO_NAME, false, false)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
if (EnableMerge)
     mergeSExts(ValToSExtendedUses, ToRemove);
 
   // Remove all instructions marked as ToRemove.
   for (Instruction *I: ToRemove)
     I->eraseFromParent();
   return LocalChange;
 }
 
+bool AArch64AddressTypePromotion::isSExtFoldableInAllUsers(
+    Instruction *SExtInst, Instructions &GEPsInOtherBlock) {
+  BasicBlock *SExtBB = SExtInst->getParent();
+  for (User *U : SExtInst->users()) {
+    // FIXME: To be simple, for now, we handle sign extensions used only by GEPs
+    // directly. We could also sink the promoted instructions forming a chain
+    // between SExt and GEP if foldable all together.
+    Instruction *GEPInst = dyn_cast<GetElementPtrInst>(U);
+    if (!GEPInst)
+      return false;
+
+    // FIXME: For foldability check in GEP, we simply see if all operands except
+    // SExtInst are constants (i.g., %base + TypeSize * sext + constant),  and
+    // check if the type size indexed by the sign extension is the foldable
+    // amount as sign extension in add/sub instruction.
+    gep_type_iterator GTI = gep_type_begin(GEPInst);
+    for (unsigned i = 1, e = GEPInst->getNumOperands(); i != e; ++i, ++GTI) {
+      if (GEPInst->getOperand(i) == SExtInst) {
+        uint64_t TypeSize = DL->getTypeAllocSize(GTI.getIndexedType());
+        if (TypeSize > 16 || !isPowerOf2_64(TypeSize))
+          return false;
+      } else if (!isa<ConstantInt>(GEPInst->getOperand(i)))
+        return false;
+    }
+    if (GEPInst->getParent() != SExtBB)
+      GEPsInOtherBlock.push_back(GEPInst);
+  }
+  return true;
+}
+
+void AArch64AddressTypePromotion::sinkFoldableSExt(
+    Instruction *SExtInst, Instructions &GEPsInOtherBlock) {
+  for (auto *GEPInst : GEPsInOtherBlock) {
+    assert(isa<GetElementPtrInst>(GEPInst) && "Expect only GEP as a user.");
+    assert(GEPInst->getParent() != SExtInst->getParent() &&
+           "Expect it to be in different block.");
+    Instruction *SExtInstSunk = SExtInst->clone();
+    SExtInstSunk->insertBefore(GEPInst);
+    GEPInst->replaceUsesOfWith(SExtInst, SExtInstSunk);
+  }
+}
+
 void AArch64AddressTypePromotion::mergeSExts(ValueToInsts &ValToSExtendedUses,
                                              SetOfInstructions &ToRemove) {
   DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
 
   for (auto &Entry : ValToSExtendedUses) {
     Instructions &Insts = Entry.second;
     Instructions CurPts;
     for (Instruction *Inst : Insts) {
       if (ToRemove.count(Inst))
         continue;
+
+      // Sink sign extension operations if foldable into address calculation.
+      Instructions GEPsInOtherBlock;
+      if (isSExtFoldableInAllUsers(Inst, GEPsInOtherBlock)) {
+        if (!GEPsInOtherBlock.empty())
+          sinkFoldableSExt(Inst, GEPsInOtherBlock);
+        // No need to be merged if foldable as a part of address calculation.
+        continue;
+      }
+
       bool inserted = false;
       for (auto &Pt : CurPts) {
         if (DT.dominates(Inst, Pt)) {
           DEBUG(dbgs() << "Replace all uses of:\n" << *Pt << "\nwith:\n"
                        << *Inst << '\n');
           Pt->replaceAllUsesWith(Inst);
           ToRemove.insert(Pt);
           Pt = Inst;

 bool AArch64AddressTypePromotion::runOnFunction(Function &F) {
   if (skipFunction(F))
     return false;
 
   if (F.isDeclaration())
     return false;
   Func = &F;
   ConsideredSExtType = Type::getInt64Ty(Func->getContext());
+  DL = &F.getParent()->getDataLayout();
 
   DEBUG(dbgs() << "*** " << getPassName() << ": " << Func->getName() << '\n');
 
   Instructions SExtInsts;
   analyzeSExtension(SExtInsts);
   return propagateSignExtension(SExtInsts);
 }

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
+}
+
+%struct.2B = type { i8, i8  }
+define i32 @func_2B(i8 %c, i8 %c2, i8* %base, i32 %i, i8 %v16, %struct.2B* %P) {
+; CHECK-LABEL: @func_2B
+
+entry:
+; CHECK-LABEL: %entry
+; CHECK: ldrb   w{{[0-9]+}}, [x[[ADDR0:[0-9]+]]]
+  %s_ext = sext i32 %i to i64
+  %addr0 = getelementptr inbounds %struct.2B, %struct.2B* %P, i64 %s_ext, i32 0
+  %cc = load i8, i8* %addr0
+  %cmp = icmp eq i8 %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: ldrb   w{{[0-9]+}}, [x{{[0-9]+}}, w{{[0-9]+}}, sxtw]
+
+  %addr1 = getelementptr inbounds i8, i8* %base, i64 %s_ext
+  %v = load i8, i8* %addr1
+  %cmp2 = icmp eq i8 %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: strb   w{{[0-9]+}}, [x[[ADDR0]]]
+; CHECK: strb   w{{[0-9]+}}, [x[[ADDR0]], #1]
+  %addr2 = getelementptr inbounds %struct.2B, %struct.2B* %P, i64 %s_ext, i32 1
+  store i8 %v16, i8* %addr0
+  store i8 %v16, i8* %addr2
+  ret i32  0
+
+out:
+  ret i32 0
+}
+
+%struct.1B = type { i8 }
+define i32 @func_1B(i8 %c, i8 %c2, i8* %base, i32 %i, i8 %v16, %struct.1B* %P) {
+; CHECK-LABEL: @func_1B
+
+entry:
+; CHECK-LABEL: %entry
+; CHECK-NOT: sxtw x{{[0-9]+}}, w{{[0-9]+}}
+; CHECK: ldrb   w{{[0-9]+}}, [x{{[0-9]+}}, w{{[0-9]+}}, sxtw]
+
+  %s_ext = sext i32 %i to i64
+  %addr0 = getelementptr inbounds %struct.1B, %struct.1B* %P, i64 %s_ext, i32 0
+  %cc = load i8, i8* %addr0
+  %cmp = icmp eq i8 %cc, %c
+  br i1 %cmp, label %if.then, label %out
+
+if.then:
+
+; CHECK-LABEL: %if.then
+; CHECK: ldrb   w{{[0-9]+}}, [x{{[0-9]+}}, w{{[0-9]+}}, sxtw]
+
+  %addr1 = getelementptr inbounds i8, i8* %base, i64 %s_ext
+  %v = load i8, i8* %addr1
+  %cmp2 = icmp eq i8 %v, %c2
+  br i1 %cmp2, label %if.then2, label %out
+
+if.then2:
+
+; CHECK-LABEL: %if.then2
+; CHECK: strb   w{{[0-9]+}}, [x{{[0-9]+}}, w{{[0-9]+}}, sxtw]
+  %addr2 = getelementptr inbounds %struct.1B, %struct.1B* %P, i64 %s_ext, i32 0
+  store i8 %v16, i8* %addr2
+  ret i32  0
+
+out:
+  ret i32 0
+}
+
+%struct.6B = type { i16, i16 , i16}
+define i32 @func_6B(i16 %c, i16 %c2, i16* %base, i32 %i, i16 %v16, %struct.6B* %P) {
+; CHECK-LABEL: @func_6B
+
+entry:
+  %s_ext = sext i32 %i to i64
+  %addr0 = getelementptr inbounds %struct.6B, %struct.6B* %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: sxtw x{{[0-9]+}}, w{{[0-9]+}}
+  %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:madd
+  %addr2 = getelementptr inbounds %struct.6B, %struct.6B* %P, i64 %s_ext, i32 1
+  store i16 %v16, i16* %addr0
+  store i16 %v16, i16* %addr2
+  ret i32  0
+
+out:
+  ret i32 0
+}
+
+%struct.nomergeup = type { i16, i16 }
+define i16 @func_no_mergeup(i32 %i, i16 %v16,  %struct.nomergeup* %P) {
+; CHECK-LABEL: @func_no_mergeup
+entry:
+; CHECK-LABEL: %entry
+; CHECK: add x[[ADDR:[0-9]+]], x{{[0-9]+}}, w{{[0-9]+}}, sxtw #2
+
+  %add1 = add nsw i32 %i, 1
+  %s_ext1 = sext i32 %add1 to i64
+  %addr2 = getelementptr inbounds %struct.nomergeup, %struct.nomergeup* %P, i64 %s_ext1, i32 0
+
+  %lv  = load i16, i16* %addr2
+  %cmp = icmp eq i16 %lv, %v16
+  br i1 %cmp, label %if.then, label %if.then2
+
+if.then:
+; CHECK-LABEL: %if.then
+; CHECK-NOT: sxtw x{{[0-9]+}}, w{{[0-9]+}}
+; CHECK: ldrh   w{{[0-9]+}}, [x[[ADDR]], #2]
+
+  %s_ext2 = sext i32 %i to i64
+  %addr = getelementptr inbounds %struct.nomergeup, %struct.nomergeup* %P, i64 %s_ext2, i32 1
+  %v = load i16, i16* %addr
+  ret i16 %v
+
+if.then2:
+  ret i16 0
+}