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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Florian Hahn (fhahn)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/108958.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp (+21-4) 
- (modified) llvm/test/Transforms/InstCombine/assume-align.ll (+4-3) 


``````````diff
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}
+;.

``````````

</details>


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