[X86] Speculatively load operands of select instruction

[lsaba]

Summary

For a select instruction where the operands are address calculations of two independent loads, the pass tries to speculate the loads and feed them into the select instruction, this allows early parallel execution of the loads and possibly memory folding into the CMOV instructions later on.
The pass currently only handles cases where the loads are elements of the same struct.

include/llvm/Target/TargetLowering.h

public:
 
   /// This method query the target whether it is beneficial for dag combiner to
   /// promote the specified node. If true, it should return the desired
   /// promotion type by reference.
   virtual bool IsDesirableToPromoteOp(SDValue /*Op*/, EVT &/*PVT*/) const {
     return false;
   }
 
+  /// Return true if it is desirable to speculatively load the operands
+  /// of a select instruction for the target.
+  virtual bool isDesirableToSpeculateSelectLoad() const {
+    return false;
+  }
+
   /// Return true if the target supports swifterror attribute. It optimizes
   /// loads and stores to reading and writing a specific register.
   virtual bool supportSwiftError() const {
     return false;
   }
 
   /// Return true if the target supports that a subset of CSRs for the given
   /// machine function is handled explicitly via copies.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (LHS.getOperand(0) != RHS.getOperand(0) ||
         // src value info, don't do the transformation if the memory
         // locations are not in the default address space.
         LLD->getPointerInfo().getAddrSpace() != 0 ||
         RLD->getPointerInfo().getAddrSpace() != 0 ||
         !TLI.isOperationLegalOrCustom(TheSelect->getOpcode(),
                                       LLD->getBasePtr().getValueType()))
       return false;
 
