[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.

Author: fhahn

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:

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