[llvm] [RISCV] Support FrameIndex operands in getMemOperandsWithOffsetWidth / getMemOperandWithOffsetWidth (PR #73802)

[Alex Bradbury via llvm-commits]

https://github.com/asb created https://github.com/llvm/llvm-project/pull/73802

I noted AArch64 happily accepts a FrameIndex operand as well as a register. This doesn't cause any changes outside of my C++ unit test for the current state of in-tree, but this will cause additional test changes if #73789 is rebased on top of it.

---

I do have some caution here that the returned Offset doesn't seem at all as meaningful if you have a FrameIndex base. Though as I say, AArch64 has been doing this for some time - see https://reviews.llvm.org/D54847. I believe this change won't harm the approach taken in shouldClusterMemOps because memOpsHaveSameBasePtr will only return true if the FrameIndex operand is the same for both operations. (Though reviewers, please double-check you agree). This patch is directly applyable, so I've opted not to base it on top of #73789.

>From 26b5b99f3b495af80b61c7b9b845fd1e8c1f90fc Mon Sep 17 00:00:00 2001
From: Alex Bradbury <asb at igalia.com>
Date: Wed, 29 Nov 2023 14:56:36 +0000
Subject: [PATCH] [RISCV] Support FrameIndex operands in
 getMemOperandsWithOffsetWidth / getMemOperandWithOffsetWidth

I noted AArch64 happily accepts a FrameIndex operand as well as a
register. This doesn't cause any changes outside of my C++ unit test for
the current state of in-tree, but this will cause additional test
changes if #73789 is rebased on top of it.
---
 llvm/lib/Target/RISCV/RISCVInstrInfo.cpp           | 2 +-
 llvm/unittests/Target/RISCV/RISCVInstrInfoTest.cpp | 9 +++++++--
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp b/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp
index 2918e5654db4f9f..900eede0c4ef24c 100644
--- a/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp
+++ b/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp
@@ -2304,7 +2304,7 @@ bool RISCVInstrInfo::getMemOperandWithOffsetWidth(
   // load/store instructions.
   if (LdSt.getNumExplicitOperands() != 3)
     return false;
-  if (!LdSt.getOperand(1).isReg() || !LdSt.getOperand(2).isImm())
+  if ((!LdSt.getOperand(1).isReg() && !LdSt.getOperand(1).isFI()) || !LdSt.getOperand(2).isImm())
     return false;
 
   if (!LdSt.hasOneMemOperand())
diff --git a/llvm/unittests/Target/RISCV/RISCVInstrInfoTest.cpp b/llvm/unittests/Target/RISCV/RISCVInstrInfoTest.cpp
index 135d7dbb426e3c2..b4c96a9c2a62ce7 100644
--- a/llvm/unittests/Target/RISCV/RISCVInstrInfoTest.cpp
+++ b/llvm/unittests/Target/RISCV/RISCVInstrInfoTest.cpp
@@ -154,7 +154,6 @@ TEST_P(RISCVInstrInfoTest, GetMemOperandsWithOffsetWidth) {
   Res = TII->getMemOperandsWithOffsetWidth(*MI, BaseOps, Offset,
                                            OffsetIsScalable, Width, TRI);
 
-  // TODO: AArch64 can handle this case, and we probably should too.
   BaseOps.clear();
   MMO = MF->getMachineMemOperand(MachinePointerInfo(),
                                  MachineMemOperand::MOStore, 4, Align(4));
@@ -165,7 +164,13 @@ TEST_P(RISCVInstrInfoTest, GetMemOperandsWithOffsetWidth) {
            .addMemOperand(MMO);
   Res = TII->getMemOperandsWithOffsetWidth(*MI, BaseOps, Offset,
                                            OffsetIsScalable, Width, TRI);
-  EXPECT_FALSE(Res);
+  ASSERT_TRUE(Res);
+  ASSERT_EQ(BaseOps.size(), 1u);
+  ASSERT_TRUE(BaseOps.front()->isFI());
+  EXPECT_EQ(BaseOps.front()->getIndex(), 2);
+  EXPECT_EQ(Offset, 4);
+  EXPECT_FALSE(OffsetIsScalable);
+  EXPECT_EQ(Width, 4u);
 }
 
 } // namespace