+    // Avoid combining the load if both loads are GEPs into elements of the same
+    // struct. TODO: handle cases where the GEP is bitcasted to another type.
+    if (TLI.isDesirableToSpeculateSelectLoad()) {
+      if (LLD->getMemOperand()->getValue() &&
+          RLD->getMemOperand()->getValue()) {
+        const GetElementPtrInst *GEPTrue =
+            dyn_cast<GetElementPtrInst>(LLD->getMemOperand()->getValue());
+        const GetElementPtrInst *GEPFalse =
+            dyn_cast<GetElementPtrInst>(RLD->getMemOperand()->getValue());
+        if (GEPTrue && GEPFalse) {
+          if (GEPTrue->getSourceElementType()->isStructTy() &&
+              GEPFalse->getSourceElementType()->isStructTy() &&
+              GEPTrue->getPointerOperand() == GEPFalse->getPointerOperand() &&
+              GEPTrue->hasAllConstantIndices() &&
+              GEPFalse->hasAllConstantIndices())
+            LLD->dump();
+          return false;
+        }
+      }
+    }
+
     // Check that the select condition doesn't reach either load.  If so,
     // folding this will induce a cycle into the DAG.  If not, this is safe to
     // xform, so create a select of the addresses.
     SDValue Addr;
     if (TheSelect->getOpcode() == ISD::SELECT) {
       SDNode *CondNode = TheSelect->getOperand(0).getNode();
       if ((LLD->hasAnyUseOfValue(1) && LLD->isPredecessorOf(CondNode)) ||
           (RLD->hasAnyUseOfValue(1) && RLD->isPredecessorOf(CondNode)))

lib/Target/X86/CMakeLists.txt

set(sources
   X86MachineFunctionInfo.cpp
   X86MacroFusion.cpp
   X86OptimizeLEAs.cpp
   X86PadShortFunction.cpp
   X86RegisterBankInfo.cpp
   X86RegisterInfo.cpp
   X86SelectionDAGInfo.cpp
   X86ShuffleDecodeConstantPool.cpp
+  X86SpeculateSelectLoad.cpp
   X86Subtarget.cpp
   X86TargetMachine.cpp
   X86TargetObjectFile.cpp
   X86TargetTransformInfo.cpp
   X86VZeroUpper.cpp
   X86WinAllocaExpander.cpp
   X86WinEHState.cpp
   X86CallingConv.cpp

lib/Target/X86/X86.h

 /// done by replacing esp-relative movs with pushes.
 FunctionPass *createX86CallFrameOptimization();
 
 /// Return an IR pass that inserts EH registration stack objects and explicit
 /// EH state updates. This pass must run after EH preparation, which does
 /// Windows-specific but architecture-neutral preparation.
 FunctionPass *createX86WinEHStatePass();
 
+/// Return an IR pass that tries to speculatively load the operands of
+/// a select instruction when profitable.
+FunctionPass *createX86SpeculateSelectLoadPass();
+
 /// Return a Machine IR pass that expands X86-specific pseudo
 /// instructions into a sequence of actual instructions. This pass
 /// must run after prologue/epilogue insertion and before lowering
 /// the MachineInstr to MC.
 FunctionPass *createX86ExpandPseudoPass();
 
 /// This pass converts X86 cmov instructions into branch when profitable.
 FunctionPass *createX86CmovConverterPass();

lib/Target/X86/X86ISelLowering.h

public:
     bool isTypeDesirableForOp(unsigned Opc, EVT VT) const override;
 
     /// Return true if the target has native support for the
     /// specified value type and it is 'desirable' to use the type. e.g. On x86
     /// i16 is legal, but undesirable since i16 instruction encodings are longer
     /// and some i16 instructions are slow.
     bool IsDesirableToPromoteOp(SDValue Op, EVT &PVT) const override;
 
+    /// Return true if it is desirable to speculatively load the operands
+    /// of a select instruction for the target.
+    bool isDesirableToSpeculateSelectLoad() const override { return true; }
+
     MachineBasicBlock *
     EmitInstrWithCustomInserter(MachineInstr &MI,
                                 MachineBasicBlock *MBB) const override;
 
     /// This method returns the name of a target specific DAG node.
     const char *getTargetNodeName(unsigned Opcode) const override;
 
     bool mergeStoresAfterLegalization() const override { return true; }

lib/Target/X86/X86SpeculateSelectLoad.cpp

+//===-- X86SpeculateSelectLoad- Speculatively load operands of a select --===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===---------------------------------------------------------------------===//
+// For a select instruction where the operands are address calculations
+// of two independent loads, the pass tries to speculate the loads and
+// feed them into the select instruction, this allows early parallel execution
+// of the loads and possibly memory folding into the CMOV instructions later on.
+// The pass currently only handles cases where the loads are elements of the
+// same struct.
+//===---------------------------------------------------------------------===//
+
+#include "X86.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/Debug.h"
+#include <deque>
+
+using namespace llvm;
+
+#define DEBUG_TYPE "x86speculateload"
+
+namespace llvm {
+void initializeX86SpeculateSelectLoadPassPass(PassRegistry &);
+}
+
+namespace {
+
+class X86SpeculateSelectLoadPass : public FunctionPass {
+public:
+  static char ID; // Pass identification, replacement for typeid.
+
+  X86SpeculateSelectLoadPass() : FunctionPass(ID) {
+    initializeX86SpeculateSelectLoadPassPass(*PassRegistry::getPassRegistry());
+  }
+
+  bool runOnFunction(Function &Fn) override;
+  StringRef getPassName() const override {
+    return "X86 Speculatively load before select instruction";
+  }
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<TargetTransformInfoWrapperPass>();
+  }
+
+private:
+  bool OptimizeSelectInst(SelectInst *SI);
+  const DataLayout *DL;
+  const TargetTransformInfo *TTI;
+  SmallVector<Instruction *, 2> InstrForRemoval;
+};
+}
+
+FunctionPass *llvm::createX86SpeculateSelectLoadPass() {
+  return new X86SpeculateSelectLoadPass();
+}
+
+char X86SpeculateSelectLoadPass::ID = 0;
+
+INITIALIZE_PASS_BEGIN(X86SpeculateSelectLoadPass, "x86-speculateload",
+                      "X86 Speculatively load before select instruction", false,
+                      false)
+INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
+INITIALIZE_PASS_END(X86SpeculateSelectLoadPass, "x86-speculateload",
+                    "X86 Speculatively load before select instruction", false,
+                    false)
+
+bool X86SpeculateSelectLoadPass::OptimizeSelectInst(SelectInst *SI) {
+  GetElementPtrInst *GEPIT = dyn_cast<GetElementPtrInst>(SI->getTrueValue());
+  GetElementPtrInst *GEPIF = dyn_cast<GetElementPtrInst>(SI->getFalseValue());
+  if (GEPIT == nullptr || GEPIF == nullptr)
+    return false;
+
+  // The pass currently only handles cases where the loads are elements of the
+  // same struct.
+  if (!GEPIT->getSourceElementType()->isStructTy() ||
+      !GEPIF->getSourceElementType()->isStructTy() ||
+      GEPIT->getPointerOperand() != GEPIF->getPointerOperand() ||
+      !GEPIT->hasAllConstantIndices() || !GEPIF->hasAllConstantIndices())
+    return false;
+
+  if (!SI->hasOneUse())
+    return false;
+
+  // Bail out if there is a good chance we'll be loading from two different
+  // cache lines instead of one.
+  // This is a shallow check of the struct's fields layout.
+  // TODO: deeper check for GEPs with more than one index.
+  if (StructType *STy = dyn_cast<StructType>(GEPIT->getSourceElementType())) {
+    // Get the indices of the elements in the struct
+
+    ConstantInt *Idx1 = dyn_cast<ConstantInt>(GEPIT->getOperand(2));
+    ConstantInt *Idx2 = dyn_cast<ConstantInt>(GEPIF->getOperand(2));
+    const StructLayout *STL = DL->getStructLayout(STy);
+    if (Idx1 && Idx2 && STL) {
+      signed offset1 = STL->getElementOffset(Idx1->getZExtValue());
+      signed offset2 = STL->getElementOffset(Idx2->getZExtValue());
+      unsigned dist = abs(offset1 - offset2);
+      assert((TTI->getCacheLineSize() > 0) &&
+             "CacheLineSize information is missing for X86");
+      if (dist > TTI->getCacheLineSize())
+        return false;
+    }
+  }
+
+  for (User *U : SI->users()) {
+    if (LoadInst *LI = dyn_cast<LoadInst>(U)) {
+      if (!LI->isSimple())
+        return false;
+      IRBuilder<> Builder(SI);
+      LoadInst *LT = Builder.CreateAlignedLoad(GEPIT, LI->getAlignment());
+      LoadInst *LF = Builder.CreateAlignedLoad(GEPIF, LI->getAlignment());
+      SelectInst *NewSI =
+          cast<SelectInst>(Builder.CreateSelect(SI->getCondition(), LT, LF));
+      LI->replaceAllUsesWith(NewSI);
+      InstrForRemoval.push_back(LI);
+      InstrForRemoval.push_back(SI);
+
+      return true;
+    }
+  }
+  return false;
+}
+
+bool X86SpeculateSelectLoadPass::runOnFunction(Function &F) {
+  if (skipFunction(F) || F.optForSize())
+    return false;
+
+  DL = &F.getParent()->getDataLayout();
+  TTI = &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
+
+  bool Changed = false;
+
+  for (auto &BB : F)
+    for (auto &I : BB)
+      if (SelectInst *SI = dyn_cast<SelectInst>(&I))
+        Changed |= OptimizeSelectInst(SI);
+
+  for (auto Instr : InstrForRemoval)
+    Instr->eraseFromParent();
+
+  InstrForRemoval.clear();
+
+  return Changed;
+}
+

lib/Target/X86/X86TargetMachine.cpp

static cl::opt<bool> EnableMachineCombinerPass("x86-machine-combiner",
                                cl::desc("Enable the machine combiner pass"),
                                cl::init(true), cl::Hidden);
 
 namespace llvm {
 
 void initializeWinEHStatePassPass(PassRegistry &);
 void initializeFixupLEAPassPass(PassRegistry &);
 void initializeX86ExecutionDepsFixPass(PassRegistry &);
+void initializeX86SpeculateSelectLoadPassPass(PassRegistry &);
 
 } // end namespace llvm
 
 extern "C" void LLVMInitializeX86Target() {
   // Register the target.
   RegisterTargetMachine<X86TargetMachine> X(getTheX86_32Target());
   RegisterTargetMachine<X86TargetMachine> Y(getTheX86_64Target());
 
   PassRegistry &PR = *PassRegistry::getPassRegistry();
   initializeGlobalISel(PR);
   initializeWinEHStatePassPass(PR);
   initializeFixupBWInstPassPass(PR);
   initializeEvexToVexInstPassPass(PR);
   initializeFixupLEAPassPass(PR);
   initializeX86ExecutionDepsFixPass(PR);
+  initializeX86SpeculateSelectLoadPassPass(PR);
 }
 
 static std::unique_ptr<TargetLoweringObjectFile> createTLOF(const Triple &TT) {
   if (TT.isOSBinFormatMachO()) {
     if (TT.getArch() == Triple::x86_64)
       return llvm::make_unique<X86_64MachoTargetObjectFile>();
     return llvm::make_unique<TargetLoweringObjectFileMachO>();
   }
TargetPassConfig *X86TargetMachine::createPassConfig(PassManagerBase &PM) {
   return new X86PassConfig(*this, PM);
 }
 
 void X86PassConfig::addIRPasses() {
   addPass(createAtomicExpandPass());
 
   TargetPassConfig::addIRPasses();
 
-  if (TM->getOptLevel() != CodeGenOpt::None)
+  if (TM->getOptLevel() != CodeGenOpt::None) {
     addPass(createInterleavedAccessPass());
+    addPass(createX86SpeculateSelectLoadPass());
+  }
 }
 
 bool X86PassConfig::addInstSelector() {
   // Install an instruction selector.
   addPass(createX86ISelDag(getX86TargetMachine(), getOptLevel()));
 
   // For ELF, cleanup any local-dynamic TLS accesses.
   if (TM->getTargetTriple().isOSBinFormatELF() &&

lib/Target/X86/X86TargetTransformInfo.h

public:
   bool isLegalMaskedLoad(Type *DataType);
   bool isLegalMaskedStore(Type *DataType);
   bool isLegalMaskedGather(Type *DataType);
   bool isLegalMaskedScatter(Type *DataType);
   bool areInlineCompatible(const Function *Caller,
                            const Function *Callee) const;
   bool expandMemCmp(Instruction *I, unsigned &MaxLoadSize);
   bool enableInterleavedAccessVectorization();
+  unsigned getCacheLineSize() const;
+
 private:
   int getGSScalarCost(unsigned Opcode, Type *DataTy, bool VariableMask,
                       unsigned Alignment, unsigned AddressSpace);
   int getGSVectorCost(unsigned Opcode, Type *DataTy, Value *Ptr,
                       unsigned Alignment, unsigned AddressSpace);
 
   /// @}
 };
 
 } // end namespace llvm
 
 #endif

lib/Target/X86/X86TargetTransformInfo.cpp

return getInterleavedMemoryOpCostAVX512(Opcode, VecTy, Factor, Indices,
                                             Alignment, AddressSpace);
   if (ST->hasAVX2())
     return getInterleavedMemoryOpCostAVX2(Opcode, VecTy, Factor, Indices,
                                           Alignment, AddressSpace);
 
   return BaseT::getInterleavedMemoryOpCost(Opcode, VecTy, Factor, Indices,
                                            Alignment, AddressSpace);
 }
+
+unsigned X86TTIImpl::getCacheLineSize() const { return 64; }

test/CodeGen/X86/speculate-select-load.ll

+; RUN: opt < %s -mcpu=x86-64 -S -x86-speculateload  | FileCheck %s
+; RUN: opt < %s -mcpu=i386 -S -x86-speculateload  | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+%struct.S = type { i64, %struct.S*, %struct.S*, i64, i16, i64, i64, i64, i64, i64 }
+
+; Function Attrs: norecurse nounwind readonly uwtable
+define %struct.S* @spec_load(i32 %x, %struct.S* nocapture readnone %A, %struct.S* nocapture readonly %B) local_unnamed_addr #0 {
+entry:
+; CHECK-LABEL:@spec_load
+; CHECK: getelementptr inbounds %struct.S, %struct.S* %B, i64 0, i32 1
+; CHECK: getelementptr inbounds %struct.S, %struct.S* %B, i64 0, i32 2
+; CHECK: [[A:%[0-9]+]] = load %struct.S*, %struct.S** %{{.*}}, align 8
+; CHECK: [[B:%[0-9]+]] = load %struct.S*, %struct.S** %{{.*}}, align 8
+; CHECK: select i1 %tobool, %struct.S* [[A]], %struct.S* [[B]]
+
+  %tobool = icmp eq i32 %x, 0
+  %b = getelementptr inbounds %struct.S, %struct.S* %B, i64 0, i32 1
+  %c = getelementptr inbounds %struct.S, %struct.S* %B, i64 0, i32 2
+  %A.addr.0.in = select i1 %tobool, %struct.S** %c, %struct.S** %b
+  %A.addr.0 = load %struct.S*, %struct.S** %A.addr.0.in, align 8
+  ret %struct.S* %A.addr.0
+}
+
+; Function Attrs: norecurse nounwind readonly uwtable
+define i64 @no_spec_load(i32 %x, i64 %A, %struct.S* nocapture readonly %B) local_unnamed_addr #0 {
+entry:
+; CHECK-LABEL:@no_spec_load
+; CHECK: [[A:%[a-z]+]] = getelementptr inbounds %struct.S, %struct.S* %B, i64 0, i32 0
+; CHECK: [[B:%[a-z]+]] = getelementptr inbounds %struct.S, %struct.S* %B, i64 0, i32 9
+; CHECK: select i1 %tobool, i64* [[B]], i64* [[A]]
+; CHECK-NEXT: load i64, i64* %{{.*}}, align 8
+
+  %tobool = icmp eq i32 %x, 0
+  %a = getelementptr inbounds %struct.S, %struct.S* %B, i64 0, i32 0
+  %j = getelementptr inbounds %struct.S, %struct.S* %B, i64 0, i32 9
+  %A.addr.0.in = select i1 %tobool, i64* %j, i64* %a
+  %A.addr.0 = load i64, i64* %A.addr.0.in, align 8
+  ret i64 %A.addr.0
+}
+
+attributes #0 = { norecurse nounwind readonly uwtable "target-cpu"="core-avx2" "target-features"="+fxsr,+mmx,+sse,+sse2,+x87"}
+