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

Florian Hahn via llvm-commits llvm-commits at lists.llvm.org
Wed Jan 15 11:17:08 PST 2025


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

>From 90f86103adc124e518a6616e92c23cf1a281a9b3 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 1/7] [EarlyCSE] Rematerialize alignment assumption.

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

diff --git a/llvm/lib/Transforms/Scalar/EarlyCSE.cpp b/llvm/lib/Transforms/Scalar/EarlyCSE.cpp
index 3a0ae6b01a1144..4f8a2e06e6b451 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"
@@ -1599,6 +1600,18 @@ 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

>From bddd57d4cee58fb0ff604140154234296da37e7e Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Thu, 19 Sep 2024 21:33:40 +0100
Subject: [PATCH 2/7] !fixup check noundef

---
 llvm/lib/Transforms/Scalar/EarlyCSE.cpp       | 20 ++++++-----
 .../EarlyCSE/materialize-align-assumptions.ll | 35 ++++++++++++++-----
 2 files changed, 39 insertions(+), 16 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/EarlyCSE.cpp b/llvm/lib/Transforms/Scalar/EarlyCSE.cpp
index 4f8a2e06e6b451..3745311e0ff307 100644
--- a/llvm/lib/Transforms/Scalar/EarlyCSE.cpp
+++ b/llvm/lib/Transforms/Scalar/EarlyCSE.cpp
@@ -1601,14 +1601,18 @@ bool EarlyCSE::processNode(DomTreeNode *Node) {
           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 the load has align and noundef 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 (Inst.hasMetadata(LLVMContext::MD_noundef)) {
+          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);
+            }
           }
         }
 
diff --git a/llvm/test/Transforms/EarlyCSE/materialize-align-assumptions.ll b/llvm/test/Transforms/EarlyCSE/materialize-align-assumptions.ll
index 628577b0975071..837a73a00d6431 100644
--- a/llvm/test/Transforms/EarlyCSE/materialize-align-assumptions.ll
+++ b/llvm/test/Transforms/EarlyCSE/materialize-align-assumptions.ll
@@ -3,6 +3,24 @@
 
 declare void @foo(ptr)
 
