PNaCl: Add support for GCC/LLVM vector extensions

lib/Transforms/NaCl/ConstantInsertExtractElementIndex.cpp

+//===- ConstantInsertExtractElementIndex.cpp - Insert/Extract element -----===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// Transform all InsertElement and ExtractElement with non-constant or
+// out-of-bounds indices into stack accesses. This moves all undefined
+// behavior to the stack, making InsertElement and ExtractElement
+// well-defined.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/IR/DataLayout.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Pass.h"
+#include "llvm/Transforms/NaCl.h"
+
+#include <algorithm>
+
+using namespace llvm;
+
+namespace {
+class ConstantInsertExtractElementIndex : public BasicBlockPass {
+public:
+  static char ID; // Pass identification, replacement for typeid
+  ConstantInsertExtractElementIndex() : BasicBlockPass(ID), M(0), TD(0) {
+    initializeConstantInsertExtractElementIndexPass(
+        *PassRegistry::getPassRegistry());
+  }
+  virtual void getAnalysisUsage(AnalysisUsage &AU) const {
+    AU.addRequired<DataLayout>();
+    BasicBlockPass::getAnalysisUsage(AU);
+  }
+  virtual bool doInitialization(Module &Mod) {
+    M = &Mod;
+    return false; // Unchanged.
+  }
+  virtual bool runOnBasicBlock(BasicBlock &BB);
+
+private:
+  typedef SmallVector<Instruction *, 8> Instructions;
+  const Module *M;
+  const DataLayout *TD;
+
+  /// Get the index of an InsertElement or ExtractElement instruction, or null.
+  static Value *getIdx(const Instruction *I);
+
+  void findInvalidInsertExtractElements(const BasicBlock &BB,
+                                        Instructions &BadVectorIdx) const;
+  void fixInvalidInsertExtractElements(BasicBlock &BB,
+                                       const Instructions &BadVectorIdx) const;
+};
+}
+
+char ConstantInsertExtractElementIndex::ID = 0;
+INITIALIZE_PASS(
+    ConstantInsertExtractElementIndex, "constant-insert-extract-element-index",
+    "Force insert and extract vector elemement to always be in bounds", false,
+    false)
+
+Value *ConstantInsertExtractElementIndex::getIdx(const Instruction *I) {
+  switch (I->getOpcode()) {
+  default: return 0;
+  case Instruction::InsertElement: return I->getOperand(2);
+  case Instruction::ExtractElement: return I->getOperand(1);
+  }
+}
+
+void ConstantInsertExtractElementIndex::findInvalidInsertExtractElements(
+    const BasicBlock &BB, Instructions &BadVectorIdx) const {
+  for (BasicBlock::const_iterator BBI = BB.begin(), BBE = BB.end(); BBI != BBE;
+       ++BBI) {
+    const Instruction *I = &*BBI;
+    Value *Idx = getIdx(I);
+    if (!Idx)
+      // Not and InsertElement or ExtractElement, ignore the instruction.
+      continue;
+    if (ConstantInt *CI = dyn_cast<ConstantInt>(Idx)) {
+      unsigned NumElements =
+          cast<VectorType>(I->getOperand(0)->getType())->getNumElements();
+      if (CI->getValue().ult(NumElements))
+        // Index is constant in-range for the vector type, keep as-is.
+        continue;
+    }
+    BadVectorIdx.push_back(const_cast<Instruction *>(I));
+  }
+}
+
+void ConstantInsertExtractElementIndex::fixInvalidInsertExtractElements(
+    BasicBlock &BB, const Instructions &BadVectorIdx) const {
+  for (Instructions::const_iterator IB = BadVectorIdx.begin(),
+                                    IE = BadVectorIdx.end();
+       IB != IE; ++IB) {
+    Instruction *I = *IB;
+    Value *Vec = I->getOperand(0);
+    Value *Idx = getIdx(I);
+    VectorType *VecTy = cast<VectorType>(Vec->getType());
+    Type *ElemTy = VecTy->getElementType();
+    unsigned ElemAlign = TD->getPrefTypeAlignment(ElemTy);
+    unsigned VecAlign = std::max(ElemAlign, TD->getPrefTypeAlignment(VecTy));
+
+    IRBuilder<> IRB(I);
+    AllocaInst *Alloca = IRB.CreateAlloca(
+        ElemTy, ConstantInt::get(Type::getInt32Ty(M->getContext()),
+                                 VecTy->getNumElements()));
+    Alloca->setAlignment(VecAlign);
+    Value *AllocaAsVec = IRB.CreateBitCast(Alloca, VecTy->getPointerTo());
+    IRB.CreateAlignedStore(Vec, AllocaAsVec, Alloca->getAlignment());
+    Value *GEP = IRB.CreateGEP(Alloca, Idx);
+
+    Value *Res;
+    switch (I->getOpcode()) {
+    default:
+      llvm_unreachable("expected InsertElement or ExtractElement");
+    case Instruction::InsertElement: {
+      IRB.CreateAlignedStore(I->getOperand(1), GEP, ElemAlign);
+      Res = IRB.CreateAlignedLoad(AllocaAsVec, Alloca->getAlignment());
+      break;
+    }
+    case Instruction::ExtractElement: {
+      Res = IRB.CreateAlignedLoad(GEP, ElemAlign);
+      break;
+    }
+    }
+
+    I->replaceAllUsesWith(Res);
+    I->eraseFromParent();
+  }
+}
+
+bool ConstantInsertExtractElementIndex::runOnBasicBlock(BasicBlock &BB) {
+  bool Changed = false;
+  if (!TD)
+    TD = &getAnalysis<DataLayout>();
+  Instructions BadVectorIdx;
+  findInvalidInsertExtractElements(BB, BadVectorIdx);
+  if (!BadVectorIdx.empty()) {
+    Changed = true;
+    fixInvalidInsertExtractElements(BB, BadVectorIdx);
+  }
+  return Changed;
+}
+
+BasicBlockPass *llvm::createConstantInsertExtractElementIndexPass() {
+  return new ConstantInsertExtractElementIndex();
+}