[TwoAddressInstruction] Check if segment was found when updating subreg LIs

llvm/lib/CodeGen/TwoAddressInstructionPass.cpp

@@ -1871,11 +1871,30 @@ bool TwoAddressInstructionPass::runOnMachineFunction(MachineFunction &Func) {
             SlotIndex Idx = LIS->getInstructionIndex(*mi);
             for (auto &S : LI.subranges()) {
               if ((S.LaneMask & LaneMask).none()) {
+                // If Idx is 160B, and we have a subrange that isn't in
+                // %reg.subidx like so:
+                //
+                // [152r,160r)[160r,256r)
+                //
+                // Merge the two segments together so the subrange becomes:
+                //
+                // [152r,256r)
                 LiveRange::iterator UseSeg = S.FindSegmentContaining(Idx);
-                LiveRange::iterator DefSeg = std::next(UseSeg);
-                S.MergeValueNumberInto(DefSeg->valno, UseSeg->valno);
+                if (UseSeg != S.end()) {
+                  LiveRange::iterator DefSeg = std::next(UseSeg);
+                  assert(DefSeg != S.end());
+                  S.MergeValueNumberInto(DefSeg->valno, UseSeg->valno);
+                }
+                // Otherwise, it should have only one segment that starts at
+                // 160r which we should remove.
+                else {
+                  assert(S.containsOneValue());
+                  assert(S.begin()->start == Idx.getRegSlot());
+                  S.removeSegment(S.begin());
+                }
               }
             }
+            LI.removeEmptySubRanges();
 
             // The COPY no longer has a use of %reg.
             LIS->shrinkToUses(&LI);

llvm/test/CodeGen/RISCV/rvv/twoaddressinstruction-subreg-liveness-update.mir

@@ -0,0 +1,42 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 3
+# RUN: llc %s -o - -mtriple=riscv64 -mattr=+v -run-pass=liveintervals,twoaddressinstruction,register-coalescer -verify-machineinstrs | FileCheck %s
+
+...
+---
+name:            f
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $v8
+    ; CHECK-LABEL: name: f
+    ; CHECK: liveins: $v8
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: KILL $v8
+    ; CHECK-NEXT: PseudoRET
+    %0:vr = COPY killed $v8
+    %2:vrm8 = INSERT_SUBREG undef %1:vrm8, killed %0, %subreg.sub_vrm1_0
+    %3:vrm4 = COPY %2.sub_vrm4_0:vrm8
+
+    PseudoRET
+
+...
+---
+name:            g
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $v8, $v16
+    ; CHECK-LABEL: name: g
+    ; CHECK: liveins: $v8, $v16
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: KILL $v8
+    ; CHECK-NEXT: dead [[COPY:%[0-9]+]]:vrm8 = COPY $v16
+    ; CHECK-NEXT: PseudoRET
+    %0:vr = COPY killed $v8
+    %1:vrm8 = COPY killed $v16
+    %2:vrm8 = INSERT_SUBREG %1:vrm8, killed %0, %subreg.sub_vrm1_0
+    %3:vrm4 = COPY %2.sub_vrm4_0:vrm8
+
+    PseudoRET
+
+...