Support placing switch tables in target-specific address spaces

[dylanmckay]

Summary

In the AVR architecture, switch tables need to be placed into a special
'program memory' address space (id=1).

In order to support this, add a hook to TargetTransformInfo that
targets can override.

If there are any more suitable reviewers I couldn't think of, please add them :)

[dylanmckay]

I have found a similar problem with a similar fix, I am currently looking into a nice way to generalise this patch.

[dylanmckay]

WIP, not ready for review yet

[hfinkel]

FYI: TTI should not be required in order to the mid-level optimizer to correctly transform a program. If this is required for correctness, it should be in DataLayout.

[dylanmckay]

Will rewrite once D36916 lands

include/llvm/Analysis/TargetTransformInfo.h

public:
   /// Returns the address space ID for a target's 'flat' address space. Note
   /// this is not necessarily the same as addrspace(0), which LLVM sometimes
   /// refers to as the generic address space. The flat address space is a
   /// generic address space that can be used access multiple segments of memory
   /// with different address spaces. Access of a memory location through a
   /// pointer with this address space is expected to be legal but slower
   /// compared to the same memory location accessed through a pointer with a
   /// different address space.
-  //
+  ///
   /// This is for for targets with different pointer representations which can
   /// be converted with the addrspacecast instruction. If a pointer is converted
   /// to this address space, optimizations should attempt to replace the access
   /// with the source address space.
   ///
   /// \returns ~0u if the target does not have such a flat address space to
   /// optimize away.
   unsigned getFlatAddressSpace() const;
 
+  /// Returns the address space ID that switch lookup tables should be placed
+  /// into.
+  ///
+  /// This is useful for Harvard architectures where you often want to place
+  /// lookup tables into a different address space (for example, in AVR, lookup
+  /// tables reside in the 'program memory' space.
+  unsigned getSwitchTableAddressSpace() const;
+
   /// \brief Test whether calls to a function lower to actual program function
   /// calls.
   ///
   /// The idea is to test whether the program is likely to require a 'call'
   /// instruction or equivalent in order to call the given function.
   ///
   /// FIXME: It's not clear that this is a good or useful query API. Client's
   /// should probably move to simpler cost metrics using the above.
public:
   virtual unsigned getEstimatedNumberOfCaseClusters(const SwitchInst &SI,
                                                     unsigned &JTSize) = 0;
   virtual int
   getUserCost(const User *U, ArrayRef<const Value *> Operands) = 0;
   virtual bool hasBranchDivergence() = 0;
   virtual bool isSourceOfDivergence(const Value *V) = 0;
   virtual bool isAlwaysUniform(const Value *V) = 0;
   virtual unsigned getFlatAddressSpace() = 0;
