[llvm] [InstCombine] Propagate invariant.load metadata across unpacked loads (PR #152186)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Vedant Paranjape (VedantParanjape)

<details>
<summary>Changes</summary>

For loads that operate on aggregate type, instcombine unpacks the loads. It does not preserve the invariant.load metadata. This patch fixes that, it looks for the metadata in the parent load and attaches the metadata to the unpacked loads.

```
%struct.double2 = type { double, double }
%struct.double1 = type { double }

define %struct.double2 @<!-- -->func1(ptr %a) {
  %1 = load %struct.double2, ptr %a, align 16, !invariant.load !1
  ret %struct.double2 %1
}

!1 = !{}
```
Reproducer: https://godbolt.org/z/hcY8MMvYh

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


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp (+10) 
- (added) llvm/test/Transforms/InstCombine/invariant-metadata-propagation.ll (+46) 


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp b/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
index 0be1034b046b6..68252b06e3a9a 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
@@ -718,6 +718,14 @@ static Instruction *combineLoadToOperationType(InstCombinerImpl &IC,
   return nullptr;
 }
 
+// Check if the aggregate load has a invariant.load metadata
+// If aggregate load has invariant.load metadata, add it to the
+// unpacked loads as well.
+static void copyInvariantLoadMetadata(LoadInst &LI, LoadInst *NewLoad) {
+  if (MDNode *MD = LI.getMetadata("invariant.load"))
+    NewLoad->setMetadata("invariant.load", MD);
+}
+
 static Instruction *unpackLoadToAggregate(InstCombinerImpl &IC, LoadInst &LI) {
   // FIXME: We could probably with some care handle both volatile and atomic
   // stores here but it isn't clear that this is important.
@@ -737,6 +745,7 @@ static Instruction *unpackLoadToAggregate(InstCombinerImpl &IC, LoadInst &LI) {
       LoadInst *NewLoad = IC.combineLoadToNewType(LI, ST->getTypeAtIndex(0U),
                                                   ".unpack");
       NewLoad->setAAMetadata(LI.getAAMetadata());
+      copyInvariantLoadMetadata(LI, NewLoad);
       return IC.replaceInstUsesWith(LI, IC.Builder.CreateInsertValue(
         PoisonValue::get(T), NewLoad, 0, Name));
     }
@@ -764,6 +773,7 @@ static Instruction *unpackLoadToAggregate(InstCombinerImpl &IC, LoadInst &LI) {
           Name + ".unpack");
       // Propagate AA metadata. It'll still be valid on the narrowed load.
       L->setAAMetadata(LI.getAAMetadata());
+      copyInvariantLoadMetadata(LI, L);
       V = IC.Builder.CreateInsertValue(V, L, i);
     }
 
diff --git a/llvm/test/Transforms/InstCombine/invariant-metadata-propagation.ll b/llvm/test/Transforms/InstCombine/invariant-metadata-propagation.ll
new file mode 100644
index 0000000000000..acc5e7ca8d2b4
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/invariant-metadata-propagation.ll
@@ -0,0 +1,46 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S < %s -passes=instcombine | FileCheck %s
+
+%struct.double2 = type { double, double }
+%struct.double1 = type { double }
+
+define %struct.double2 @func1(ptr addrspace(1) %a) {
+; CHECK-LABEL: define %struct.double2 @func1(
+; CHECK-SAME: ptr addrspace(1) [[A:%.*]]) {
+; CHECK-NEXT:    [[DOTUNPACK:%.*]] = load double, ptr addrspace(1) [[A]], align 16, !invariant.load [[META0:![0-9]+]]
+; CHECK-NEXT:    [[TMP1:%.*]] = insertvalue [[STRUCT_DOUBLE2:%.*]] poison, double [[DOTUNPACK]], 0
+; CHECK-NEXT:    [[DOTELT1:%.*]] = getelementptr inbounds nuw i8, ptr addrspace(1) [[A]], i64 8
+; CHECK-NEXT:    [[DOTUNPACK2:%.*]] = load double, ptr addrspace(1) [[DOTELT1]], align 8, !invariant.load [[META0]]
+; CHECK-NEXT:    [[TMP2:%.*]] = insertvalue [[STRUCT_DOUBLE2]] [[TMP1]], double [[DOTUNPACK2]], 1
+; CHECK-NEXT:    ret [[STRUCT_DOUBLE2]] [[TMP2]]
+;
+  %1 = load %struct.double2, ptr addrspace(1) %a, align 16, !invariant.load !1
+  ret %struct.double2 %1
+}
+
+define %struct.double2 @func2(ptr %a) {
+; CHECK-LABEL: define %struct.double2 @func2(
+; CHECK-SAME: ptr [[A:%.*]]) {
+; CHECK-NEXT:    [[DOTUNPACK:%.*]] = load double, ptr [[A]], align 16, !invariant.load [[META0]]
+; CHECK-NEXT:    [[TMP1:%.*]] = insertvalue [[STRUCT_DOUBLE2:%.*]] poison, double [[DOTUNPACK]], 0
+; CHECK-NEXT:    [[DOTELT1:%.*]] = getelementptr inbounds nuw i8, ptr [[A]], i64 8
+; CHECK-NEXT:    [[DOTUNPACK2:%.*]] = load double, ptr [[DOTELT1]], align 8, !invariant.load [[META0]]
+; CHECK-NEXT:    [[TMP2:%.*]] = insertvalue [[STRUCT_DOUBLE2]] [[TMP1]], double [[DOTUNPACK2]], 1
+; CHECK-NEXT:    ret [[STRUCT_DOUBLE2]] [[TMP2]]
+;
+  %1 = load %struct.double2, ptr %a, align 16, !invariant.load !1
+  ret %struct.double2 %1
+}
+
+define %struct.double1 @func3(ptr %a) {
+; CHECK-LABEL: define %struct.double1 @func3(
+; CHECK-SAME: ptr [[A:%.*]]) {
+; CHECK-NEXT:    [[DOTUNPACK:%.*]] = load double, ptr [[A]], align 16, !invariant.load [[META0]]
+; CHECK-NEXT:    [[TMP1:%.*]] = insertvalue [[STRUCT_DOUBLE1:%.*]] poison, double [[DOTUNPACK]], 0
+; CHECK-NEXT:    ret [[STRUCT_DOUBLE1]] [[TMP1]]
+;
+  %1 = load %struct.double1, ptr %a, align 16, !invariant.load !1
+  ret %struct.double1 %1
+}
+
+!1 = !{}

``````````

</details>


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