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

[Vedant Paranjape via llvm-commits]

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

>From a3d0e64801b72f84e73fde09eb037535106a86fb Mon Sep 17 00:00:00 2001
From: Vedant Paranjape <vedantvinaya at nvidia.com>
Date: Tue, 5 Aug 2025 00:04:44 +0000
Subject: [PATCH] [InstCombine] Propagate invariant.load metadata across
 unpacked loads

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.
---
 .../InstCombineLoadStoreAlloca.cpp            | 10 ++++
 .../invariant-metadata-propagation.ll         | 46 +++++++++++++++++++
 2 files changed, 56 insertions(+)
 create mode 100644 llvm/test/Transforms/InstCombine/invariant-metadata-propagation.ll

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 = !{}