[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

Author: actinks

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,

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);
       }
 

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: {{.*}}