[NFC] [HWASan] [MTE] factor out threadlong increment

[fmayer]

No description provided.

[llvmbot]

@llvm/pr-subscribers-compiler-rt-sanitizer
@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-backend-aarch64

Author: Florian Mayer (fmayer)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/110340.diff

4 Files Affected:

• (modified) llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h (+2)
• (modified) llvm/lib/Target/AArch64/AArch64StackTagging.cpp (+2-13)
• (modified) llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp (+1-29)
• (modified) llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp (+34)

diff --git a/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h b/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h
index 1401c5fcde5f77..60255d57364099 100644
--- a/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h
+++ b/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h
@@ -99,6 +99,8 @@ Value *getPC(const Triple &TargetTriple, IRBuilder<> &IRB);
 Value *getAndroidSlotPtr(IRBuilder<> &IRB, int Slot);
 
 void annotateDebugRecords(AllocaInfo &Info, unsigned int Tag);
+Value *incrementThreadLong(IRBuilder<> &IRB, Value *ThreadLong,
+                           unsigned int Inc);
 
 } // namespace memtag
 } // namespace llvm
diff --git a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
index 72823fdcd858df..72266e261a1f6c 100644
--- a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
+++ b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
@@ -507,19 +507,8 @@ Instruction *AArch64StackTagging::insertBaseTaggedPointer(
     Value *RecordPtr = IRB.CreateIntToPtr(ThreadLong, IRB.getPtrTy(0));
     IRB.CreateStore(PC, RecordPtr);
     IRB.CreateStore(TaggedFP, IRB.CreateConstGEP1_64(IntptrTy, RecordPtr, 1));
-    // Update the ring buffer. Top byte of ThreadLong defines the size of the
-    // buffer in pages, it must be a power of two, and the start of the buffer
-    // must be aligned by twice that much. Therefore wrap around of the ring
-    // buffer is simply Addr &= ~((ThreadLong >> 56) << 12).
-    // The use of AShr instead of LShr is due to
-    //   https://bugs.llvm.org/show_bug.cgi?id=39030
-    // Runtime library makes sure not to use the highest bit.
-    Value *WrapMask = IRB.CreateXor(
-        IRB.CreateShl(IRB.CreateAShr(ThreadLong, 56), 12, "", true, true),
-        ConstantInt::get(IntptrTy, (uint64_t)-1));
-    Value *ThreadLongNew = IRB.CreateAnd(
-        IRB.CreateAdd(ThreadLong, ConstantInt::get(IntptrTy, 16)), WrapMask);
-    IRB.CreateStore(ThreadLongNew, SlotPtr);
+
+    IRB.CreateStore(memtag::incrementThreadLong(IRB, ThreadLong, 16), SlotPtr);
   }
   return Base;
 }
diff --git a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
index 669b63343e994e..cc7f20cffea771 100644
--- a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
@@ -1421,35 +1421,7 @@ void HWAddressSanitizer::emitPrologue(IRBuilder<> &IRB, bool WithFrameRecord) {
           IRB.CreateIntToPtr(ThreadLongMaybeUntagged, IRB.getPtrTy(0));
       IRB.CreateStore(FrameRecordInfo, RecordPtr);
 
-      // Update the ring buffer. Top byte of ThreadLong defines the size of the
-      // buffer in pages, it must be a power of two, and the start of the buffer
-      // must be aligned by twice that much. Therefore wrap around of the ring
-      // buffer is simply Addr &= ~((ThreadLong >> 56) << 12).
-      // The use of AShr instead of LShr is due to
-      //   https://bugs.llvm.org/show_bug.cgi?id=39030
-      // Runtime library makes sure not to use the highest bit.
-      //
-      // Mechanical proof of this address calculation can be found at:
-      // https://github.com/google/sanitizers/blob/master/hwaddress-sanitizer/prove_hwasanwrap.smt2
-      //
-      // Example of the wrap case for N = 1
-      // Pointer:   0x01AAAAAAAAAAAFF8
-      //                     +
-      //            0x0000000000000008
-      //                     =
-      //            0x01AAAAAAAAAAB000
-      //                     &
-      // WrapMask:  0xFFFFFFFFFFFFF000
-      //                     =
-      //            0x01AAAAAAAAAAA000
-      //
-      // Then the WrapMask will be a no-op until the next wrap case.
-      Value *WrapMask = IRB.CreateXor(
-          IRB.CreateShl(IRB.CreateAShr(ThreadLong, 56), 12, "", true, true),
-          ConstantInt::get(IntptrTy, (uint64_t)-1));
-      Value *ThreadLongNew = IRB.CreateAnd(
-          IRB.CreateAdd(ThreadLong, ConstantInt::get(IntptrTy, 8)), WrapMask);
-      IRB.CreateStore(ThreadLongNew, SlotPtr);
+      IRB.CreateStore(memtag::incrementThreadLong(IRB, ThreadLong, 8), SlotPtr);
       break;
     }
     case none: {
diff --git a/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp b/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp
index c5d22fbcc4e639..1cb1a7b396badc 100644
--- a/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp
+++ b/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp
@@ -338,5 +338,39 @@ void annotateDebugRecords(AllocaInfo &Info, unsigned int Tag) {
   llvm::for_each(Info.DbgVariableRecords, AnnotateDbgRecord);
 }
 
+Value *incrementThreadLong(IRBuilder<> &IRB, Value *ThreadLong,
+                           unsigned int Inc) {
+  // Update the ring buffer. Top byte of ThreadLong defines the size of the
+  // buffer in pages, it must be a power of two, and the start of the buffer
+  // must be aligned by twice that much. Therefore wrap around of the ring
+  // buffer is simply Addr &= ~((ThreadLong >> 56) << 12).
+  // The use of AShr instead of LShr is due to
+  //   https://bugs.llvm.org/show_bug.cgi?id=39030
+  // Runtime library makes sure not to use the highest bit.
+  //
+  // Mechanical proof of this address calculation can be found at:
+  // https://github.com/google/sanitizers/blob/master/hwaddress-sanitizer/prove_hwasanwrap.smt2
+  //
+  // Example of the wrap case for N = 1
+  // Pointer:   0x01AAAAAAAAAAAFF8
+  //                     +
+  //            0x0000000000000008
+  //                     =
+  //            0x01AAAAAAAAAAB000
+  //                     &
+  // WrapMask:  0xFFFFFFFFFFFFF000
+  //                     =
+  //            0x01AAAAAAAAAAA000
+  //
+  // Then the WrapMask will be a no-op until the next wrap case.
+  assert((4096 % Inc) == 0);
+  Value *WrapMask = IRB.CreateXor(
+      IRB.CreateShl(IRB.CreateAShr(ThreadLong, 56), 12, "", true, true),
+      ConstantInt::get(ThreadLong->getType(), (uint64_t)-1));
+  return IRB.CreateAnd(
+      IRB.CreateAdd(ThreadLong, ConstantInt::get(ThreadLong->getType(), Inc)),
+      WrapMask);
+}
+
 } // namespace memtag
 } // namespace llvm