[llvm] [InstCombine] Fold align assume into load's !align metadata if possible. (PR #108958)

Florian Hahn via llvm-commits llvm-commits at lists.llvm.org
Wed Sep 18 03:38:29 PDT 2024


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

>From 8b65faa8305425de29cbd1e0098a44b05df82ccc Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Tue, 17 Sep 2024 10:50:09 +0100
Subject: [PATCH 1/7] [InstCombine] Fold align assume into load's !align
 metadata if possible.

If an alignment assumption is valid in the context of a corresponding
load of the pointer the assumption applies to, the assumption can be
replaced !align metadata on the load.

The benefits of folding it into !align are that existing code makes
better use of !align and it allows removing the now-redundant call
instructions.
---
 .../InstCombine/InstCombineCalls.cpp          | 25 ++++++++++++++++---
 .../Transforms/InstCombine/assume-align.ll    |  7 +++---
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index 61011d55227e7b..596de10e20b6de 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3076,12 +3076,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) {
@@ -3095,6 +3096,22 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
         MaybeSimplifyHint(OBU.Inputs[0]);
         MaybeSimplifyHint(OBU.Inputs[1]);
       }
+
+      // 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) {
+        auto *LI = dyn_cast<LoadInst>(OBU.Inputs[0]);
+        if (!LI || !isValidAssumeForContext(II, LI, &DT, /*AllowEphemerals=*/true))
+          continue;
+        auto *Align = cast<ConstantInt>(OBU.Inputs[1]);
+        if (!isPowerOf2_64(Align->getZExtValue()))
+          continue;
+        LI->setMetadata(LLVMContext::MD_align,
+                        MDNode::get(II->getContext(),
+                                    ValueAsMetadata::getConstant(
+                                        Align)));
+        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 2b8ca5d25fd1a8..65256377696a59 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
@@ -171,3 +169,6 @@ 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
 }
+;.
+; CHECK: [[META0]] = !{i64 8}
+;.

>From a92e56cee7e58ce6290b019367ae97482da211e0 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Tue, 17 Sep 2024 12:09:07 +0100
Subject: [PATCH 2/7] !fix formatting.

---
 .../InstCombine/InstCombineCalls.cpp          | 21 ++++++++++---------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index 596de10e20b6de..bff44e11c3ab8d 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3079,10 +3079,10 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
     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.
+      // 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) {
@@ -3097,18 +3097,19 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
         MaybeSimplifyHint(OBU.Inputs[1]);
       }
 
-      // Try to fold alignment assumption into a load's !align metadata, if the assumption is valid in the load's context.
+      // 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) {
         auto *LI = dyn_cast<LoadInst>(OBU.Inputs[0]);
-        if (!LI || !isValidAssumeForContext(II, LI, &DT, /*AllowEphemerals=*/true))
+        if (!LI ||
+            !isValidAssumeForContext(II, LI, &DT, /*AllowEphemerals=*/true))
           continue;
         auto *Align = cast<ConstantInt>(OBU.Inputs[1]);
         if (!isPowerOf2_64(Align->getZExtValue()))
           continue;
-        LI->setMetadata(LLVMContext::MD_align,
-                        MDNode::get(II->getContext(),
-                                    ValueAsMetadata::getConstant(
-                                        Align)));
+        LI->setMetadata(
+            LLVMContext::MD_align,
+            MDNode::get(II->getContext(), ValueAsMetadata::getConstant(Align)));
         auto *New = CallBase::removeOperandBundle(II, OBU.getTagID());
         return New;
       }

>From 164823218af45ebc435df859ae438e7d5d3b035e Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Tue, 17 Sep 2024 13:04:56 +0100
Subject: [PATCH 3/7] !fixup check alignment is ConstantInt.

---
 .../lib/Transforms/InstCombine/InstCombineCalls.cpp |  4 ++--
 llvm/test/Transforms/InstCombine/assume-align.ll    | 13 +++++++++++++
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index bff44e11c3ab8d..66ceda9aa42643 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3104,8 +3104,8 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
         if (!LI ||
             !isValidAssumeForContext(II, LI, &DT, /*AllowEphemerals=*/true))
           continue;
-        auto *Align = cast<ConstantInt>(OBU.Inputs[1]);
-        if (!isPowerOf2_64(Align->getZExtValue()))
+        auto *Align = dyn_cast<ConstantInt>(OBU.Inputs[1]);
+        if (!Align || !isPowerOf2_64(Align->getZExtValue()))
           continue;
         LI->setMetadata(
             LLVMContext::MD_align,
diff --git a/llvm/test/Transforms/InstCombine/assume-align.ll b/llvm/test/Transforms/InstCombine/assume-align.ll
index 65256377696a59..9d9f18068bb3e8 100644
--- a/llvm/test/Transforms/InstCombine/assume-align.ll
+++ b/llvm/test/Transforms/InstCombine/assume-align.ll
@@ -169,6 +169,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
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P2]], i64 [[ALIGN:%.*]]) ]
+; 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}
 ;.

