Strip invalid TBAA when reading bitcode

[mehdi_amini]

Summary

This ensure backward compatibility on bitcode loading.

[sanjoy]

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

51

I'd have mildly preferred doing the NFC header sorting in a separate change.

157

Why are you skipping these? Would be nice to add a one-liner, unless it is patently obvious.

476

I'd have named this as TBAAStripped since it really indicates if we've already stripped out TBAA metadata (and hence there is no need to check this in for newer functions). Right now this seems imperative, and it sounds like setting this to true will cause TBAA to be stripped.

llvm/lib/Bitcode/Reader/MetadataLoader.cpp

1367

Should this not continue (otherwise you'll drop all metadata once you've seen !tbaa)?

Actually, I'd rather write this out separately as:

if (I->second == LLVMContext::MD_tbaa && StripTBAA)
  continue;

right after the if (!MD) check.

llvm/lib/Bitcode/Reader/MetadataLoader.h

58

End comment with period. Also, why not /// for doxygen?

[mehdi_amini]

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

51

That was not intentional, clang-format did it for me.

157

A function that is Materializable hasn't been materialized, so it does not have a body. It is not a declaration either as its body may be parsed later. Trying to iterate it results in an assertion.

476

Actually it will also cause TBAA to be stripped :)
The first time you encounter an invalid TBAA, you set this boolean to request that from now on when parsing the body of a function, it has to ignore TBAA MD.
It is also used line 4469 to see if we're in a mode where we strip TBAA during reading, so we don't need to check every instruction in the function.
Make sense?

llvm/lib/Bitcode/Reader/MetadataLoader.cpp

1367

It is breaking out of the switch. I'll move the check above.

[sanjoy]

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

157

I think this use of dropUnknownNonDebugMetadata is wrong (and so is the usage in CodeGenPrepare::stripInvariantGroupMetadata).

The argument to dropUnknownNonDebugMetadata is the set of metadata IDs that are "known" and should not be dropped. All non-debug metadata *not* in that set is dropped.

The reason why this works is that TBAA->getMetadataID() is an entirely different thing. It is the IDs that are used to distinguish between MDString, MDNode etc. (the subclass id, basically, which also feeds into dyn_cast etc.). Since that ID is not equal to LLVMContext::MD_tbaa, we end up dropping MD_tbaa.

The right usage here is: I.setMetadata(LLVMContext::MD_tbaa, nullptr).

476

The first time you encounter an invalid TBAA, you set this boolean to request that from now on when parsing the body of a function, it has to ignore TBAA MD.

But that is not this flag is it? I thought you were dropping future TBAA by setting the flag on MetadataLoaderImpl.

From what I can gather locally, if I manually bool StripTBAA = false; to bool StripTBAA = true; then I would not drop TBAA at all, and you're using this flag solely to remember if you've asked MDLoader to not load new TBAA, and so you don't have to walk instructions(F). In that sense, it isn't imperative, it just "remembers" if a certain action was done in the past.

However, I'm not cognizant of the inter-relationships between these classes that you changed, so I may be missing the forest for the trees.

llvm/lib/Bitcode/Reader/MetadataLoader.cpp

1367

Right, but isn't it also breaking out of the innermost for?

[sanjoy]

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

157

I've fixed the bad usage in CGP in https://reviews.llvm.org/rL289973

[mehdi_amini]

Address comments!

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

157

Uh... I took indeed example on the usage in CodeGenPrepare!
I'll switch to I.setMetadata(LLVMContext::MD_tbaa, nullptr).

476

Indeed it is redundant with the state on the MetadataLoader. I think I can remove it from here and add a getter on the MetadataLoader. You'll tell me with the next update if it helps :)

llvm/lib/Bitcode/Reader/MetadataLoader.cpp

1367

Right, I missed the inner for! I was looking at the break that follows a few lines below, but missed the fact that we had an enclosing for.

[sanjoy]

lgtm

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

159

The checking is redundant -- you can just do I.setMetadata(LLVMContext::MD_tbaa, nullptr) unconditionally (that should not do anything if MD_tbaa was not present on the instruction).

[mehdi_amini]

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

159

Good point.

I also modified the test to check that we only strip the TBAA attachment on instructions.

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

 #include "llvm/IR/GlobalAlias.h"
 #include "llvm/IR/GlobalIFunc.h"
 #include "llvm/IR/GlobalIndirectSymbol.h"
 #include "llvm/IR/GlobalObject.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/GVMaterializer.h"
 #include "llvm/IR/InlineAsm.h"