+  virtual unsigned getSwitchTableAddressSpace() const = 0;
+
   virtual bool isLoweredToCall(const Function *F) = 0;
   virtual void getUnrollingPreferences(Loop *L, ScalarEvolution &,
                                        UnrollingPreferences &UP) = 0;
   virtual bool isLegalAddImmediate(int64_t Imm) = 0;
   virtual bool isLegalICmpImmediate(int64_t Imm) = 0;
   virtual bool isLegalAddressingMode(Type *Ty, GlobalValue *BaseGV,
                                      int64_t BaseOffset, bool HasBaseReg,
                                      int64_t Scale,
public:
   bool isAlwaysUniform(const Value *V) override {
     return Impl.isAlwaysUniform(V);
   }
 
   unsigned getFlatAddressSpace() override {
     return Impl.getFlatAddressSpace();
   }
 
+  unsigned getSwitchTableAddressSpace() const override {
+    return Impl.getSwitchTableAddressSpace();
+  }
+
   bool isLoweredToCall(const Function *F) override {
     return Impl.isLoweredToCall(F);
   }
   void getUnrollingPreferences(Loop *L, ScalarEvolution &SE,
                                UnrollingPreferences &UP) override {
     return Impl.getUnrollingPreferences(L, SE, UP);
   }
   bool isLegalAddImmediate(int64_t Imm) override {

include/llvm/Analysis/TargetTransformInfoImpl.h

bool isConstantStridedAccessLessThan(ScalarEvolution *SE, const SCEV *Ptr,
     APInt StrideVal = Step->getAPInt();
     if (StrideVal.getBitWidth() > 64)
       return false;
     // FIXME: need to take absolute value for negtive stride case  
     return StrideVal.getSExtValue() < MergeDistance;
   }
 };
 
+/// A memory architecture where the program and the data live in
+/// separate address spaces.
+template<unsigned ProgramSpace = 1, unsigned DataSpace = 0>
+class HarvardArchitecture {
+public:
+  static unsigned getSwitchTableAddressSpace() { return ProgramSpace; }
+};
+
+/// A memory architecture where the program and the data live in the
+/// same address space.
+class VonNeumannArchitecture {
+public:
+  static unsigned getSwitchTableAddressSpace() { return 0; }
+};
+
 /// \brief CRTP base class for use as a mix-in that aids implementing
 /// a TargetTransformInfo-compatible class.
-template <typename T>
+template <typename T, typename MemArch>
 class TargetTransformInfoImplCRTPBase : public TargetTransformInfoImplBase {
 private:
   typedef TargetTransformInfoImplBase BaseT;
 
 protected:
   explicit TargetTransformInfoImplCRTPBase(const DataLayout &DL) : BaseT(DL) {}
 
 public:
if (const CastInst *CI = dyn_cast<CastInst>(U)) {
       if (isa<CmpInst>(CI->getOperand(0)))
         return TTI::TCC_Free;
     }
 
     return static_cast<T *>(this)->getOperationCost(
         Operator::getOpcode(U), U->getType(),
         U->getNumOperands() == 1 ? U->getOperand(0)->getType() : nullptr);
   }
+
+  unsigned getSwitchTableAddressSpace() const {
+    return MemArch::getSwitchTableAddressSpace();
+  }
 };
 }
 
 #endif

