Only strip invariant.load from special pointers (#57386)

Other backends (in this case NVPTX) require that `invariant.load`
metadata is maintained to generate non-coherent loads.

Currently, we unconditionally strip that metadata from all loads,
since our other uses of it may have become invalid.

x-ref: llvm/llvm-project#112834
JuliaGPU/CUDA.jl#2531

---------

Co-authored-by: Gabriel Baraldi <[email protected]>

Author: vchuravy

src/llvm-late-gc-lowering.cpp

@@ -2011,8 +2011,10 @@ bool LateLowerGCFrame::CleanupIR(Function &F, State *S, bool *CFGModified) {
             // strip all constant alias information, as it might depend on the gc having
             // preserved a gc root, which stops being true after this pass (#32215)
             // similar to RewriteStatepointsForGC::stripNonValidData, but less aggressive
-            if (I->getMetadata(LLVMContext::MD_invariant_load))
-                I->setMetadata(LLVMContext::MD_invariant_load, NULL);
+            if (auto *LI = dyn_cast<LoadInst>(I)){
+                if (isSpecialPtr(LI->getPointerOperand()->getType()) && LI->getMetadata(LLVMContext::MD_invariant_load))
+                    LI->setMetadata(LLVMContext::MD_invariant_load, NULL);
+            }
             if (MDNode *TBAA = I->getMetadata(LLVMContext::MD_tbaa)) {
                 if (TBAA->getNumOperands() == 4 && isTBAA(TBAA, {"jtbaa_const", "jtbaa_memoryptr", "jtbaa_memorylen", "tbaa_memoryown"})) {
                     MDNode *MutableTBAA = createMutableTBAAAccessTag(TBAA);

test/llvmpasses/late-lower-gc.ll

@@ -90,6 +90,20 @@ top:
     ret void
 }
 
+; Confirm that `invariant.load` on other loads survive
+define void @gc_keep_invariant(float addrspace(1)* %0) {
+top:
+; CHECK-LABEL: @gc_keep_invariant
+    %pgcstack = call {}*** @julia.get_pgcstack()
+    %1 = bitcast {}*** %pgcstack to {}**
+    %current_task = getelementptr inbounds {}*, {}** %1, i64 -12
+
+; CHECK: %current_task = getelementptr inbounds ptr, ptr %1, i64 -12
+    %2 = load float, ptr addrspace(1) %0, align 4, !invariant.load !1
+; CHECK-NEXT: %2 = load float, ptr addrspace(1) %0, align 4, !invariant.load
+    ret void
+}
+
 define i32 @callee_root({} addrspace(10)* %v0, {} addrspace(10)* %v1) {
 top:
 ; CHECK-LABEL: @callee_root