[RISCV] Teach RISCVMergeBaseOffset to handle inline asm (#78945)

For inline asm with memory operands, we can merge the offset into
the second operand of memory constraint operands.

Differential Revision: https://reviews.llvm.org/D158062

Author: wangpc-pp

llvm/include/llvm/CodeGen/MachineOperand.h

@@ -782,6 +782,10 @@ class MachineOperand {
   void ChangeToGA(const GlobalValue *GV, int64_t Offset,
                   unsigned TargetFlags = 0);
 
+  /// ChangeToBA - Replace this operand with a new block address operand.
+  void ChangeToBA(const BlockAddress *BA, int64_t Offset,
+                  unsigned TargetFlags = 0);
+
   /// ChangeToMCSymbol - Replace this operand with a new MC symbol operand.
   void ChangeToMCSymbol(MCSymbol *Sym, unsigned TargetFlags = 0);
 

llvm/lib/CodeGen/MachineOperand.cpp

@@ -206,6 +206,19 @@ void MachineOperand::ChangeToGA(const GlobalValue *GV, int64_t Offset,
   setTargetFlags(TargetFlags);
 }
 
+void MachineOperand::ChangeToBA(const BlockAddress *BA, int64_t Offset,
+                                unsigned TargetFlags) {
+  assert((!isReg() || !isTied()) &&
+         "Cannot change a tied operand into a block address");
+
+  removeRegFromUses();
+
+  OpKind = MO_BlockAddress;
+  Contents.OffsetedInfo.Val.BA = BA;
+  setOffset(Offset);
+  setTargetFlags(TargetFlags);
+}
+
 void MachineOperand::ChangeToMCSymbol(MCSymbol *Sym, unsigned TargetFlags) {
   assert((!isReg() || !isTied()) &&
          "Cannot change a tied operand into an MCSymbol");

llvm/lib/Target/RISCV/RISCVAsmPrinter.cpp

@@ -323,7 +323,8 @@ bool RISCVAsmPrinter::PrintAsmMemoryOperand(const MachineInstr *MI,
   // RISCVDAGToDAGISel::SelectInlineAsmMemoryOperand).
   if (!AddrReg.isReg())
     return true;
-  if (!Offset.isImm() && !Offset.isGlobal() && !Offset.isBlockAddress())
+  if (!Offset.isImm() && !Offset.isGlobal() && !Offset.isBlockAddress() &&
+      !Offset.isMCSymbol())
     return true;
 
   MCOperand MCO;
@@ -332,7 +333,7 @@ bool RISCVAsmPrinter::PrintAsmMemoryOperand(const MachineInstr *MI,
 
   if (Offset.isImm())
     OS << MCO.getImm();
-  else if (Offset.isGlobal() || Offset.isBlockAddress())
+  else if (Offset.isGlobal() || Offset.isBlockAddress() || Offset.isMCSymbol())
     OS << *MCO.getExpr();
   OS << "(" << RISCVInstPrinter::getRegisterName(AddrReg.getReg()) << ")";
   return false;

llvm/lib/Target/RISCV/RISCVMergeBaseOffset.cpp

@@ -362,6 +362,8 @@ bool RISCVMergeBaseOffsetOpt::foldIntoMemoryOps(MachineInstr &Hi,
   // Tail: lw vreg3, 8(vreg2)
 
   std::optional<int64_t> CommonOffset;
+  DenseMap<const MachineInstr *, SmallVector<unsigned>>
+      InlineAsmMemoryOpIndexesMap;
   for (const MachineInstr &UseMI : MRI->use_instructions(DestReg)) {
     switch (UseMI.getOpcode()) {
     default:
@@ -396,6 +398,49 @@ bool RISCVMergeBaseOffsetOpt::foldIntoMemoryOps(MachineInstr &Hi,
       if (CommonOffset && Offset != CommonOffset)
         return false;
       CommonOffset = Offset;
+      break;
+    }
+    case RISCV::INLINEASM:
+    case RISCV::INLINEASM_BR: {
+      SmallVector<unsigned> InlineAsmMemoryOpIndexes;
+      unsigned NumOps = 0;
+      for (unsigned I = InlineAsm::MIOp_FirstOperand;
+           I < UseMI.getNumOperands(); I += 1 + NumOps) {
+        const MachineOperand &FlagsMO = UseMI.getOperand(I);
+        // Should be an imm.
+        if (!FlagsMO.isImm())
+          continue;
+
+        const InlineAsm::Flag Flags(FlagsMO.getImm());
+        NumOps = Flags.getNumOperandRegisters();
+
+        // Memory constraints have two operands.
+        if (NumOps != 2 || !Flags.isMemKind())
+          continue;
+
+        // We can't do this for constraint A because AMO instructions don't have
+        // an immediate offset field.
+        if (Flags.getMemoryConstraintID() == InlineAsm::ConstraintCode::A)
+          return false;
+
+        const MachineOperand &AddrMO = UseMI.getOperand(I + 1);
+        if (!AddrMO.isReg() || AddrMO.getReg() != DestReg)
+          continue;
+
+        const MachineOperand &OffsetMO = UseMI.getOperand(I + 2);
+        if (!OffsetMO.isImm())
+          continue;
+
+        // All inline asm memory operands must use the same offset.
+        int64_t Offset = OffsetMO.getImm();
+        if (CommonOffset && Offset != CommonOffset)
+          return false;
+        CommonOffset = Offset;
+        InlineAsmMemoryOpIndexes.push_back(I + 1);
+      }
+      InlineAsmMemoryOpIndexesMap.insert(
+          std::make_pair(&UseMI, InlineAsmMemoryOpIndexes));
+      break;
     }
     }
   }
@@ -420,13 +465,36 @@ bool RISCVMergeBaseOffsetOpt::foldIntoMemoryOps(MachineInstr &Hi,
   // Update the immediate in the load/store instructions to add the offset.
   for (MachineInstr &UseMI :
        llvm::make_early_inc_range(MRI->use_instructions(DestReg))) {
-    UseMI.removeOperand(2);
-    UseMI.addOperand(ImmOp);
-    // Update the base reg in the Tail instruction to feed from LUI.
-    // Output of Hi is only used in Lo, no need to use MRI->replaceRegWith().
-    UseMI.getOperand(1).setReg(Hi.getOperand(0).getReg());
+    if (UseMI.getOpcode() == RISCV::INLINEASM ||
+        UseMI.getOpcode() == RISCV::INLINEASM_BR) {
+      auto &InlineAsmMemoryOpIndexes = InlineAsmMemoryOpIndexesMap[&UseMI];
+      for (unsigned I : InlineAsmMemoryOpIndexes) {
+        MachineOperand &MO = UseMI.getOperand(I + 1);
+        switch (ImmOp.getType()) {
+        case MachineOperand::MO_GlobalAddress:
+          MO.ChangeToGA(ImmOp.getGlobal(), ImmOp.getOffset(),
+                        ImmOp.getTargetFlags());
+          break;
+        case MachineOperand::MO_MCSymbol:
+          MO.ChangeToMCSymbol(ImmOp.getMCSymbol(), ImmOp.getTargetFlags());
+          MO.setOffset(ImmOp.getOffset());
+          break;
+        case MachineOperand::MO_BlockAddress:
+          MO.ChangeToBA(ImmOp.getBlockAddress(), ImmOp.getOffset(),
+                        ImmOp.getTargetFlags());
+          break;
+        default:
+          report_fatal_error("unsupported machine operand type");
+          break;
+        }
+      }
+    } else {
+      UseMI.removeOperand(2);
+      UseMI.addOperand(ImmOp);
+    }
   }
 
+  MRI->replaceRegWith(Lo.getOperand(0).getReg(), Hi.getOperand(0).getReg());
   Lo.eraseFromParent();
   return true;
 }