[HWASAN] Update the tag info for X86_64

[xiangzhangllvm]

Summary

In LAM model X86_64 will use bits 57-62 (of 0-63) as HWASAN tag.
So here we make sure the tag shift position and tag mask is correct for x86-64.

And support the HWASAN for setjmp longjmp in compiler-rt

[xiangzhangllvm]

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

188

I add back untag point, (disable it in default), it is useful to test some simple test with HWASAN.

[xiangzhangllvm]

Update clang-format

[pcc]

It doesn't look like this was reviewed by anyone. Could you revert this please until it gets reviewed?

[xiangzhangllvm]

In D102472#2770228, @pcc wrote:

It doesn't look like this was reviewed by anyone. Could you revert this please until it gets reviewed?

Done,
revert it at 02f2d739e074782ac47531196739a2e003603e62

Could you help review it ?

[xiangzhangllvm]

Hello @eugenis @vitalybuka @morehouse I see you change the related code at compiler-rt, Could you help review this patch? This patch just limited the tag for x86_64.

[morehouse]

Hi Xiang, thanks for the patch.

I haven't looked too close yet, but a couple questions:

1. There seems to be multiple things going on in this patch: adding stack support on x86, adding a (debug?) flag ClUntagPointer, adding x86_64 assembly to the runtime. Could you break this up into smaller separate patches for independent review?
2. How are you testing the patch? I am currently setting up a QEMU+LAM buildbot for CI, and I would prefer if we waited to land any more LAM stuff until the buildbot is setup to ensure things don't break.

[xiangzhangllvm]

Hello @morehouse, thanks for your work too:
1 OK, I'll break it up into 2 patches, one focus on changes in compiler, the other focus on changes in compiler-rt.

2 In fact, We can divide it into two steps, Step 1 is "Tag correctness checking", Step 2 is "Hardware supporting point with tag".

 The option "ClUntagPointer" can help us first testing the Step 1.
 For example: we write 7 in to int *p, for HWASAN, we do 2 things:
 Step1:   Checking tag of p :
              call void @__hwasan_store4 (p_tagged)     // Checking tag is correct or not
 Step2:   Write 7 into mem:
              *p_tagged = 7          // Need Hardware supporting.  But if we use "ClUntagPointer" change "*p_tagged = 7" --> "*p = 7", we can run the program, and test Step 1 for HWASAN.

For HWASAN, Most code in compiler and compiler-rt is doing work for Step 1. So we can use "ClUntagPointer" first test it.
After we make sure Step 1 is OK, when the Hardware/OS is ready, we just need to totally/really enable HWASAN by just, I think, updating the system call in InitPrctl.

[xiangzhangllvm]

Hello, @morehouse, I first spilt the change of compiler to https://reviews.llvm.org/D102901
I'll update this link after that small patch in.

[morehouse]

In D102472#2772657, @xiangzhangllvm wrote:

2 In fact, We can divide it into two steps, Step 1 is "Tag correctness checking", Step 2 is "Hardware supporting point with tag".

 The option "ClUntagPointer" can help us first testing the Step 1.
 For example: we write 7 in to int *p, for HWASAN, we do 2 things:
 Step1:   Checking tag of p :
              call void @__hwasan_store4 (p_tagged)     // Checking tag is correct or not
 Step2:   Write 7 into mem:
              *p_tagged = 7          // Need Hardware supporting.  But if we use "ClUntagPointer" change "*p_tagged = 7" --> "*p = 7", we can run the program, and test Step 1 for HWASAN.

For HWASAN, Most code in compiler and compiler-rt is doing work for Step 1. So we can use "ClUntagPointer" first test it.
After we make sure Step 1 is OK, when the Hardware/OS is ready, we just need to totally/really enable HWASAN by just, I think, updating the system call in InitPrctl.

I recently removed a bunch of untagging stuff from our tests, and I'd like to avoid adding it back. This is why I'm working on a buildbot to run tests in QEMU with LAM enabled.

I've also been testing locally with LAM in QEMU.

[xiangzhangllvm]

