[StackColoring] Don't merge slots with differing StackIDs

[frasercrmck]

Summary

The documentation for this specifically mentions that this should not
happen. We could think about adding target hooks to permit it (and how
to merge IDs) in the future if that is desirable.

This specific test case was merging a scalable-vector slot into a
non-scalable one and dropping the notion of scalability, meaning we
failed to allocate enough stack space for the object.

[sdesmalen]

We have a similar fix downstream and I actually thought we already upstreamed this. In any case, good fix!

llvm/lib/CodeGen/StackColoring.cpp

1322

minor nit: if you move this code two lines below, you can use FirstSlot and SecondSlot.

[arsenm]

I feel like I fixed this one before

[MaskRay]

With sdesmalen's suggestion applied.

[frasercrmck]

apply suggestion

[sdesmalen]

In D125699#3517128, @arsenm wrote:

I feel like I fixed this one before

I think that may have been in StackSlotColoring as opposed the StackColoring pass changed here :)

[frasercrmck]

llvm/lib/CodeGen/StackColoring.cpp

1322

Ach of course, thank you. And I'm glad you got in touch as I woke up this morning feeling guilty I hadn't notified any AArch64 people when this might (and indeed does) affect them :)

[sdesmalen]

llvm/lib/CodeGen/StackColoring.cpp

1322

No worries, I've got a Herald rule set up for patches related to scalable vectors, hence why I spotted it. I appreciate the (after)thought though :)

llvm/lib/CodeGen/StackColoring.cpp

for (unsigned I = 0; I < NumSlots; ++I) {
       if (SortedSlots[I] == -1)
         continue;
 
       for (unsigned J=I+1; J < NumSlots; ++J) {
         if (SortedSlots[J] == -1)
           continue;
 
         int FirstSlot = SortedSlots[I];
         int SecondSlot = SortedSlots[J];
+
+        // Objects with different stack IDs cannot be merged.
+        if (MFI->getStackID(FirstSlot) != MFI->getStackID(SecondSlot))
+          continue;
+
         LiveInterval *First = &*Intervals[FirstSlot];
         LiveInterval *Second = &*Intervals[SecondSlot];
         auto &FirstS = LiveStarts[FirstSlot];
         auto &SecondS = LiveStarts[SecondSlot];
         assert(!First->empty() && !Second->empty() && "Found an empty range");
 
         // Merge disjoint slots. This is a little bit tricky - see the
         // Implementation Notes section for an explanation.

llvm/test/CodeGen/RISCV/rvv/stack-coloring-scalablevec.mir

+# RUN: llc -mtriple riscv64 -mattr=+m,+v -run-pass=stack-coloring \
+# RUN:     -riscv-v-vector-bits-min=512 -o - %s | FileCheck %s
+
+# Test that a scalable slot (%stack.1) is not merged into a non-scalable one
+# (%stack.0)
+
+# CHECK:    {{^}}stack:
+# CHECK-NEXT: - { id: 0,
+# CHECK:      - { id: 1,
+
+--- |
+  define dso_local void @dont_merge() {
+  entry:
+    %buf1 = alloca <4 x i32>
+    %buf2 = alloca <vscale x 4 x i32>
+    ret void
+  }
+
+...
+---
+name:            dont_merge
+tracksRegLiveness: true
+stack:
+  - { id: 0, name: buf1, size: 16, alignment: 16 }
+  - { id: 1, name: buf2, size: 16, alignment: 16, stack-id: scalable-vector }
+body:             |
+  bb.0.entry:
+    liveins: $v8, $v10, $x10, $x11
+
+    LIFETIME_START %stack.0
+    PseudoVSPILL_M1 killed renamable $v8, %stack.0 :: (store 16 into %stack.0, align 16)
+    renamable $v8 = PseudoVRELOAD_M1 killed $x10 :: (load 16 from %stack.0, align 16)
+    LIFETIME_END %stack.0
+    LIFETIME_START %stack.1
+    PseudoVSPILL_M2 killed renamable $v10m2, %stack.1 :: (store unknown-size into %stack.1, align 16)
+    renamable $v10m2 = PseudoVRELOAD_M2 killed $x11 :: (load unknown-size from %stack.1, align 16)
+    LIFETIME_END %stack.1
+    PseudoRET
+...