[CodeGen] Enhance `MachineInstrSpan` to allow the end of MBB to be used.

[hliao]

Summary

• Explicitly specify the parent MBB to allow the end iterator to be used.

[arsenm]

LGTM

[dstenb]

The tests introduced in this commit currently fail when running UBSan, due to invoking getRegInfo() with RegInfo being null:

[ RUN      ] MachineInstrSpan.DistanceBegin
/b/sanitizer-x86_64-linux-fast/build/llvm/include/llvm/CodeGen/MachineFunction.h:474:46: runtime error: reference binding to null pointer of type 'llvm::MachineRegisterInfo'
    #0 0x1bed93c in getRegInfo /b/sanitizer-x86_64-linux-fast/build/llvm/include/llvm/CodeGen/MachineFunction.h:474:39
    #1 0x1bed93c in llvm::ilist_traits<llvm::MachineInstr>::addNodeToList(llvm::MachineInstr*) /b/sanitizer-x86_64-linux-fast/build/llvm/lib/CodeGen/MachineBasicBlock.cpp:111
    #2 0x48939c in llvm::iplist_impl<llvm::simple_ilist<llvm::MachineInstr, llvm::ilist_sentinel_tracking<true> >, llvm::ilist_traits<llvm::MachineInstr> >::insert(llvm::ilist_iterator<llvm::ilist_detail::node_options<llvm::MachineInstr, true, true, void>, false, false>, llvm::MachineInstr*) /b/sanitizer-x86_64-linux-fast/build/llvm/include/llvm/ADT/ilist.h:227:11
    #3 0x488e2a in llvm::MachineBasicBlock::insert(llvm::MachineInstrBundleIterator<llvm::MachineInstr, false>, llvm::MachineInstr*) /b/sanitizer-x86_64-linux-fast/build/llvm/include/llvm/CodeGen/MachineBasicBlock.h:627:18
    #4 0x488a93 in (anonymous namespace)::MachineInstrSpan_DistanceBegin_Test::TestBody() /b/sanitizer-x86_64-linux-fast/build/llvm/unittests/CodeGen/MachineInstrTest.cpp:289:8
    #5 0x2506006 in testing::Test::Run() /b/sanitizer-x86_64-linux-fast/build/llvm/utils/unittest/googletest/src/gtest.cc:2474:5
    #6 0x2506cf3 in testing::TestInfo::Run() /b/sanitizer-x86_64-linux-fast/build/llvm/utils/unittest/googletest/src/gtest.cc:2656:11
    #7 0x25076d2 in testing::TestCase::Run() /b/sanitizer-x86_64-linux-fast/build/llvm/utils/unittest/googletest/src/gtest.cc:2774:28
    #8 0x250ee02 in testing::internal::UnitTestImpl::RunAllTests() /b/sanitizer-x86_64-linux-fast/build/llvm/utils/unittest/googletest/src/gtest.cc:4649:43
    #9 0x250e846 in testing::UnitTest::Run() /b/sanitizer-x86_64-linux-fast/build/llvm/utils/unittest/googletest/src/gtest.cc:4257:10
    #10 0x492765 in RUN_ALL_TESTS /b/sanitizer-x86_64-linux-fast/build/llvm/utils/unittest/googletest/include/gtest/gtest.h:2233:46
    #11 0x492765 in main /b/sanitizer-x86_64-linux-fast/build/llvm/unittests/CodeGen/TargetOptionsTest.cpp:75
    #12 0x7f0cc584c2e0 in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x202e0)
    #13 0x44d429 in _start (/b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/unittests/CodeGen/CodeGenTests+0x44d429)

Log: http://lab.llvm.org:8011/builders/sanitizer-x86_64-linux-fast/builds/33393/steps/check-llvm%20ubsan/logs/stdio

Is someone looking into that? I did not find any indication of that on llvm-commits, llvm-dev, or in bugzilla, so I'm writing on this review just to make sure.

(I have not been able to run this myself, so I have not actually verified that the tests start breaking with rL365240 rather than some later commit.)

llvm/include/llvm/CodeGen/MachineBasicBlock.h

 /// containing the instruction it was initialized with, along with all
 /// those instructions inserted prior to or following that instruction
 /// at some point after the MachineInstrSpan is constructed.
 class MachineInstrSpan {
   MachineBasicBlock &MBB;
   MachineBasicBlock::iterator I, B, E;
 
 public:
-  MachineInstrSpan(MachineBasicBlock::iterator I)
-    : MBB(*I->getParent()),
-      I(I),
-      B(I == MBB.begin() ? MBB.end() : std::prev(I)),
-      E(std::next(I)) {}
+  MachineInstrSpan(MachineBasicBlock::iterator I, MachineBasicBlock *BB)
+      : MBB(*BB), I(I), B(I == MBB.begin() ? MBB.end() : std::prev(I)),
+        E(std::next(I)) {
+    assert(I == BB->end() || I->getParent() == BB);
+  }
 
   MachineBasicBlock::iterator begin() {
     return B == MBB.end() ? MBB.begin() : std::next(B);
   }
   MachineBasicBlock::iterator end() { return E; }
   bool empty() { return begin() == end(); }
 
   MachineBasicBlock::iterator getInitial() { return I; }

llvm/lib/CodeGen/InlineSpiller.cpp

if (!MI->isRegTiedToDefOperand(Idx))
       FoldOps.push_back(Idx);
   }
 
   // If we only have implicit uses, we won't be able to fold that.
   // Moreover, TargetInstrInfo::foldMemoryOperand will assert if we try!
   if (FoldOps.empty())
     return false;
 