Hello @morehouse , I saw the change of compiler for X86 ( https://reviews.llvm.org/D102901 ) has been merged.
Do you have plan to support the Set/longjmp on compiler-rt for X86 (the rest of current patch did) ?

[morehouse]

In D102472#2889460, @xiangzhangllvm wrote:

Hello @morehouse , I saw the change of compiler for X86 ( https://reviews.llvm.org/D102901 ) has been merged.
Do you have plan to support the Set/longjmp on compiler-rt for X86 (the rest of current patch did) ?

Eventually. The existing tests at the time passed without it, so I haven't done anything with this patch.

[xiangzhangllvm]

Eventually. The existing tests at the time passed without it, so I haven't done anything with this patch.

Does the "existing tests" just means LIT tests about HWASAN in compiler-rt ? thanks.

[morehouse]

Yes.

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

const char kHwasanShadowMemoryDynamicAddress[] =
     "__hwasan_shadow_memory_dynamic_address";
 
 // Accesses sizes are powers of two: 1, 2, 4, 8, 16.
 static const size_t kNumberOfAccessSizes = 5;
 
 static const size_t kDefaultShadowScale = 4;
 static const uint64_t kDynamicShadowSentinel =
     std::numeric_limits<uint64_t>::max();
-static const unsigned kPointerTagShift = 56;
 
 static const unsigned kShadowBaseAlignment = 32;
 
-static cl::opt<std::string> ClMemoryAccessCallbackPrefix(
-    "hwasan-memory-access-callback-prefix",
-    cl::desc("Prefix for memory access callbacks"), cl::Hidden,
-    cl::init("__hwasan_"));
-
-static cl::opt<bool>
-    ClInstrumentWithCalls("hwasan-instrument-with-calls",
-                cl::desc("instrument reads and writes with callbacks"),
-                cl::Hidden, cl::init(false));
+static cl::opt<std::string>
+    ClMemoryAccessCallbackPrefix("hwasan-memory-access-callback-prefix",
+                                 cl::desc("Prefix for memory access callbacks"),
+                                 cl::Hidden, cl::init("__hwasan_"));
+
+static cl::opt<bool> ClInstrumentWithCalls(
+    "hwasan-instrument-with-calls",
+    cl::desc("instrument reads and writes with callbacks"), cl::Hidden,
+    cl::init(false));
 
 static cl::opt<bool> ClInstrumentReads("hwasan-instrument-reads",
                                        cl::desc("instrument read instructions"),
                                        cl::Hidden, cl::init(true));
 
-static cl::opt<bool> ClInstrumentWrites(
-    "hwasan-instrument-writes", cl::desc("instrument write instructions"),
-    cl::Hidden, cl::init(true));
+static cl::opt<bool>
+    ClInstrumentWrites("hwasan-instrument-writes",
+                       cl::desc("instrument write instructions"), cl::Hidden,
+                       cl::init(true));
 
 static cl::opt<bool> ClInstrumentAtomics(
     "hwasan-instrument-atomics",
     cl::desc("instrument atomic instructions (rmw, cmpxchg)"), cl::Hidden,
     cl::init(true));
 
 static cl::opt<bool> ClInstrumentByval("hwasan-instrument-byval",
                                        cl::desc("instrument byval arguments"),
                                        cl::Hidden, cl::init(true));
 
-static cl::opt<bool> ClRecover(
-    "hwasan-recover",
+static cl::opt<bool>
+    ClRecover("hwasan-recover",
               cl::desc("Enable recovery mode (continue-after-error)."),
               cl::Hidden, cl::init(false));
 
 static cl::opt<bool> ClInstrumentStack("hwasan-instrument-stack",
                                        cl::desc("instrument stack (allocas)"),
                                        cl::Hidden, cl::init(true));
 
 static cl::opt<bool> ClUARRetagToZero(
     "hwasan-uar-retag-to-zero",
     cl::desc("Clear alloca tags before returning from the function to allow "

 static cl::opt<bool> ClGlobals("hwasan-globals", cl::desc("Instrument globals"),
                                cl::Hidden, cl::init(false), cl::ZeroOrMore);
 
 static cl::opt<int> ClMatchAllTag(
     "hwasan-match-all-tag",
     cl::desc("don't report bad accesses via pointers with this tag"),
     cl::Hidden, cl::init(-1));
 
-static cl::opt<bool> ClEnableKhwasan(
-    "hwasan-kernel",
+static cl::opt<bool>
+    ClEnableKhwasan("hwasan-kernel",
                     cl::desc("Enable KernelHWAddressSanitizer instrumentation"),
                     cl::Hidden, cl::init(false));
 
 // These flags allow to change the shadow mapping and control how shadow memory
 // is accessed. The shadow mapping looks like:
 //    Shadow = (Mem >> scale) + offset
 
 static cl::opt<uint64_t>
     ClMappingOffset("hwasan-mapping-offset",
                     cl::desc("HWASan shadow mapping offset [EXPERIMENTAL]"),
static cl::opt<bool> ClInstrumentPersonalityFunctions(
     "hwasan-instrument-personality-functions",
     cl::desc("instrument personality functions"), cl::Hidden, cl::init(false),
     cl::ZeroOrMore);
 
 static cl::opt<bool> ClInlineAllChecks("hwasan-inline-all-checks",
                                        cl::desc("inline all checks"),
                                        cl::Hidden, cl::init(false));
 
+static cl::opt<bool> ClUntagPointer("hwasan-untag-mem-operation",
+                                    cl::desc("untag mem operate"), cl::Hidden,
+                                    cl::init(false));
+
+// TODO: Need to co-related with clang option.
+static cl::opt<bool> ClUsePageAlias("hwasan-use-page-alias",
+                                    cl::desc("hwasan use page alias"),
+                                    cl::Hidden, cl::init(true));
+
 namespace {
 
 /// An instrumentation pass implementing detection of addressability bugs
 /// using tagged pointers.
 class HWAddressSanitizer {
 public:
   explicit HWAddressSanitizer(Module &M, bool CompileKernel = false,
-                              bool Recover = false) : M(M) {
+                              bool Recover = false)
+      : M(M) {
     this->Recover = ClRecover.getNumOccurrences() > 0 ? ClRecover : Recover;
-    this->CompileKernel = ClEnableKhwasan.getNumOccurrences() > 0 ?
-        ClEnableKhwasan : CompileKernel;
+    this->CompileKernel = ClEnableKhwasan.getNumOccurrences() > 0
+                              ? ClEnableKhwasan
+                              : CompileKernel;
 
     initializeModule();
   }
 
   bool sanitizeFunction(Function &F);
   void initializeModule();
   void createHwasanCtorComdat();
 
public:
   void instrumentMemIntrinsic(MemIntrinsic *MI);
   bool instrumentMemAccess(InterestingMemoryOperand &O);
   bool ignoreAccess(Value *Ptr);
   void getInterestingMemoryOperands(
       Instruction *I, SmallVectorImpl<InterestingMemoryOperand> &Interesting);
 
   bool isInterestingAlloca(const AllocaInst &AI);
   bool tagAlloca(IRBuilder<> &IRB, AllocaInst *AI, Value *Tag, size_t Size);
+  Value *getTargetTagPtr(IRBuilder<> &IRB, Value *PtrLong, Value *Tag);
   Value *tagPointer(IRBuilder<> &IRB, Type *Ty, Value *PtrLong, Value *Tag);
   Value *untagPointer(IRBuilder<> &IRB, Value *PtrLong);
   bool instrumentStack(
       SmallVectorImpl<AllocaInst *> &Allocas,
       DenseMap<AllocaInst *, std::vector<DbgVariableIntrinsic *>> &AllocaDbgMap,
       SmallVectorImpl<Instruction *> &RetVec, Value *StackTag);
   Value *readRegister(IRBuilder<> &IRB, StringRef Name);
   bool instrumentLandingPads(SmallVectorImpl<Instruction *> &RetVec);
   Value *getNextTagWithCall(IRBuilder<> &IRB);
   Value *getStackBaseTag(IRBuilder<> &IRB);
   Value *getAllocaTag(IRBuilder<> &IRB, Value *StackTag, AllocaInst *AI,
                       unsigned AllocaNo);
   Value *getUARTag(IRBuilder<> &IRB, Value *StackTag);
 
   Value *getHwasanThreadSlotPtr(IRBuilder<> &IRB, Type *Ty);
+  Value *retagTargetTag(IRBuilder<> &IRB, Value *OldTag);
   void emitPrologue(IRBuilder<> &IRB, bool WithFrameRecord);
 
   void instrumentGlobal(GlobalVariable *GV, uint8_t Tag);
   void instrumentGlobals();
 
   void instrumentPersonalityFunctions();
 
+  unsigned getTargetTagShift() {
+    if (TargetTriple.getArch() == Triple::x86_64)
+      return 57;
+    else
+      return 56;
+  }
+
+  uint64_t getTargetTagMask() {
+    if (TargetTriple.getArch() == Triple::x86_64)
+      return 0x3FLL;
+    else
+      return 0xFFLL;
+  }
+
+  unsigned kPointerTagShift;
+  uint64_t TagMaskByte;
+
 private:
   LLVMContext *C;
   Module &M;
   Triple TargetTriple;
   FunctionCallee HWAsanMemmove, HWAsanMemcpy, HWAsanMemset;
   FunctionCallee HWAsanHandleVfork;
 
   /// This struct defines the shadow mapping using the rule:

 void HWAddressSanitizer::initializeModule() {
   LLVM_DEBUG(dbgs() << "Init " << M.getName() << "\n");
   auto &DL = M.getDataLayout();
 
   TargetTriple = Triple(M.getTargetTriple());
 
   // x86_64 uses userspace pointer aliases, currently heap-only with callback
   // instrumentation only.
-  UsePageAliases = TargetTriple.getArch() == Triple::x86_64;
+  UsePageAliases = ClUsePageAlias && (TargetTriple.getArch() == Triple::x86_64);
   InstrumentWithCalls = UsePageAliases ? true : ClInstrumentWithCalls;
   InstrumentStack = UsePageAliases ? false : ClInstrumentStack;
 
   Mapping.init(TargetTriple, InstrumentWithCalls);
 
   C = &(M.getContext());
   IRBuilder<> IRB(*C);
   IntptrTy = IRB.getIntPtrTy(DL);
   Int8PtrTy = IRB.getInt8PtrTy();
   Int8Ty = IRB.getInt8Ty();
   Int32Ty = IRB.getInt32Ty();
 
   HwasanCtorFunction = nullptr;
+  kPointerTagShift = getTargetTagShift();
+  TagMaskByte = getTargetTagMask();
 
   // Older versions of Android do not have the required runtime support for
   // short granules, global or personality function instrumentation. On other
   // platforms we currently require using the latest version of the runtime.
   bool NewRuntime =
       !TargetTriple.isAndroid() || !TargetTriple.isAndroidVersionLT(30);
 
   UseShortGranules =
void HWAddressSanitizer::instrumentMemAccessInline(Value *Ptr, bool IsWrite,
   SplitBlockAndInsertIfThen(InlineTagMismatch, CheckTerm, false,
                             MDBuilder(*C).createBranchWeights(1, 100000),
                             (DomTreeUpdater *)nullptr, nullptr,
                             CheckFailTerm->getParent());
 
   IRB.SetInsertPoint(CheckFailTerm);
   InlineAsm *Asm;
   switch (TargetTriple.getArch()) {
-    case Triple::x86_64:
-      // The signal handler will find the data address in rdi.
-      Asm = InlineAsm::get(
-          FunctionType::get(IRB.getVoidTy(), {PtrLong->getType()}, false),
-          "int3\nnopl " +
-              itostr(0x40 + (AccessInfo & HWASanAccessInfo::RuntimeMask)) +
-              "(%rax)",
-          "{rdi}",
-          /*hasSideEffects=*/true);
-      break;
-    case Triple::aarch64:
-    case Triple::aarch64_be:
-      // The signal handler will find the data address in x0.
-      Asm = InlineAsm::get(
-          FunctionType::get(IRB.getVoidTy(), {PtrLong->getType()}, false),
-          "brk #" +
-              itostr(0x900 + (AccessInfo & HWASanAccessInfo::RuntimeMask)),
-          "{x0}",
-          /*hasSideEffects=*/true);
-      break;
-    default:
-      report_fatal_error("unsupported architecture");
+        "brk #" + itostr(0x900 + (AccessInfo & HWASanAccessInfo::RuntimeMask)),
   }
   IRB.CreateCall(Asm, PtrLong);
   if (Recover)
     cast<BranchInst>(CheckFailTerm)->setSuccessor(0, CheckTerm->getParent());
 }
 
 void HWAddressSanitizer::instrumentMemIntrinsic(MemIntrinsic *MI) {
   IRBuilder<> IRB(MI);

 }
 
 bool HWAddressSanitizer::instrumentMemAccess(InterestingMemoryOperand &O) {
   Value *Addr = O.getPtr();
 
   LLVM_DEBUG(dbgs() << "Instrumenting: " << O.getInsn() << "\n");
 
   if (O.MaybeMask)
-    return false; //FIXME
+    return false; // FIXME
 
   IRBuilder<> IRB(O.getInsn());
   if (isPowerOf2_64(O.TypeSize) &&
       (O.TypeSize / 8 <= (1ULL << (kNumberOfAccessSizes - 1))) &&
       (!O.Alignment || *O.Alignment >= (1ULL << Mapping.Scale) ||
        *O.Alignment >= O.TypeSize / 8)) {
     size_t AccessSizeIndex = TypeSizeToSizeIndex(O.TypeSize);
     if (InstrumentWithCalls) {
       IRB.CreateCall(HwasanMemoryAccessCallback[O.IsWrite][AccessSizeIndex],
                      IRB.CreatePointerCast(Addr, IntptrTy));
     } else {
       instrumentMemAccessInline(Addr, O.IsWrite, AccessSizeIndex, O.getInsn());
     }
   } else {
     IRB.CreateCall(HwasanMemoryAccessCallbackSized[O.IsWrite],
                    {IRB.CreatePointerCast(Addr, IntptrTy),
                     ConstantInt::get(IntptrTy, O.TypeSize / 8)});
   }
+  if (ClUntagPointer)
     untagPointerOperand(O.getInsn(), Addr);
 
   return true;
 }
 
 static uint64_t getAllocaSizeInBytes(const AllocaInst &AI) {
   uint64_t ArraySize = 1;
   if (AI.isArrayAllocation()) {
     const ConstantInt *CI = dyn_cast<ConstantInt>(AI.getArraySize());
     assert(CI && "non-constant array size");
     ArraySize = CI->getZExtValue();
   }
   Type *Ty = AI.getAllocatedType();
   uint64_t SizeInBytes = AI.getModule()->getDataLayout().getTypeAllocSize(Ty);
   return SizeInBytes * ArraySize;
 }
 
-bool HWAddressSanitizer::tagAlloca(IRBuilder<> &IRB, AllocaInst *AI,
-                                   Value *Tag, size_t Size) {
+bool HWAddressSanitizer::tagAlloca(IRBuilder<> &IRB, AllocaInst *AI, Value *Tag,
+                                   size_t Size) {
   size_t AlignedSize = alignTo(Size, Mapping.getObjectAlignment());
   if (!UseShortGranules)
     Size = AlignedSize;
 
   Value *JustTag = IRB.CreateTrunc(Tag, IRB.getInt8Ty());
   if (InstrumentWithCalls) {
     IRB.CreateCall(HwasanTagMemoryFunc,
                    {IRB.CreatePointerCast(AI, Int8PtrTy), JustTag,
if (Size != AlignedSize) {
       IRB.CreateStore(JustTag, IRB.CreateConstGEP1_32(
                                    Int8Ty, IRB.CreateBitCast(AI, Int8PtrTy),
                                    AlignedSize - 1));
     }
   }
   return true;
 }
 
-static unsigned RetagMask(unsigned AllocaNo) {
+static unsigned retagMaskArmv8(unsigned AllocaNo) {
   // A list of 8-bit numbers that have at most one run of non-zero bits.
   // x = x ^ (mask << 56) can be encoded as a single armv8 instruction for these
   // masks.
   // The list does not include the value 255, which is used for UAR.
   //
   // Because we are more likely to use earlier elements of this list than later
   // ones, it is sorted in increasing order of probability of collision with a
   // mask allocated (temporally) nearby. The program that generated this list
   // can be found at:
   // https://github.com/google/sanitizers/blob/master/hwaddress-sanitizer/sort_masks.py
   static unsigned FastMasks[] = {0,  128, 64,  192, 32,  96,  224, 112, 240,
                                  48, 16,  120, 248, 56,  24,  8,   124, 252,
                                  60, 28,  12,  4,   126, 254, 62,  30,  14,
                                  6,  2,   127, 63,  31,  15,  7,   3,   1};
   return FastMasks[AllocaNo % (sizeof(FastMasks) / sizeof(FastMasks[0]))];
 }
 
+static unsigned retagX86_64(unsigned AllocaNo, uint64_t TagMaskByte) {
+  return AllocaNo & TagMaskByte;
+}
+
+static unsigned retagMask(unsigned AllocaNo, Triple TargetTriple,
+                          uint64_t TagMaskByte) {
+  if (TargetTriple.getArch() == Triple::x86_64)
+    return retagX86_64(AllocaNo, TagMaskByte);
+  return retagMaskArmv8(AllocaNo);
+}
+
+// The whole pass imply the Tag is 8-bits size, because it is first
+// implemented by aarch64 whose tag is happen 1 byte. We should diff
+// it frome other targets now.
+Value *HWAddressSanitizer::retagTargetTag(IRBuilder<> &IRB, Value *OldTag) {
+  if (TargetTriple.getArch() == Triple::x86_64) {
+    Constant *TagMask = ConstantInt::get(IntptrTy, TagMaskByte);
+    Value *NewTag = IRB.CreateAnd(OldTag, TagMask);
+    return NewTag;
+  } else
+    return OldTag;
+}
+
 Value *HWAddressSanitizer::getNextTagWithCall(IRBuilder<> &IRB) {
   return IRB.CreateZExt(IRB.CreateCall(HwasanGenerateTagFunc), IntptrTy);
 }
 
 Value *HWAddressSanitizer::getStackBaseTag(IRBuilder<> &IRB) {
   if (ClGenerateTagsWithCalls)
     return getNextTagWithCall(IRB);
   if (StackBaseTag)
Value *HWAddressSanitizer::getStackBaseTag(IRBuilder<> &IRB) {
 
   // Extract some entropy from the stack pointer for the tags.
   // Take bits 20..28 (ASLR entropy) and xor with bits 0..8 (these differ
   // between functions).
   Value *StackPointerLong = IRB.CreatePointerCast(StackPointer, IntptrTy);
   Value *StackTag =
       IRB.CreateXor(StackPointerLong, IRB.CreateLShr(StackPointerLong, 20),
                     "hwasan.stack.base.tag");
+
+  StackTag = retagTargetTag(IRB, StackTag);
+
   return StackTag;
 }
 
 Value *HWAddressSanitizer::getAllocaTag(IRBuilder<> &IRB, Value *StackTag,
                                         AllocaInst *AI, unsigned AllocaNo) {
   if (ClGenerateTagsWithCalls)
     return getNextTagWithCall(IRB);
-  return IRB.CreateXor(StackTag,
-                       ConstantInt::get(IntptrTy, RetagMask(AllocaNo)));
+
+  Value *Tag = IRB.CreateXor(
+      StackTag, ConstantInt::get(
+                    IntptrTy, retagMask(AllocaNo, TargetTriple, TagMaskByte)));
+  return retagTargetTag(IRB, Tag);
 }
 
 Value *HWAddressSanitizer::getUARTag(IRBuilder<> &IRB, Value *StackTag) {
   if (ClUARRetagToZero)
     return ConstantInt::get(IntptrTy, 0);
+
   if (ClGenerateTagsWithCalls)
     return getNextTagWithCall(IRB);
-  return IRB.CreateXor(StackTag, ConstantInt::get(IntptrTy, 0xFFU));
+
+  Value *Tag = IRB.CreateXor(StackTag, ConstantInt::get(IntptrTy, 0xFFU));
+  return retagTargetTag(IRB, Tag);
 }
 
-// Add a tag to an address.
-Value *HWAddressSanitizer::tagPointer(IRBuilder<> &IRB, Type *Ty,
-                                      Value *PtrLong, Value *Tag) {
-  assert(!UsePageAliases);
+Value *HWAddressSanitizer::getTargetTagPtr(IRBuilder<> &IRB, Value *PtrLong,
+                                           Value *Tag) {
   Value *TaggedPtrLong;
-  if (CompileKernel) {
-    // Kernel addresses have 0xFF in the most significant byte.
-    Value *ShiftedTag = IRB.CreateOr(
-        IRB.CreateShl(Tag, kPointerTagShift),
-        ConstantInt::get(IntptrTy, (1ULL << kPointerTagShift) - 1));
+  Value *Mask;
+  Value *ShiftedTag = IRB.CreateShl(Tag, kPointerTagShift);
+  if (CompileKernel) {
+    // X86_64 can not tag the highest bit.
+    if (TargetTriple.getArch() == Triple::x86_64)
+      Mask = ConstantInt::get(IntptrTy, ((1ULL << kPointerTagShift) - 1) |
+                                            (1ULL << (64 - 1)));
+    else
+      Mask = ConstantInt::get(IntptrTy, (1ULL << kPointerTagShift) - 1);
+    ShiftedTag = IRB.CreateOr(ShiftedTag, Mask);
     TaggedPtrLong = IRB.CreateAnd(PtrLong, ShiftedTag);
   } else {
-    // Userspace can simply do OR (tag << 56);
-    Value *ShiftedTag = IRB.CreateShl(Tag, kPointerTagShift);
+    // Userspace can simply do OR (tag << kPointerTagShift);
     TaggedPtrLong = IRB.CreateOr(PtrLong, ShiftedTag);
   }
+  return TaggedPtrLong;
+}
+
+// Add a tag to an address.
+Value *HWAddressSanitizer::tagPointer(IRBuilder<> &IRB, Type *Ty,
+                                      Value *PtrLong, Value *Tag) {
+  assert(!UsePageAliases);
+  Value *TaggedPtrLong = getTargetTagPtr(IRB, PtrLong, Tag);
   return IRB.CreateIntToPtr(TaggedPtrLong, Ty);
 }
 
 // Remove tag from an address.
 Value *HWAddressSanitizer::untagPointer(IRBuilder<> &IRB, Value *PtrLong) {
   assert(!UsePageAliases);
   Value *UntaggedPtrLong;
   if (CompileKernel) {
     // Kernel addresses have 0xFF in the most significant byte.
-    UntaggedPtrLong = IRB.CreateOr(PtrLong,
-        ConstantInt::get(PtrLong->getType(), 0xFFULL << kPointerTagShift));
+    UntaggedPtrLong =
+        IRB.CreateOr(PtrLong, ConstantInt::get(PtrLong->getType(),
+                                               0xFFULL << kPointerTagShift));
   } else {
     // Userspace addresses have 0x00.
-    UntaggedPtrLong = IRB.CreateAnd(PtrLong,
+    UntaggedPtrLong = IRB.CreateAnd(
+        PtrLong,
         ConstantInt::get(PtrLong->getType(), ~(0xFFULL << kPointerTagShift)));
   }
   return UntaggedPtrLong;
 }
 
 Value *HWAddressSanitizer::getHwasanThreadSlotPtr(IRBuilder<> &IRB, Type *Ty) {
   Module *M = IRB.GetInsertBlock()->getParent()->getParent();
   if (TargetTriple.isAArch64() && TargetTriple.isAndroid()) {
     // Android provides a fixed TLS slot for sanitizers. See TLS_SLOT_SANITIZER
     // in Bionic's libc/private/bionic_tls.h.
     Function *ThreadPointerFunc =
         Intrinsic::getDeclaration(M, Intrinsic::thread_pointer);
     Value *SlotPtr = IRB.CreatePointerCast(
         IRB.CreateConstGEP1_32(IRB.getInt8Ty(),
                                IRB.CreateCall(ThreadPointerFunc), 0x30),
         Ty->getPointerTo(0));
     return SlotPtr;
   }
   if (ThreadPtrGlobal)
     return ThreadPtrGlobal;
 
-
   return nullptr;
 }
 
 void HWAddressSanitizer::emitPrologue(IRBuilder<> &IRB, bool WithFrameRecord) {
   if (!Mapping.InTls) {
     ShadowBase = getShadowNonTls(IRB);
     return;
   }
for (unsigned N = 0; N < Allocas.size(); ++N) {
 
     AI->replaceUsesWithIf(Replacement,
                           [AILong](Use &U) { return U.getUser() != AILong; });
 
     for (auto *DDI : AllocaDbgMap.lookup(AI)) {
       // Prepend "tag_offset, N" to the dwarf expression.
       // Tag offset logically applies to the alloca pointer, and it makes sense
       // to put it at the beginning of the expression.
-      SmallVector<uint64_t, 8> NewOps = {dwarf::DW_OP_LLVM_tag_offset,
-                                         RetagMask(N)};
+      SmallVector<uint64_t, 8> NewOps = {
+          dwarf::DW_OP_LLVM_tag_offset,
+          retagMask(N, TargetTriple, TagMaskByte)};
       auto Locations = DDI->location_ops();
       unsigned LocNo = std::distance(Locations.begin(), find(Locations, AI));
       DDI->setExpression(
           DIExpression::appendOpsToArg(DDI->getExpression(), NewOps, LocNo));
     }
 
     size_t Size = getAllocaSizeInBytes(*AI);
     tagAlloca(IRB, AI, Tag, Size);
bool HWAddressSanitizer::sanitizeFunction(Function &F) {
 
   if (!F.hasFnAttribute(Attribute::SanitizeHWAddress))
     return false;
 
   LLVM_DEBUG(dbgs() << "Function: " << F.getName() << "\n");
 
   SmallVector<InterestingMemoryOperand, 16> OperandsToInstrument;
   SmallVector<MemIntrinsic *, 16> IntrinToInstrument;
-  SmallVector<AllocaInst*, 8> AllocasToInstrument;
-  SmallVector<Instruction*, 8> RetVec;
-  SmallVector<Instruction*, 8> LandingPadVec;
+  SmallVector<AllocaInst *, 8> AllocasToInstrument;
+  SmallVector<Instruction *, 8> RetVec;
+  SmallVector<Instruction *, 8> LandingPadVec;
   DenseMap<AllocaInst *, std::vector<DbgVariableIntrinsic *>> AllocaDbgMap;
   for (auto &BB : F) {
     for (auto &Inst : BB) {
       if (InstrumentStack)
         if (AllocaInst *AI = dyn_cast<AllocaInst>(&Inst)) {
           if (isInterestingAlloca(*AI))
             AllocasToInstrument.push_back(AI);
           continue;
void HWAddressSanitizer::instrumentGlobals() {
   MD5 Hasher;
   Hasher.update(M.getSourceFileName());
   MD5::MD5Result Hash;
   Hasher.final(Hash);
   uint8_t Tag = Hash[0];
 
   for (GlobalVariable *GV : Globals) {
     // Skip tag 0 in order to avoid collisions with untagged memory.
+    Tag &= TagMaskByte;
     if (Tag == 0)
       Tag = 1;
+
     instrumentGlobal(GV, Tag++);
   }
 }
 
 void HWAddressSanitizer::instrumentPersonalityFunctions() {
   // We need to untag stack frames as we unwind past them. That is the job of
   // the personality function wrapper, which either wraps an existing
   // personality function or acts as a personality function on its own. Each

llvm/test/Instrumentation/HWAddressSanitizer/X86/stack.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -hwasan -hwasan-instrument-with-calls -hwasan-instrument-stack -hwasan-use-page-alias=false -S | FileCheck %s
+
+source_filename = "stack.c"
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; Function Attrs: noinline nounwind optnone uwtable
+define dso_local i32 @main() sanitize_hwaddress {
+; CHECK-LABEL: @main(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[DOTHWASAN_SHADOW:%.*]] = call i8* asm "", "=r,0"(i8* null)
+; CHECK-NEXT:    [[TMP0:%.*]] = call i8* @llvm.frameaddress.p0i8(i32 0)
+; CHECK-NEXT:    [[TMP1:%.*]] = ptrtoint i8* [[TMP0]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = lshr i64 [[TMP1]], 20
+; CHECK-NEXT:    [[HWASAN_STACK_BASE_TAG:%.*]] = xor i64 [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[TMP3:%.*]] = and i64 [[HWASAN_STACK_BASE_TAG]], 63
+; CHECK-NEXT:    [[RETVAL:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[LV0:%.*]] = alloca { i32, [12 x i8] }, align 16
+; CHECK-NEXT:    [[TMP4:%.*]] = bitcast { i32, [12 x i8] }* [[LV0]] to i32*
+; CHECK-NEXT:    [[TMP5:%.*]] = xor i64 [[TMP3]], 0
+; CHECK-NEXT:    [[TMP6:%.*]] = and i64 [[TMP5]], 63
+; CHECK-NEXT:    [[TMP7:%.*]] = ptrtoint i32* [[TMP4]] to i64
+; CHECK-NEXT:    [[TMP8:%.*]] = shl i64 [[TMP6]], 57
+; CHECK-NEXT:    [[TMP9:%.*]] = or i64 [[TMP7]], [[TMP8]]
+; CHECK-NEXT:    [[LV0_HWASAN:%.*]] = inttoptr i64 [[TMP9]] to i32*
+; CHECK-NEXT:    [[TMP10:%.*]] = trunc i64 [[TMP6]] to i8
+; CHECK-NEXT:    [[TMP11:%.*]] = bitcast i32* [[TMP4]] to i8*
+; CHECK-NEXT:    call void @__hwasan_tag_memory(i8* [[TMP11]], i8 [[TMP10]], i64 16)
+; CHECK-NEXT:    [[TMP12:%.*]] = ptrtoint i32* [[RETVAL]] to i64
+; CHECK-NEXT:    call void @__hwasan_store4(i64 [[TMP12]])
+; CHECK-NEXT:    store i32 0, i32* [[RETVAL]], align 4
+; CHECK-NEXT:    [[TMP13:%.*]] = ptrtoint i32* [[LV0_HWASAN]] to i64
+; CHECK-NEXT:    call void @__hwasan_store4(i64 [[TMP13]])
+; CHECK-NEXT:    store i32 12345, i32* [[LV0_HWASAN]], align 4
+; CHECK-NEXT:    call void @foo(i32* [[LV0_HWASAN]])
+; CHECK-NEXT:    [[TMP14:%.*]] = ptrtoint i32* [[LV0_HWASAN]] to i64
+; CHECK-NEXT:    call void @__hwasan_load4(i64 [[TMP14]])
+; CHECK-NEXT:    [[TMP15:%.*]] = load i32, i32* [[LV0_HWASAN]], align 4
+; CHECK-NEXT:    [[TMP16:%.*]] = bitcast i32* [[TMP4]] to i8*
+; CHECK-NEXT:    call void @__hwasan_tag_memory(i8* [[TMP16]], i8 0, i64 16)
+; CHECK-NEXT:    ret i32 [[TMP15]]
+;
+entry:
+  %retval = alloca i32, align 4
+  %lv0 = alloca i32, align 4
+  store i32 0, i32* %retval, align 4
+  store i32 12345, i32* %lv0, align 4
+  call void @foo(i32* %lv0)
+  %0 = load i32, i32* %lv0, align 4
+  ret i32 %0
+}
+
+declare dso_local void @foo(i32*)