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

public:
   bool isSourceOfDivergence(const Value *V) { return false; }
 
   bool isAlwaysUniform(const Value *V) { return false; }
 
   unsigned getFlatAddressSpace () {
     return -1;
   }
 
+  unsigned getSwitchTableAddressSpace() const { return 0; }
+
   bool isLoweredToCall(const Function *F) {
     // FIXME: These should almost certainly not be handled here, and instead
     // handled with the help of TLI or the target itself. This was largely
     // ported from existing analysis heuristics here so that such refactorings
     // can take place in the future.
 
     if (F->isIntrinsic())
       return false;

include/llvm/CodeGen/BasicTTIImpl.h

public:
 
   bool isAlwaysUniform(const Value *V) { return false; }
 
   unsigned getFlatAddressSpace() {
     // Return an invalid address space.
     return -1;
   }
 
+  unsigned getSwitchTableAddressSpace() const { return 0; }
+
   bool isLegalAddImmediate(int64_t imm) {
     return getTLI()->isLegalAddImmediate(imm);
   }
 
   bool isLegalICmpImmediate(int64_t imm) {
     return getTLI()->isLegalICmpImmediate(imm);
   }
 

lib/Analysis/TargetTransformInfo.cpp

 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/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> {
+  typedef BasicTTIImplBase<AVRTTIImpl> 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/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
+}
+