[RISCV] Teach RISCVMergeBaseOffset to handle inline asm

For inline asm with memory operands, we can merge the offset into the second operand of memory constraint operands. Reviewed By: craig.topper Differential Revision: https://reviews.llvm.org/D158062
smeenai · Aug 31, 2023 · f281543 · f281543
1 parent 0d73259
commit f281543
Show file tree

Hide file tree

Showing 3 changed files with 130 additions and 115 deletions.
diff --git a/llvm/lib/Target/RISCV/RISCVAsmPrinter.cpp b/llvm/lib/Target/RISCV/RISCVAsmPrinter.cpp
@@ -237,7 +237,7 @@ bool RISCVAsmPrinter::PrintAsmMemoryOperand(const MachineInstr *MI,
   // RISCVDAGToDAGISel::SelectInlineAsmMemoryOperand).
   if (!AddrReg.isReg())
     return true;
-  if (!Offset.isImm() && !Offset.isGlobal())
+  if (!Offset.isImm() && !Offset.isGlobal() && !Offset.isMCSymbol())
     return true;
 
   MCOperand MCO;
@@ -246,7 +246,7 @@ bool RISCVAsmPrinter::PrintAsmMemoryOperand(const MachineInstr *MI,
 
   if (Offset.isImm())
     OS << MCO.getImm();
-  else if (Offset.isGlobal())
+  else if (Offset.isGlobal() || Offset.isMCSymbol())
     OS << *MCO.getExpr();
   OS << "(" << RISCVInstPrinter::getRegisterName(AddrReg.getReg()) << ")";
   return false;

diff --git a/llvm/lib/Target/RISCV/RISCVMergeBaseOffset.cpp b/llvm/lib/Target/RISCV/RISCVMergeBaseOffset.cpp
@@ -357,6 +357,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:
@@ -391,6 +393,44 @@ 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;
+
+        unsigned Flags = FlagsMO.getImm();
+        NumOps = InlineAsm::getNumOperandRegisters(Flags);
+
+        // Memory constraints have two operands.
+        if (NumOps != 2 || !InlineAsm::isMemKind(Flags))
+          continue;
+
+        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;
     }
     }
   }
@@ -415,13 +455,32 @@ 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 llvm::MachineOperand::MachineOperandType::MO_MCSymbol:
+          MO.ChangeToMCSymbol(ImmOp.getMCSymbol(), ImmOp.getTargetFlags());
+          MO.setOffset(ImmOp.getOffset());
+          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;
 }