include/llvm/CodeGen/BasicTTIImpl.h

 #include "llvm/Support/CommandLine.h"
 #include "llvm/Target/TargetLowering.h"
 #include "llvm/Target/TargetSubtargetInfo.h"
 
 namespace llvm {
 
 extern cl::opt<unsigned> PartialUnrollingThreshold;
 
+/// The default memory architecture.
+typedef VonNeumannArchitecture DefaultMemArch;
+
 /// \brief Base class which can be used to help build a TTI implementation.
 ///
 /// This class provides as much implementation of the TTI interface as is
 /// possible using the target independent parts of the code generator.
 ///
 /// In order to subclass it, your class must implement a getST() method to
 /// return the subtarget, and a getTLI() method to return the target lowering.
 /// We need these methods implemented in the derived class so that this class
 /// doesn't have to duplicate storage for them.
-template <typename T>
-class BasicTTIImplBase : public TargetTransformInfoImplCRTPBase<T> {
+template <typename T, typename MemArch = DefaultMemArch>
+class BasicTTIImplBase : public TargetTransformInfoImplCRTPBase<T, MemArch> {
 private:
-  typedef TargetTransformInfoImplCRTPBase<T> BaseT;
+  typedef TargetTransformInfoImplCRTPBase<T, MemArch> BaseT;
   typedef TargetTransformInfo TTI;
 
   /// Estimate a cost of shuffle as a sequence of extract and insert
   /// operations.
   unsigned getPermuteShuffleOverhead(Type *Ty) {
     assert(Ty->isVectorTy() && "Can only shuffle vectors");
     unsigned Cost = 0;
     // Shuffle cost is equal to the cost of extracting element from its argument

 
   unsigned getVectorSplitCost() { return 1; }
 
   /// @}
 };
 
 /// \brief Concrete BasicTTIImpl that can be used if no further customization
 /// is needed.
-class BasicTTIImpl : public BasicTTIImplBase<BasicTTIImpl> {
-  typedef BasicTTIImplBase<BasicTTIImpl> BaseT;
-  friend class BasicTTIImplBase<BasicTTIImpl>;
+template<typename MemArch = DefaultMemArch>
+class BasicTTIImpl : public BasicTTIImplBase<BasicTTIImpl<MemArch>, MemArch> {
+  typedef BasicTTIImplBase<BasicTTIImpl, MemArch> BaseT;
+  friend class BasicTTIImplBase<BasicTTIImpl, MemArch>;
 
   const TargetSubtargetInfo *ST;
   const TargetLoweringBase *TLI;
 
   const TargetSubtargetInfo *getST() const { return ST; }
   const TargetLoweringBase *getTLI() const { return TLI; }
 
 public:
-  explicit BasicTTIImpl(const TargetMachine *ST, const Function &F);
+  explicit BasicTTIImpl(const TargetMachine *TM, const Function &F)
+   : BaseT(TM, F.getParent()->getDataLayout()), ST(TM->getSubtargetImpl(F)),
+     TLI(ST->getTargetLowering()) {}
 };
 
 }
 
 #endif

lib/Analysis/TargetTransformInfo.cpp

static cl::opt<bool> UseWideMemcpyLoopLowering(
     cl::desc("Enables the new wide memcpy loop lowering in Transforms/Utils."),
     cl::Hidden);
 
 namespace {
 /// \brief No-op implementation of the TTI interface using the utility base
 /// classes.
 ///
 /// This is used when no target specific information is available.
-struct NoTTIImpl : TargetTransformInfoImplCRTPBase<NoTTIImpl> {
+struct NoTTIImpl : TargetTransformInfoImplCRTPBase<NoTTIImpl, VonNeumannArchitecture> {
   explicit NoTTIImpl(const DataLayout &DL)
-      : TargetTransformInfoImplCRTPBase<NoTTIImpl>(DL) {}
+      : TargetTransformInfoImplCRTPBase<NoTTIImpl, VonNeumannArchitecture>(DL) {}
 };
 }
 
 TargetTransformInfo::TargetTransformInfo(const DataLayout &DL)
     : TTIImpl(new Model<NoTTIImpl>(NoTTIImpl(DL))) {}
 
 TargetTransformInfo::~TargetTransformInfo() {}
 

 bool llvm::TargetTransformInfo::isAlwaysUniform(const Value *V) const {
   return TTIImpl->isAlwaysUniform(V);
 }
 
 unsigned TargetTransformInfo::getFlatAddressSpace() const {
   return TTIImpl->getFlatAddressSpace();
 }
 
+unsigned TargetTransformInfo::getSwitchTableAddressSpace() const {
+  return TTIImpl->getSwitchTableAddressSpace();
+}
+
 bool TargetTransformInfo::isLoweredToCall(const Function *F) const {
   return TTIImpl->isLoweredToCall(F);
 }
 
 void TargetTransformInfo::getUnrollingPreferences(
     Loop *L, ScalarEvolution &SE, UnrollingPreferences &UP) const {
   return TTIImpl->getUnrollingPreferences(L, SE, UP);
 }

lib/CodeGen/BasicTargetTransformInfo.cpp

 
 // This flag is used by the template base class for BasicTTIImpl, and here to
 // provide a definition.
 cl::opt<unsigned>
     llvm::PartialUnrollingThreshold("partial-unrolling-threshold", cl::init(0),
                                     cl::desc("Threshold for partial unrolling"),
                                     cl::Hidden);
 
-BasicTTIImpl::BasicTTIImpl(const TargetMachine *TM, const Function &F)
-    : BaseT(TM, F.getParent()->getDataLayout()), ST(TM->getSubtargetImpl(F)),
-      TLI(ST->getTargetLowering()) {}

lib/CodeGen/LLVMTargetMachine.cpp

LLVMTargetMachine::LLVMTargetMachine(const Target &T,
   T.adjustCodeGenOpts(TT, RM, CM);
   this->RM = RM;
   this->CMModel = CM;
   this->OptLevel = OL;
 }
 
 TargetIRAnalysis LLVMTargetMachine::getTargetIRAnalysis() {
   return TargetIRAnalysis([this](const Function &F) {
-    return TargetTransformInfo(BasicTTIImpl(this, F));
+    return TargetTransformInfo(BasicTTIImpl<>(this, F));
   });
 }
 
 /// addPassesToX helper drives creation and initialization of TargetPassConfig.
 static MCContext *
 addPassesToGenerateCode(LLVMTargetMachine *TM, PassManagerBase &PM,
                         bool DisableVerify, AnalysisID StartBefore,
                         AnalysisID StartAfter, AnalysisID StopBefore,

lib/Target/AVR/AVRTargetMachine.h

 public:
   AVRTargetMachine(const Target &T, const Triple &TT, StringRef CPU,
                    StringRef FS, const TargetOptions &Options, Optional<Reloc::Model> RM,
                    CodeModel::Model CM, CodeGenOpt::Level OL);
 
   const AVRSubtarget *getSubtargetImpl() const;
   const AVRSubtarget *getSubtargetImpl(const Function &) const override;
 
+  TargetIRAnalysis getTargetIRAnalysis() override;
+
   TargetLoweringObjectFile *getObjFileLowering() const override {
     return this->TLOF.get();
   }
 
   TargetPassConfig *createPassConfig(PassManagerBase &PM) override;
 
   bool isMachineVerifierClean() const override {
     return false;

lib/Target/AVR/AVRTargetMachine.cpp

 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/TargetRegistry.h"
 
 #include "AVR.h"
 #include "AVRTargetObjectFile.h"
+#include "AVRTargetTransformInfo.h"
 #include "MCTargetDesc/AVRMCTargetDesc.h"
 
 namespace llvm {
 
 static const char *AVRDataLayout = "e-p:16:16:16-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-n8";
 
 /// Processes a CPU name.
 static StringRef getCPU(StringRef CPU) {
AVRTargetMachine::AVRTargetMachine(const Target &T, const Triple &TT,
     : LLVMTargetMachine(
           T, AVRDataLayout, TT,
           getCPU(CPU), FS, Options, getEffectiveRelocModel(RM), CM, OL),
       SubTarget(TT, getCPU(CPU), FS, *this) {
   this->TLOF = make_unique<AVRTargetObjectFile>();
   initAsmInfo();
 }
 
+TargetIRAnalysis AVRTargetMachine::getTargetIRAnalysis() {
+  return TargetIRAnalysis([this](const Function &F) {
+    return TargetTransformInfo(AVRTTIImpl(this, F));
+  });
+}
+
 namespace {
 /// AVR Code Generator Pass Configuration Options.
 class AVRPassConfig : public TargetPassConfig {
 public:
   AVRPassConfig(AVRTargetMachine &TM, PassManagerBase &PM)
       : TargetPassConfig(TM, PM) {}
 
   AVRTargetMachine &getAVRTargetMachine() const {

lib/Target/AVR/AVRTargetTransformInfo.h

+//===---- AVRTargetTransformInfo.h - AVR specific TTI -----------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file a TargetTransformInfo::Concept conforming object specific to the
+// AVR target machine. It uses the target's detailed information to
+// provide more precise answers to certain TTI queries, while letting the
+// target independent and default TTI implementations handle the rest.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_TARGET_AVR_AVRTARGETTRANSFORMINFO_H
+#define LLVM_LIB_TARGET_AVR_AVRTARGETTRANSFORMINFO_H
+
+#include "AVR.h"
+#include "AVRSubtarget.h"
+#include "AVRTargetMachine.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/CodeGen/BasicTTIImpl.h"
+#include "llvm/Target/TargetLowering.h"
+
+namespace llvm {
+class AVRTTIImpl : public BasicTTIImplBase<AVRTTIImpl, HarvardArchitecture<>> {
+  typedef BasicTTIImplBase<AVRTTIImpl, HarvardArchitecture<>> BaseT;
+  typedef TargetTransformInfo TTI;
+  friend BaseT;
+
+  const AVRSubtarget *ST;
+  const AVRTargetLowering *TLI;
+
+  const AVRSubtarget *getST() const { return ST; }
+  const AVRTargetLowering *getTLI() const { return TLI; }
+
+public:
+  explicit AVRTTIImpl(const AVRTargetMachine *TM, const Function &F)
+      : BaseT(TM, F.getParent()->getDataLayout()), ST(TM->getSubtargetImpl(F)),
+        TLI(ST->getTargetLowering()) {}
+
+  unsigned getSwitchTableAddressSpace() const { return AVR::ProgramMemory; }
+};
+
+} // end namespace llvm
+
+#endif // LLVM_LIB_TARGET_AVR_AVRTARGETTRANSFORMINFO_H

lib/Target/AVR/LLVMBuild.txt

 has_asmprinter = 1
 has_asmparser = 1
 has_disassembler = 1
 
 [component_1]
 type = Library
 name = AVRCodeGen
 parent = AVR
-required_libraries = AsmPrinter CodeGen Core MC AVRAsmPrinter AVRDesc AVRInfo SelectionDAG Support Target
+required_libraries =
+ Analysis
+ AsmPrinter
+ CodeGen
+ Core
+ MC
+ AVRAsmPrinter
+ AVRDesc
+ AVRInfo
+ SelectionDAG
+ Support
+ Target
 add_to_library_groups = AVR
 

lib/Target/Mips/MipsTargetMachine.cpp

TargetIRAnalysis MipsTargetMachine::getTargetIRAnalysis() {
   return TargetIRAnalysis([this](const Function &F) {
     if (Subtarget->allowMixed16_32()) {
       DEBUG(errs() << "No Target Transform Info Pass Added\n");
       // FIXME: This is no longer necessary as the TTI returned is per-function.
       return TargetTransformInfo(F.getParent()->getDataLayout());
     }
 
     DEBUG(errs() << "Target Transform Info Pass Added\n");
-    return TargetTransformInfo(BasicTTIImpl(this, F));
+    return TargetTransformInfo(BasicTTIImpl<>(this, F));
   });
 }
 
 // Implemented by targets that want to run passes immediately before
 // machine code is emitted. return true if -print-machineinstrs should
 // print out the code after the passes.
 void MipsPassConfig::addPreEmitPass() {
   addPass(createMicroMipsSizeReductionPass());

lib/Transforms/Utils/SimplifyCFG.cpp

 /// This class represents a lookup table that can be used to replace a switch.
 class SwitchLookupTable {
 public:
   /// Create a lookup table to use as a switch replacement with the contents
   /// of Values, using DefaultValue to fill any holes in the table.
   SwitchLookupTable(
       Module &M, uint64_t TableSize, ConstantInt *Offset,
       const SmallVectorImpl<std::pair<ConstantInt *, Constant *>> &Values,
-      Constant *DefaultValue, const DataLayout &DL, const StringRef &FuncName);
+      Constant *DefaultValue, const DataLayout &DL, const StringRef &FuncName,
+      const TargetTransformInfo &TTI);
 
   /// Build instructions with Builder to retrieve the value at
   /// the position given by Index in the lookup table.
   Value *BuildLookup(Value *Index, IRBuilder<> &Builder);
 
   /// Return true if a table with TableSize elements of
   /// type ElementType would fit in a target-legal register.
   static bool WouldFitInRegister(const DataLayout &DL, uint64_t TableSize,
private:
   GlobalVariable *Array;
 };
 
 } // end anonymous namespace
 
 SwitchLookupTable::SwitchLookupTable(
     Module &M, uint64_t TableSize, ConstantInt *Offset,
     const SmallVectorImpl<std::pair<ConstantInt *, Constant *>> &Values,
-    Constant *DefaultValue, const DataLayout &DL, const StringRef &FuncName)
+    Constant *DefaultValue, const DataLayout &DL, const StringRef &FuncName,
+    const TargetTransformInfo &TTI)
     : SingleValue(nullptr), BitMap(nullptr), BitMapElementTy(nullptr),
       LinearOffset(nullptr), LinearMultiplier(nullptr), Array(nullptr) {
   assert(Values.size() && "Can't build lookup table without values!");
   assert(TableSize >= Values.size() && "Can't fit values in table!");
 
   // If all values in the table are equal, this is that value.
   SingleValue = Values.begin()->second;
 
SwitchLookupTable::SwitchLookupTable(
   }
 
   // Store the table in an array.
   ArrayType *ArrayTy = ArrayType::get(ValueType, TableSize);
   Constant *Initializer = ConstantArray::get(ArrayTy, TableContents);
 
   Array = new GlobalVariable(M, ArrayTy, /*constant=*/true,
                              GlobalVariable::PrivateLinkage, Initializer,
-                             "switch.table." + FuncName);
+                             "switch.table." + FuncName,
+                             /*InsertBefore=*/nullptr,
+                             GlobalValue::NotThreadLocal,
+                             TTI.getSwitchTableAddressSpace());
   Array->setUnnamedAddr(GlobalValue::UnnamedAddr::Global);
   Kind = ArrayKind;
 }
 
 Value *SwitchLookupTable::BuildLookup(Value *Index, IRBuilder<> &Builder) {
   switch (Kind) {
   case SingleValueKind:
     return SingleValue;
static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
   for (size_t I = 0, E = PHIs.size(); I != E; ++I) {
     PHINode *PHI = PHIs[I];
     const ResultListTy &ResultList = ResultLists[PHI];
 
     // If using a bitmask, use any value to fill the lookup table holes.
     Constant *DV = NeedMask ? ResultLists[PHI][0].second : DefaultResults[PHI];
     StringRef FuncName = SI->getParent()->getParent()->getName();
     SwitchLookupTable Table(Mod, TableSize, MinCaseVal, ResultList, DV, DL,
-                            FuncName);
+                            FuncName, TTI);
 
     Value *Result = Table.BuildLookup(TableIndex, Builder);
 
     // If the result is used to return immediately from the function, we want to
     // do that right here.
     if (PHI->hasOneUse() && isa<ReturnInst>(*PHI->user_begin()) &&
         PHI->user_back() == CommonDest->getFirstNonPHIOrDbg()) {
       Builder.CreateRet(Result);

test/CodeGen/AVR/switch-lookup-table-in-progmem-space.ll

+; RUN: opt -S -mtriple=avr -latesimplifycfg < %s | FileCheck %s
+
+; Switch lookup tables should be placed in program memory (address space 1).
+
+%MyType = type { i8, [0 x i8], [0 x i8] }
+
+define i8 @foo(i8) {
+start:
+  %_0 = alloca %MyType
+
+; CHECK: @switch.table.foo = {{.*}}addrspace(1) constant [15 x i8] c"\01\01\01\01\01\00\00\00\00\00\00\00\00\00\01"
+  switch i8 %0, label %bb7 [
+    i8 0, label %bb1
+    i8 1, label %bb2
+    i8 2, label %bb3
+    i8 3, label %bb4
+    i8 4, label %bb5
+    i8 14, label %bb6
+  ]
+
+bb1:
+  %1 = getelementptr inbounds %MyType, %MyType* %_0, i32 0, i32 0
+  store i8 1, i8* %1
+  br label %bb8
+
+bb2:
+  %2 = getelementptr inbounds %MyType, %MyType* %_0, i32 0, i32 0
+  store i8 1, i8* %2
+  br label %bb8
+
+bb3:
+  %3 = getelementptr inbounds %MyType, %MyType* %_0, i32 0, i32 0
+  store i8 1, i8* %3
+  br label %bb8
+
+bb4:
+  %4 = getelementptr inbounds %MyType, %MyType* %_0, i32 0, i32 0
+  store i8 1, i8* %4
+  br label %bb8
+
+bb5:
+  %5 = getelementptr inbounds %MyType, %MyType* %_0, i32 0, i32 0
+  store i8 1, i8* %5
+  br label %bb8
+
+bb6:
+  %6 = getelementptr inbounds %MyType, %MyType* %_0, i32 0, i32 0
+  store i8 1, i8* %6
+  br label %bb8
+
+bb7:
+  %7 = getelementptr inbounds %MyType, %MyType* %_0, i32 0, i32 0
+  store i8 0, i8* %7
+  br label %bb8
+
+bb8:
+  %8 = bitcast %MyType* %_0 to i8*
+  %9 = load i8, i8* %8, align 1
+  ret i8 %9
+}
+