[llvm] [CodeGen] Don't schedule across frame register save/restore for sanitized functions (PR #221714)

[via llvm-commits]

https://github.com/heiher created https://github.com/llvm/llvm-project/pull/221714

Even if the target's Post-RA Scheduler can reorder instructions across the frame register save/restore, it must not do so for sanitized functions. A crash can happen at any point, and sanitizers need to be able to unwind from the PC at the point of the crash using a frame register that still points to the current frame.

This is the same class of issue fixed for shrink-wrapping in 2cdcfd23cd01 ("[ShrinkWrapping] Disable the optimization for functions with sanitize like attribute.").

>From 083a8b3442a5814fd7fdd9c560dc4df491ac8ebd Mon Sep 17 00:00:00 2001
From: WANG Rui <wangrui at loongson.cn>
Date: Mon, 7 Sep 2026 16:50:47 +0800
Subject: [PATCH] [CodeGen] Don't schedule across frame register save/restore
 for sanitized functions

Even if the target's Post-RA Scheduler can reorder instructions across
the frame register save/restore, it must not do so for sanitized
functions. A crash can happen at any point, and sanitizers need to be
able to unwind from the PC at the point of the crash using a frame
register that still points to the current frame.

This is the same class of issue fixed for shrink-wrapping in
2cdcfd23cd01 ("[ShrinkWrapping] Disable the optimization for functions
with sanitize like attribute.").
---
 llvm/lib/CodeGen/TargetInstrInfo.cpp              | 15 +++++++++++++++
 .../CodeGen/LoongArch/postra-sched-sanitize.ll    | 15 +++++++++++++++
 llvm/test/CodeGen/RISCV/postra-sched-sanitize.ll  | 15 +++++++++++++++
 3 files changed, 45 insertions(+)
 create mode 100644 llvm/test/CodeGen/LoongArch/postra-sched-sanitize.ll
 create mode 100644 llvm/test/CodeGen/RISCV/postra-sched-sanitize.ll

diff --git a/llvm/lib/CodeGen/TargetInstrInfo.cpp b/llvm/lib/CodeGen/TargetInstrInfo.cpp
index 92fc628e888e5..c3533a652a5f4 100644
--- a/llvm/lib/CodeGen/TargetInstrInfo.cpp
+++ b/llvm/lib/CodeGen/TargetInstrInfo.cpp
@@ -1711,6 +1711,21 @@ bool TargetInstrInfo::isSchedulingBoundary(const MachineInstr &MI,
   if (MI.getOpcode() == TargetOpcode::INLINEASM_BR)
     return true;
 
+  // Sanitizers need a correct stack frame at the point of a crash,
+  // which can happen anywhere. Treat frame register modifications
+  // as scheduling boundaries so the scheduler can't reorder other
+  // instructions across the frame register save/restore.
+  if (MF.getFunction().hasFnAttribute(Attribute::SanitizeAddress) ||
+      MF.getFunction().hasFnAttribute(Attribute::SanitizeThread) ||
+      MF.getFunction().hasFnAttribute(Attribute::SanitizeMemory) ||
+      MF.getFunction().hasFnAttribute(Attribute::SanitizeType) ||
+      MF.getFunction().hasFnAttribute(Attribute::SanitizeHWAddress) ||
+      MF.getFunction().hasFnAttribute(Attribute::SanitizeMemTag)) {
+    Register FrameReg = TRI.getFrameRegister(MF);
+    if (MI.modifiesRegister(FrameReg, &TRI))
+      return true;
+  }
+
   // Don't attempt to schedule around any instruction that defines
   // a stack-oriented pointer, as it's unlikely to be profitable. This
   // saves compile time, because it doesn't require every single
diff --git a/llvm/test/CodeGen/LoongArch/postra-sched-sanitize.ll b/llvm/test/CodeGen/LoongArch/postra-sched-sanitize.ll
new file mode 100644
index 0000000000000..3b24985e1a19e
--- /dev/null
+++ b/llvm/test/CodeGen/LoongArch/postra-sched-sanitize.ll
@@ -0,0 +1,15 @@
+; RUN: llc --mtriple=loongarch32 --frame-pointer=all < %s | FileCheck %s
+; RUN: llc --mtriple=loongarch64 --frame-pointer=all < %s | FileCheck %s
+
+; CHECK-LABEL: foo:
+; CHECK: addi.{{w|d}}	[[REG1:\$[a-z0-9]+]], [[REG1]], 1
+; CHECK: st.w	[[REG1]], {{\$[a-z0-9]+}}, {{[0-9]+}}
+; CHECK: ld.{{w|d}}	$fp, $sp, {{[0-9]+}}
+define void @foo() nounwind sanitize_address {
+entry:
+  %1 = load ptr, ptr inttoptr (i32 64 to ptr), align 64
+  %2 = load i32, ptr %1, align 8
+  %3 = add nsw i32 %2, 1
+  store i32 %3, ptr %1, align 8
+  ret void
+}
diff --git a/llvm/test/CodeGen/RISCV/postra-sched-sanitize.ll b/llvm/test/CodeGen/RISCV/postra-sched-sanitize.ll
new file mode 100644
index 0000000000000..e4eff62099b0c
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/postra-sched-sanitize.ll
@@ -0,0 +1,15 @@
+; RUN: llc --mtriple=riscv32 --mattr=+use-postra-scheduler --frame-pointer=all < %s | FileCheck %s
+; RUN: llc --mtriple=riscv64 --mattr=+use-postra-scheduler --frame-pointer=all < %s | FileCheck %s
+
+; CHECK-LABEL: foo:
+; CHECK: addi	[[REG1:[a-z0-9]+]], [[REG1]], 1
+; CHECK: sw	[[REG1]], {{[0-9]+}}({{[a-z0-9]+}})
+; CHECK: {{lw|ld}}	s0, {{[0-9]+}}(sp)
+define void @foo() nounwind sanitize_address {
+entry:
+  %1 = load ptr, ptr inttoptr (i32 64 to ptr), align 64
+  %2 = load i32, ptr %1, align 8
+  %3 = add nsw i32 %2, 1
+  store i32 %3, ptr %1, align 8
+  ret void
+}