+#include "llvm/IR/InstIterator.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/ModuleSummaryIndex.h"
 #include "llvm/IR/OperandTraits.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/TrackingMDRef.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/ValueHandle.h"
+#include "llvm/IR/Verifier.h"
 #include "llvm/Support/AtomicOrdering.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/ManagedStatic.h"
static bool convertToString(ArrayRef<uint64_t> Record, unsigned Idx,
   if (Idx > Record.size())
     return true;
 
   for (unsigned i = Idx, e = Record.size(); i != e; ++i)
     Result += (char)Record[i];
   return false;
 }
 
+// Strip all the TBAA attachment for the module.
+void stripTBAA(Module *M) {
+  for (auto &F : *M) {
+    if (F.isMaterializable())
+      continue;
+    for (auto &I : instructions(F))
+      if (MDNode *TBAA = I.getMetadata(LLVMContext::MD_tbaa))
+        I.setMetadata(LLVMContext::MD_tbaa, nullptr);
+  }
+}
+
 /// Read the "IDENTIFICATION_BLOCK_ID" block, do some basic enforcement on the
 /// "epoch" encoded in the bitcode, and return the producer name if any.
 Expected<std::string> readIdentificationBlock(BitstreamCursor &Stream) {
   if (Stream.EnterSubBlock(bitc::IDENTIFICATION_BLOCK_ID))
     return error("Invalid record");
 
   // Read all the records.
   SmallVector<uint64_t, 64> Record;
class BitcodeReader : public BitcodeReaderBase, public GVMaterializer {
   /// not need this flag.
   bool UseRelativeIDs = false;
 
   /// True if all functions will be materialized, negating the need to process
   /// (e.g.) blockaddress forward references.
   bool WillMaterializeAllForwardRefs = false;
 
   bool StripDebugInfo = false;
+  TBAAVerifier TBAAVerifyHelper;
 
   std::vector<std::string> BundleTags;
 
 public:
   BitcodeReader(BitstreamCursor Stream, StringRef ProducerIdentification,
                 LLVMContext &Context);
 
   Error materializeForwardReferencedFunctions();
for (auto UI = I.first->materialized_user_begin(), UE = I.first->user_end();
          UI != UE;)
       // Don't expect any other users than call sites
       CallSite(*UI++).setCalledFunction(I.second);
 
   // Finish fn->subprogram upgrade for materialized functions.
   if (DISubprogram *SP = MDLoader->lookupSubprogramForFunction(F))
     F->setSubprogram(SP);
 
+  // Check if the TBAA Metadata are valid, otherwise we will need to strip them.
+  if (!MDLoader->isStrippingTBAA()) {
+    for (auto &I : instructions(F)) {
+      MDNode *TBAA = I.getMetadata(LLVMContext::MD_tbaa);
+      if (!TBAA || TBAAVerifyHelper.visitTBAAMetadata(I, TBAA))
+        continue;
+      MDLoader->setStripTBAA(true);
+      stripTBAA(F->getParent());
+    }
+  }
+
   // Bring in any functions that this function forward-referenced via
   // blockaddresses.
   return materializeForwardReferencedFunctions();
 }
 
 Error BitcodeReader::materializeModule() {
   if (Error Err = materializeMetadata(/* OnDemand = */ false))
     return Err;

llvm/lib/Bitcode/Reader/MetadataLoader.h

public:
   Error parseModuleMetadata(bool OnDemand) {
     return parseMetadata(true, OnDemand);
   }
   // Parse a function metadata block
   Error parseFunctionMetadata() {
     return parseMetadata(false, /* OnDemand = */ false);
   }
 
+  /// Set the mode to strip TBAA metadata on load.
+  void setStripTBAA(bool StripTBAA = true);
+
+  /// Return true if the Loader is stripping TBAA metadata.
+  bool isStrippingTBAA();
+
   // Return true there are remaining unresolved forward references.
   bool hasFwdRefs() const;
 
   /// Return the given metadata, creating a replaceable forward reference if
   /// necessary.
   Metadata *getMetadataFwdRef(unsigned Idx);
 
   MDNode *getMDNodeFwdRefOrNull(unsigned Idx);

llvm/lib/Bitcode/Reader/MetadataLoader.cpp

class MetadataLoader::MetadataLoaderImpl {
 
   /// Functions that need to be matched with subprograms when upgrading old
   /// metadata.
   SmallDenseMap<Function *, DISubprogram *, 16> FunctionsWithSPs;
 
   // Map the bitcode's custom MDKind ID to the Module's MDKind ID.
   DenseMap<unsigned, unsigned> MDKindMap;
 
+  bool StripTBAA = false;
   bool HasSeenOldLoopTags = false;
   Error parseOneMetadata(SmallVectorImpl<uint64_t> &Record, unsigned Code,
                          PlaceholderQueue &Placeholders, StringRef Blob,
                          unsigned &NextMetadataNo);
   Error parseMetadataStrings(ArrayRef<uint64_t> Record, StringRef Blob,
                              unsigned &NextMetadataNo);
   Error parseGlobalObjectAttachment(GlobalObject &GO,
                                     ArrayRef<uint64_t> Record);
public:
 
   bool hasSeenOldLoopTags() { return HasSeenOldLoopTags; }
 
   Error parseMetadataAttachment(
       Function &F, const SmallVectorImpl<Instruction *> &InstructionList);
 
   Error parseMetadataKinds();
 
+  void setStripTBAA(bool Value) { StripTBAA = Value; }
+  bool isStrippingTBAA() { return StripTBAA; }
+
   unsigned size() const { return MetadataList.size(); }
   void shrinkTo(unsigned N) { MetadataList.shrinkTo(N); }
 };
 
 Error error(const Twine &Message) {
   return make_error<StringError>(
       Message, make_error_code(BitcodeError::CorruptedBitcode));
 }
case bitc::METADATA_ATTACHMENT: {
 
       // An instruction attachment.
       Instruction *Inst = InstructionList[Record[0]];
       for (unsigned i = 1; i != RecordLength; i = i + 2) {
         unsigned Kind = Record[i];
         DenseMap<unsigned, unsigned>::iterator I = MDKindMap.find(Kind);
         if (I == MDKindMap.end())
           return error("Invalid ID");
+        if (I->second == LLVMContext::MD_tbaa && StripTBAA)
+          continue;
+
         Metadata *Node = MetadataList.getMetadataFwdRef(Record[i + 1]);
         if (isa<LocalAsMetadata>(Node))
           // Drop the attachment.  This used to be legal, but there's no
           // upgrade path.
           break;
         MDNode *MD = dyn_cast_or_null<MDNode>(Node);
         if (!MD)
           return error("Invalid metadata attachment");
 
         if (HasSeenOldLoopTags && I->second == LLVMContext::MD_loop)
           MD = upgradeInstructionLoopAttachment(*MD);
 
         if (I->second == LLVMContext::MD_tbaa) {
           assert(!MD->isTemporary() && "should load MDs before attachments");
           MD = UpgradeTBAANode(*MD);
         }
         Inst->setMetadata(I->second, MD);
       }
       break;
     }
     }
   }
 }
