[llvm] [CodeExtractor] Add align metadata to extracted pointers (PR #131131)

[Dominik Adamski via llvm-commits]

https://github.com/DominikAdamski updated https://github.com/llvm/llvm-project/pull/131131

>From 802fef5f5e76dcdbaed87c1b665d74bce7c9c514 Mon Sep 17 00:00:00 2001
From: Dominik Adamski <dominik.adamski at amd.com>
Date: Tue, 11 Mar 2025 07:41:50 -0500
Subject: [PATCH 1/2] [CodeExtractor] Add align metadata to extracted pointers

Moving code to another function can lead to missed optimization opportunities,
because function passes operate on smaller chunks of code,
and they cannot figure out all details.

One example of missed optimization opportunities after code extraction
is information about pointer alignment. The instruction combine pass
adds information about pointer alignment to LLVM intrinsic memcpy calls
if it can deduce it from the code or if align metadata is added.
If this information is not present, then further optimization passes
can generate inefficient code.

If we add align metadata to extracted pointers, then the instruction
combine pass can add the align attribute to the LLVM intrinsic memcpy call
and unblock further optimization.
---
 llvm/lib/Transforms/Utils/CodeExtractor.cpp | 22 +++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
index 7277603b3ec2b..61d70a1500028 100644
--- a/llvm/lib/Transforms/Utils/CodeExtractor.cpp
+++ b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
@@ -1604,8 +1604,26 @@ void CodeExtractor::emitFunctionBody(
       Idx[1] = ConstantInt::get(Type::getInt32Ty(header->getContext()), aggIdx);
       GetElementPtrInst *GEP = GetElementPtrInst::Create(
           StructArgTy, AggArg, Idx, "gep_" + inputs[i]->getName(), newFuncRoot);
-      RewriteVal = new LoadInst(StructArgTy->getElementType(aggIdx), GEP,
-                                "loadgep_" + inputs[i]->getName(), newFuncRoot);
+      LoadInst *LoadGEP =
+          new LoadInst(StructArgTy->getElementType(aggIdx), GEP,
+                       "loadgep_" + inputs[i]->getName(), newFuncRoot);
+      PointerType *ItemType =
+          dyn_cast<PointerType>(StructArgTy->getElementType(aggIdx));
+      if (ItemType && !LoadGEP->getMetadata(LLVMContext::MD_align)) {
+        unsigned AddressSpace = ItemType->getAddressSpace();
+        unsigned AlignmentValue = oldFunction->getDataLayout()
+                                      .getPointerPrefAlignment(AddressSpace)
+                                      .value();
+
+        MDBuilder MDB(header->getContext());
+        LoadGEP->setMetadata(
+            LLVMContext::MD_align,
+            MDNode::get(
+                header->getContext(),
+                MDB.createConstant(ConstantInt::get(
+                    Type::getInt64Ty(header->getContext()), AlignmentValue))));
+      }
+      RewriteVal = LoadGEP;
       ++aggIdx;
     } else
       RewriteVal = &*ScalarAI++;

>From 3a41608e17f24df357e2d0131c353a182b659024 Mon Sep 17 00:00:00 2001
From: Dominik Adamski <dominik.adamski at amd.com>
Date: Fri, 14 Mar 2025 04:49:40 -0500
Subject: [PATCH 2/2] Applied remarks

---
 llvm/lib/Transforms/Utils/CodeExtractor.cpp | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
index 61d70a1500028..9ec2ba6dc3df3 100644
--- a/llvm/lib/Transforms/Utils/CodeExtractor.cpp
+++ b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
@@ -1609,11 +1609,11 @@ void CodeExtractor::emitFunctionBody(
                        "loadgep_" + inputs[i]->getName(), newFuncRoot);
       PointerType *ItemType =
           dyn_cast<PointerType>(StructArgTy->getElementType(aggIdx));
-      if (ItemType && !LoadGEP->getMetadata(LLVMContext::MD_align)) {
+      if (ItemType) {
         unsigned AddressSpace = ItemType->getAddressSpace();
-        unsigned AlignmentValue = oldFunction->getDataLayout()
-                                      .getPointerPrefAlignment(AddressSpace)
-                                      .value();
+        // Use the same alignment as the one used for struct allocation.
+        unsigned AlignmentValue =
+            oldFunction->getDataLayout().getPrefTypeAlign(StructArgTy).value();
 
         MDBuilder MDB(header->getContext());
         LoadGEP->setMetadata(
@@ -1817,7 +1817,8 @@ CallInst *CodeExtractor::emitReplacerCall(
 
     AllocaInst *alloca = new AllocaInst(
         output->getType(), DL.getAllocaAddrSpace(), nullptr,
-        output->getName() + ".loc", AllocaBlock->getFirstInsertionPt());
+        DL.getPrefTypeAlign(output->getType()), output->getName() + ".loc",
+        AllocaBlock->getFirstInsertionPt());
     params.push_back(alloca);
     ReloadOutputs.push_back(alloca);
   }
@@ -1825,7 +1826,8 @@ CallInst *CodeExtractor::emitReplacerCall(
   AllocaInst *Struct = nullptr;
   if (!StructValues.empty()) {
     Struct = new AllocaInst(StructArgTy, DL.getAllocaAddrSpace(), nullptr,
-                            "structArg", AllocaBlock->getFirstInsertionPt());
+                            DL.getPrefTypeAlign(StructArgTy), "structArg",
+                            AllocaBlock->getFirstInsertionPt());
     if (ArgsInZeroAddressSpace && DL.getAllocaAddrSpace() != 0) {
       auto *StructSpaceCast = new AddrSpaceCastInst(
           Struct, PointerType ::get(Context, 0), "structArg.ascast");