diff --git a/compiler-rt/lib/hwasan/hwasan.h b/compiler-rt/lib/hwasan/hwasan.h
--- a/compiler-rt/lib/hwasan/hwasan.h
+++ b/compiler-rt/lib/hwasan/hwasan.h
@@ -37,8 +37,14 @@
 
 // TBI (Top Byte Ignore) feature of AArch64: bits [63:56] are ignored in address
 // translation and can be used to store a tag.
-const unsigned kAddressTagShift = 56;
-const uptr kAddressTagMask = 0xFFUL << kAddressTagShift;
+constexpr unsigned kAddressTagShift = 56;
+constexpr unsigned kTagBits = 8;
+
+// Mask for extracting tag bits from the lower 8 bits.
+constexpr uptr kTagMask = (1UL << kTagBits) - 1;
+
+// Masks for extracting and removing tags from full pointers.
+constexpr uptr kAddressTagMask = kTagMask << kAddressTagShift;
 
 // Minimal alignment of the shadow base address. Determines the space available
 // for threads and stack histories. This is an ABI constant.
@@ -50,7 +56,7 @@
 const unsigned kRecordFPModulus = 1 << (64 - kRecordFPShift + kRecordFPLShift);
 
 static inline tag_t GetTagFromPointer(uptr p) {
-  return p >> kAddressTagShift;
+  return (p >> kAddressTagShift) & kTagMask;
 }
 
 static inline uptr UntagAddr(uptr tagged_addr) {
diff --git a/compiler-rt/lib/hwasan/hwasan_thread.h b/compiler-rt/lib/hwasan/hwasan_thread.h
--- a/compiler-rt/lib/hwasan/hwasan_thread.h
+++ b/compiler-rt/lib/hwasan/hwasan_thread.h
@@ -42,7 +42,7 @@
   HeapAllocationsRingBuffer *heap_allocations() { return heap_allocations_; }
   StackAllocationsRingBuffer *stack_allocations() { return stack_allocations_; }
 
-  tag_t GenerateRandomTag();
+  tag_t GenerateRandomTag(uptr num_bits = kTagBits);
 
   void DisableTagging() { tagging_disabled_++; }
   void EnableTagging() { tagging_disabled_--; }
diff --git a/compiler-rt/lib/hwasan/hwasan_thread.cpp b/compiler-rt/lib/hwasan/hwasan_thread.cpp
--- a/compiler-rt/lib/hwasan/hwasan_thread.cpp
+++ b/compiler-rt/lib/hwasan/hwasan_thread.cpp
@@ -113,18 +113,21 @@
 }
 
 // Generate a (pseudo-)random non-zero tag.
-tag_t Thread::GenerateRandomTag() {
+tag_t Thread::GenerateRandomTag(uptr num_bits) {
+  DCHECK_GT(num_bits, 0);
   if (tagging_disabled_) return 0;
   tag_t tag;
+  const uptr tag_mask = (1ULL << num_bits) - 1;
   do {
     if (flags()->random_tags) {
       if (!random_buffer_)
         random_buffer_ = random_state_ = xorshift(random_state_);
       CHECK(random_buffer_);
-      tag = random_buffer_ & 0xFF;
-      random_buffer_ >>= 8;
+      tag = random_buffer_ & tag_mask;
+      random_buffer_ >>= num_bits;
     } else {
-      tag = random_state_ = (random_state_ + 1) & 0xFF;
+      random_state_ += 1;
+      tag = random_state_ & tag_mask;
     }
   } while (!tag);
   return tag;