[SandboxIR] Implement AllocaInst (#102027)

This patch implements sandboxir::AllocaInst which mirrors
llvm::AllocaInst.

Author: vporpo

llvm/include/llvm/SandboxIR/SandboxIR.h

@@ -124,6 +124,7 @@ class GetElementPtrInst;
 class CastInst;
 class PtrToIntInst;
 class BitCastInst;
+class AllocaInst;
 
 /// Iterator for the `Use` edges of a User's operands.
 /// \Returns the operand `Use` when dereferenced.
@@ -240,6 +241,7 @@ class Value {
   friend class InvokeInst;        // For getting `Val`.
   friend class CallBrInst;        // For getting `Val`.
   friend class GetElementPtrInst; // For getting `Val`.
+  friend class AllocaInst;        // For getting `Val`.
   friend class CastInst;          // For getting `Val`.
   friend class PHINode;           // For getting `Val`.
 
@@ -633,6 +635,7 @@ class Instruction : public sandboxir::User {
   friend class InvokeInst;        // For getTopmostLLVMInstruction().
   friend class CallBrInst;        // For getTopmostLLVMInstruction().
   friend class GetElementPtrInst; // For getTopmostLLVMInstruction().
+  friend class AllocaInst;        // For getTopmostLLVMInstruction().
   friend class CastInst;          // For getTopmostLLVMInstruction().
   friend class PHINode;           // For getTopmostLLVMInstruction().
 
@@ -1393,6 +1396,103 @@ class GetElementPtrInst final : public Instruction {
 #endif
 };
 
+class AllocaInst final : public UnaryInstruction {
+  Use getOperandUseInternal(unsigned OpIdx, bool Verify) const final {
+    return getOperandUseDefault(OpIdx, Verify);
+  }
+  SmallVector<llvm::Instruction *, 1> getLLVMInstrs() const final {
+    return {cast<llvm::Instruction>(Val)};
+  }
+
+  AllocaInst(llvm::AllocaInst *AI, Context &Ctx)
+      : UnaryInstruction(ClassID::Alloca, Instruction::Opcode::Alloca, AI,
+                         Ctx) {}
+  friend class Context; // For constructor.
+
+public:
+  static AllocaInst *create(Type *Ty, unsigned AddrSpace, BBIterator WhereIt,
+                            BasicBlock *WhereBB, Context &Ctx,
+                            Value *ArraySize = nullptr, const Twine &Name = "");
+  static AllocaInst *create(Type *Ty, unsigned AddrSpace,
+                            Instruction *InsertBefore, Context &Ctx,
+                            Value *ArraySize = nullptr, const Twine &Name = "");
+  static AllocaInst *create(Type *Ty, unsigned AddrSpace,
+                            BasicBlock *InsertAtEnd, Context &Ctx,
+                            Value *ArraySize = nullptr, const Twine &Name = "");
+
+  unsigned getUseOperandNo(const Use &Use) const final {
+    return getUseOperandNoDefault(Use);
+  }
+  unsigned getNumOfIRInstrs() const final { return 1u; }
+
+  /// Return true if there is an allocation size parameter to the allocation
+  /// instruction that is not 1.
+  bool isArrayAllocation() const {
+    return cast<llvm::AllocaInst>(Val)->isArrayAllocation();
+  }
+  /// Get the number of elements allocated. For a simple allocation of a single
+  /// element, this will return a constant 1 value.
+  Value *getArraySize();
+  const Value *getArraySize() const {
+    return const_cast<AllocaInst *>(this)->getArraySize();
+  }
+  /// Overload to return most specific pointer type.
+  PointerType *getType() const {
+    return cast<llvm::AllocaInst>(Val)->getType();
+  }
+  /// Return the address space for the allocation.
+  unsigned getAddressSpace() const {
+    return cast<llvm::AllocaInst>(Val)->getAddressSpace();
+  }
+  /// Get allocation size in bytes. Returns std::nullopt if size can't be
+  /// determined, e.g. in case of a VLA.
+  std::optional<TypeSize> getAllocationSize(const DataLayout &DL) const {
+    return cast<llvm::AllocaInst>(Val)->getAllocationSize(DL);
+  }
+  /// Get allocation size in bits. Returns std::nullopt if size can't be
+  /// determined, e.g. in case of a VLA.
+  std::optional<TypeSize> getAllocationSizeInBits(const DataLayout &DL) const {
+    return cast<llvm::AllocaInst>(Val)->getAllocationSizeInBits(DL);
+  }
+  /// Return the type that is being allocated by the instruction.
+  Type *getAllocatedType() const {
+    return cast<llvm::AllocaInst>(Val)->getAllocatedType();
+  }
+  /// for use only in special circumstances that need to generically
+  /// transform a whole instruction (eg: IR linking and vectorization).
+  void setAllocatedType(Type *Ty);
+  /// Return the alignment of the memory that is being allocated by the
+  /// instruction.
+  Align getAlign() const { return cast<llvm::AllocaInst>(Val)->getAlign(); }
+  void setAlignment(Align Align);
+  /// Return true if this alloca is in the entry block of the function and is a
+  /// constant size. If so, the code generator will fold it into the
+  /// prolog/epilog code, so it is basically free.
+  bool isStaticAlloca() const {
+    return cast<llvm::AllocaInst>(Val)->isStaticAlloca();
+  }
+  /// Return true if this alloca is used as an inalloca argument to a call. Such
+  /// allocas are never considered static even if they are in the entry block.
+  bool isUsedWithInAlloca() const {
+    return cast<llvm::AllocaInst>(Val)->isUsedWithInAlloca();
+  }
+  /// Specify whether this alloca is used to represent the arguments to a call.
+  void setUsedWithInAlloca(bool V);
+
+  static bool classof(const Value *From) {
+    if (auto *I = dyn_cast<Instruction>(From))
+      return I->getSubclassID() == Instruction::ClassID::Alloca;
+    return false;
+  }
+#ifndef NDEBUG
+  void verify() const final {
+    assert(isa<llvm::AllocaInst>(Val) && "Expected AllocaInst!");
+  }
+  void dump(raw_ostream &OS) const override;
+  LLVM_DUMP_METHOD void dump() const override;
+#endif
+};
+
 class CastInst : public UnaryInstruction {
   static Opcode getCastOpcode(llvm::Instruction::CastOps CastOp) {
     switch (CastOp) {
@@ -1726,6 +1826,8 @@ class Context {
   friend CallBrInst; // For createCallBrInst()
   GetElementPtrInst *createGetElementPtrInst(llvm::GetElementPtrInst *I);
   friend GetElementPtrInst; // For createGetElementPtrInst()
+  AllocaInst *createAllocaInst(llvm::AllocaInst *I);
+  friend AllocaInst; // For createAllocaInst()
   CastInst *createCastInst(llvm::CastInst *I);
   friend CastInst; // For createCastInst()
   PHINode *createPHINode(llvm::PHINode *I);

llvm/include/llvm/SandboxIR/SandboxIRValues.def

@@ -43,6 +43,7 @@ DEF_INSTR(Call,          OP(Call),          CallInst)
 DEF_INSTR(Invoke,        OP(Invoke),        InvokeInst)
 DEF_INSTR(CallBr,        OP(CallBr),        CallBrInst)
 DEF_INSTR(GetElementPtr, OP(GetElementPtr), GetElementPtrInst)
+DEF_INSTR(Alloca,        OP(Alloca),        AllocaInst)
 DEF_INSTR(Cast,  OPCODES(\
                          OP(ZExt)          \
                          OP(SExt)          \

llvm/include/llvm/SandboxIR/Tracker.h

@@ -58,6 +58,7 @@ class LoadInst;
 class StoreInst;
 class Instruction;
 class Tracker;
+class AllocaInst;
 
 /// The base class for IR Change classes.
 class IRChangeBase {
@@ -255,6 +256,57 @@ class CallBrInstSetDefaultDest : public IRChangeBase {
 #endif
 };
 
+class AllocaSetAllocatedType final : public IRChangeBase {
+  AllocaInst *Alloca;
+  Type *OrigType;
+
+public:
+  AllocaSetAllocatedType(AllocaInst *Alloca, Tracker &Tracker);
+  void revert() final;
+  void accept() final {}
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const final {
+    dumpCommon(OS);
+    OS << "AllocaSetAllocatedType";
+  }
+  LLVM_DUMP_METHOD void dump() const final;
+#endif
+};
+
+class AllocaSetAlignment final : public IRChangeBase {
+  AllocaInst *Alloca;
+  Align OrigAlign;
+
+public:
+  AllocaSetAlignment(AllocaInst *Alloca, Tracker &Tracker);
+  void revert() final;
+  void accept() final {}
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const final {
+    dumpCommon(OS);
+    OS << "AllocaSetAlignment";
+  }
+  LLVM_DUMP_METHOD void dump() const final;
+#endif
+};
+
+class AllocaSetUsedWithInAlloca final : public IRChangeBase {
+  AllocaInst *Alloca;
+  bool Orig;
+
+public:
+  AllocaSetUsedWithInAlloca(AllocaInst *Alloca, Tracker &Tracker);
+  void revert() final;
+  void accept() final {}
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const final {
+    dumpCommon(OS);
+    OS << "AllocaSetUsedWithInAlloca";
+  }
+  LLVM_DUMP_METHOD void dump() const final;
+#endif
+};
+
 class CallBrInstSetIndirectDest : public IRChangeBase {
   CallBrInst *CallBr;
   unsigned Idx;

llvm/lib/SandboxIR/SandboxIR.cpp

@@ -1232,6 +1232,69 @@ static llvm::Instruction::CastOps getLLVMCastOp(Instruction::Opcode Opc) {
   }
 }
 
+AllocaInst *AllocaInst::create(Type *Ty, unsigned AddrSpace, BBIterator WhereIt,
+                               BasicBlock *WhereBB, Context &Ctx,
+                               Value *ArraySize, const Twine &Name) {
+  auto &Builder = Ctx.getLLVMIRBuilder();
+  if (WhereIt == WhereBB->end())
+    Builder.SetInsertPoint(cast<llvm::BasicBlock>(WhereBB->Val));
+  else
+    Builder.SetInsertPoint((*WhereIt).getTopmostLLVMInstruction());
+  auto *NewAlloca = Builder.CreateAlloca(Ty, AddrSpace, ArraySize->Val, Name);
+  return Ctx.createAllocaInst(NewAlloca);
+}
+
+AllocaInst *AllocaInst::create(Type *Ty, unsigned AddrSpace,
+                               Instruction *InsertBefore, Context &Ctx,
+                               Value *ArraySize, const Twine &Name) {
+  return create(Ty, AddrSpace, InsertBefore->getIterator(),
+                InsertBefore->getParent(), Ctx, ArraySize, Name);
+}
+
+AllocaInst *AllocaInst::create(Type *Ty, unsigned AddrSpace,
+                               BasicBlock *InsertAtEnd, Context &Ctx,
+                               Value *ArraySize, const Twine &Name) {
+  return create(Ty, AddrSpace, InsertAtEnd->end(), InsertAtEnd, Ctx, ArraySize,
+                Name);
+}
+
+void AllocaInst::setAllocatedType(Type *Ty) {
+  auto &Tracker = Ctx.getTracker();
+  if (Tracker.isTracking())
+    Tracker.track(std::make_unique<AllocaSetAllocatedType>(this, Tracker));
+  cast<llvm::AllocaInst>(Val)->setAllocatedType(Ty);
+}
+
+void AllocaInst::setAlignment(Align Align) {
+  auto &Tracker = Ctx.getTracker();
+  if (Tracker.isTracking())
+    Tracker.track(std::make_unique<AllocaSetAlignment>(this, Tracker));
+  cast<llvm::AllocaInst>(Val)->setAlignment(Align);
+}
+
+void AllocaInst::setUsedWithInAlloca(bool V) {
+  auto &Tracker = Ctx.getTracker();
+  if (Tracker.isTracking())
+    Tracker.track(std::make_unique<AllocaSetUsedWithInAlloca>(this, Tracker));
+  cast<llvm::AllocaInst>(Val)->setUsedWithInAlloca(V);
+}
+
+Value *AllocaInst::getArraySize() {
+  return Ctx.getValue(cast<llvm::AllocaInst>(Val)->getArraySize());
+}
+
+#ifndef NDEBUG
+void AllocaInst::dump(raw_ostream &OS) const {
+  dumpCommonPrefix(OS);
+  dumpCommonSuffix(OS);
+}
+
+void AllocaInst::dump() const {
+  dump(dbgs());
+  dbgs() << "\n";
+}
+#endif // NDEBUG
+
 Value *CastInst::create(Type *DestTy, Opcode Op, Value *Operand,
                         BBIterator WhereIt, BasicBlock *WhereBB, Context &Ctx,
                         const Twine &Name) {
@@ -1472,6 +1535,11 @@ Value *Context::getOrCreateValueInternal(llvm::Value *LLVMV, llvm::User *U) {
         new GetElementPtrInst(LLVMGEP, *this));
     return It->second.get();
   }
+  case llvm::Instruction::Alloca: {
+    auto *LLVMAlloca = cast<llvm::AllocaInst>(LLVMV);
+    It->second = std::unique_ptr<AllocaInst>(new AllocaInst(LLVMAlloca, *this));
+    return It->second.get();
+  }
   case llvm::Instruction::ZExt:
   case llvm::Instruction::SExt:
   case llvm::Instruction::FPToUI:
@@ -1558,7 +1626,10 @@ Context::createGetElementPtrInst(llvm::GetElementPtrInst *I) {
       std::unique_ptr<GetElementPtrInst>(new GetElementPtrInst(I, *this));
   return cast<GetElementPtrInst>(registerValue(std::move(NewPtr)));
 }
-
+AllocaInst *Context::createAllocaInst(llvm::AllocaInst *I) {
+  auto NewPtr = std::unique_ptr<AllocaInst>(new AllocaInst(I, *this));
+  return cast<AllocaInst>(registerValue(std::move(NewPtr)));
+}
 CastInst *Context::createCastInst(llvm::CastInst *I) {
   auto NewPtr = std::unique_ptr<CastInst>(new CastInst(I, *this));
   return cast<CastInst>(registerValue(std::move(NewPtr)));

llvm/lib/SandboxIR/Tracker.cpp

@@ -204,6 +204,46 @@ void RemoveFromParent::dump() const {
 }
 #endif
 
+AllocaSetAllocatedType::AllocaSetAllocatedType(AllocaInst *Alloca,
+                                               Tracker &Tracker)
+    : IRChangeBase(Tracker), Alloca(Alloca),
+      OrigType(Alloca->getAllocatedType()) {}
+
+void AllocaSetAllocatedType::revert() { Alloca->setAllocatedType(OrigType); }
+
+#ifndef NDEBUG
+void AllocaSetAllocatedType::dump() const {
+  dump(dbgs());
+  dbgs() << "\n";
+}
+#endif // NDEBUG
+
+AllocaSetAlignment::AllocaSetAlignment(AllocaInst *Alloca, Tracker &Tracker)
+    : IRChangeBase(Tracker), Alloca(Alloca), OrigAlign(Alloca->getAlign()) {}
+
+void AllocaSetAlignment::revert() { Alloca->setAlignment(OrigAlign); }
+
+#ifndef NDEBUG
+void AllocaSetAlignment::dump() const {
+  dump(dbgs());
+  dbgs() << "\n";
+}
+#endif // NDEBUG
+
+AllocaSetUsedWithInAlloca::AllocaSetUsedWithInAlloca(AllocaInst *Alloca,
+                                                     Tracker &Tracker)
+    : IRChangeBase(Tracker), Alloca(Alloca),
+      Orig(Alloca->isUsedWithInAlloca()) {}
+
+void AllocaSetUsedWithInAlloca::revert() { Alloca->setUsedWithInAlloca(Orig); }
+
+#ifndef NDEBUG
+void AllocaSetUsedWithInAlloca::dump() const {
+  dump(dbgs());
+  dbgs() << "\n";
+}
+#endif // NDEBUG
+
 CallBrInstSetDefaultDest::CallBrInstSetDefaultDest(CallBrInst *CallBr,
                                                    Tracker &Tracker)
     : IRChangeBase(Tracker), CallBr(CallBr) {