[llvm] [llvm] Optimize MachineMemOperand::getAlign with KnownBits information (PR #143872)

[Acthinks Yang via llvm-commits]

https://github.com/Acthinks created https://github.com/llvm/llvm-project/pull/143872

When using the `align_value` attribute, the original getAlign calculation (BaseAlign + Offset) cannot capture more precise alignment information after the load/store legalization phases (specifically after `expandUnalignedLoad`/`expandUnalignedStore` generate new load/store instructions).

Fixes#143215

>From 3bd2d07687b6953749cbb6e1a2b82793f7863cd3 Mon Sep 17 00:00:00 2001
From: Acthinks <yangzhh at mail.ustc.edu.cn>
Date: Thu, 12 Jun 2025 18:16:59 +0800
Subject: [PATCH] [llvm] Optimize MachineMemOperand::getAlign with KnownBits
 information

When using the `align_value` attribute, the original getAlign calculation
(BaseAlign + Offset) cannot capture more precise alignment information
after the load/store legalization phases (specifically after
`expandUnalignedLoad`/`expandUnalignedStore` generate new load/store
instructions).

Fixes#143215
---
 llvm/lib/CodeGen/MachineOperand.cpp           | 16 +++++++++++++++-
 .../SelectionDAG/SelectionDAGBuilder.cpp      |  4 ++++
 .../CodeGen/RISCV/unaligned-load-store.ll     | 19 +++++++++++++++++++
 3 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/MachineOperand.cpp b/llvm/lib/CodeGen/MachineOperand.cpp
index 0d251697f2567..0a5663a7e4ab2 100644
--- a/llvm/lib/CodeGen/MachineOperand.cpp
+++ b/llvm/lib/CodeGen/MachineOperand.cpp
@@ -14,6 +14,7 @@
 #include "llvm/ADT/StableHashing.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Analysis/Loads.h"
+#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/CodeGen/MIRFormatter.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineJumpTableInfo.h"
@@ -1136,7 +1137,20 @@ void MachineMemOperand::refineAlignment(const MachineMemOperand *MMO) {
 /// getAlign - Return the minimum known alignment in bytes of the
 /// actual memory reference.
 Align MachineMemOperand::getAlign() const {
-  return commonAlignment(getBaseAlign(), getOffset());
+  Align RetAlign = commonAlignment(getBaseAlign(), getOffset());
+  if (const Value *V = getValue()) {
+    if (auto *I = dyn_cast<Instruction>(V)) {
+      DataLayout DL = I->getDataLayout();
+      unsigned PtrWidth = DL.getPointerTypeSizeInBits(V->getType());
+      KnownBits Known(PtrWidth);
+      llvm::computeKnownBits(V, Known, DL);
+      Known = KnownBits::add(
+          Known, KnownBits::makeConstant(APInt(PtrWidth, getOffset())));
+      unsigned AlignBits = Known.countMinTrailingZeros();
+      RetAlign = std::max(Align(1ull << std::min(31U, AlignBits)), RetAlign);
+    }
+  }
+  return RetAlign;
 }
 
 void MachineMemOperand::print(raw_ostream &OS, ModuleSlotTracker &MST,
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index e6a1dc930685c..278a3a12bfbc6 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -11848,6 +11848,10 @@ void SelectionDAGISel::LowerArguments(const Function &F) {
           OutVal = DAG.getNode(ISD::AssertNoFPClass, dl, OutVal.getValueType(),
                                OutVal, SDNoFPClass);
         }
+        if (Arg.hasAttribute(Attribute::Alignment)) {
+          OutVal =
+              DAG.getAssertAlign(dl, OutVal, Arg.getParamAlign().valueOrOne());
+        }
         ArgValues.push_back(OutVal);
       }
 
diff --git a/llvm/test/CodeGen/RISCV/unaligned-load-store.ll b/llvm/test/CodeGen/RISCV/unaligned-load-store.ll
index c9c49e8f7f532..af83f434171c1 100644
--- a/llvm/test/CodeGen/RISCV/unaligned-load-store.ll
+++ b/llvm/test/CodeGen/RISCV/unaligned-load-store.ll
@@ -578,5 +578,24 @@ define void @store_large_constant(ptr %x) {
   store i64 18364758544493064720, ptr %x, align 1
   ret void
 }
+
+define void @store_const_with_align_attribute(ptr align 2 %p) {
+; SLOW-LABEL: store_const_with_align_attribute:
+; SLOW:       # %bb.0: # %entry
+; SLOW-NEXT:    sb zero, 3(a0)
+; SLOW-NEXT:    sh zero, 4(a0)
+; SLOW-NEXT:    sb zero, 6(a0)
+; SLOW-NEXT:    ret
+;
+; FAST-LABEL: store_const_with_align_attribute:
+; FAST:       # %bb.0: # %entry
+; FAST-NEXT:    sw zero, 3(a0)
+; FAST-NEXT:    ret
+entry:
+  %len = getelementptr inbounds nuw i8, ptr %p, i32 3
+  store i32 0, ptr %len, align 1
+  ret void
+}
+
 ;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
 ; SLOWZBKB: {{.*}}