>From 1e33607d2121280256b2d641a29130b8086f28be Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Tue, 17 Sep 2024 13:10:16 +0100
Subject: [PATCH 4/7] !fixup check align type size.

---
 llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp |  2 +-
 llvm/test/Transforms/InstCombine/assume-align.ll     | 11 +++++++++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index 66ceda9aa42643..40da0288eccb16 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3105,7 +3105,7 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
             !isValidAssumeForContext(II, LI, &DT, /*AllowEphemerals=*/true))
           continue;
         auto *Align = dyn_cast<ConstantInt>(OBU.Inputs[1]);
-        if (!Align || !isPowerOf2_64(Align->getZExtValue()))
+        if (!Align || !isPowerOf2_64(Align->getZExtValue()) || Align->getType()->getScalarSizeInBits() != 64)
           continue;
         LI->setMetadata(
             LLVMContext::MD_align,
diff --git a/llvm/test/Transforms/InstCombine/assume-align.ll b/llvm/test/Transforms/InstCombine/assume-align.ll
index 9d9f18068bb3e8..c80275b6e4d784 100644
--- a/llvm/test/Transforms/InstCombine/assume-align.ll
+++ b/llvm/test/Transforms/InstCombine/assume-align.ll
@@ -133,6 +133,17 @@ define ptr @fold_assume_align_pow2_of_loaded_pointer_into_align_metadata(ptr %p)
   ret ptr %p2
 }
 
+define ptr @dont_fold_assume_align_i32_pow2_of_loaded_pointer_into_align_metadata(ptr %p) {
+; CHECK-LABEL: @dont_fold_assume_align_i32_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]], i32 8) ]
+; 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

>From 990095c33825c7accc75a27e24a4e4340517c8e9 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Tue, 17 Sep 2024 13:17:56 +0100
Subject: [PATCH 5/7] !fixup use getKnowledgeFromBundle

---
 .../lib/Transforms/InstCombine/InstCombineCalls.cpp | 13 +++++++++----
 llvm/test/Transforms/InstCombine/assume-align.ll    |  9 ++++-----
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index 40da0288eccb16..e9e476b10ff340 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3100,16 +3100,21 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
       // 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;
-        auto *Align = dyn_cast<ConstantInt>(OBU.Inputs[1]);
-        if (!Align || !isPowerOf2_64(Align->getZExtValue()) || Align->getType()->getScalarSizeInBits() != 64)
-          continue;
+
         LI->setMetadata(
             LLVMContext::MD_align,
-            MDNode::get(II->getContext(), ValueAsMetadata::getConstant(Align)));
+            MDNode::get(II->getContext(), ValueAsMetadata::getConstant(
+                                              Builder.getInt64(RK.ArgValue))));
         auto *New = CallBase::removeOperandBundle(II, OBU.getTagID());
         return New;
       }
diff --git a/llvm/test/Transforms/InstCombine/assume-align.ll b/llvm/test/Transforms/InstCombine/assume-align.ll
index c80275b6e4d784..a434bb9a247709 100644
--- a/llvm/test/Transforms/InstCombine/assume-align.ll
+++ b/llvm/test/Transforms/InstCombine/assume-align.ll
@@ -133,10 +133,9 @@ define ptr @fold_assume_align_pow2_of_loaded_pointer_into_align_metadata(ptr %p)
   ret ptr %p2
 }
 
-define ptr @dont_fold_assume_align_i32_pow2_of_loaded_pointer_into_align_metadata(ptr %p) {
+define ptr @fold_assume_align_i32_pow2_of_loaded_pointer_into_align_metadata(ptr %p) {
 ; CHECK-LABEL: @dont_fold_assume_align_i32_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]], i32 8) ]
+; CHECK-NEXT:    [[P2:%.*]] = load ptr, ptr [[P:%.*]], align 8, !align [[META0]]
 ; CHECK-NEXT:    ret ptr [[P2]]
 ;
   %p2 = load ptr, ptr %p
@@ -184,8 +183,7 @@ define ptr @dont_fold_assume_align_zero_of_loaded_pointer_into_align_metadata(pt
 ; !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
-; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P2]], i64 [[ALIGN:%.*]]) ]
+; CHECK-NEXT:    [[P2:%.*]] = load ptr, ptr [[P:%.*]], align 8, !align [[META1:![0-9]+]]
 ; CHECK-NEXT:    ret ptr [[P2]]
 ;
   %p2 = load ptr, ptr %p
