[RS4GC] Drop invalid metadata after pointers are relocated

[anna]

Summary

After RS4GC, we should drop metadata that is no longer valid. These metadata
is used by optimizations scheduled after RS4GC, and can cause a miscompile.
One such metadata is invariant.load which is used by LICM sinking transform.
After rewriting statepoints, the address of a load maybe relocated. With
invariant.load metadata on a load instruction, LICM sinking assumes the
loaded value (from a dererenceable address) to be invariant, and
rematerializes the load operand and the load at the exit block.
This transforms the IR to have an unrelocated use of the
address after a statepoint, which is incorrect.
Other metadata we conservatively remove are related to
dereferenceability and noalias metadata.

This patch drops such metadata on store and load instructions after
rewriting statepoints.

[anna]

ping

[dneilson]

lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

2318 ↗

(On Diff #100962)

Can you order these values in the order provided in IR/LLVMContext.h? Makes it easier to see which ones are missing if the ordering is the same.

[anna]

Ordered the metadata according to LLVMContext.h and updated test comment.

[dneilson]

lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

2318 ↗

(On Diff #101927)

Why drop MD_nontemporal?

[reames]

lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

2351 ↗

(On Diff #101927)

I think that a better factoring on this would be to factor out something from RemoveNonValidAttrAtIndex which works on a load, call it on a load, then add !invariant.load to that list. This is not specific to invariant load. Really, we'd just completely failed to think about loads in implementing this.

[anna]

lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

2351 ↗

(On Diff #101927)

I don't think we can factor out anything from RemoveNonValidAttrAtIndex. We are trying to remove metadata and not attributes. Also, there is no way to specify something like I.removeMetadata(invariant_load), and such a function should *not* be added in Instruction class. The reason being we need to keep auditing code for correctness to keep track of all new metadata that gets added to langref. The dropUnknownNonDebugMetadata(knownIDs) works *correctly* without any changes even when new metadata gets added.

Also, I.dropUnknownNonDebugMetadata updates private members in Instruction once no metadata is left behind in that instruction.

I can introduce a function like DropInvalidMetadataOnInstruction and call it on a load instruction. Does that work?

[anna]

lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

2351 ↗

(On Diff #101927)

Talked to Philip offline, and he agreed with my analysis (since we are dealing with metadata and not attributes).

I'm going to update the comments and code to make this distinction clearer.

[anna]

Updated code to remove invalid metadata on load and stores.
Now, we remove also noalias and dereferenceability metadata (which are analogous to their attribute counterparts).
Added some comments and renamed function names.

[anna]

Updated comments (NFC).

[reames]

LGTM and thanks for the clarifications.

llvm/trunk/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

struct RewriteStatepointsForGC : public ModulePass {
   }
   bool runOnFunction(Function &F);
   bool runOnModule(Module &M) override {
     bool Changed = false;
     for (Function &F : M)
       Changed |= runOnFunction(F);
 
     if (Changed) {
-      // stripNonValidAttributes asserts that shouldRewriteStatepointsIn
+      // stripNonValidAttributesAndMetadata asserts that shouldRewriteStatepointsIn
       // returns true for at least one function in the module.  Since at least
       // one function changed, we know that the precondition is satisfied.
-      stripNonValidAttributes(M);
+      stripNonValidAttributesAndMetadata(M);
     }
 
     return Changed;
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     // We add and rewrite a bunch of instructions, but don't really do much
     // else.  We could in theory preserve a lot more analyses here.
     AU.addRequired<DominatorTreeWrapperPass>();
     AU.addRequired<TargetTransformInfoWrapperPass>();
   }
 
-  /// The IR fed into RewriteStatepointsForGC may have had attributes implying
-  /// dereferenceability that are no longer valid/correct after
-  /// RewriteStatepointsForGC has run.  This is because semantically, after
+  /// The IR fed into RewriteStatepointsForGC may have had attributes and
+  /// metadata implying dereferenceability that are no longer valid/correct after
+  /// RewriteStatepointsForGC has run. This is because semantically, after
   /// RewriteStatepointsForGC runs, all calls to gc.statepoint "free" the entire
-  /// heap.  stripNonValidAttributes (conservatively) restores correctness
-  /// by erasing all attributes in the module that externally imply
-  /// dereferenceability.
-  /// Similar reasoning also applies to the noalias attributes. gc.statepoint
-  /// can touch the entire heap including noalias objects.
-  void stripNonValidAttributes(Module &M);
+  /// heap. stripNonValidAttributesAndMetadata (conservatively) restores
+  /// correctness by erasing all attributes in the module that externally imply
+  /// dereferenceability. Similar reasoning also applies to the noalias
+  /// attributes and metadata. gc.statepoint can touch the entire heap including
+  /// noalias objects.
+  void stripNonValidAttributesAndMetadata(Module &M);
 
-  // Helpers for stripNonValidAttributes
-  void stripNonValidAttributesFromBody(Function &F);
+  // Helpers for stripNonValidAttributesAndMetadata
+  void stripNonValidAttributesAndMetadataFromBody(Function &F);
   void stripNonValidAttributesFromPrototype(Function &F);
+  // Certain metadata on instructions are invalid after running RS4GC.
+  // Optimizations that run after RS4GC can incorrectly use this metadata to
+  // optimize functions. We drop such metadata on the instruction.
+  void stripInvalidMetadataFromInstruction(Instruction &I);
 };
 } // namespace
 
 char RewriteStatepointsForGC::ID = 0;
 
 ModulePass *llvm::createRewriteStatepointsForGCPass() {
   return new RewriteStatepointsForGC();
 }
for (Argument &A : F.args())
     if (isa<PointerType>(A.getType()))
       RemoveNonValidAttrAtIndex(Ctx, F,
                                 A.getArgNo() + AttributeList::FirstArgIndex);
 
   if (isa<PointerType>(F.getReturnType()))
     RemoveNonValidAttrAtIndex(Ctx, F, AttributeList::ReturnIndex);
 }
 
-void RewriteStatepointsForGC::stripNonValidAttributesFromBody(Function &F) {
+void RewriteStatepointsForGC::stripInvalidMetadataFromInstruction(Instruction &I) {
+
+  if (!isa<LoadInst>(I) && !isa<StoreInst>(I))
+    return;
+  // These are the attributes that are still valid on loads and stores after
+  // RS4GC.
+  // The metadata implying dereferenceability and noalias are (conservatively)
+  // dropped.  This is because semantically, after RewriteStatepointsForGC runs,
+  // all calls to gc.statepoint "free" the entire heap. Also, gc.statepoint can
+  // touch the entire heap including noalias objects. Note: The reasoning is
+  // same as stripping the dereferenceability and noalias attributes that are
+  // analogous to the metadata counterparts.
+  // We also drop the invariant.load metadata on the load because that metadata
+  // implies the address operand to the load points to memory that is never
+  // changed once it became dereferenceable. This is no longer true after RS4GC.
+  // Similar reasoning applies to invariant.group metadata, which applies to
+  // loads within a group.
+  unsigned ValidMetadataAfterRS4GC[] = {LLVMContext::MD_tbaa,
+                         LLVMContext::MD_range,
+                         LLVMContext::MD_alias_scope,
+                         LLVMContext::MD_nontemporal,
+                         LLVMContext::MD_nonnull,
+                         LLVMContext::MD_align,
+                         LLVMContext::MD_type};
+
+  // Drops all metadata on the instruction other than ValidMetadataAfterRS4GC.
+  I.dropUnknownNonDebugMetadata(ValidMetadataAfterRS4GC);
+
+}
+
+void RewriteStatepointsForGC::stripNonValidAttributesAndMetadataFromBody(Function &F) {
   if (F.empty())
     return;
 
   LLVMContext &Ctx = F.getContext();
   MDBuilder Builder(Ctx);
 
+
   for (Instruction &I : instructions(F)) {
     if (const MDNode *MD = I.getMetadata(LLVMContext::MD_tbaa)) {
       assert(MD->getNumOperands() < 5 && "unrecognized metadata shape!");
       bool IsImmutableTBAA =
           MD->getNumOperands() == 4 &&
           mdconst::extract<ConstantInt>(MD->getOperand(3))->getValue() == 1;
 
       if (!IsImmutableTBAA)
         continue; // no work to do, MD_tbaa is already marked mutable
 
       MDNode *Base = cast<MDNode>(MD->getOperand(0));
       MDNode *Access = cast<MDNode>(MD->getOperand(1));
       uint64_t Offset =
           mdconst::extract<ConstantInt>(MD->getOperand(2))->getZExtValue();
 
       MDNode *MutableTBAA =
           Builder.createTBAAStructTagNode(Base, Access, Offset);
       I.setMetadata(LLVMContext::MD_tbaa, MutableTBAA);
     }
 
+    stripInvalidMetadataFromInstruction(I);
+
     if (CallSite CS = CallSite(&I)) {
       for (int i = 0, e = CS.arg_size(); i != e; i++)
         if (isa<PointerType>(CS.getArgument(i)->getType()))
           RemoveNonValidAttrAtIndex(Ctx, CS, i + AttributeList::FirstArgIndex);
       if (isa<PointerType>(CS.getType()))
         RemoveNonValidAttrAtIndex(Ctx, CS, AttributeList::ReturnIndex);
     }
   }
 }
 
 /// Returns true if this function should be rewritten by this pass.  The main
 /// point of this function is as an extension point for custom logic.
 static bool shouldRewriteStatepointsIn(Function &F) {
   // TODO: This should check the GCStrategy
   if (F.hasGC()) {
     const auto &FunctionGCName = F.getGC();
     const StringRef StatepointExampleName("statepoint-example");
     const StringRef CoreCLRName("coreclr");
     return (StatepointExampleName == FunctionGCName) ||
            (CoreCLRName == FunctionGCName);
   } else
     return false;
 }
 
-void RewriteStatepointsForGC::stripNonValidAttributes(Module &M) {
+void RewriteStatepointsForGC::stripNonValidAttributesAndMetadata(Module &M) {
 #ifndef NDEBUG
   assert(any_of(M, shouldRewriteStatepointsIn) && "precondition!");
 #endif
 
   for (Function &F : M)
     stripNonValidAttributesFromPrototype(F);
 
   for (Function &F : M)
-    stripNonValidAttributesFromBody(F);
+    stripNonValidAttributesAndMetadataFromBody(F);
 }
 
 bool RewriteStatepointsForGC::runOnFunction(Function &F) {
   // Nothing to do for declarations.
   if (F.isDeclaration() || F.empty())
     return false;
 
   // Policy choice says not to rewrite - the most common reason is that we're

llvm/trunk/test/Transforms/RewriteStatepointsForGC/drop-invalid-metadata.ll

+; RUN: opt -S -rewrite-statepoints-for-gc < %s | FileCheck %s
+
+; This test checks that metadata that's invalid after RS4GC is dropped. 
+; We can miscompile if optimizations scheduled after RS4GC uses the
+; metadata that's infact invalid.
+
+declare void @bar()
+
+declare void @baz(i32)
+; Confirm that loadedval instruction does not contain invariant.load metadata.
+; but contains the range metadata.
+; Since loadedval is not marked invariant, it will prevent incorrectly sinking
+; %loadedval in LICM and avoid creation of an unrelocated use of %baseaddr.
+define void @test_invariant_load() gc "statepoint-example" {
+; CHECK-LABEL: @test_invariant_load
+; CHECK: %loadedval = load i32, i32 addrspace(1)* %baseaddr, align 8, !range !0
+bb:
+  br label %outerloopHdr
+
+outerloopHdr:                                              ; preds = %bb6, %bb
+  %baseaddr = phi i32 addrspace(1)* [ undef, %bb ], [ %tmp4, %bb6 ]
+; LICM may sink this load to exit block after RS4GC because it's tagged invariant.
+  %loadedval = load i32, i32 addrspace(1)* %baseaddr, align 8, !range !0, !invariant.load !1
+  br label %innerloopHdr
+
+innerloopHdr:                                              ; preds = %innerlooplatch, %outerloopHdr
+  %tmp4 = phi i32 addrspace(1)* [ %baseaddr, %outerloopHdr ], [ %gep, %innerlooplatch ]
+  br label %innermostloophdr
+
+innermostloophdr:                                              ; preds = %bb6, %innerloopHdr
+  br i1 undef, label %exitblock, label %bb6
+
+bb6:                                              ; preds = %innermostloophdr
+  switch i32 undef, label %innermostloophdr [
+    i32 0, label %outerloopHdr
+    i32 1, label %innerlooplatch
+  ]
+
+innerlooplatch:                                              ; preds = %bb6
+  call void @bar()
+  %gep = getelementptr inbounds i32, i32 addrspace(1)* %tmp4, i64 8
+  br label %innerloopHdr
+
+exitblock:                                             ; preds = %innermostloophdr
+  %tmp13 = add i32 42, %loadedval
+  call void @baz(i32 %tmp13)
+  unreachable
+}
+
+; drop the noalias metadata.
+define void @test_noalias(i32 %x, i32 addrspace(1)* %p, i32 addrspace(1)* %q) gc "statepoint-example" {
+; CHECK-LABEL: test_noalias
+; CHECK: %y = load i32, i32 addrspace(1)* %q, align 16
+; CHECK: gc.statepoint
+; CHECK: %p.relocated
+; CHECK-NEXT: %p.relocated.casted = bitcast i8 addrspace(1)* %p.relocated to i32 addrspace(1)*
+; CHECK-NEXT: store i32 %x, i32 addrspace(1)* %p.relocated.casted, align 16
+entry:
+  %y = load i32, i32 addrspace(1)* %q, align 16, !noalias !3
+  call void @baz(i32 %x)
+  store i32 %x, i32 addrspace(1)* %p, align 16, !noalias !4
+  ret void
+}
+
+; drop the dereferenceable metadata
+define void @test_dereferenceable(i32 addrspace(1)* addrspace(1)* %p, i32 %x, i32 addrspace(1)* %q) gc "statepoint-example" {
+; CHECK-LABEL: test_dereferenceable
+; CHECK: %v1 = load i32 addrspace(1)*, i32 addrspace(1)* addrspace(1)* %p
+; CHECK-NEXT: %v2 = load i32, i32 addrspace(1)* %v1
+; CHECK: gc.statepoint
+  %v1 = load i32 addrspace(1)*, i32 addrspace(1)* addrspace(1)* %p, !dereferenceable !5
+  %v2 = load i32, i32 addrspace(1)* %v1
+  call void @baz(i32 %x)
+  store i32 %v2, i32 addrspace(1)* %q, align 16
+  ret void
+}
+
+declare token @llvm.experimental.gc.statepoint.p0f_isVoidi32f(i64, i32, void (i32)*, i32, i32, ...)
+
+; Function Attrs: nounwind readonly
+declare i8 addrspace(1)* @llvm.experimental.gc.relocate.p1i8(token, i32, i32) #0
+
+declare token @llvm.experimental.gc.statepoint.p0f_isVoidf(i64, i32, void ()*, i32, i32, ...)
+
+attributes #0 = { nounwind readonly }
+
+!0 = !{i32 0, i32 2147483647}
+!1 = !{}
+!2 = !{i32 10, i32 1}
+!3 = !{!3}
+!4 = !{!4}
+!5 = !{i64 8}