-  MachineInstrSpan MIS(MI);
+  MachineInstrSpan MIS(MI, MI->getParent());
 
   MachineInstr *FoldMI =
       LoadMI ? TII.foldMemoryOperand(*MI, FoldOps, *LoadMI, &LIS)
              : TII.foldMemoryOperand(*MI, FoldOps, StackSlot, &LIS, &VRM);
   if (!FoldMI)
     return false;
 
   // Remove LIS for any dead defs in the original MI not in FoldMI.
foldMemoryOperand(ArrayRef<std::pair<MachineInstr *, unsigned>> Ops,
   return true;
 }
 
 void InlineSpiller::insertReload(unsigned NewVReg,
                                  SlotIndex Idx,
                                  MachineBasicBlock::iterator MI) {
   MachineBasicBlock &MBB = *MI->getParent();
 
-  MachineInstrSpan MIS(MI);
+  MachineInstrSpan MIS(MI, &MBB);
   TII.loadRegFromStackSlot(MBB, MI, NewVReg, StackSlot,
                            MRI.getRegClass(NewVReg), &TRI);
 
   LIS.InsertMachineInstrRangeInMaps(MIS.begin(), MI);
 
   LLVM_DEBUG(dumpMachineInstrRangeWithSlotIndex(MIS.begin(), MI, LIS, "reload",
                                                 NewVReg));
   ++NumReloads;
static bool isFullUndefDef(const MachineInstr &Def) {
   return !Def.getOperand(0).getSubReg();
 }
 
 /// insertSpill - Insert a spill of NewVReg after MI.
 void InlineSpiller::insertSpill(unsigned NewVReg, bool isKill,
                                  MachineBasicBlock::iterator MI) {
   MachineBasicBlock &MBB = *MI->getParent();
 
-  MachineInstrSpan MIS(MI);
+  MachineInstrSpan MIS(MI, &MBB);
   bool IsRealSpill = true;
   if (isFullUndefDef(*MI)) {
     // Don't spill undef value.
     // Anything works for undef, in particular keeping the memory
     // uninitialized is a viable option and it saves code size and
     // run time.
     BuildMI(MBB, std::next(MI), MI->getDebugLoc(), TII.get(TargetOpcode::KILL))
         .addReg(NewVReg, getKillRegState(isKill));

llvm/lib/Target/AMDGPU/SILowerSGPRSpills.cpp

static void insertCSRSaves(MachineBasicBlock &SaveBlock,
   const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
 
   MachineBasicBlock::iterator I = SaveBlock.begin();
   if (!TFI->spillCalleeSavedRegisters(SaveBlock, I, CSI, TRI)) {
     for (const CalleeSavedInfo &CS : CSI) {
       // Insert the spill to the stack frame.
       unsigned Reg = CS.getReg();
 
-      MachineInstrSpan MIS(I);
+      MachineInstrSpan MIS(I, &SaveBlock);
       const TargetRegisterClass *RC = TRI->getMinimalPhysRegClass(Reg);
 
       TII.storeRegToStackSlot(SaveBlock, I, Reg, true, CS.getFrameIdx(), RC,
                               TRI);
 
       if (LIS) {
         assert(std::distance(MIS.begin(), I) == 1);
         MachineInstr &Inst = *std::prev(I);

llvm/unittests/CodeGen/MachineInstrTest.cpp

 //===- MachineInstrTest.cpp -----------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/TargetFrameLowering.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/DebugInfoMetadata.h"
TEST(MachineInstrPrintingTest, DebugLocPrinting) {
   MI->print(OS, /*IsStandalone*/true, /*SkipOpers*/false, /*SkipDebugLoc*/false,
             /*AddNewLine*/false);
   ASSERT_TRUE(
       StringRef(OS.str()).startswith("$noreg = UNKNOWN debug-location "));
   ASSERT_TRUE(
       StringRef(OS.str()).endswith("filename:1:5"));
 }
 
+TEST(MachineInstrSpan, DistanceBegin) {
+  auto MF = createMachineFunction();
+  auto MBB = MF->CreateMachineBasicBlock();
+
+  MCInstrDesc MCID = {0, 0,       0,       0,       0, 0,
+                      0, nullptr, nullptr, nullptr, 0, nullptr};
+
+  auto MII = MBB->begin();
+  MachineInstrSpan MIS(MII, MBB);
+  ASSERT_TRUE(MIS.empty());
+
+  auto MI = MF->CreateMachineInstr(MCID, DebugLoc());
+  MBB->insert(MII, MI);
+  ASSERT_TRUE(std::distance(MIS.begin(), MII) == 1);
+}
+
+TEST(MachineInstrSpan, DistanceEnd) {
+  auto MF = createMachineFunction();
+  auto MBB = MF->CreateMachineBasicBlock();
+
+  MCInstrDesc MCID = {0, 0,       0,       0,       0, 0,
+                      0, nullptr, nullptr, nullptr, 0, nullptr};
+
+  auto MII = MBB->end();
+  MachineInstrSpan MIS(MII, MBB);
+  ASSERT_TRUE(MIS.empty());
+
+  auto MI = MF->CreateMachineInstr(MCID, DebugLoc());
+  MBB->insert(MII, MI);
+  ASSERT_TRUE(std::distance(MIS.begin(), MII) == 1);
+}
+
 static_assert(is_trivially_copyable<MCOperand>::value, "trivially copyable");
 
 } // end namespace