[AArch64] Add Defs=[NZCV] to MTE loop pseudos.

The `STGloop` family of pseudo-instructions all expand to a loop which
iterates over a region of memory setting all its MTE tags to a given
value. The loop writes to the flags in order to check termination. But
the unexpanded pseudo-instructions were not marked as modifying the
flags. Therefore it was possible for one to end up in a location where
the flags were live, and then the loop would corrupt them.

We spotted the effect of this in a libc++ test involving a lot of
complicated inlining, and haven't been able to construct a smaller
test case that demonstrates actual incorrect output code. So my test
here is just checking that `implicit-def $nzcv` shows up on the
pseudo-instructions as they're output from isel.

Reviewed By: DavidSpickett

Differential Revision: https://reviews.llvm.org/D158262

(cherry picked from commit b09c575975b691e988a0f2e31d632c5f1038ab1d)

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
index 0f3d346..9e72d37 100644 (file)
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.td
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
@@ -2253,7 +2253,7 @@ def : Pat<(int_aarch64_irg_sp i64:$Rm), (IRGstack SP, i64:$Rm)>;
 
 // Large STG to be expanded into a loop. $sz is the size, $Rn is start address.
 // $Rn_wback is one past the end of the range. $Rm is the loop counter.
-let isCodeGenOnly=1, mayStore=1 in {
+let isCodeGenOnly=1, mayStore=1, Defs=[NZCV] in {
 def STGloop_wback
     : Pseudo<(outs GPR64common:$Rm, GPR64sp:$Rn_wback), (ins i64imm:$sz, GPR64sp:$Rn),
              [], "$Rn = $Rn_wback,@earlyclobber $Rn_wback,@earlyclobber $Rm" >,

diff --git a/llvm/test/CodeGen/AArch64/memtag-loop-nzcv.ll b/llvm/test/CodeGen/AArch64/memtag-loop-nzcv.ll
new file mode 100644 (file)
index 0000000..86bafd1
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/memtag-loop-nzcv.ll
@@ -0,0 +1,59 @@
+; RUN: llc -O2 -print-after-isel -mtriple=aarch64-linux-gnu %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=CHECK
+
+; This test function includes a 256-byte buffer. We expect it to require its
+; MTE tags to be set to a useful value on entry, and cleared again on exit. At
+; the time of writing this test, the pseudo-instructions chosen are
+; STGloop_wback and STGloop respectively, but if different pseudos are selected
+; in future, that's not a problem. The important thing is that both should
+; include that implicit-def of $nzcv, because these pseudo-instructions will
+; expand into loops that use the flags for their termination tests.
+
+; CHECK: STGloop_wback 256, {{.*}}, implicit-def dead $nzcv
+; CHECK: STGloop       256, {{.*}}, implicit-def dead $nzcv
+
+define i32 @foo(i32 noundef %0) #0 {
+  %2 = alloca i32, align 4
+  %3 = alloca [256 x i8], align 1
+  %4 = alloca i64, align 8
+  %5 = alloca i32, align 4
+  %6 = alloca i64, align 8
+  store i32 %0, ptr %2, align 4
+  %7 = load i32, ptr %2, align 4
+  %8 = getelementptr inbounds [256 x i8], ptr %3, i64 0, i64 0
+  %9 = call i64 @read(i32 noundef %7, ptr noundef %8, i64 noundef 256)
+  store i64 %9, ptr %4, align 8
+  store i32 0, ptr %5, align 4
+  store i64 0, ptr %6, align 8
+  br label %10
+
+10:                                               ; preds = %21, %1
+  %11 = load i64, ptr %6, align 8
+  %12 = load i64, ptr %4, align 8
+  %13 = icmp ult i64 %11, %12
+  br i1 %13, label %14, label %24
+
+14:                                               ; preds = %10
+  %15 = load i64, ptr %6, align 8
+  %16 = getelementptr inbounds [256 x i8], ptr %3, i64 0, i64 %15
+  %17 = load i8, ptr %16, align 1
+  %18 = zext i8 %17 to i32
+  %19 = load i32, ptr %5, align 4
+  %20 = add nsw i32 %19, %18
+  store i32 %20, ptr %5, align 4
+  br label %21
+
+21:                                               ; preds = %14
+  %22 = load i64, ptr %6, align 8
+  %23 = add i64 %22, 1
+  store i64 %23, ptr %6, align 8
+  br label %10
+
+24:                                               ; preds = %10
+  %25 = load i32, ptr %5, align 4
+  %26 = srem i32 %25, 251
+  ret i32 %26
+}
+
+declare i64 @read(i32 noundef, ptr noundef, i64 noundef)
+
+attributes #0 = { sanitize_memtag "target-features"="+mte" }