[llvm] [EarlyCSE] Rematerialize alignment assumption. (PR #109131)

[Florian Hahn via llvm-commits]

https://github.com/fhahn updated https://github.com/llvm/llvm-project/pull/109131

>From b4f54c4b9451fb54a1a3037b754009d414ef2501 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Wed, 18 Sep 2024 13:38:52 +0100
Subject: [PATCH] [EarlyCSE] Rematerialize alignment assumption.

---
 llvm/lib/Transforms/Scalar/EarlyCSE.cpp               | 11 +++++++++++
 .../EarlyCSE/materialize-align-assumptions.ll         |  3 +++
 2 files changed, 14 insertions(+)

diff --git a/llvm/lib/Transforms/Scalar/EarlyCSE.cpp b/llvm/lib/Transforms/Scalar/EarlyCSE.cpp
index cf11f5bc885a75..da66ed49ebe8b2 100644
--- a/llvm/lib/Transforms/Scalar/EarlyCSE.cpp
+++ b/llvm/lib/Transforms/Scalar/EarlyCSE.cpp
@@ -31,6 +31,7 @@
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
@@ -1588,6 +1589,16 @@ bool EarlyCSE::processNode(DomTreeNode *Node) {
         if (InVal.IsLoad)
           if (auto *I = dyn_cast<Instruction>(Op))
             combineMetadataForCSE(I, &Inst, false);
+
+        // If the load has align metadata, preserve it via an alignment assumption. Note that this doesn't use salavageKnowledge, as we need to create the assumption for the value we replaced the load with.
+        if (auto *AlignMD = Inst.getMetadata(LLVMContext::MD_align)) {
+          auto *A = mdconst::extract<ConstantInt>(AlignMD->getOperand(0));
+          if (Op->getPointerAlignment(SQ.DL).value() % A->getZExtValue() != 0) {
+            IRBuilder B(&Inst);
+            B.CreateAlignmentAssumption(SQ.DL, Op, A);
+          }
+        }
+
         if (!Inst.use_empty())
           Inst.replaceAllUsesWith(Op);
         salvageKnowledge(&Inst, &AC);
diff --git a/llvm/test/Transforms/EarlyCSE/materialize-align-assumptions.ll b/llvm/test/Transforms/EarlyCSE/materialize-align-assumptions.ll
index ea63376957162b..628577b0975071 100644
--- a/llvm/test/Transforms/EarlyCSE/materialize-align-assumptions.ll
+++ b/llvm/test/Transforms/EarlyCSE/materialize-align-assumptions.ll
@@ -10,6 +10,7 @@ define ptr @align_replacement_does_not_have_align_metadata(ptr noalias %p) {
 ; CHECK-NEXT:    call void @foo(ptr [[L_1]])
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr [[L_1]], i64 4
 ; CHECK-NEXT:    store ptr [[GEP]], ptr [[P]], align 8
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[GEP]], i64 4) ]
 ; CHECK-NEXT:    ret ptr [[GEP]]
 ;
   %l.1 = load ptr, ptr %p, align 8
@@ -27,6 +28,7 @@ define ptr @align_replacement_does_not_have_align_metadata2(ptr noalias %p) {
 ; CHECK-NEXT:    [[L_1:%.*]] = load ptr, ptr [[P]], align 8
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr [[L_1]], i64 4
 ; CHECK-NEXT:    store ptr [[GEP]], ptr [[P]], align 8
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[GEP]], i64 4) ]
 ; CHECK-NEXT:    ret ptr [[GEP]]
 ;
   %l.1 = load ptr, ptr %p, align 8
@@ -54,6 +56,7 @@ define ptr @align_replacement_has_smaller_alignment(ptr noalias %p) {
 ; CHECK-SAME: ptr noalias [[P:%.*]]) {
 ; CHECK-NEXT:    [[L_1:%.*]] = load ptr, ptr [[P]], align 8, !align [[META0]]
 ; CHECK-NEXT:    call void @foo(ptr [[L_1]])
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[L_1]], i64 8) ]
 ; CHECK-NEXT:    ret ptr [[L_1]]
 ;
   %l.1 = load ptr, ptr %p, align 8, !align !0