Fix assertion failure in getMemOperandWithOffsetWidth

[kristof.beyls]

Summary

This fixes an assertion failure that triggers inside
getMemOperandWithOffset when Machine Sinking calls it on a MachineInstr
that is not a memory operation.

Different backends implement getMemOperandWithOffset differently: some
return false on non-memory MachineInstrs, others assert.

The Machine Sinking pass in at least SinkingPreventsImplicitNullCheck
relies on getMemOperandWithOffset to return false on non-memory
MachineInstrs, instead of asserting.

This patch updates the documentation on getMemOperandWithOffset that it
should return false on any MachineInstr it cannot handle, instead of
asserting. It also adapts the in-tree backends accordingly where
necessary.

[Jim]

Could you remove "[PATCH]" from the title?

[sdesmalen]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

2031–2032

Should this return false as well, instead of asserting? And if so, is there a way to construct a test for that?

llvm/test/CodeGen/AArch64/machine-sink-getmemoperandwithoffset.mir

35

nit: Can you add a comment clarifying which instruction (that is not a load or store) would have caused the assert to fail?

[kristof.beyls]

Addressed Sander's comments.

[kristof.beyls]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

2031–2032

Thanks - good catch!

llvm/test/CodeGen/AArch64/machine-sink-getmemoperandwithoffset.mir

35

I now added a comment that hopefully is useful for people reading this test in the future, see a few lines below.

[sdesmalen]

LGTM, thanks @kristof.beyls!

llvm/test/CodeGen/AArch64/machine-sink-getmemoperandwithoffset.mir

35

Thanks, that comment seems useful to me.

llvm/include/llvm/CodeGen/TargetInstrInfo.h

public:
   virtual bool shouldScheduleLoadsNear(SDNode *Load1, SDNode *Load2,
                                        int64_t Offset1, int64_t Offset2,
                                        unsigned NumLoads) const {
     return false;
   }
 
   /// Get the base operand and byte offset of an instruction that reads/writes
   /// memory.
+  /// It returns false if MI does not read/write memory.
+  /// It returns false if no base operand and offset was found.
+  /// It is not guaranteed to always recognize base operand and offsets in all
+  /// cases.
   virtual bool getMemOperandWithOffset(const MachineInstr &MI,
                                        const MachineOperand *&BaseOp,
                                        int64_t &Offset,
                                        const TargetRegisterInfo *TRI) const {
     return false;
   }
 
   /// Return true if the instruction contains a base register and offset. If

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