+define ptr @align_replacement_does_not_have_align_metadata_missing_noundef(ptr noalias %p) {
+; CHECK-LABEL: define ptr @align_replacement_does_not_have_align_metadata_missing_noundef(
+; CHECK-SAME: ptr noalias [[P:%.*]]) {
+; CHECK-NEXT:    [[L_1:%.*]] = load ptr, ptr [[P]], align 8
+; 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:    ret ptr [[GEP]]
+;
+  %l.1 = load ptr, ptr %p, align 8
+  call void @foo(ptr %l.1)
+  %l.2 = load ptr, ptr %p, align 8
+  %gep = getelementptr i8, ptr %l.2, i64 4
+  store ptr %gep, ptr %p, align 8
+  %l.3 = load ptr, ptr %p, align 8, !align !0
+  ret ptr %l.3
+}
+
 define ptr @align_replacement_does_not_have_align_metadata(ptr noalias %p) {
 ; CHECK-LABEL: define ptr @align_replacement_does_not_have_align_metadata(
 ; CHECK-SAME: ptr noalias [[P:%.*]]) {
@@ -18,7 +36,7 @@ define ptr @align_replacement_does_not_have_align_metadata(ptr noalias %p) {
   %l.2 = load ptr, ptr %p, align 8
   %gep = getelementptr i8, ptr %l.2, i64 4
   store ptr %gep, ptr %p, align 8
-  %l.3 = load ptr, ptr %p, align 8, !align !0
+  %l.3 = load ptr, ptr %p, align 8, !align !0, !noundef !{}
   ret ptr %l.3
 }
 
@@ -34,7 +52,7 @@ define ptr @align_replacement_does_not_have_align_metadata2(ptr noalias %p) {
   %l.1 = load ptr, ptr %p, align 8
   %gep = getelementptr i8, ptr %l.1, i64 4
   store ptr %gep, ptr %p, align 8
-  %l.2 = load ptr, ptr %p, align 8, !align !0
+  %l.2 = load ptr, ptr %p, align 8, !align !0, !noundef !{}
   ret ptr %l.2
 }
 
@@ -61,7 +79,7 @@ define ptr @align_replacement_has_smaller_alignment(ptr noalias %p) {
 ;
   %l.1 = load ptr, ptr %p, align 8, !align !0
   call void @foo(ptr %l.1)
-  %l.2 = load ptr, ptr %p, align 8, !align !1
+  %l.2 = load ptr, ptr %p, align 8, !align !1, !noundef !{}
   ret ptr %l.2
 }
 
@@ -70,12 +88,12 @@ define ptr @align_replacement_has_larger_alignment(ptr %p) {
 ; CHECK-SAME: ptr [[P:%.*]]) {
 ; CHECK-NEXT:    [[L_1:%.*]] = load ptr, ptr [[P]], align 8, !align [[META1:![0-9]+]]
 ; CHECK-NEXT:    call void @foo(ptr [[L_1]])
-; CHECK-NEXT:    [[L_2:%.*]] = load ptr, ptr [[P]], align 8, !align [[META0]]
+; CHECK-NEXT:    [[L_2:%.*]] = load ptr, ptr [[P]], align 8, !align [[META0]], !noundef [[META2:![0-9]+]]
 ; CHECK-NEXT:    ret ptr [[L_2]]
 ;
   %l.1 = load ptr, ptr %p, align 8, !align !1
   call void @foo(ptr %l.1)
-  %l.2 = load ptr, ptr %p, align 8, !align !0
+  %l.2 = load ptr, ptr %p, align 8, !align !0, !noundef !{}
   ret ptr %l.2
 }
 
@@ -84,12 +102,12 @@ define ptr @align_1(ptr %p) {
 ; CHECK-SAME: ptr [[P:%.*]]) {
 ; CHECK-NEXT:    [[L_1:%.*]] = load ptr, ptr [[P]], align 8
 ; CHECK-NEXT:    call void @foo(ptr [[L_1]])
-; CHECK-NEXT:    [[L_2:%.*]] = load ptr, ptr [[P]], align 8, !align [[META2:![0-9]+]]
+; CHECK-NEXT:    [[L_2:%.*]] = load ptr, ptr [[P]], align 8, !align [[META3:![0-9]+]], !noundef [[META2]]
 ; CHECK-NEXT:    ret ptr [[L_2]]
 ;
   %l.1 = load ptr, ptr %p, align 8
   call void @foo(ptr %l.1)
-  %l.2 = load ptr, ptr %p, align 8, !align !2
+  %l.2 = load ptr, ptr %p, align 8, !align !2, !noundef !{}
   ret ptr %l.2
 }
 
@@ -99,5 +117,6 @@ define ptr @align_1(ptr %p) {
 ;.
 ; CHECK: [[META0]] = !{i64 4}
 ; CHECK: [[META1]] = !{i64 8}
-; CHECK: [[META2]] = !{i64 1}
+; CHECK: [[META2]] = !{}
+; CHECK: [[META3]] = !{i64 1}
 ;.

>From 8b5312148bee68e3545c876c986bf2766845f707 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Mon, 13 Jan 2025 21:31:42 +0000
Subject: [PATCH 3/7] Use computeKnownBits.

---
 llvm/lib/Transforms/Scalar/EarlyCSE.cpp | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/EarlyCSE.cpp b/llvm/lib/Transforms/Scalar/EarlyCSE.cpp
index 3745311e0ff307..378c16f126f4de 100644
--- a/llvm/lib/Transforms/Scalar/EarlyCSE.cpp
+++ b/llvm/lib/Transforms/Scalar/EarlyCSE.cpp
@@ -16,6 +16,7 @@
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopedHashTable.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AssumptionCache.h"
@@ -1607,9 +1608,11 @@ bool EarlyCSE::processNode(DomTreeNode *Node) {
         // load with.
         if (Inst.hasMetadata(LLVMContext::MD_noundef)) {
           if (auto *AlignMD = Inst.getMetadata(LLVMContext::MD_align)) {
+            Inst.setMetadata(LLVMContext::MD_align, nullptr);
             auto *A = mdconst::extract<ConstantInt>(AlignMD->getOperand(0));
-            if (Op->getPointerAlignment(SQ.DL).value() % A->getZExtValue() !=
-                0) {
+            auto KB = computeKnownBits(Op, SQ.DL);
+            unsigned AlignFromKB = 1 << KB.countMinTrailingZeros();
+            if (AlignFromKB < A->getZExtValue()) {
               IRBuilder B(&Inst);
               B.CreateAlignmentAssumption(SQ.DL, Op, A);
             }

>From b182b4d1edd3083a8750bff731f2c4bd295900cb Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Mon, 13 Jan 2025 21:32:09 +0000
Subject: [PATCH 4/7] Remove align assumptions if possible.

---
 .../InstCombine/InstCombineCalls.cpp          | 36 +++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index dd5a4ba5a4724a..9af24c18b9531c 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3226,6 +3226,42 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
         MaybeSimplifyHint(OBU.Inputs[0]);
         MaybeSimplifyHint(OBU.Inputs[1]);
       }
+
+      if (OBU.getTagName() == "align" && OBU.Inputs.size() == 2) {
+        RetainedKnowledge RK = getKnowledgeFromBundle(
+            *cast<AssumeInst>(II), II->bundle_op_info_begin()[Idx]);
+        if (!RK || RK.AttrKind != Attribute::Alignment ||
+            !isPowerOf2_64(RK.ArgValue))
+          continue;
+        SetVector<const Instruction *> WorkList;
+        bool AlignNeeded = false;
+        WorkList.insert(II);
+        for (unsigned I = 0; I != WorkList.size(); ++I) {
+          if (auto *LI = dyn_cast<LoadInst>(WorkList[I])) {
+            if (auto *AlignMD = LI->getMetadata(LLVMContext::MD_align)) {
+              auto *A = mdconst::extract<ConstantInt>(AlignMD->getOperand(0));
+
+              if (A->getZExtValue() % RK.ArgValue != 0) {
+                AlignNeeded = true;
+                break;
+              }
+            }
+          }
+          if (isa<ICmpInst>(WorkList[I])) {
+            AlignNeeded = true;
+            break;
+          }
+          if (WorkList.size() > 16) {
+            AlignNeeded = true;
+            break;
+          }
+
+          for (const User *U : WorkList[I]->users())
+            WorkList.insert(cast<Instruction>(U));
+        }
+        auto *New = CallBase::removeOperandBundle(II, OBU.getTagID());
+        return New;
+      }
     }
 
     // Convert nonnull assume like:

>From 7b2f38bb5bc441aa2858b0f683e0b8c929634ed7 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Tue, 14 Jan 2025 14:08:44 +0000
Subject: [PATCH 5/7] step

---
 .../InstCombine/InstCombineCalls.cpp          | 36 ---------------
 llvm/lib/Transforms/Scalar/EarlyCSE.cpp       | 44 +++++++++++++++++--
 2 files changed, 40 insertions(+), 40 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index 9af24c18b9531c..dd5a4ba5a4724a 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3226,42 +3226,6 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
         MaybeSimplifyHint(OBU.Inputs[0]);
         MaybeSimplifyHint(OBU.Inputs[1]);
       }
-
-      if (OBU.getTagName() == "align" && OBU.Inputs.size() == 2) {
-        RetainedKnowledge RK = getKnowledgeFromBundle(
-            *cast<AssumeInst>(II), II->bundle_op_info_begin()[Idx]);
-        if (!RK || RK.AttrKind != Attribute::Alignment ||
-            !isPowerOf2_64(RK.ArgValue))
-          continue;
-        SetVector<const Instruction *> WorkList;
-        bool AlignNeeded = false;
-        WorkList.insert(II);
-        for (unsigned I = 0; I != WorkList.size(); ++I) {
-          if (auto *LI = dyn_cast<LoadInst>(WorkList[I])) {
-            if (auto *AlignMD = LI->getMetadata(LLVMContext::MD_align)) {
-              auto *A = mdconst::extract<ConstantInt>(AlignMD->getOperand(0));
-
-              if (A->getZExtValue() % RK.ArgValue != 0) {
-                AlignNeeded = true;
-                break;
-              }
-            }
-          }
-          if (isa<ICmpInst>(WorkList[I])) {
-            AlignNeeded = true;
-            break;
-          }
-          if (WorkList.size() > 16) {
-            AlignNeeded = true;
-            break;
-          }
-
-          for (const User *U : WorkList[I]->users())
-            WorkList.insert(cast<Instruction>(U));
-        }
-        auto *New = CallBase::removeOperandBundle(II, OBU.getTagID());
-        return New;
-      }
     }
 
     // Convert nonnull assume like:
diff --git a/llvm/lib/Transforms/Scalar/EarlyCSE.cpp b/llvm/lib/Transforms/Scalar/EarlyCSE.cpp
index 378c16f126f4de..fd7f0eb1ff1fc4 100644
--- a/llvm/lib/Transforms/Scalar/EarlyCSE.cpp
+++ b/llvm/lib/Transforms/Scalar/EarlyCSE.cpp
@@ -1609,12 +1609,48 @@ bool EarlyCSE::processNode(DomTreeNode *Node) {
         if (Inst.hasMetadata(LLVMContext::MD_noundef)) {
           if (auto *AlignMD = Inst.getMetadata(LLVMContext::MD_align)) {
             Inst.setMetadata(LLVMContext::MD_align, nullptr);
-            auto *A = mdconst::extract<ConstantInt>(AlignMD->getOperand(0));
+            auto *B = mdconst::extract<ConstantInt>(AlignMD->getOperand(0));
             auto KB = computeKnownBits(Op, SQ.DL);
             unsigned AlignFromKB = 1 << KB.countMinTrailingZeros();
-            if (AlignFromKB < A->getZExtValue()) {
-              IRBuilder B(&Inst);
-              B.CreateAlignmentAssumption(SQ.DL, Op, A);
+            if (AlignFromKB < B->getZExtValue()) {
+              SetVector<const Value *> WorkList;
+              bool AlignNeeded = false;
+                for (const User *U : Op->users())
+                  if (auto *I = dyn_cast<Instruction>(U))
+                    WorkList.insert(I);
+
+              for (unsigned I = 0; I != WorkList.size(); ++I) {
+                auto *Curr = WorkList[I];
+                if (auto *LI = dyn_cast<LoadInst>(Curr)) {
+                  if (LI->getAlign().value() < B->getZExtValue()) {
+                    AlignNeeded = true;
+                    break;
+                  }
+                  continue;
+                }
+                if (auto *SI = dyn_cast<StoreInst>(Curr)) {
+                  if (SI->getAlign().value() < B->getZExtValue()) {
+                    AlignNeeded = true;
+                    break;
+                  }
+                  continue;
+                }
+                if (isa<ICmpInst>(Curr) && !isa<Constant>(cast<Instruction>(Curr)->getOperand(0)) && !isa<Constant>(cast<Instruction>(Curr)->getOperand(1))) {
+                  AlignNeeded = true;
+                  break;
+                }
+                if (WorkList.size() > 16) {
+                  AlignNeeded = true;
+                  break;
+                }
+
+                for (const User *U : Curr->users())
+                  WorkList.insert(cast<Instruction>(U));
+              }
+              if (AlignNeeded) {
+                  IRBuilder Builder(&Inst);
+                  Builder.CreateAlignmentAssumption(SQ.DL, Op, B);
+              }
             }
           }
         }

>From ed33c3d23dc09b2d1640ec0bb8d2bfcda0e42486 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Sun, 22 Dec 2024 11:16:58 +0000
Subject: [PATCH 6/7] [InstCombine] Add option to clean up some assumptions.

---
 .../llvm/Transforms/InstCombine/InstCombine.h |  6 +++++
 llvm/lib/Passes/PassBuilder.cpp               |  2 ++
 llvm/lib/Passes/PassBuilderPipelines.cpp      | 25 ++++++++++++-------
 llvm/lib/Passes/PassRegistry.def              |  2 +-
 .../InstCombine/InstCombineCalls.cpp          |  7 ++++++
 .../InstCombine/InstCombineInternal.h         |  8 ++++--
 .../InstCombine/InstructionCombining.cpp      |  4 +--
 7 files changed, 40 insertions(+), 14 deletions(-)

diff --git a/llvm/include/llvm/Transforms/InstCombine/InstCombine.h b/llvm/include/llvm/Transforms/InstCombine/InstCombine.h
index c12d749709cd25..dd9d05cbbcf7c9 100644
--- a/llvm/include/llvm/Transforms/InstCombine/InstCombine.h
+++ b/llvm/include/llvm/Transforms/InstCombine/InstCombine.h
@@ -31,6 +31,7 @@ struct InstCombineOptions {
   // Verify that a fix point has been reached after MaxIterations.
   bool VerifyFixpoint = false;
   unsigned MaxIterations = InstCombineDefaultMaxIterations;
+  bool CleanupAssumptions = false;
 
   InstCombineOptions() = default;
 
@@ -43,6 +44,11 @@ struct InstCombineOptions {
     MaxIterations = Value;
     return *this;
   }
+
+  InstCombineOptions &setCleanupAssumptions(bool Value) {
+    CleanupAssumptions = Value;
+    return *this;
+  }
 };
 
 class InstCombinePass : public PassInfoMixin<InstCombinePass> {
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index f698a3df08ef78..ebf91c2ff1ca90 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -918,6 +918,8 @@ Expected<InstCombineOptions> parseInstCombineOptions(StringRef Params) {
                     ParamName).str(),
             inconvertibleErrorCode());
       Result.setMaxIterations((unsigned)MaxIterations.getZExtValue());
+    } else if (ParamName == "cleanup-assumptions") {
+      Result.setCleanupAssumptions(Enable);
     } else {
       return make_error<StringError>(
           formatv("invalid InstCombine pass parameter '{0}' ", ParamName).str(),
diff --git a/llvm/lib/Passes/PassBuilderPipelines.cpp b/llvm/lib/Passes/PassBuilderPipelines.cpp
index 4ec0fb8fc81ea4..5d5408aceeb449 100644
--- a/llvm/lib/Passes/PassBuilderPipelines.cpp
+++ b/llvm/lib/Passes/PassBuilderPipelines.cpp
@@ -1305,7 +1305,8 @@ void PassBuilder::addVectorPasses(OptimizationLevel Level,
     FPM.addPass(LoopLoadEliminationPass());
   }
   // Cleanup after the loop optimization passes.
-  FPM.addPass(InstCombinePass());
+  FPM.addPass(
+      InstCombinePass(InstCombineOptions()));
 
   if (Level.getSpeedupLevel() > 1 && ExtraVectorizerPasses) {
     ExtraFunctionPassManager<ShouldRunExtraVectorPasses> ExtraPasses;
@@ -1317,7 +1318,8 @@ void PassBuilder::addVectorPasses(OptimizationLevel Level,
     // dead (or speculatable) control flows or more combining opportunities.
     ExtraPasses.addPass(EarlyCSEPass());
     ExtraPasses.addPass(CorrelatedValuePropagationPass());
-    ExtraPasses.addPass(InstCombinePass());
+    ExtraPasses.addPass(
+        InstCombinePass(InstCombineOptions()));
     LoopPassManager LPM;
     LPM.addPass(LICMPass(PTO.LicmMssaOptCap, PTO.LicmMssaNoAccForPromotionCap,
                          /*AllowSpeculation=*/true));
@@ -1328,7 +1330,8 @@ void PassBuilder::addVectorPasses(OptimizationLevel Level,
                                         /*UseBlockFrequencyInfo=*/true));
     ExtraPasses.addPass(
         SimplifyCFGPass(SimplifyCFGOptions().convertSwitchRangeToICmp(true)));
-    ExtraPasses.addPass(InstCombinePass());
+    ExtraPasses.addPass(
+        InstCombinePass(InstCombineOptions()));
     FPM.addPass(std::move(ExtraPasses));
   }
 
@@ -1351,7 +1354,8 @@ void PassBuilder::addVectorPasses(OptimizationLevel Level,
 
   if (IsFullLTO) {
     FPM.addPass(SCCPPass());
-    FPM.addPass(InstCombinePass());
+    FPM.addPass(
+        InstCombinePass(InstCombineOptions()));
     FPM.addPass(BDCEPass());
   }
 
@@ -1366,7 +1370,8 @@ void PassBuilder::addVectorPasses(OptimizationLevel Level,
   FPM.addPass(VectorCombinePass());
 
   if (!IsFullLTO) {
-    FPM.addPass(InstCombinePass());
+    FPM.addPass(
+        InstCombinePass(InstCombineOptions()));
     // Unroll small loops to hide loop backedge latency and saturate any
     // parallel execution resources of an out-of-order processor. We also then
     // need to clean up redundancies and loop invariant code.
@@ -1392,7 +1397,12 @@ void PassBuilder::addVectorPasses(OptimizationLevel Level,
   }
 
   FPM.addPass(InferAlignmentPass());
-  FPM.addPass(InstCombinePass());
+  // Now that we've vectorized and unrolled loops, we may have more refined
+  // alignment information, try to re-derive it here.
+  FPM.addPass(AlignmentFromAssumptionsPass());
+
+  FPM.addPass(
+      InstCombinePass(InstCombineOptions().setCleanupAssumptions(true)));
 
   // This is needed for two reasons:
   //   1. It works around problems that instcombine introduces, such as sinking
@@ -1405,9 +1415,6 @@ void PassBuilder::addVectorPasses(OptimizationLevel Level,
                /*AllowSpeculation=*/true),
       /*UseMemorySSA=*/true, /*UseBlockFrequencyInfo=*/false));
 
-  // Now that we've vectorized and unrolled loops, we may have more refined
-  // alignment information, try to re-derive it here.
-  FPM.addPass(AlignmentFromAssumptionsPass());
 }
 
 ModulePassManager
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index a93a995655a147..fd985e46c71667 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -539,7 +539,7 @@ FUNCTION_PASS_WITH_PARAMS(
     [](InstCombineOptions Opts) { return InstCombinePass(Opts); },
     parseInstCombineOptions,
     "no-use-loop-info;use-loop-info;no-verify-fixpoint;verify-fixpoint;"
-    "max-iterations=N")
+    "max-iterations=N;cleanup-assumptions")
 FUNCTION_PASS_WITH_PARAMS(
     "loop-unroll", "LoopUnrollPass",
     [](LoopUnrollOptions Opts) { return LoopUnrollPass(Opts); },
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index dd5a4ba5a4724a..4aaa2549df265d 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3226,6 +3226,13 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
         MaybeSimplifyHint(OBU.Inputs[0]);
         MaybeSimplifyHint(OBU.Inputs[1]);
       }
+
+      // Try to clean up some assumption that are not very useful after this
+      // point.
+      if (CleanupAssumptions && OBU.getTagName() == "align") {
+        auto *New = CallBase::removeOperandBundle(II, OBU.getTagID());
+        return New;
+      }
     }
 
     // Convert nonnull assume like:
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
index 83e1da98deeda0..b365a41a87d7ac 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
+++ b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
@@ -59,6 +59,8 @@ class User;
 class LLVM_LIBRARY_VISIBILITY InstCombinerImpl final
     : public InstCombiner,
       public InstVisitor<InstCombinerImpl, Instruction *> {
+  bool CleanupAssumptions = false;
+
 public:
   InstCombinerImpl(InstructionWorklist &Worklist, BuilderTy &Builder,
                    bool MinimizeSize, AAResults *AA, AssumptionCache &AC,
@@ -66,9 +68,11 @@ class LLVM_LIBRARY_VISIBILITY InstCombinerImpl final
                    DominatorTree &DT, OptimizationRemarkEmitter &ORE,
                    BlockFrequencyInfo *BFI, BranchProbabilityInfo *BPI,
                    ProfileSummaryInfo *PSI, const DataLayout &DL,
-                   ReversePostOrderTraversal<BasicBlock *> &RPOT)
+                   ReversePostOrderTraversal<BasicBlock *> &RPOT,
+                   bool CleanupAssumptions)
       : InstCombiner(Worklist, Builder, MinimizeSize, AA, AC, TLI, TTI, DT, ORE,
-                     BFI, BPI, PSI, DL, RPOT) {}
+                     BFI, BPI, PSI, DL, RPOT),
+        CleanupAssumptions(CleanupAssumptions) {}
 
   virtual ~InstCombinerImpl() = default;
 
diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 2fb60ef11499c7..d3a44e06eb61fa 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -5545,7 +5545,7 @@ static bool combineInstructionsOverFunction(
                       << F.getName() << "\n");
 
     InstCombinerImpl IC(Worklist, Builder, F.hasMinSize(), AA, AC, TLI, TTI, DT,
-                        ORE, BFI, BPI, PSI, DL, RPOT);
+                        ORE, BFI, BPI, PSI, DL, RPOT, Opts.CleanupAssumptions);
     IC.MaxArraySizeForCombine = MaxArraySize;
     bool MadeChangeInThisIteration = IC.prepareWorklist(F);
     MadeChangeInThisIteration |= IC.run();
@@ -5594,7 +5594,7 @@ PreservedAnalyses InstCombinePass::run(Function &F,
                                        FunctionAnalysisManager &AM) {
   auto &LRT = AM.getResult<LastRunTrackingAnalysis>(F);
   // No changes since last InstCombine pass, exit early.
-  if (LRT.shouldSkip(&ID))
+  if (LRT.shouldSkip(&ID) && !Options.CleanupAssumptions)
     return PreservedAnalyses::all();
 
   auto &AC = AM.getResult<AssumptionAnalysis>(F);

>From 82398385bd25fe4e49a6537ed8e767dcc5f87281 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Wed, 15 Jan 2025 19:16:26 +0000
Subject: [PATCH 7/7] [InstCombine] Fold align assume into load's !align
 metadata if possible

https://github.com/llvm/llvm-project/pull/108958/
---
 .../InstCombine/InstCombineCalls.cpp          | 31 ++++++++++++++++---
 .../Transforms/InstCombine/assume-align.ll    | 30 ++++++++++++++++--
 2 files changed, 54 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index 4aaa2549df265d..a66db4a3cad54f 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3207,12 +3207,13 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
       // TODO: apply range metadata for range check patterns?
     }
 
-    // Separate storage assumptions apply to the underlying allocations, not any
-    // particular pointer within them. When evaluating the hints for AA purposes
-    // we getUnderlyingObject them; by precomputing the answers here we can
-    // avoid having to do so repeatedly there.
     for (unsigned Idx = 0; Idx < II->getNumOperandBundles(); Idx++) {
       OperandBundleUse OBU = II->getOperandBundleAt(Idx);
+
+      // Separate storage assumptions apply to the underlying allocations, not
+      // any particular pointer within them. When evaluating the hints for AA
+      // purposes we getUnderlyingObject them; by precomputing the answers here
+      // we can avoid having to do so repeatedly there.
       if (OBU.getTagName() == "separate_storage") {
         assert(OBU.Inputs.size() == 2);
         auto MaybeSimplifyHint = [&](const Use &U) {
@@ -3233,6 +3234,28 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
         auto *New = CallBase::removeOperandBundle(II, OBU.getTagID());
         return New;
       }
+
+      // Try to fold alignment assumption into a load's !align metadata, if the
+      // assumption is valid in the load's context.
+      if (OBU.getTagName() == "align" && OBU.Inputs.size() == 2) {
+        RetainedKnowledge RK = getKnowledgeFromBundle(
+            *cast<AssumeInst>(II), II->bundle_op_info_begin()[Idx]);
+        if (!RK || RK.AttrKind != Attribute::Alignment ||
+            !isPowerOf2_64(RK.ArgValue))
+          continue;
+
+        auto *LI = dyn_cast<LoadInst>(OBU.Inputs[0]);
+        if (!LI ||
+            !isValidAssumeForContext(II, LI, &DT, /*AllowEphemerals=*/true))
+          continue;
+
+        LI->setMetadata(
+            LLVMContext::MD_align,
+            MDNode::get(II->getContext(), ValueAsMetadata::getConstant(
+                                              Builder.getInt64(RK.ArgValue))));
+        auto *New = CallBase::removeOperandBundle(II, OBU.getTagID());
+        return New;
+      }
     }
 
     // Convert nonnull assume like:
diff --git a/llvm/test/Transforms/InstCombine/assume-align.ll b/llvm/test/Transforms/InstCombine/assume-align.ll
index 47659ff8c84909..549821802fe674 100644
--- a/llvm/test/Transforms/InstCombine/assume-align.ll
+++ b/llvm/test/Transforms/InstCombine/assume-align.ll
@@ -123,11 +123,9 @@ define i8 @assume_align_non_pow2(ptr %p) {
   ret i8 %v
 }
 
-; TODO: Can fold alignment assumption into !align metadata on load.
 define ptr @fold_assume_align_pow2_of_loaded_pointer_into_align_metadata(ptr %p) {
 ; CHECK-LABEL: @fold_assume_align_pow2_of_loaded_pointer_into_align_metadata(
-; CHECK-NEXT:    [[P2:%.*]] = load ptr, ptr [[P:%.*]], align 8
-; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P2]], i64 8) ]
+; CHECK-NEXT:    [[P2:%.*]] = load ptr, ptr [[P:%.*]], align 8, !align [[META0:![0-9]+]]
 ; CHECK-NEXT:    ret ptr [[P2]]
 ;
   %p2 = load ptr, ptr %p
@@ -135,6 +133,16 @@ define ptr @fold_assume_align_pow2_of_loaded_pointer_into_align_metadata(ptr %p)
   ret ptr %p2
 }
 
+define ptr @fold_assume_align_i32_pow2_of_loaded_pointer_into_align_metadata(ptr %p) {
+; CHECK-LABEL: @fold_assume_align_i32_pow2_of_loaded_pointer_into_align_metadata(
+; CHECK-NEXT:    [[P2:%.*]] = load ptr, ptr [[P:%.*]], align 8, !align [[META0]]
+; CHECK-NEXT:    ret ptr [[P2]]
+;
+  %p2 = load ptr, ptr %p
+  call void @llvm.assume(i1 true) [ "align"(ptr %p2, i32 8) ]
+  ret ptr %p2
+}
+
 define ptr @dont_fold_assume_align_pow2_of_loaded_pointer_into_align_metadata_due_to_call(ptr %p) {
 ; CHECK-LABEL: @dont_fold_assume_align_pow2_of_loaded_pointer_into_align_metadata_due_to_call(
 ; CHECK-NEXT:    [[P2:%.*]] = load ptr, ptr [[P:%.*]], align 8
@@ -171,3 +179,19 @@ define ptr @dont_fold_assume_align_zero_of_loaded_pointer_into_align_metadata(pt
   call void @llvm.assume(i1 true) [ "align"(ptr %p2, i64 0) ]
   ret ptr %p2
 }
+
+; !align must have a constant integer alignment.
+define ptr @dont_fold_assume_align_not_constant_of_loaded_pointer_into_align_metadata(ptr %p, i64 %align) {
+; CHECK-LABEL: @dont_fold_assume_align_not_constant_of_loaded_pointer_into_align_metadata(
+; CHECK-NEXT:    [[P2:%.*]] = load ptr, ptr [[P:%.*]], align 8, !align [[META1:![0-9]+]]
+; CHECK-NEXT:    ret ptr [[P2]]
+;
+  %p2 = load ptr, ptr %p
+  call void @llvm.assume(i1 true) [ "align"(ptr %p2, i64 %align) ]
+  ret ptr %p2
+}
+
+;.
+; CHECK: [[META0]] = !{i64 8}
+; CHECK: [[META1]] = !{i64 1}
+;.



More information about the llvm-commits mailing list