[RS4GC] Strip off invariant.start because memory locations arent invariant

[anna]

Summary

Invariant.start on memory locations has the property that the memory
location is unchanging. However, this is not true in the face of
rewriting statepoints for GC.
Teach RS4GC about removing invariant.start so that optimizations after
RS4GC does not incorrect sink a load from the memory location past a
statepoint.

Added test showcasing the issue.

[dneilson]

Is is possible to also remove the corresponding invariant.end calls to the removed invariant.start calls? This patch will leave orphaned calls to invariant.end sitting around in the IR that have no corresponding start.

[anna]

In D39388#911770, @dneilson wrote:

Is is possible to also remove the corresponding invariant.end calls to the removed invariant.start calls? This patch will leave orphaned calls to invariant.end sitting around in the IR that have no corresponding start.

Yes, we could remove it (which was what I initially did), but I didn't see any benefit for doing so. It would be special casing just one type of use of invariant.start. As testcase points out, we can have any instruction using invariant.start, and we cannot just remove those uses, because those uses may have other uses...

Just to be clear: having orphaned invariant.end is perfectly valid IR. On the other hand, as an IR cleanup, that's a good thing to do, especially because invariant.end cannot have any uses. I'll update the patch with special casing invariant.end

[anna]

addressed review comments.

[anna]

just a note: We haven't yet added support to forward loads/stores across calls when the calls are dominated by invariant.start. However, that's a perfectly legal transform, and this patch makes sure we don't miscompile when the transformation is supported.

[dneilson]

lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

131

Note that this method is called if *any* function in the entire module is changed. This is also the only call of this method.

2390

It seems really odd & counterintuitive to me that a method that purports to strip attributes & metadata will now also be removing instructions. I'd never expect that from the name of the method.

I'm not a fan of introducing another loop over every instruction, but also not a fan of instructions being removed in a method with a name like this. Not sure how to resolve that.

Also, due to the way that this method is invoked the result of putting the invariant.start removal in this method is that if RS4GC changes any function in the module (even if it's only one) then we will remove the invariant.start/end pairs from *every* function in the module. That doesn't seem correct to me.

2480

Note: Indiscriminately called for every function in the module.

[anna]

lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

131

yes, that is what we want, we have to be conservative for correctness. We do not know at which stage RS4GC is done in the pipeline and we cannot rely on that information.

If we have one more stage of inlining after RS4GC, followed by other optimizations, there is a possibility of the metadata being incorrect. Just as an example, consider invariant.start at the end of function F and F has no calls within it. So, RS4GC does nothing for this function. If we didn't strip off invariant.start and then inlined F into G after RS4GC, we can still face this same problem.

That is why we remove all incorrect metadata, attributes and instructions from all functions.

2390

Regarding the name, I'll change it to stripNonValidDataFromBody.

Regarding the second concern, I've addressed it in the first comment response. We actually need to remove the invariant.start for correctness from *every* function.

2480

yes, that''s what we want.

[anna]

Updated function names per review comment. NFC wrt previous diff.

[dneilson]

LGTM

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) {
-      // stripNonValidAttributesAndMetadata asserts that shouldRewriteStatepointsIn
+      // stripNonValidData 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.
-      stripNonValidAttributesAndMetadata(M);
+      stripNonValidData(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>();
     AU.addRequired<TargetLibraryInfoWrapperPass>();
   }
 
   /// 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. stripNonValidAttributesAndMetadata (conservatively) restores
+  /// heap. stripNonValidData (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);
+  /// Apart from attributes and metadata, we also remove instructions that imply
+  /// constant physical memory: llvm.invariant.start.
+  void stripNonValidData(Module &M);
 
-  // Helpers for stripNonValidAttributesAndMetadata
-  void stripNonValidAttributesAndMetadataFromBody(Function &F);
+  // Helpers for stripNonValidData
+  void stripNonValidDataFromBody(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);
 };
 
