Recommit "Add a heap alloc site marker field to the ExtraInfo in MachineInstrs"

Summary:
Fixes some things from original commit at https://reviews.llvm.org/D69136. The main
change is that the heap alloc marker is always stored as ExtraInfo in the machine
instruction instead of in the PointerSumType because it cannot hold more than
4 pointer types.

Add instruction marker to MachineInstr ExtraInfo. This does almost the
same thing as Pre/PostInstrSymbols, except that it doesn't create a label until
printing instructions. This allows for labels to be put around instructions that
are deleted/duplicated somewhere.
Use this marker to track heap alloc site call instructions.

Reviewers: rnk

Subscribers: MatzeB, hiraditya, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D69536

1 files changed, 140 insertions, 0 deletions

diff --git a/llvm/unittests/CodeGen/MachineInstrTest.cpp b/llvm/unittests/CodeGen/MachineInstrTest.cpp
index d48297734e1..9e455548e15 100644
--- a/llvm/unittests/CodeGen/MachineInstrTest.cpp
+++ b/llvm/unittests/CodeGen/MachineInstrTest.cpp
@@ -9,6 +9,7 @@
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineMemOperand.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/TargetFrameLowering.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
@@ -16,6 +17,8 @@
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/ModuleSlotTracker.h"
+#include "llvm/MC/MCAsmInfo.h"
+#include "llvm/MC/MCSymbol.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Target/TargetMachine.h"
@@ -125,6 +128,7 @@ public:
       : LLVMTargetMachine(Target(), "", Triple(""), "", "", TargetOptions(),
                           Reloc::Static, CodeModel::Small, CodeGenOpt::Default),
         ST(*this) {}
+
   ~BogusTargetMachine() override {}
 
   const TargetSubtargetInfo *getSubtargetImpl(const Function &) const override {
@@ -135,6 +139,13 @@ private:
   BogusSubtarget ST;
 };
 
+static MCAsmInfo AsmInfo = MCAsmInfo();
+
+std::unique_ptr<MCContext> createMCContext() {
+  return std::make_unique<MCContext>(
+      &AsmInfo, nullptr, nullptr, nullptr, nullptr, false);
+}
+
 std::unique_ptr<BogusTargetMachine> createTargetMachine() {
   return std::make_unique<BogusTargetMachine>();
 }
@@ -361,6 +372,135 @@ TEST(MachineInstrSpan, DistanceEnd) {
   ASSERT_TRUE(std::distance(MIS.begin(), MII) == 1);
 }
 