bool AArch64InstrInfo::isCandidateToMergeOrPair(const MachineInstr &MI) const {
 
   return true;
 }
 
 bool AArch64InstrInfo::getMemOperandWithOffset(const MachineInstr &LdSt,
                                           const MachineOperand *&BaseOp,
                                           int64_t &Offset,
                                           const TargetRegisterInfo *TRI) const {
+  if (!LdSt.mayLoadOrStore())
+    return false;
+
   unsigned Width;
   return getMemOperandWithOffsetWidth(LdSt, BaseOp, Offset, Width, TRI);
 }
 
 bool AArch64InstrInfo::getMemOperandWithOffsetWidth(
     const MachineInstr &LdSt, const MachineOperand *&BaseOp, int64_t &Offset,
     unsigned &Width, const TargetRegisterInfo *TRI) const {
   assert(LdSt.mayLoadOrStore() && "Expected a memory operation.");
bool AArch64InstrInfo::getMemOperandWithOffsetWidth(
   if (LdSt.getNumExplicitOperands() == 3) {
     BaseOp = &LdSt.getOperand(1);
     Offset = LdSt.getOperand(2).getImm() * Scale;
   } else {
     assert(LdSt.getNumExplicitOperands() == 4 && "invalid number of operands");
     BaseOp = &LdSt.getOperand(2);
     Offset = LdSt.getOperand(3).getImm() * Scale;
   }
 
-  assert((BaseOp->isReg() || BaseOp->isFI()) &&
-         "getMemOperandWithOffset only supports base "
-         "operands of type register or frame index.");
+  if (!BaseOp->isReg() && !BaseOp->isFI())
+    return false;
 
   return true;
 }
 
 MachineOperand &
 AArch64InstrInfo::getMemOpBaseRegImmOfsOffsetOperand(MachineInstr &LdSt) const {
   assert(LdSt.mayLoadOrStore() && "Expected a memory operation.");
   MachineOperand &OfsOp = LdSt.getOperand(LdSt.getNumExplicitOperands() - 1);

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

default:
     return false;
   }
 }
 
 bool SIInstrInfo::getMemOperandWithOffset(const MachineInstr &LdSt,
                                           const MachineOperand *&BaseOp,
                                           int64_t &Offset,
                                           const TargetRegisterInfo *TRI) const {
+  if (!LdSt.mayLoadOrStore())
+    return false;
+
   unsigned Opc = LdSt.getOpcode();
 
   if (isDS(LdSt)) {
     const MachineOperand *OffsetImm =
         getNamedOperand(LdSt, AMDGPU::OpName::offset);
     if (OffsetImm) {
       // Normal, single offset LDS instruction.
       BaseOp = getNamedOperand(LdSt, AMDGPU::OpName::addr);
       // TODO: ds_consume/ds_append use M0 for the base address. Is it safe to
       // report that here?
-      if (!BaseOp)
+      if (!BaseOp || !BaseOp->isReg())
         return false;
 
       Offset = OffsetImm->getImm();
-      assert(BaseOp->isReg() && "getMemOperandWithOffset only supports base "
-                                "operands of type register.");
+
       return true;
     }
 
     // The 2 offset instructions use offset0 and offset1 instead. We can treat
     // these as a load with a single offset if the 2 offsets are consecutive. We
     // will use this for some partially aligned loads.
     const MachineOperand *Offset0Imm =
         getNamedOperand(LdSt, AMDGPU::OpName::offset0);
if (Offset1 > Offset0 && Offset1 - Offset0 == 1) {
         int Data0Idx = AMDGPU::getNamedOperandIdx(Opc, AMDGPU::OpName::data0);
         EltSize = TRI->getRegSizeInBits(*getOpRegClass(LdSt, Data0Idx)) / 8;
       }
 
       if (isStride64(Opc))
         EltSize *= 64;
 
       BaseOp = getNamedOperand(LdSt, AMDGPU::OpName::addr);
+      if (!BaseOp->isReg())
+        return false;
+
       Offset = EltSize * Offset0;
-      assert(BaseOp->isReg() && "getMemOperandWithOffset only supports base "
-                                "operands of type register.");
+
       return true;
     }
 
     return false;
   }
 
   if (isMUBUF(LdSt) || isMTBUF(LdSt)) {
     const MachineOperand *SOffset = getNamedOperand(LdSt, AMDGPU::OpName::soffset);
if (isMUBUF(LdSt) || isMTBUF(LdSt)) {
     const MachineOperand *AddrReg = getNamedOperand(LdSt, AMDGPU::OpName::vaddr);
     if (!AddrReg)
       return false;
 
     const MachineOperand *OffsetImm =
         getNamedOperand(LdSt, AMDGPU::OpName::offset);
     BaseOp = AddrReg;
     Offset = OffsetImm->getImm();
-
     if (SOffset) // soffset can be an inline immediate.
       Offset += SOffset->getImm();
 
-    assert(BaseOp->isReg() && "getMemOperandWithOffset only supports base "
-                              "operands of type register.");
+    if (!BaseOp->isReg())
+      return false;
+
     return true;
   }
 
   if (isSMRD(LdSt)) {
     const MachineOperand *OffsetImm =
         getNamedOperand(LdSt, AMDGPU::OpName::offset);
     if (!OffsetImm)
       return false;
 
     const MachineOperand *SBaseReg = getNamedOperand(LdSt, AMDGPU::OpName::sbase);
     BaseOp = SBaseReg;
     Offset = OffsetImm->getImm();
-    assert(BaseOp->isReg() && "getMemOperandWithOffset only supports base "
-                              "operands of type register.");
+    if (!BaseOp->isReg())
+      return false;
+
     return true;
   }
 
   if (isFLAT(LdSt)) {
     const MachineOperand *VAddr = getNamedOperand(LdSt, AMDGPU::OpName::vaddr);
     if (VAddr) {
       // Can't analyze 2 offsets.
       if (getNamedOperand(LdSt, AMDGPU::OpName::saddr))
         return false;
 
       BaseOp = VAddr;
     } else {
       // scratch instructions have either vaddr or saddr.
       BaseOp = getNamedOperand(LdSt, AMDGPU::OpName::saddr);
     }
 
     Offset = getNamedOperand(LdSt, AMDGPU::OpName::offset)->getImm();
-    assert(BaseOp->isReg() && "getMemOperandWithOffset only supports base "
-                              "operands of type register.");
+    if (!BaseOp->isReg())
+      return false;
     return true;
   }
 
   return false;
 }
 
 static bool memOpsHaveSameBasePtr(const MachineInstr &MI1,
                                   const MachineOperand &BaseOp1,

llvm/lib/Target/Hexagon/HexagonInstrInfo.cpp

 }
 
 /// Get the base register and byte offset of a load/store instr.
 bool HexagonInstrInfo::getMemOperandWithOffset(
     const MachineInstr &LdSt, const MachineOperand *&BaseOp, int64_t &Offset,
     const TargetRegisterInfo *TRI) const {
   unsigned AccessSize = 0;
   BaseOp = getBaseAndOffset(LdSt, Offset, AccessSize);
-  assert((!BaseOp || BaseOp->isReg()) &&
-         "getMemOperandWithOffset only supports base "
-         "operands of type register.");
-  return BaseOp != nullptr;
+  return BaseOp != nullptr && BaseOp->isReg();
 }
 
 /// Can these instructions execute at the same time in a bundle.
 bool HexagonInstrInfo::canExecuteInBundle(const MachineInstr &First,
       const MachineInstr &Second) const {
   if (Second.mayStore() && First.getOpcode() == Hexagon::S2_allocframe) {
     const MachineOperand &Op = Second.getOperand(0);
     if (Op.isReg() && Op.isUse() && Op.getReg() == Hexagon::R29)

llvm/lib/Target/Lanai/LanaiInstrInfo.cpp

bool LanaiInstrInfo::getMemOperandWithOffsetWidth(
   case Lanai::LDBz_RI:
   case Lanai::STB_RI:
     Width = 1;
     break;
   }
 
   BaseOp = &LdSt.getOperand(1);
   Offset = LdSt.getOperand(2).getImm();
-  assert(BaseOp->isReg() && "getMemOperandWithOffset only supports base "
-                            "operands of type register.");
+
+  if (!BaseOp->isReg())
+    return false;
+
   return true;
 }
 
 bool LanaiInstrInfo::getMemOperandWithOffset(const MachineInstr &LdSt,
                                         const MachineOperand *&BaseOp,
                                         int64_t &Offset,
                                         const TargetRegisterInfo *TRI) const {
   switch (LdSt.getOpcode()) {

llvm/lib/Target/PowerPC/PPCInstrInfo.cpp

for (auto &I : PreHeader.instrs())
     if (I.getOpcode() == LOOPi)
       return &I;
   return nullptr;
 }
 
 // Return true if get the base operand, byte offset of an instruction and the
 // memory width. Width is the size of memory that is being loaded/stored.
 bool PPCInstrInfo::getMemOperandWithOffsetWidth(
-  const MachineInstr &LdSt,
-  const MachineOperand *&BaseReg,
-  int64_t &Offset,
-  unsigned &Width,
-  const TargetRegisterInfo *TRI) const {
-  assert(LdSt.mayLoadOrStore() && "Expected a memory operation.");
+    const MachineInstr &LdSt, const MachineOperand *&BaseReg, int64_t &Offset,
+    unsigned &Width, const TargetRegisterInfo *TRI) const {
+  if (!LdSt.mayLoadOrStore())
+    return false;
 
   // Handle only loads/stores with base register followed by immediate offset.
   if (LdSt.getNumExplicitOperands() != 3)
     return false;
   if (!LdSt.getOperand(1).isImm() || !LdSt.getOperand(2).isReg())
     return false;
 
   if (!LdSt.hasOneMemOperand())

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

bool RISCVInstrInfo::verifyInstruction(const MachineInstr &MI,
   return true;
 }
 
 // Return true if get the base operand, byte offset of an instruction and the
 // memory width. Width is the size of memory that is being loaded/stored.
 bool RISCVInstrInfo::getMemOperandWithOffsetWidth(
     const MachineInstr &LdSt, const MachineOperand *&BaseReg, int64_t &Offset,
     unsigned &Width, const TargetRegisterInfo *TRI) const {
-  assert(LdSt.mayLoadOrStore() && "Expected a memory operation.");
+  if (!LdSt.mayLoadOrStore())
+    return false;
 
   // Here we assume the standard RISC-V ISA, which uses a base+offset
   // addressing mode. You'll need to relax these conditions to support custom
   // load/stores instructions.
   if (LdSt.getNumExplicitOperands() != 3)
     return false;
   if (!LdSt.getOperand(1).isReg() || !LdSt.getOperand(2).isImm())
     return false;

llvm/lib/Target/X86/X86InstrInfo.cpp

bool X86InstrInfo::getMemOperandWithOffset(
   const MachineOperand &DispMO = MemOp.getOperand(MemRefBegin + X86::AddrDisp);
 
   // Displacement can be symbolic
   if (!DispMO.isImm())
     return false;
 
   Offset = DispMO.getImm();
 
-  assert(BaseOp->isReg() && "getMemOperandWithOffset only supports base "
-                            "operands of type register.");
+  if (!BaseOp->isReg())
+    return false;
+
   return true;
 }
 
 static unsigned getStoreRegOpcode(unsigned SrcReg,
                                   const TargetRegisterClass *RC,
                                   bool isStackAligned,
                                   const X86Subtarget &STI) {
   return getLoadStoreRegOpcode(SrcReg, RC, isStackAligned, STI, false);

llvm/test/CodeGen/AArch64/machine-sink-getmemoperandwithoffset.mir

+# RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass machine-sink -o - %s | FileCheck %s
+--- |
+  define i8 @g() {
+  else.7:
+    br i1 undef, label %then.8, label %else.8, !make.implicit !0
+
+  then.8:                                           ; preds = %else.8, %else.7
+    %merge = phi i8 [ 1, %else.7 ], [ 0, %else.8 ]
+    ret i8 %merge ;1 ;%merge
+
+  else.8:                                           ; preds = %else.7
+    %icmp.8 = icmp eq i64 undef, undef
+  ;  ret i8 0 ; added
+    br i1 %icmp.8, label %else.11, label %then.8
+
+  else.11:                                          ; preds = %else.8
+    ret i8 undef
+  }
+
+  !0 = !{}
+...
+---
+name:            g
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: gpr32all }
+  - { id: 1, class: gpr32all }
+  - { id: 2, class: gpr32 }
+  - { id: 3, class: gpr32 }
+  - { id: 4, class: gpr32all }
+  - { id: 5, class: gpr32 }
+  - { id: 6, class: gpr32all }
+body:             |
+  ; Just check that the pass didn't crash/assert.
+  ; CHECK-LABEL: name: g
+  bb.0.else.7:
+    successors: %bb.1, %bb.2
+
+    %2:gpr32 = MOVi32imm 1
+  ; Sinking the below COPY instruction caused an assert to trigger before
+  ; requiring getMemOperandWithOffset to return false rather than assert
+  ; when handling non-memory operations.
+    %1:gpr32all = COPY %2
+    %3:gpr32 = COPY $wzr
+    CBNZW %3, %bb.2
+    B %bb.1
+
+  bb.1.then.8:
+    %0:gpr32all = PHI %1, %bb.0, %4, %bb.2
+    $w0 = COPY %0
+    RET_ReallyLR implicit $w0
+
+  bb.2.else.8:
+    successors: %bb.3, %bb.1
+
+    %5:gpr32 = COPY $wzr
+    %4:gpr32all = COPY %5
+    CBNZW %5, %bb.1
+    B %bb.3
+
+  bb.3.else.11:
+    %6:gpr32all = IMPLICIT_DEF
+    $w0 = COPY %6
+    RET_ReallyLR implicit $w0
+...