[llvm] [hwasan] Optimize outlined memaccess for fixed shadow on Aarch64 (PR #88544)

[Thurston Dang via llvm-commits]

https://github.com/thurstond updated https://github.com/llvm/llvm-project/pull/88544

>From ea9ae978a691d8eaf179016209aad0ebc4fe30b3 Mon Sep 17 00:00:00 2001
From: Thurston Dang <thurston at google.com>
Date: Tue, 23 Apr 2024 00:40:55 +0000
Subject: [PATCH] [hwasan] Optimize outlined memaccess for fixed shadow on
 Aarch64

The HWASan transform currently always uses x20 to pass the shadow base to hwasan_check_memaccess_shortgranules, even
if the shadow base is a constant known at compile time (via -hwasan-mapping-offset). This patch uses the fixed
shadow variant of the hwasan_check_memaccess_shortgranules intrinsic (introduced in
https://github.com/llvm/llvm-project/commit/365bddf634993d5ea357e9715d8aacd7ee40c4b5), allowing the shadow base to
be materialized inside the memaccess callee.

We currently only support this optimization for AArch64. It is a no-op on other platforms, or if -hwasan-mapping-offset is not specified.

Note: when -hwasan-mapping-offset is specified, it is necessary to specify HWASAN_OPTIONS=fixed_shadow_base=... (see ea991a1) to ensure that the runtime will map the shadow appropriately.
---
 .../Instrumentation/HWAddressSanitizer.cpp    | 31 ++++++++++++++++---
 1 file changed, 26 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
index a35f24447cc39b..322beb8f0d58f5 100644
--- a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
@@ -930,11 +930,32 @@ void HWAddressSanitizer::instrumentMemAccessOutline(Value *Ptr, bool IsWrite,
 
   IRBuilder<> IRB(InsertBefore);
   Module *M = IRB.GetInsertBlock()->getParent()->getParent();
-  IRB.CreateCall(Intrinsic::getDeclaration(
-                     M, UseShortGranules
-                            ? Intrinsic::hwasan_check_memaccess_shortgranules
-                            : Intrinsic::hwasan_check_memaccess),
-                 {ShadowBase, Ptr, ConstantInt::get(Int32Ty, AccessInfo)});
+  bool useFixedShadowIntrinsic = false;
+  // The memaccess fixed shadow intrinsic is only supported on AArch64,
+  // which allows a 16-bit immediate to be left-shifted by 32.
+  // Since kShadowBaseAlignment == 32, and Linux by default will not
+  // mmap above 48-bits, practically any valid shadow offset is
+  // representable.
+  // In particular, an offset of 4TB (1024 << 32) is representable, and
+  // ought to be good enough for anybody.
+  if (TargetTriple.isAArch64() && ClMappingOffset.getNumOccurrences() > 0) {
+    uint16_t offset_shifted = Mapping.Offset >> 32;
+    useFixedShadowIntrinsic = (uint64_t)offset_shifted << 32 == Mapping.Offset;
+  }
+
+  if (useFixedShadowIntrinsic)
+    IRB.CreateCall(
+        Intrinsic::getDeclaration(
+            M, UseShortGranules
+                   ? Intrinsic::hwasan_check_memaccess_shortgranules_fixedshadow
+                   : Intrinsic::hwasan_check_memaccess_fixedshadow),
+        {Ptr, ConstantInt::get(Int32Ty, AccessInfo), ConstantInt::get(Int64Ty, Mapping.Offset)});
+  else
+    IRB.CreateCall(Intrinsic::getDeclaration(
+                       M, UseShortGranules
+                              ? Intrinsic::hwasan_check_memaccess_shortgranules
+                              : Intrinsic::hwasan_check_memaccess),
+                   {ShadowBase, Ptr, ConstantInt::get(Int32Ty, AccessInfo)});
 }
 
 void HWAddressSanitizer::instrumentMemAccessInline(Value *Ptr, bool IsWrite,