@@ -195,4 +193,5 @@ define ptr @dont_fold_assume_align_not_constant_of_loaded_pointer_into_align_met
 
 ;.
 ; CHECK: [[META0]] = !{i64 8}
+; CHECK: [[META1]] = !{i64 1}
 ;.

>From eb129775b8a268d0067ccc40c5bb0543b78a8ded Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Wed, 18 Sep 2024 11:37:09 +0100
Subject: [PATCH 6/7] !fixup update test check

---
 llvm/test/Transforms/InstCombine/assume-align.ll | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/test/Transforms/InstCombine/assume-align.ll b/llvm/test/Transforms/InstCombine/assume-align.ll
index a434bb9a247709..7020e5920f328c 100644
--- a/llvm/test/Transforms/InstCombine/assume-align.ll
+++ b/llvm/test/Transforms/InstCombine/assume-align.ll
@@ -134,7 +134,7 @@ define ptr @fold_assume_align_pow2_of_loaded_pointer_into_align_metadata(ptr %p)
 }
 
 define ptr @fold_assume_align_i32_pow2_of_loaded_pointer_into_align_metadata(ptr %p) {
-; CHECK-LABEL: @dont_fold_assume_align_i32_pow2_of_loaded_pointer_into_align_metadata(
+; 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]]
 ;

>From a7ad32d7bcb0a17b9e9686aa87d70d5bf7e8c990 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Wed, 18 Sep 2024 11:37:30 +0100
Subject: [PATCH 7/7] !fixup materialize assumption for !align of load

---
 llvm/lib/Transforms/Scalar/EarlyCSE.cpp           |  2 +-
 llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp | 13 +++++++++++--
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/EarlyCSE.cpp b/llvm/lib/Transforms/Scalar/EarlyCSE.cpp
index cf11f5bc885a75..fe65bbb3e14048 100644
--- a/llvm/lib/Transforms/Scalar/EarlyCSE.cpp
+++ b/llvm/lib/Transforms/Scalar/EarlyCSE.cpp
@@ -1588,9 +1588,9 @@ bool EarlyCSE::processNode(DomTreeNode *Node) {
         if (InVal.IsLoad)
           if (auto *I = dyn_cast<Instruction>(Op))
             combineMetadataForCSE(I, &Inst, false);
+        salvageKnowledge(&Inst, &AC);
         if (!Inst.use_empty())
           Inst.replaceAllUsesWith(Op);
-        salvageKnowledge(&Inst, &AC);
         removeMSSA(Inst);
         Inst.eraseFromParent();
         Changed = true;
diff --git a/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp b/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp
index 3cf68e07da5be2..941078b2305603 100644
--- a/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp
+++ b/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp
@@ -254,7 +254,7 @@ struct AssumeBuilderState {
                              ->getDataLayout()
                              .getTypeStoreSize(AccType)
                              .getKnownMinValue();
-    if (DerefSize != 0) {
+    if (EnableKnowledgeRetention && DerefSize != 0) {
       addKnowledge({Attribute::Dereferenceable, DerefSize, Pointer});
       if (!NullPointerIsDefined(MemInst->getFunction(),
                                 Pointer->getType()->getPointerAddressSpace()))
@@ -265,6 +265,15 @@ struct AssumeBuilderState {
   }
 
   void addInstruction(Instruction *I) {
+    if (auto *L = dyn_cast<LoadInst>(I))
+      if (auto *Align = L->getMetadata(LLVMContext::MD_align)) {
+        addKnowledge({Attribute::Alignment,
+                      mdconst::extract<ConstantInt>(Align->getOperand(0))
+                          ->getZExtValue(),
+                      I});
+      }
+    if (!EnableKnowledgeRetention)
+      return;
     if (auto *Call = dyn_cast<CallBase>(I))
       return addCall(Call);
     if (auto *Load = dyn_cast<LoadInst>(I))
@@ -291,7 +300,7 @@ AssumeInst *llvm::buildAssumeFromInst(Instruction *I) {
 
 bool llvm::salvageKnowledge(Instruction *I, AssumptionCache *AC,
                             DominatorTree *DT) {
-  if (!EnableKnowledgeRetention || I->isTerminator())
+  if (I->isTerminator())
     return false;
   bool Changed = false;
   AssumeBuilderState Builder(I->getModule(), I, AC, DT);



More information about the llvm-commits mailing list