Error MetadataLoader::parseMetadataAttachment(
     Function &F, const SmallVectorImpl<Instruction *> &InstructionList) {
   return Pimpl->parseMetadataAttachment(F, InstructionList);
 }
 
 Error MetadataLoader::parseMetadataKinds() {
   return Pimpl->parseMetadataKinds();
 }
 
+void MetadataLoader::setStripTBAA(bool StripTBAA) {
+  return Pimpl->setStripTBAA(StripTBAA);
+}
+
+bool MetadataLoader::isStrippingTBAA() { return Pimpl->isStrippingTBAA(); }
+
 unsigned MetadataLoader::size() const { return Pimpl->size(); }
 void MetadataLoader::shrinkTo(unsigned N) { return Pimpl->shrinkTo(N); }

llvm/test/Verifier/tbaa.ll

 ; RUN: not llvm-as < %s 2>&1 | FileCheck %s
+; RUN: llvm-as -disable-verify < %s 2>&1 | opt -verify -S | FileCheck %s  --check-prefix=STRIP
+
+; STRIP-NOT: tbaa
+; STRIP: @f_0
 
 define void @f_0(i32* %ptr) {
 ; This part checks for the easy syntactic verifier rules.
 
 ; CHECK: Struct tag metadata must have either 3 or 4 operands
 ; CHECK-NEXT:  store i32 0, i32* %ptr, !tbaa !{{[0-9]+}}
 
 ; CHECK: Immutability tag on struct tag metadata must be a constant