From 1f411252d06566153109e47cce6f425746a43e7f Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa@amd.com>
Date: Mon, 19 Aug 2024 21:07:57 +0530
Subject: [PATCH] [AMDGPU] Update instrumentAddress method to support aligned
 size and unsusual size access.

---
 .../AMDGPU/AMDGPUAsanInstrumentation.cpp      | 51 ++++++++++++++++---
 .../Target/AMDGPU/AMDGPUAsanInstrumentation.h |  8 +--
 2 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUAsanInstrumentation.cpp b/llvm/lib/Target/AMDGPU/AMDGPUAsanInstrumentation.cpp
index 593fca5bc3ed6..4c8ddbd9aabd5 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUAsanInstrumentation.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUAsanInstrumentation.cpp
@@ -147,11 +147,13 @@ static Value *memToShadow(Module &M, IRBuilder<> &IRB, Type *IntptrTy,
   return IRB.CreateAdd(Shadow, ShadowBase);
 }
 
-void instrumentAddress(Module &M, IRBuilder<> &IRB, Instruction *OrigIns,
-                       Instruction *InsertBefore, Value *Addr,
-                       MaybeAlign Alignment, uint32_t TypeStoreSize,
-                       bool IsWrite, Value *SizeArgument, bool UseCalls,
-                       bool Recover, int AsanScale, int AsanOffset) {
+static void instrumentAddressImpl(Module &M, IRBuilder<> &IRB,
+                                  Instruction *OrigIns,
+                                  Instruction *InsertBefore, Value *Addr,
+                                  Align Alignment, uint32_t TypeStoreSize,
+                                  bool IsWrite, Value *SizeArgument,
+                                  bool UseCalls, bool Recover, int AsanScale,
+                                  int AsanOffset) {
   Type *AddrTy = Addr->getType();
   Type *IntptrTy = M.getDataLayout().getIntPtrType(
       M.getContext(), AddrTy->getPointerAddressSpace());
@@ -164,7 +166,7 @@ void instrumentAddress(Module &M, IRBuilder<> &IRB, Instruction *OrigIns,
   Value *ShadowPtr =
       memToShadow(M, IRB, IntptrTy, AddrLong, AsanScale, AsanOffset);
   const uint64_t ShadowAlign =
-      std::max<uint64_t>(Alignment.valueOrOne().value() >> AsanScale, 1);
+      std::max<uint64_t>(Alignment.value() >> AsanScale, 1);
   Value *ShadowValue = IRB.CreateAlignedLoad(
       ShadowTy, IRB.CreateIntToPtr(ShadowPtr, ShadowPtrTy), Align(ShadowAlign));
   Value *Cmp = IRB.CreateIsNotNull(ShadowValue);
@@ -179,6 +181,43 @@ void instrumentAddress(Module &M, IRBuilder<> &IRB, Instruction *OrigIns,
   return;
 }
 
+void instrumentAddress(Module &M, IRBuilder<> &IRB, Instruction *OrigIns,
+                       Instruction *InsertBefore, Value *Addr, Align Alignment,
+                       TypeSize TypeStoreSize, bool IsWrite,
+                       Value *SizeArgument, bool UseCalls, bool Recover,
+                       int AsanScale, int AsanOffset) {
+  if (!TypeStoreSize.isScalable()) {
+    unsigned Granularity = 1 << AsanScale;
+    const auto FixedSize = TypeStoreSize.getFixedValue();
+    switch (FixedSize) {
+    case 8:
+    case 16:
+    case 32:
+    case 64:
+    case 128:
+      if (Alignment.value() >= Granularity ||
+          Alignment.value() >= FixedSize / 8)
+        return instrumentAddressImpl(
+            M, IRB, OrigIns, InsertBefore, Addr, Alignment, FixedSize, IsWrite,
+            SizeArgument, UseCalls, Recover, AsanScale, AsanOffset);
+    }
+  }
+  // Instrument unusual size or unusual alignment.
+  IRB.SetInsertPoint(InsertBefore);
+  Type *AddrTy = Addr->getType();
+  Type *IntptrTy = M.getDataLayout().getIntPtrType(AddrTy);
+  Value *NumBits = IRB.CreateTypeSize(IntptrTy, TypeStoreSize);
+  Value *Size = IRB.CreateLShr(NumBits, ConstantInt::get(IntptrTy, 3));
+  Value *AddrLong = IRB.CreatePtrToInt(Addr, IntptrTy);
+  Value *SizeMinusOne = IRB.CreateAdd(Size, ConstantInt::get(IntptrTy, -1));
+  Value *LastByte =
+      IRB.CreateIntToPtr(IRB.CreateAdd(AddrLong, SizeMinusOne), AddrTy);
+  instrumentAddressImpl(M, IRB, OrigIns, InsertBefore, Addr, {}, 8, IsWrite,
+                        SizeArgument, UseCalls, Recover, AsanScale, AsanOffset);
+  instrumentAddressImpl(M, IRB, OrigIns, InsertBefore, LastByte, {}, 8, IsWrite,
+                        SizeArgument, UseCalls, Recover, AsanScale, AsanOffset);
+}
+
 void getInterestingMemoryOperands(
     Module &M, Instruction *I,
     SmallVectorImpl<InterestingMemoryOperand> &Interesting) {
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUAsanInstrumentation.h b/llvm/lib/Target/AMDGPU/AMDGPUAsanInstrumentation.h
index 64d78c4aeb692..f28952b431410 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUAsanInstrumentation.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPUAsanInstrumentation.h
@@ -43,10 +43,10 @@ uint64_t getRedzoneSizeForGlobal(int Scale, uint64_t SizeInBytes);
 /// Instrument the memory operand Addr.
 /// Generates report blocks that catch the addressing errors.
 void instrumentAddress(Module &M, IRBuilder<> &IRB, Instruction *OrigIns,
-                       Instruction *InsertBefore, Value *Addr,
-                       MaybeAlign Alignment, uint32_t TypeStoreSize,
-                       bool IsWrite, Value *SizeArgument, bool UseCalls,
-                       bool Recover, int Scale, int Offset);
+                       Instruction *InsertBefore, Value *Addr, Align Alignment,
+                       TypeSize TypeStoreSize, bool IsWrite,
+                       Value *SizeArgument, bool UseCalls, bool Recover,
+                       int Scale, int Offset);
 
 /// Get all the memory operands from the instruction
 /// that needs to be instrumented