[llvm] Handled the uimm9 offset while FrameIndex folding. (PR #149303)

[via llvm-commits]

https://github.com/ukalappa-mips updated https://github.com/llvm/llvm-project/pull/149303

>From 16e7516a04a1f93ed9a891720e022ed5e7388b99 Mon Sep 17 00:00:00 2001
From: Umesh Kalappa <ukalappa.mips at gmail.com>
Date: Thu, 17 Jul 2025 13:14:54 +0000
Subject: [PATCH] Handle and honour  the uimm9 offset for mips prefetch.

---
 llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp |  8 ++--
 llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp |  8 ++++
 llvm/test/CodeGen/RISCV/xmips-cbop.ll       | 44 +++++++++++++++++++++
 3 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp b/llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp
index 0f948b22759fe..ddba47b1975c7 100644
--- a/llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp
@@ -2942,8 +2942,8 @@ bool RISCVDAGToDAGISel::SelectAddrRegImm(SDValue Addr, SDValue &Base,
 /// Similar to SelectAddrRegImm, except that the offset is restricted to uimm9.
 bool RISCVDAGToDAGISel::SelectAddrRegImm9(SDValue Addr, SDValue &Base,
                                           SDValue &Offset) {
-  // FIXME: Support FrameIndex. Need to teach eliminateFrameIndex that only
-  // a 9-bit immediate can be folded.
+  if (SelectAddrFrameIndex(Addr, Base, Offset))
+    return true;
 
   SDLoc DL(Addr);
   MVT VT = Addr.getSimpleValueType();
@@ -2953,8 +2953,8 @@ bool RISCVDAGToDAGISel::SelectAddrRegImm9(SDValue Addr, SDValue &Base,
     if (isUInt<9>(CVal)) {
       Base = Addr.getOperand(0);
 
-      // FIXME: Support FrameIndex. Need to teach eliminateFrameIndex that only
-      // a 9-bit immediate can be folded.
+      if (auto *FIN = dyn_cast<FrameIndexSDNode>(Base))
+        Base = CurDAG->getTargetFrameIndex(FIN->getIndex(), VT);
       Offset = CurDAG->getSignedTargetConstant(CVal, DL, VT);
       return true;
     }
diff --git a/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp b/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp
index 540412366026b..4b9ba0f33276e 100644
--- a/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp
+++ b/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp
@@ -573,6 +573,9 @@ bool RISCVRegisterInfo::eliminateFrameIndex(MachineBasicBlock::iterator II,
     int64_t Val = Offset.getFixed();
     int64_t Lo12 = SignExtend64<12>(Val);
     unsigned Opc = MI.getOpcode();
+    int64_t Imm9Val = SignExtend64<9>(Val);
+    auto &Subtarget = MF.getSubtarget<RISCVSubtarget>();
+
     if (Opc == RISCV::ADDI && !isInt<12>(Val)) {
       // We chose to emit the canonical immediate sequence rather than folding
       // the offset into the using add under the theory that doing so doesn't
@@ -585,6 +588,11 @@ bool RISCVRegisterInfo::eliminateFrameIndex(MachineBasicBlock::iterator II,
                (Lo12 & 0b11111) != 0) {
       // Prefetch instructions require the offset to be 32 byte aligned.
       MI.getOperand(FIOperandNum + 1).ChangeToImmediate(0);
+    } else if ((Opc == RISCV::PREFETCH_I || Opc == RISCV::PREFETCH_R ||
+                Opc == RISCV::PREFETCH_W) &&
+               Subtarget.hasVendorXMIPSCBOP() && !isUInt<9>(Imm9Val)) {
+      // MIPS Prefetch instructions require the offset to be 9 bits encoded.
+      MI.getOperand(FIOperandNum + 1).ChangeToImmediate(0);
     } else if ((Opc == RISCV::PseudoRV32ZdinxLD ||
                 Opc == RISCV::PseudoRV32ZdinxSD) &&
                Lo12 >= 2044) {
diff --git a/llvm/test/CodeGen/RISCV/xmips-cbop.ll b/llvm/test/CodeGen/RISCV/xmips-cbop.ll
index cbbd1de13192c..0cfe3e3f0ab7c 100644
--- a/llvm/test/CodeGen/RISCV/xmips-cbop.ll
+++ b/llvm/test/CodeGen/RISCV/xmips-cbop.ll
@@ -49,3 +49,47 @@ define void @prefetch_inst_read(ptr noundef %ptr) nounwind  {
   tail call void @llvm.prefetch.p0(ptr nonnull %arrayidx, i32 0, i32 0, i32 0)
   ret void
 }
+
+define void @prefetch_frameindex_test_neg() nounwind {
+; RV32XMIPSPREFETCH-LABEL: prefetch_frameindex_test_neg:
+; RV32XMIPSPREFETCH:       # %bb.0:
+; RV32XMIPSPREFETCH-NEXT:    addi sp, sp, -512
+; RV32XMIPSPREFETCH-NEXT:    addi a0, sp, -32
+; RV32XMIPSPREFETCH-NEXT:    mips.pref 8, 0(a0)
+; RV32XMIPSPREFETCH-NEXT:    addi sp, sp, 512
+; RV32XMIPSPREFETCH-NEXT:    ret
+;
+; RV64XMIPSPREFETCH-LABEL: prefetch_frameindex_test_neg:
+; RV64XMIPSPREFETCH:       # %bb.0:
+; RV64XMIPSPREFETCH-NEXT:    addi sp, sp, -512
+; RV64XMIPSPREFETCH-NEXT:    addi a0, sp, -32
+; RV64XMIPSPREFETCH-NEXT:    mips.pref 8, 0(a0)
+; RV64XMIPSPREFETCH-NEXT:    addi sp, sp, 512
+; RV64XMIPSPREFETCH-NEXT:    ret
+ %data = alloca [128 x i32], align 4
+  %base = bitcast ptr %data to ptr
+  %ptr = getelementptr [128 x i32], ptr %base, i32 0, i32 -8
+  call void @llvm.prefetch(ptr %ptr, i32 0, i32 0, i32 1)
+  ret void
+}
+
+define void @prefetch_frameindex_test() nounwind {
+; RV32XMIPSPREFETCH-LABEL: prefetch_frameindex_test:
+; RV32XMIPSPREFETCH:       # %bb.0:
+; RV32XMIPSPREFETCH-NEXT:    addi sp, sp, -512
+; RV32XMIPSPREFETCH-NEXT:    mips.pref 8, 32(sp)
+; RV32XMIPSPREFETCH-NEXT:    addi sp, sp, 512
+; RV32XMIPSPREFETCH-NEXT:    ret
+;
+; RV64XMIPSPREFETCH-LABEL: prefetch_frameindex_test:
+; RV64XMIPSPREFETCH:       # %bb.0:
+; RV64XMIPSPREFETCH-NEXT:    addi sp, sp, -512
+; RV64XMIPSPREFETCH-NEXT:    mips.pref 8, 32(sp)
+; RV64XMIPSPREFETCH-NEXT:    addi sp, sp, 512
+; RV64XMIPSPREFETCH-NEXT:    ret
+ %data = alloca [128 x i32], align 4
+  %base = bitcast ptr %data to ptr
+  %ptr = getelementptr [128 x i32], ptr %base, i32 0, i32 8
+  call void @llvm.prefetch(ptr %ptr, i32 0, i32 0, i32 1)
+  ret void
+}