From 599ff247dee8f462638731c3b2fdd3d1e332a2fe Mon Sep 17 00:00:00 2001
From: Fraser Cormack <fraser@codeplay.com>
Date: Mon, 16 May 2022 16:38:52 +0100
Subject: [PATCH] [StackColoring] Don't merge slots with differing StackIDs

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.

Reviewed By: arsenm, MaskRay, sdesmalen

Differential Revision: https://reviews.llvm.org/D125699
---
 llvm/lib/CodeGen/StackColoring.cpp            |  5 +++
 .../RISCV/rvv/stack-coloring-scalablevec.mir  | 39 +++++++++++++++++++
 2 files changed, 44 insertions(+)
 create mode 100644 llvm/test/CodeGen/RISCV/rvv/stack-coloring-scalablevec.mir

diff --git a/llvm/lib/CodeGen/StackColoring.cpp b/llvm/lib/CodeGen/StackColoring.cpp
index de549eff1787b..11c6bdc699568 100644
--- a/llvm/lib/CodeGen/StackColoring.cpp
+++ b/llvm/lib/CodeGen/StackColoring.cpp
@@ -1320,6 +1320,11 @@ bool StackColoring::runOnMachineFunction(MachineFunction &Func) {
 
         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];
diff --git a/llvm/test/CodeGen/RISCV/rvv/stack-coloring-scalablevec.mir b/llvm/test/CodeGen/RISCV/rvv/stack-coloring-scalablevec.mir
new file mode 100644
index 0000000000000..201014ae035f1
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/rvv/stack-coloring-scalablevec.mir
@@ -0,0 +1,39 @@
+# 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
+...