[asan] Fix shadow load alignment for unaligned 128-bit load/store

When a 128-bit load/store is aligned by 8, we incorrectly emit `load i16, ptr ..., align 2`
while the shadow memory address may not be aligned by 2.

This manifests as possibly-misaligned shadow memory load with `-mstrict-align`,
e.g. `clang --target=aarch64-linux -O2 -mstrict-align -fsanitize=address`
```
__attribute__((noinline)) void foo(unsigned long *ptr) {
  ptr[0] = 3;
  ptr[1] = 3;
}
// ldrh    w8, [x9, x8]  // the shadow memory load may not be aligned by 2
```

Infer the shadow memory alignment from the load/store alignment to set the
correct alignment. The generated code now uses two ldrb and one orr.

Fix llvm#63258

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

Author: MaskRay

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

@@ -684,7 +684,8 @@ struct AddressSanitizer {
                      const DataLayout &DL);
   void instrumentPointerComparisonOrSubtraction(Instruction *I);
   void instrumentAddress(Instruction *OrigIns, Instruction *InsertBefore,
-                         Value *Addr, uint32_t TypeStoreSize, bool IsWrite,
+                         Value *Addr, MaybeAlign Alignment,
+                         uint32_t TypeStoreSize, bool IsWrite,
                          Value *SizeArgument, bool UseCalls, uint32_t Exp);
   Instruction *instrumentAMDGPUAddress(Instruction *OrigIns,
                                        Instruction *InsertBefore, Value *Addr,
@@ -1480,8 +1481,9 @@ static void doInstrumentAddress(AddressSanitizer *Pass, Instruction *I,
     case 128:
       if (!Alignment || *Alignment >= Granularity ||
           *Alignment >= FixedSize / 8)
-        return Pass->instrumentAddress(I, InsertBefore, Addr, FixedSize,
-                                       IsWrite, nullptr, UseCalls, Exp);
+        return Pass->instrumentAddress(I, InsertBefore, Addr, Alignment,
+                                       FixedSize, IsWrite, nullptr, UseCalls,
+                                       Exp);
     }
   }
   Pass->instrumentUnusualSizeOrAlignment(I, InsertBefore, Addr, TypeStoreSize,
@@ -1681,6 +1683,7 @@ Instruction *AddressSanitizer::instrumentAMDGPUAddress(
 
 void AddressSanitizer::instrumentAddress(Instruction *OrigIns,
                                          Instruction *InsertBefore, Value *Addr,
+                                         MaybeAlign Alignment,
                                          uint32_t TypeStoreSize, bool IsWrite,
                                          Value *SizeArgument, bool UseCalls,
                                          uint32_t Exp) {
@@ -1720,8 +1723,10 @@ void AddressSanitizer::instrumentAddress(Instruction *OrigIns,
       IntegerType::get(*C, std::max(8U, TypeStoreSize >> Mapping.Scale));
   Type *ShadowPtrTy = PointerType::get(ShadowTy, 0);
   Value *ShadowPtr = memToShadow(AddrLong, IRB);
-  Value *ShadowValue =
-      IRB.CreateLoad(ShadowTy, IRB.CreateIntToPtr(ShadowPtr, ShadowPtrTy));
+  const uint64_t ShadowAlign =
+      std::max<uint64_t>(Alignment.valueOrOne().value() >> Mapping.Scale, 1);
+  Value *ShadowValue = IRB.CreateAlignedLoad(
+      ShadowTy, IRB.CreateIntToPtr(ShadowPtr, ShadowPtrTy), Align(ShadowAlign));
 
   Value *Cmp = IRB.CreateIsNotNull(ShadowValue);
   size_t Granularity = 1ULL << Mapping.Scale;
@@ -1778,8 +1783,9 @@ void AddressSanitizer::instrumentUnusualSizeOrAlignment(
     Value *LastByte = IRB.CreateIntToPtr(
         IRB.CreateAdd(AddrLong, SizeMinusOne),
         Addr->getType());
-    instrumentAddress(I, InsertBefore, Addr, 8, IsWrite, Size, false, Exp);
-    instrumentAddress(I, InsertBefore, LastByte, 8, IsWrite, Size, false, Exp);
+    instrumentAddress(I, InsertBefore, Addr, {}, 8, IsWrite, Size, false, Exp);
+    instrumentAddress(I, InsertBefore, LastByte, {}, 8, IsWrite, Size, false,
+                      Exp);
   }
 }
 

llvm/test/Instrumentation/AddressSanitizer/basic.ll

@@ -135,6 +135,27 @@ define void @i64test_align1(ptr %b) nounwind uwtable sanitize_address {
 ; CHECK: __asan_report_store_n{{.*}}, i64 8)
 ; CHECK: ret void
 
+define void @i128test_align8(ptr %a) nounwind uwtable sanitize_address {
+entry:
+  store i128 0, ptr %a, align 8
+  ret void
+}
+; CHECK-LABEL: define {{[^@]+}}@i128test_align8(
+; CHECK-S3:      load i16, ptr %[[#]], align 1
+; CHECK-S3-NEXT: icmp ne i16 %[[#]], 0
+; CHECK-S5:      load i8, ptr %[[#]], align 1
+; CHECK-S5:      load i8, ptr %[[#]], align 1
+
+define void @i128test_align16(ptr %a) nounwind uwtable sanitize_address {
+entry:
+  store i128 0, ptr %a, align 16
+  ret void
+}
+; CHECK-LABEL: define {{[^@]+}}@i128test_align16(
+; CHECK-S3:      load i16, ptr %[[#]], align 2
+; CHECK-S3-NEXT: icmp ne i16 %[[#]], 0
+; CHECK-S5:      load i8, ptr %[[#]], align 1
+; CHECK-S5-NEXT: icmp ne i8 %[[#]], 0
 
 define void @i80test(ptr %a, ptr %b) nounwind uwtable sanitize_address {
   entry:

llvm/test/Instrumentation/AddressSanitizer/vector-load-store.ll

@@ -375,6 +375,31 @@ define void @store.v16i32(ptr %p) sanitize_address {
   ret void
 }
 
+define void @store.v2i32.align8(ptr %p) sanitize_address {
+; CHECK-LABEL: @store.v2i32.align8(
+; CHECK-NEXT:    [[TMP1:%.*]] = ptrtoint ptr [[P:%.*]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = lshr i64 [[TMP1]], 3
+; CHECK-NEXT:    [[TMP3:%.*]] = or i64 [[TMP2]], 17592186044416
+; CHECK-NEXT:    [[TMP4:%.*]] = inttoptr i64 [[TMP3]] to ptr
+; CHECK-NEXT:    [[TMP5:%.*]] = load i8, ptr [[TMP4]], align 1
+; CHECK-NEXT:    [[TMP6:%.*]] = icmp ne i8 [[TMP5]], 0
+; CHECK-NEXT:    br i1 [[TMP6]], label [[TMP7:%.*]], label [[TMP8:%.*]]
+; CHECK:       7:
+; CHECK-NEXT:    call void @__asan_report_store8(i64 [[TMP1]]) #[[ATTR4]]
+; CHECK-NEXT:    unreachable
+; CHECK:       8:
+; CHECK-NEXT:    store <2 x i32> zeroinitializer, ptr [[P]], align 8
+; CHECK-NEXT:    ret void
+;
+; CALLS-LABEL: @store.v2i32.align8(
+; CALLS-NEXT:    [[TMP1:%.*]] = ptrtoint ptr [[P:%.*]] to i64
+; CALLS-NEXT:    call void @__asan_store8(i64 [[TMP1]])
+; CALLS-NEXT:    store <2 x i32> zeroinitializer, ptr [[P]], align 8
+; CALLS-NEXT:    ret void
+;
+  store <2 x i32> zeroinitializer, ptr %p, align 8
+  ret void
+}
 
 define void @load.nxv1i32(ptr %p) sanitize_address {
 ; CHECK-LABEL: @load.nxv1i32(