unsigned ValidMetadataAfterRS4GC[] = {LLVMContext::MD_tbaa,
                          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) {
+void RewriteStatepointsForGC::stripNonValidDataFromBody(Function &F) {
   if (F.empty())
     return;
 
   LLVMContext &Ctx = F.getContext();
   MDBuilder Builder(Ctx);
 
+  // Set of invariantstart instructions that we need to remove.
+  // Use this to avoid invalidating the instruction iterator.
+  SmallVector<IntrinsicInst*, 12> InvariantStartInstructions;
+
   for (Instruction &I : instructions(F)) {
+    // invariant.start on memory location implies that the referenced memory
+    // location is constant and unchanging. This is no longer true after
+    // RewriteStatepointsForGC runs because there can be calls to gc.statepoint
+    // which frees the entire heap and the presence of invariant.start allows
+    // the optimizer to sink the load of a memory location past a statepoint,
+    // which is incorrect.
+    if (auto *II = dyn_cast<IntrinsicInst>(&I))
+      if (II->getIntrinsicID() == Intrinsic::invariant_start) {
+        InvariantStartInstructions.push_back(II);
+        continue;
+      }
+
     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
for (Instruction &I : instructions(F)) {
     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);
     }
   }
+
+  // Delete the invariant.start instructions and any corresponding uses that
+  // don't have further uses, for example invariant.end.
+  for (auto *II : InvariantStartInstructions) {
+    for (auto *U : II->users())
+      if (auto *I = dyn_cast<Instruction>(U))
+       if (U->hasNUses(0))
+         I->eraseFromParent();
+    // We cannot just delete the remaining uses of II, so we RAUW undef.
+    II->replaceAllUsesWith(UndefValue::get(II->getType()));
+    II->eraseFromParent();
+  }
 }
 
 /// 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::stripNonValidAttributesAndMetadata(Module &M) {
+void RewriteStatepointsForGC::stripNonValidData(Module &M) {
 #ifndef NDEBUG
   assert(llvm::any_of(M, shouldRewriteStatepointsIn) && "precondition!");
 #endif
 
   for (Function &F : M)
     stripNonValidAttributesFromPrototype(F);
 
   for (Function &F : M)
-    stripNonValidAttributesAndMetadataFromBody(F);
+    stripNonValidDataFromBody(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

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

 ; 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
 }
 
+; invariant.start allows us to sink the load past the baz statepoint call into taken block, which is
+; incorrect. remove the invariant.start and RAUW undef.
+define void @test_inv_start(i1 %cond, i32 addrspace(1)* addrspace(1)* %p, i32 %x, i32 addrspace(1)* %q) gc "statepoint-example" {
+; CHECK-LABEL: test_inv_start
+; CHECK-NOT: invariant.start
+; CHECK: gc.statepoint
+  %v1 = load i32 addrspace(1)*, i32 addrspace(1)* addrspace(1)* %p
+  %invst = call {}* @llvm.invariant.start.p1i32(i64 1, i32 addrspace(1)* %v1)
+  %v2 = load i32, i32 addrspace(1)* %v1
+  call void @baz(i32 %x)
+  br i1 %cond, label %taken, label %untaken
+
+; CHECK-LABEL: taken:
+; CHECK-NOT: llvm.invariant.end
+taken:
+  store i32 %v2, i32 addrspace(1)* %q, align 16
+  call void @llvm.invariant.end.p1i32({}* %invst, i64 4, i32 addrspace(1)* %v1)
+  ret void
+
+; CHECK-LABEL: untaken:
+; CHECK: gc.statepoint
+untaken:
+  %foo = call i32 @escaping.invariant.start({}* %invst)
+  call void @dummy(i32 %foo)
+  ret void
+}
+
+; invariant.start and end is removed. No other uses.
+define void @test_inv_start2(i1 %cond, i32 addrspace(1)* addrspace(1)* %p, i32 %x, i32 addrspace(1)* %q) gc "statepoint-example" {
+; CHECK-LABEL: test_inv_start2
+; CHECK-NOT: invariant.start
+; CHECK: gc.statepoint
+  %v1 = load i32 addrspace(1)*, i32 addrspace(1)* addrspace(1)* %p
+  %invst = call {}* @llvm.invariant.start.p1i32(i64 1, i32 addrspace(1)* %v1)
+  %v2 = load i32, i32 addrspace(1)* %v1
+  call void @baz(i32 %x)
+  br i1 %cond, label %taken, label %untaken
+
+; CHECK-LABEL: taken:
+; CHECK-NOT: llvm.invariant.end
+taken:
+  store i32 %v2, i32 addrspace(1)* %q, align 16
+  call void @llvm.invariant.end.p1i32({}* %invst, i64 4, i32 addrspace(1)* %v1)
+  ret void
+
+; CHECK-LABEL: untaken:
+untaken:
+  ret void
+}
+declare {}* @llvm.invariant.start.p1i32(i64, i32 addrspace(1)*  nocapture) nounwind readonly
+declare void @llvm.invariant.end.p1i32({}*, i64, i32 addrspace(1)* nocapture) nounwind
+declare i32 @escaping.invariant.start({}*) nounwind
+declare void @dummy(i32)
 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}