[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.

2388

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.

2474

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.

2388

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.

2474

yes, that''s what we want.

[anna]

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

[dneilson]

LGTM

lib/Transforms/Scalar/RewriteStatepointsForGC.cpp

bool runOnModule(Module &M) override {
     bool Changed = false;
     for (Function &F : M)
       Changed |= runOnFunction(F);
 
     if (Changed) {
       // 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.
       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.
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) {
   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.
+  for (auto *II : InvariantStartInstructions) {
+    // We cannot just delete the uses of invariant.start (it could be used
+    // as an argument to another call), 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();

 #ifndef NDEBUG
   assert(llvm::any_of(M, shouldRewriteStatepointsIn) && "precondition!");
 #endif
 
   for (Function &F : M)
     stripNonValidAttributesFromPrototype(F);
 
   for (Function &F : M)
     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

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: call void @llvm.invariant.end.p1i32({}* undef, i64 4, i32 addrspace(1)* %v1.relocated.casted)
+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:
+  call void @escaping.invariant.start({}* %invst)
+  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 void @escaping.invariant.start({}*) nounwind
 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}