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

[MaskRay]

Summary

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 https://github.com/llvm/llvm-project/issues/63258

[MaskRay]

improve code

[vit9696]

Hi @MaskRay, I applied your patch on top of 15.0.7 and can confirm that it does fix the immediate issue. I will perform more extensive tests as time permits.

Thank you for such a instant patch :-)

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

struct AddressSanitizer {
   void getInterestingMemoryOperands(
       Instruction *I, SmallVectorImpl<InterestingMemoryOperand> &Interesting);
 
   void instrumentMop(ObjectSizeOffsetVisitor &ObjSizeVis,
                      InterestingMemoryOperand &O, bool UseCalls,
                      const DataLayout &DL);
   void instrumentPointerComparisonOrSubtraction(Instruction *I);
   void instrumentAddress(Instruction *OrigIns, Instruction *InsertBefore,
-                         Value *Addr, uint32_t TypeStoreSize, bool IsWrite,
-                         Value *SizeArgument, bool UseCalls, uint32_t Exp);
+                         Value *Addr, MaybeAlign Align, uint32_t TypeStoreSize,
+                         bool IsWrite, Value *SizeArgument, bool UseCalls,
+                         uint32_t Exp);
   Instruction *instrumentAMDGPUAddress(Instruction *OrigIns,
                                        Instruction *InsertBefore, Value *Addr,
                                        uint32_t TypeStoreSize, bool IsWrite,
                                        Value *SizeArgument);
   void instrumentUnusualSizeOrAlignment(Instruction *I,
                                         Instruction *InsertBefore, Value *Addr,
                                         TypeSize TypeStoreSize, bool IsWrite,
                                         Value *SizeArgument, bool UseCalls,
if (!TypeStoreSize.isScalable()) {
     switch (FixedSize) {
     case 8:
     case 16:
     case 32:
     case 64:
     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,
                                          IsWrite, nullptr, UseCalls, Exp);
 }
 
 void AddressSanitizer::instrumentMaskedLoadOrStore(
     AddressSanitizer *Pass, const DataLayout &DL, Type *IntptrTy, Value *Mask,
Instruction *AddressSanitizer::instrumentAMDGPUAddress(
   Value *AddrSpaceZeroLanding =
       SplitBlockAndInsertIfThen(Cmp, InsertBefore, false);
   InsertBefore = cast<Instruction>(AddrSpaceZeroLanding);
   return InsertBefore;
 }
 
 void AddressSanitizer::instrumentAddress(Instruction *OrigIns,
                                          Instruction *InsertBefore, Value *Addr,
+                                         MaybeAlign Alignment,
                                          uint32_t TypeStoreSize, bool IsWrite,
                                          Value *SizeArgument, bool UseCalls,
                                          uint32_t Exp) {
   if (TargetTriple.isAMDGPU()) {
     InsertBefore = instrumentAMDGPUAddress(OrigIns, InsertBefore, Addr,
                                            TypeStoreSize, IsWrite, SizeArgument);
     if (!InsertBefore)
       return;
else
                      {AddrLong, ConstantInt::get(IRB.getInt32Ty(), Exp)});
     return;
   }
 
   Type *ShadowTy =
       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));
+  int NewAlignScale =
+      std::max(int(Log2(Alignment.valueOrOne()) - Mapping.Scale), 0);
+  Value *ShadowValue = IRB.CreateAlignedLoad(
+      ShadowTy, IRB.CreateIntToPtr(ShadowPtr, ShadowPtrTy),
+      Align(1 << NewAlignScale));
 
   Value *Cmp = IRB.CreateIsNotNull(ShadowValue);
   size_t Granularity = 1ULL << Mapping.Scale;
   Instruction *CrashTerm = nullptr;
 
   if (ClAlwaysSlowPath || (TypeStoreSize < 8 * Granularity)) {
     // We use branch weights for the slow path check, to indicate that the slow
     // path is rarely taken. This seems to be the case for SPEC benchmarks.
if (UseCalls) {
     else
       IRB.CreateCall(AsanMemoryAccessCallbackSized[IsWrite][1],
                      {AddrLong, Size, ConstantInt::get(IRB.getInt32Ty(), Exp)});
   } else {
     Value *SizeMinusOne = IRB.CreateSub(Size, ConstantInt::get(IntptrTy, 1));
     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);
   }
 }
 
 void ModuleAddressSanitizer::poisonOneInitializer(Function &GlobalInit,
                                                   GlobalValue *ModuleName) {
   // Set up the arguments to our poison/unpoison functions.
   IRBuilder<> IRB(&GlobalInit.front(),
                   GlobalInit.front().getFirstInsertionPt());

llvm/test/Instrumentation/AddressSanitizer/basic.ll

define void @i64test_align1(ptr %b) nounwind uwtable sanitize_address {
   ret void
 }
 
 ; CHECK-LABEL: i64test_align1
 ; CHECK: __asan_report_store_n{{.*}}, i64 8)
 ; 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:
   %t = load i80, ptr %a
   store i80 %t, ptr %b, align 8
   ret void
 }
 

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

 ; CALLS-NEXT:    call void @__asan_storeN(i64 [[TMP1]], i64 64)
 ; CALLS-NEXT:    store <16 x i32> zeroinitializer, ptr [[P]], align 64
 ; CALLS-NEXT:    ret void
 ;
   store <16 x i32> zeroinitializer, ptr %p
   ret void
 }
 
+define void @store.v2i64.align8(ptr %p) sanitize_address {
+; CHECK-LABEL: @store.v2i64.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.v2i64.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(
 ; CHECK-NEXT:    [[TMP1:%.*]] = call i64 @llvm.vscale.i64()
 ; CHECK-NEXT:    [[TMP2:%.*]] = mul i64 [[TMP1]], 32
 ; CHECK-NEXT:    [[TMP3:%.*]] = lshr i64 [[TMP2]], 3
 ; CHECK-NEXT:    [[TMP4:%.*]] = ptrtoint ptr [[P:%.*]] to i64
 ; CHECK-NEXT:    [[TMP5:%.*]] = sub i64 [[TMP3]], 1