[llvm] [StackColoring] Do not drop AA metadata when not doing remappings (PR #71958)

Fri Nov 10 08:48:48 PST 2023

https://github.com/bjope created https://github.com/llvm/llvm-project/pull/71958

In the StackColoring pass we first scan for possible stack slot merges. A SlotRemap map is setup with the remappings that should be performed. Then the main work is done by calling remapInstructions and providing that map.

Most of the work in remapInstructions would just be a waste of time in situations when the SlotRemap map is empty, but it turns out that the part that adjusts Alias Analysis information could end up dropping AA metadata even when there are no stack slot merges being done. This happens since all instructions machine memory operands are considered, and if we can't determine the underlying object that is accessed (using getUnderlyingObjectsForCodeGen) then we conservatively drop AA metadata.

This patch simply avoids calling remapInstructions if we don't intend to do any remappings (i.e. if SlotRemap is empty). That avoid touching AA metadata when all we do is to remove lifetime markers. That seems like a safe thing to do, as it is the same thing as happens when we bail out early due to other reasons (e.g. when only having one lifetime marker).

For targets that do not care about Alias Analysis information after the StackColoring pass this shouldn't have any impact, except that it might improve compile time slightly as we now skip spending time in remapInstructions when not doing any stack merges.

>From 47ea7f7c1c2794be5ae4a43086883be07ecf5e5c Mon Sep 17 00:00:00 2001
From: Bjorn Pettersson <bjorn.a.pettersson at ericsson.com>
Date: Fri, 10 Nov 2023 17:40:51 +0100
Subject: [PATCH] [StackColoring] Do not drop AA metadata when not doing
 remappings

In the StackColoring pass we first scan for possible stack slot
merges. A SlotRemap map is setup with the remappings that should
be performed. Then the main work is done by calling remapInstructions
and providing that map.

Most of the work in remapInstructions would just be a waste of time
in situations when the SlotRemap map is empty, but it turns out that
the part that adjusts Alias Analysis information could end up
dropping AA metadata even when there are no stack slot merges being
done. This happens since all instructions machine memory operands
are considered, and if we can't determine the underlying object that
is accessed (using getUnderlyingObjectsForCodeGen) then we
conservatively drop AA metadata.

This patch simply avoids calling remapInstructions if we don't
intend to do any remappings (i.e. if SlotRemap is empty). That
avoid touching AA metadata when all we do is to remove lifetime
markers. That seems like a safe thing to do, as it is the same
thing as happens when we bail out early due to other reasons (e.g.
when only having one lifetime marker).

For targets that do not care about Alias Analysis information after
the StackColoring pass this shouldn't have any impact, except that
it might improve compile time slightly as we now skip spending time
in remapInstructions when not doing any stack merges.
---
 llvm/lib/CodeGen/StackColoring.cpp           |  6 +-
 llvm/test/CodeGen/X86/StackColoring-tbaa.mir | 68 ++++++++++++++++++++
 2 files changed, 72 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/X86/StackColoring-tbaa.mir

diff --git a/llvm/lib/CodeGen/StackColoring.cpp b/llvm/lib/CodeGen/StackColoring.cpp
index 3d261688fa8c817..a06172ef99939fd 100644
--- a/llvm/lib/CodeGen/StackColoring.cpp
+++ b/llvm/lib/CodeGen/StackColoring.cpp
@@ -1338,8 +1338,10 @@ bool StackColoring::runOnMachineFunction(MachineFunction &Func) {
 
   // Scan the entire function and update all machine operands that use frame
   // indices to use the remapped frame index.
-  expungeSlotMap(SlotRemap, NumSlots);
-  remapInstructions(SlotRemap);
+  if (!SlotRemap.empty()) {
+    expungeSlotMap(SlotRemap, NumSlots);
+    remapInstructions(SlotRemap);
+  }
 
   return removeAllMarkers();
 }
diff --git a/llvm/test/CodeGen/X86/StackColoring-tbaa.mir b/llvm/test/CodeGen/X86/StackColoring-tbaa.mir
new file mode 100644
index 000000000000000..6d7f294549d7913
--- /dev/null
+++ b/llvm/test/CodeGen/X86/StackColoring-tbaa.mir
@@ -0,0 +1,68 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 3
+# RUN: llc -run-pass=stack-coloring %s -o - | FileCheck %s
+
+# We do not expect any stack coloring remappings in this test case.
+# And then there should be no reason to drop the tbaa metadata on the
+# MOV8rm instruction, so we check that the tbaa info is kept.
+
+--- |
+  source_filename = "test_case.cc"
+  target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+  target triple = "x86_64-unknown-linux-gnu"
+
+  %struct.Agg = type { [3 x i8], [3 x i8] }
+
+  define i8 @main() {
+    %padding = alloca %struct.Agg, align 8
+    %agg = alloca %struct.Agg, align 8
+    %a2 = getelementptr inbounds %struct.Agg, ptr %agg, i64 0, i32 1
+    %a22 = getelementptr inbounds [3 x i8], ptr %a2, i64 0, i32 1
+    %t0 = load i8, ptr %a22, align 1, !tbaa !2
+    %tobool = icmp slt i8 %t0, 0
+    %t1 = load ptr, ptr %a2, align 8
+    %cond = select i1 %tobool, ptr %t1, ptr %a2
+    %add.ptr.i = getelementptr inbounds i8, ptr %cond, i64 16
+    %t2 = load i8, ptr %add.ptr.i, align 1, !tbaa !2
+    ret i8 %t2
+  }
+
+  !llvm.module.flags = !{!0}
+  !llvm.ident = !{!1}
+
+  !0 = !{i32 1, !"wchar_size", i32 4}
+  !1 = !{!"clang version 9.0.0"}
+  !2 = !{!3, !3, i64 0}
+  !3 = !{!"omnipotent char", !4, i64 0}
+  !4 = !{!"Simple C++ TBAA"}
+
+...
+---
+name:            main
+alignment:       16
+tracksRegLiveness: true
+stack:
+  - { id: 0, name: padding, type: default, offset: 0, size: 24, alignment: 16,
+      stack-id: default, callee-saved-register: '', callee-saved-restored: true,
+      debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
+  - { id: 1, name: agg, type: default, offset: 0, size: 48, alignment: 16,
+      stack-id: default, callee-saved-register: '', callee-saved-restored: true,
+      debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: main
+    ; CHECK: [[LEA64r:%[0-9]+]]:gr64 = nuw LEA64r %stack.1.agg, 1, $noreg, 24, $noreg
+    ; CHECK-NEXT: CMP8mi %stack.1.agg, 1, $noreg, 47, $noreg, 0, implicit-def $eflags :: (dereferenceable load (s8) from %ir.a22, !tbaa !2)
+    ; CHECK-NEXT: [[CMOV64rm:%[0-9]+]]:gr64 = CMOV64rm [[LEA64r]], %stack.1.agg, 1, $noreg, 24, $noreg, 8, implicit $eflags :: (dereferenceable load (s64) from %ir.a2)
+    ; CHECK-NEXT: [[MOV8rm:%[0-9]+]]:gr8 = MOV8rm killed [[CMOV64rm]], 1, $noreg, 16, $noreg :: (load (s8) from %ir.add.ptr.i, !tbaa !2)
+    ; CHECK-NEXT: $al = COPY [[MOV8rm]]
+    ; CHECK-NEXT: RET 0, $al
+    LIFETIME_START %stack.0.padding
+    LIFETIME_START %stack.1.agg
+    %8:gr64 = nuw LEA64r %stack.1.agg, 1, $noreg, 24, $noreg
+    CMP8mi %stack.1.agg, 1, $noreg, 47, $noreg, 0, implicit-def $eflags :: (dereferenceable load (s8) from %ir.a22, !tbaa !2)
+    %13:gr64 = CMOV64rm %8, %stack.1.agg, 1, $noreg, 24, $noreg, 8, implicit $eflags :: (dereferenceable load (s64) from %ir.a2)
+    %14:gr8 = MOV8rm killed %13, 1, $noreg, 16, $noreg :: (load (s8) from %ir.add.ptr.i, !tbaa !2)
+    $al = COPY %14
+    RET 0, $al
+
+...

