[llvm] [InstCombine] Preserve all load metadata when folding load of select (PR #206663)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: ayrai-gb

<details>
<summary>Changes</summary>

`InstCombinerImpl::visitLoadInst` folds a load of a select-of-pointers into a select of loads:

```llvm
load (select c, p1, p2)
    =>
select c, (load p1), (load p2)
```

The transform only fires when both pointers are provably safe to load unconditionally (`isSafeToLoadUnconditionally` on each arm).

When creating the two new loads, the transform currently copies only poison-generating metadata via

```cpp
copyMetadata(LI, Metadata::PoisonGeneratingIDs)
```

which preserves only `!range`, `!nonnull`, `!align`, and `!nofpclass`. All other load metadata is silently dropped, including `!tbaa`, `!invariant.group`, `!invariant.load`, `!nontemporal`, `!dereferenceable`, `!noalias`/`!alias.scope`, `!llvm.access.group`, and `!noundef`.

Each new load performs the same access as the original load on one arm of the select and executes only when that arm is chosen, so the original load's metadata remains valid for the new loads. There is therefore no reason to restrict preservation to poison-generating metadata.

Replace the two `copyMetadata(..., Metadata::PoisonGeneratingIDs)` calls with `copyMetadataForLoad`, the existing helper already used by the load-retype path, so that the full set of load metadata is preserved.

This is a missed optimization rather than a correctness issue—the loss of metadata is the conservative direction—but it is observable. For example, losing `!invariant.group` can inhibit devirtualization, while losing `!tbaa` and `!noalias` pessimizes later alias-analysis-driven optimizations.

The existing `preserve_load_metadata_after_select_transform*` tests in `test/Transforms/InstCombine/loadstore-metadata.ll` already exercise this transform; their `CHECK` lines previously reflected the limited (poison-only) copying. Regenerate them to show the full metadata preservation, and extend `transform1` to also cover `!nontemporal` and `!invariant.group`.

Found via @<!-- -->jlebar's X86 LLVM bug-hunt / FuzzX effort:

* https://github.com/SemiAnalysisAI/FuzzX/tree/master/x86/bugs/179-instcombine-load-of-select-drops-invariant-group-tbaa

cc @<!-- -->jlebar


---
Full diff: https://github.com/llvm/llvm-project/pull/206663.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp (+7-4) 
- (modified) llvm/test/Transforms/InstCombine/loadstore-metadata.ll (+5-5) 


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp b/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
index e6481c89c0265..12f3caf3a9a20 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
@@ -1170,10 +1170,13 @@ Instruction *InstCombinerImpl::visitLoadInst(LoadInst &LI) {
         V1->setAtomic(LI.getOrdering(), LI.getSyncScopeID());
         V2->setAlignment(Alignment);
         V2->setAtomic(LI.getOrdering(), LI.getSyncScopeID());
-        // It is safe to copy any metadata that does not trigger UB. Copy any
-        // poison-generating metadata.
-        V1->copyMetadata(LI, Metadata::PoisonGeneratingIDs);
-        V2->copyMetadata(LI, Metadata::PoisonGeneratingIDs);
+        // Each new load performs the same access as the original load in one
+        // arm of the select, and executes only when that arm is selected, so
+        // every fact the original load carried still holds. Clone all load
+        // metadata (not just the poison-generating kinds) to preserve !tbaa,
+        // !invariant.group, !nontemporal, !dereferenceable, etc.
+        copyMetadataForLoad(*V1, LI);
+        copyMetadataForLoad(*V2, LI);
         return SelectInst::Create(SI->getCondition(), V1, V2, "", nullptr,
                                   ProfcheckDisableMetadataFixes ? nullptr : SI);
       }
diff --git a/llvm/test/Transforms/InstCombine/loadstore-metadata.ll b/llvm/test/Transforms/InstCombine/loadstore-metadata.ll
index 02552df831196..add4e8145aa21 100644
--- a/llvm/test/Transforms/InstCombine/loadstore-metadata.ll
+++ b/llvm/test/Transforms/InstCombine/loadstore-metadata.ll
@@ -204,14 +204,14 @@ define ptr @preserve_load_metadata_after_select_transform1(i1 %c, ptr dereferenc
 ; CHECK-LABEL: define ptr @preserve_load_metadata_after_select_transform1(
 ; CHECK-SAME: i1 [[C:%.*]], ptr dereferenceable(8) [[A:%.*]], ptr dereferenceable(8) [[B:%.*]]) {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[B_VAL:%.*]] = load ptr, ptr [[B]], align 1, !nonnull [[META6]], !align [[META8]]
-; CHECK-NEXT:    [[A_VAL:%.*]] = load ptr, ptr [[A]], align 1, !nonnull [[META6]], !align [[META8]]
+; CHECK-NEXT:    [[B_VAL:%.*]] = load ptr, ptr [[B]], align 1, !tbaa [[SCALAR_TYPE_TBAA0]], !invariant.load [[META6]], !nontemporal [[META7]], !nonnull [[META6]], !dereferenceable [[META8]], !invariant.group [[META6]], !align [[META8]], !llvm.access.group [[META6]], !noundef [[META6]]
+; CHECK-NEXT:    [[A_VAL:%.*]] = load ptr, ptr [[A]], align 1, !tbaa [[SCALAR_TYPE_TBAA0]], !invariant.load [[META6]], !nontemporal [[META7]], !nonnull [[META6]], !dereferenceable [[META8]], !invariant.group [[META6]], !align [[META8]], !llvm.access.group [[META6]], !noundef [[META6]]
 ; CHECK-NEXT:    [[L_SEL:%.*]] = select i1 [[C]], ptr [[B_VAL]], ptr [[A_VAL]]
 ; CHECK-NEXT:    ret ptr [[L_SEL]]
 ;
 entry:
   %ptr.sel = select i1 %c, ptr %b, ptr %a
-  %l.sel = load ptr, ptr %ptr.sel, align 1, !tbaa !0, !llvm.access.group !7, !dereferenceable !9, !noundef !{}, !invariant.load !7, !align !9, !nonnull !{}
+  %l.sel = load ptr, ptr %ptr.sel, align 1, !tbaa !0, !llvm.access.group !7, !dereferenceable !9, !noundef !{}, !invariant.load !7, !align !9, !nonnull !{}, !nontemporal !8, !invariant.group !7
   ret ptr %l.sel
 }
 
@@ -220,8 +220,8 @@ define i32 @preserve_load_metadata_after_select_transform_range(i1 %c, ptr deref
 ; CHECK-LABEL: define i32 @preserve_load_metadata_after_select_transform_range(
 ; CHECK-SAME: i1 [[C:%.*]], ptr dereferenceable(8) [[A:%.*]], ptr dereferenceable(8) [[B:%.*]]) {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
-; CHECK-NEXT:    [[B_VAL:%.*]] = load i32, ptr [[B]], align 1, !range [[RNG11:![0-9]+]]
-; CHECK-NEXT:    [[A_VAL:%.*]] = load i32, ptr [[A]], align 1, !range [[RNG11]]
+; CHECK-NEXT:    [[B_VAL:%.*]] = load i32, ptr [[B]], align 1, !tbaa [[SCALAR_TYPE_TBAA0]], !range [[RNG11:![0-9]+]], !invariant.load [[META6]], !llvm.access.group [[META6]], !noundef [[META6]]
+; CHECK-NEXT:    [[A_VAL:%.*]] = load i32, ptr [[A]], align 1, !tbaa [[SCALAR_TYPE_TBAA0]], !range [[RNG11]], !invariant.load [[META6]], !llvm.access.group [[META6]], !noundef [[META6]]
 ; CHECK-NEXT:    [[L_SEL:%.*]] = select i1 [[C]], i32 [[B_VAL]], i32 [[A_VAL]]
 ; CHECK-NEXT:    ret i32 [[L_SEL]]
 ;

``````````

</details>


https://github.com/llvm/llvm-project/pull/206663