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

[simon_tatham]

Summary

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.

[DavidSpickett]

llvm/test/CodeGen/AArch64/memtag-loop-nzcv.ll

2

Any reason not to do the usual ... | FileCheck %s ...?

10

Add the reason here? of $nzcv because...

15

Might not be needed due to -mtriple above.

[simon_tatham]

Applied review suggestions in the test.

llvm/test/CodeGen/AArch64/memtag-loop-nzcv.ll

2

I cribbed the RUN line from an existing test using -print-after-isel. It did that because -print-after-isel outputs to stderr, not stdout. But I suppose if we don't also need the stdout, we can use 2>&1 |. I'll do that.

[DavidSpickett]

I thought maybe the test input could be smaller, but it would look less like the actual use case, be more confusing and be prone to breaking randomly. So keep it as is.

This all LGTM.

[DavidSpickett]

Though, should you check for the STZGloops as well? Not sure if you can reliably trigger that, maybe the clearing of tags here uses it but equally it could use STGloop with a tag value of 0.

They're all in the same group so not covering the Z variant isn't that big of a risk.

[simon_tatham]

I agree that it would be nice to test all four pseudos if possible, but I don't know how to generate the other two. The clearing of the tags on function exit in this example is done by the STGloop shown in my CHECK line, so if that didn't cause an STZGloop of any kind then I don't know what would.

I also agree that it's not vital, since the same code change affects all four pseudos in any case.

[DavidSpickett]

Closest I found was llvm/test/CodeGen/AArch64/settag.ll but that's just one set not a loop.

Well anyway, all my comments have been addressed.

[simon_tatham]

Cross-reference: this bug was already reported as https://github.com/llvm/llvm-project/issues/64309 (but I only just found out about that).

llvm/lib/Target/AArch64/AArch64InstrInfo.td

 // Explicit SP in the first operand prevents ShrinkWrap optimization
 // from leaving this instruction out of the stack frame. When IRGstack
 // is transformed into IRG, this operand is replaced with the actual
 // register / expression for the tagged base pointer of the current function.
 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" >,
       Sched<[WriteAdr, WriteST]>;
 
 def STZGloop_wback
     : Pseudo<(outs GPR64common:$Rm, GPR64sp:$Rn_wback), (ins i64imm:$sz, GPR64sp:$Rn),
              [], "$Rn = $Rn_wback,@earlyclobber $Rn_wback,@earlyclobber $Rm" >,

llvm/test/CodeGen/AArch64/memtag-loop-nzcv.ll

+; 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" }