[llvm] [TwoAddressInstruction] Check if segment was found when updating subreg LIs (PR #65916)

[Luke Lau via llvm-commits]

https://github.com/lukel97 created https://github.com/llvm/llvm-project/pull/65916:

This patch fixes a crash I was seeing when working with subregisters in a
separate upcoming patch. Some of the RVV tests have -early-live-intervals
enabled, which triggers this subregister liveness updating code in
TwoAddressInstruction.

The INSERT_SUBREG in the test case gets converted to a COPY, so it tries to
update the subranges for %1 below:

********** INTERVALS **********
V8 [0B,16r:0) 0 at 0B-phi
%0 [16r,32r:0) 0 at 16r  weight:0.000000e+00
%1 [32r,48r:0) 0 at 32r  L000000000000003C [32r,48r:0) 0 at 32r  L00000000000003C0 [32r,32d:0) 0 at 32r  weight:0.000000e+00
%2 EMPTY  weight:0.000000e+00
%3 [48r,48d:0) 0 at 48r  weight:0.000000e+00

0B	bb.0:
	  liveins: $v8
16B	  %0:vr = COPY $v8
32B	  undef %1.sub_vrm1_0:vrm8 = COPY %0:vr
48B	  dead %3:vrm4 = COPY %1.sub_vrm4_0:vrm8
64B	  PseudoRET

The subrange for %1 L00000000000003C0 [32r,32d:0) doesn't overlap with
sub_vrm1_0's lane mask, so it tries to merge 0 at 32r. It calls
FindSegmentContaining(32B), 32B being the instruction index for the
INSERT_SUBREG-turned-COPY. But 32B isn't actually in [32r,32d:0) so it returns
end(). The code wasn't checking for end() though so it crashed when
dereferencing DefSeg.

I'm not familiar with this area in the slightest, but I presume we wanted to
find the segment at 32r? And in this test case there's no other segment in the
subrange anyway, so we need to check if DefSeg exists.

I also did a quick check, nothing in llvm-check-codegen seems to trigger this
code. Not sure if that will change if -early-live-intervals is enabled by
default.


>From 5da784840c5f773e21eb224ee5887e82df7b6d50 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Mon, 11 Sep 2023 00:44:33 +0100
Subject: [PATCH] [TwoAddressInstruction] Check if segment was found when
 updating subreg LIs

This patch fixes a crash I was seeing when working with subregisters in a
separate upcoming patch. Some of the RVV tests have -early-live-intervals
enabled, which triggers this subregister liveness updating code in
TwoAddressInstruction.

The INSERT_SUBREG in the test case gets converted to a COPY, so it tries to
update the subranges for %1 below:

********** INTERVALS **********
V8 [0B,16r:0) 0 at 0B-phi
%0 [16r,32r:0) 0 at 16r  weight:0.000000e+00
%1 [32r,48r:0) 0 at 32r  L000000000000003C [32r,48r:0) 0 at 32r  L00000000000003C0 [32r,32d:0) 0 at 32r  weight:0.000000e+00
%2 EMPTY  weight:0.000000e+00
%3 [48r,48d:0) 0 at 48r  weight:0.000000e+00

0B	bb.0:
	  liveins: $v8
16B	  %0:vr = COPY $v8
32B	  undef %1.sub_vrm1_0:vrm8 = COPY %0:vr
48B	  dead %3:vrm4 = COPY %1.sub_vrm4_0:vrm8
64B	  PseudoRET

The subrange for %1 L00000000000003C0 [32r,32d:0) doesn't overlap with
sub_vrm1_0's lane mask, so it tries to merge 0 at 32r. It calls
FindSegmentContaining(32B), 32B being the instruction index for the
INSERT_SUBREG-turned-COPY. But 32B isn't actually in [32r,32d:0) so it returns
end(). The code wasn't checking for end() though so it crashed when
dereferencing DefSeg.

I'm not familiar with this area in the slightest, but I presume we wanted to
find the segment at 32r? And in this test case there's no other segment in the
subrange anyway, so we need to check if DefSeg exists.

I also did a quick check, nothing in llvm-check-codegen seems to trigger this
code. Not sure if that will change if -early-live-intervals is enabled by
default.
---
 .../lib/CodeGen/TwoAddressInstructionPass.cpp |  5 +++--
 ...ressinstruction-subreg-liveness-update.mir | 22 +++++++++++++++++++
 2 files changed, 25 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/RISCV/rvv/twoaddressinstruction-subreg-liveness-update.mir

diff --git a/llvm/lib/CodeGen/TwoAddressInstructionPass.cpp b/llvm/lib/CodeGen/TwoAddressInstructionPass.cpp
index 45f61262faf9391..6ceb4b43522f00b 100644
--- a/llvm/lib/CodeGen/TwoAddressInstructionPass.cpp
+++ b/llvm/lib/CodeGen/TwoAddressInstructionPass.cpp
@@ -1868,12 +1868,13 @@ bool TwoAddressInstructionPass::runOnMachineFunction(MachineFunction &Func) {
             // %reg.subidx.
             LaneBitmask LaneMask =
                 TRI->getSubRegIndexLaneMask(mi->getOperand(0).getSubReg());
-            SlotIndex Idx = LIS->getInstructionIndex(*mi);
+            SlotIndex Idx = LIS->getInstructionIndex(*mi).getRegSlot();
             for (auto &S : LI.subranges()) {
               if ((S.LaneMask & LaneMask).none()) {
                 LiveRange::iterator UseSeg = S.FindSegmentContaining(Idx);
                 LiveRange::iterator DefSeg = std::next(UseSeg);
-                S.MergeValueNumberInto(DefSeg->valno, UseSeg->valno);
+                if (DefSeg != S.end())
+                  S.MergeValueNumberInto(DefSeg->valno, UseSeg->valno);
               }
             }
 
diff --git a/llvm/test/CodeGen/RISCV/rvv/twoaddressinstruction-subreg-liveness-update.mir b/llvm/test/CodeGen/RISCV/rvv/twoaddressinstruction-subreg-liveness-update.mir
new file mode 100644
index 000000000000000..abafc580a968f9e
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/rvv/twoaddressinstruction-subreg-liveness-update.mir
@@ -0,0 +1,22 @@
+# 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 | FileCheck %s
+
+...
+---
+name:            insert_v2i64_nxv16i64
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $v8
+    ; CHECK-LABEL: name: insert_v2i64_nxv16i64
+    ; 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
+
+...