+TEST(MachineInstrExtraInfo, AddExtraInfo) {
+  auto MF = createMachineFunction();
+  MCInstrDesc MCID = {0, 0,       0,       0,       0, 0,
+                      0, nullptr, nullptr, nullptr, 0, nullptr};
+
+  auto MI = MF->CreateMachineInstr(MCID, DebugLoc());
+  auto MC = createMCContext();
+  auto MMO = MF->getMachineMemOperand(MachinePointerInfo(),
+                                      MachineMemOperand::MOLoad, 8, 8);
+  SmallVector<MachineMemOperand *, 2> MMOs;
+  MMOs.push_back(MMO);
+  MCSymbol *Sym1 = MC->createTempSymbol("pre_label", false);
+  MCSymbol *Sym2 = MC->createTempSymbol("post_label", false);
+  LLVMContext Ctx;
+  MDNode *MDN = MDNode::getDistinct(Ctx, None);
+
+  ASSERT_TRUE(MI->memoperands_empty());
+  ASSERT_FALSE(MI->getPreInstrSymbol());
+  ASSERT_FALSE(MI->getPostInstrSymbol());
+  ASSERT_FALSE(MI->getHeapAllocMarker());
+
+  MI->setMemRefs(*MF, MMOs);
+  ASSERT_TRUE(MI->memoperands().size() == 1);
+  ASSERT_FALSE(MI->getPreInstrSymbol());
+  ASSERT_FALSE(MI->getPostInstrSymbol());
+  ASSERT_FALSE(MI->getHeapAllocMarker());
+
+  MI->setPreInstrSymbol(*MF, Sym1);
+  ASSERT_TRUE(MI->memoperands().size() == 1);
+  ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
+  ASSERT_FALSE(MI->getPostInstrSymbol());
+  ASSERT_FALSE(MI->getHeapAllocMarker());
+
+  MI->setPostInstrSymbol(*MF, Sym2);
+  ASSERT_TRUE(MI->memoperands().size() == 1);
+  ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
+  ASSERT_TRUE(MI->getPostInstrSymbol() == Sym2);
+  ASSERT_FALSE(MI->getHeapAllocMarker());
+
+  MI->setHeapAllocMarker(*MF, MDN);
+  ASSERT_TRUE(MI->memoperands().size() == 1);
+  ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
+  ASSERT_TRUE(MI->getPostInstrSymbol() == Sym2);
+  ASSERT_TRUE(MI->getHeapAllocMarker() == MDN);
+}
+
+TEST(MachineInstrExtraInfo, ChangeExtraInfo) {
+  auto MF = createMachineFunction();
+  MCInstrDesc MCID = {0, 0,       0,       0,       0, 0,
+                      0, nullptr, nullptr, nullptr, 0, nullptr};
+
+  auto MI = MF->CreateMachineInstr(MCID, DebugLoc());
+  auto MC = createMCContext();
+  auto MMO = MF->getMachineMemOperand(MachinePointerInfo(),
+                                      MachineMemOperand::MOLoad, 8, 8);
+  SmallVector<MachineMemOperand *, 2> MMOs;
+  MMOs.push_back(MMO);
+  MCSymbol *Sym1 = MC->createTempSymbol("pre_label", false);
+  MCSymbol *Sym2 = MC->createTempSymbol("post_label", false);
+  LLVMContext Ctx;
+  MDNode *MDN = MDNode::getDistinct(Ctx, None);
+
+  MI->setMemRefs(*MF, MMOs);
+  MI->setPreInstrSymbol(*MF, Sym1);
+  MI->setPostInstrSymbol(*MF, Sym2);
+  MI->setHeapAllocMarker(*MF, MDN);
+
+  MMOs.push_back(MMO);
+
+  MI->setMemRefs(*MF, MMOs);
+  ASSERT_TRUE(MI->memoperands().size() == 2);
+  ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
+  ASSERT_TRUE(MI->getPostInstrSymbol() == Sym2);
+  ASSERT_TRUE(MI->getHeapAllocMarker() == MDN);
+
+  MI->setPostInstrSymbol(*MF, Sym1);
+  ASSERT_TRUE(MI->memoperands().size() == 2);
+  ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
+  ASSERT_TRUE(MI->getPostInstrSymbol() == Sym1);
+  ASSERT_TRUE(MI->getHeapAllocMarker() == MDN);
+}
+
+TEST(MachineInstrExtraInfo, RemoveExtraInfo) {
+  auto MF = createMachineFunction();
+  MCInstrDesc MCID = {0, 0,       0,       0,       0, 0,
+                      0, nullptr, nullptr, nullptr, 0, nullptr};
+
+  auto MI = MF->CreateMachineInstr(MCID, DebugLoc());
+  auto MC = createMCContext();
+  auto MMO = MF->getMachineMemOperand(MachinePointerInfo(),
+                                      MachineMemOperand::MOLoad, 8, 8);
+  SmallVector<MachineMemOperand *, 2> MMOs;
+  MMOs.push_back(MMO);
+  MMOs.push_back(MMO);
+  MCSymbol *Sym1 = MC->createTempSymbol("pre_label", false);
+  MCSymbol *Sym2 = MC->createTempSymbol("post_label", false);
+  LLVMContext Ctx;
+  MDNode *MDN = MDNode::getDistinct(Ctx, None);
+
+  MI->setMemRefs(*MF, MMOs);
+  MI->setPreInstrSymbol(*MF, Sym1);
+  MI->setPostInstrSymbol(*MF, Sym2);
+  MI->setHeapAllocMarker(*MF, MDN);
+
+  MI->setPostInstrSymbol(*MF, nullptr);
+  ASSERT_TRUE(MI->memoperands().size() == 2);
+  ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
+  ASSERT_FALSE(MI->getPostInstrSymbol());
+  ASSERT_TRUE(MI->getHeapAllocMarker() == MDN);
+
+  MI->setHeapAllocMarker(*MF, nullptr);
+  ASSERT_TRUE(MI->memoperands().size() == 2);
+  ASSERT_TRUE(MI->getPreInstrSymbol() == Sym1);
+  ASSERT_FALSE(MI->getPostInstrSymbol());
+  ASSERT_FALSE(MI->getHeapAllocMarker());
+
+  MI->setPreInstrSymbol(*MF, nullptr);
+  ASSERT_TRUE(MI->memoperands().size() == 2);
+  ASSERT_FALSE(MI->getPreInstrSymbol());
+  ASSERT_FALSE(MI->getPostInstrSymbol());
+  ASSERT_FALSE(MI->getHeapAllocMarker());
+
+  MI->setMemRefs(*MF, {});
+  ASSERT_TRUE(MI->memoperands_empty());
+  ASSERT_FALSE(MI->getPreInstrSymbol());
+  ASSERT_FALSE(MI->getPostInstrSymbol());
+  ASSERT_FALSE(MI->getHeapAllocMarker());
+}
+
 static_assert(is_trivially_copyable<MCOperand>::value, "trivially copyable");
 
 } // end namespace