[llvm] [MacroFusion][RISCV] Allocate same register for second instruction of fusible pair (PR #77461)

[Wang Pengcheng via llvm-commits]

https://github.com/wangpc-pp updated https://github.com/llvm/llvm-project/pull/77461

>From 627718939edefad76fe209bf9f03bb97bc900b1b Mon Sep 17 00:00:00 2001
From: wangpc <wangpengcheng.pp at bytedance.com>
Date: Tue, 9 Jan 2024 20:53:18 +0800
Subject: [PATCH] [MacroFusion][RISCV] Allocate same register for second
 instruction of fusible pair

We add a MI flag to indicate the constraint and set this flag to
true for the second instruction of fusible pairs in pre-regalloc
macrofusion.

Then, we add register allocation hints for it.

During regalloc, the allocator will choose the same register
according to the hint.

This is a PoC currently.
---
 llvm/include/llvm/CodeGen/MachineInstr.h    |  4 +++
 llvm/lib/CodeGen/MacroFusion.cpp            |  6 ++++
 llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp |  3 ++
 llvm/test/CodeGen/RISCV/pr76779.ll          | 37 +++++++++++++++++++++
 4 files changed, 50 insertions(+)
 create mode 100644 llvm/test/CodeGen/RISCV/pr76779.ll

diff --git a/llvm/include/llvm/CodeGen/MachineInstr.h b/llvm/include/llvm/CodeGen/MachineInstr.h
index fcdd73d8b65fdd..4f79c3c1872c2d 100644
--- a/llvm/include/llvm/CodeGen/MachineInstr.h
+++ b/llvm/include/llvm/CodeGen/MachineInstr.h
@@ -115,6 +115,7 @@ class MachineInstr
                              // this instruction.
     Unpredictable = 1 << 16, // Instruction with unpredictable condition.
     NoConvergent = 1 << 17,  // Call does not require convergence guarantees.
+    Fusible = 1 << 18,       // Instruction is the second of a fusible pair.
   };
 
 private:
@@ -1031,6 +1032,9 @@ class MachineInstr
     return hasProperty(MCID::Convergent, Type);
   }
 
+  /// Return true if this instruction is fusible.
+  bool isFusible() const { return getFlag(Fusible); }
+
   /// Returns true if the specified instruction has a delay slot
   /// which must be filled by the code generator.
   bool hasDelaySlot(QueryType Type = AnyInBundle) const {
diff --git a/llvm/lib/CodeGen/MacroFusion.cpp b/llvm/lib/CodeGen/MacroFusion.cpp
index 5bd6ca0978a4b1..3efae951dba53e 100644
--- a/llvm/lib/CodeGen/MacroFusion.cpp
+++ b/llvm/lib/CodeGen/MacroFusion.cpp
@@ -128,6 +128,12 @@ bool llvm::fuseInstructionPair(ScheduleDAGInstrs &DAG, SUnit &FirstSU,
     }
   }
 
+  // Mark the second instruction of fusible pair as MachineInstr::Fusible if
+  // this mutation is running in pre-ra scheduler.
+  if (!DAG.MF.getProperties().hasProperty(
+          MachineFunctionProperties::Property::NoVRegs))
+    SecondSU.getInstr()->setFlag(MachineInstr::Fusible);
+
   ++NumFused;
   return true;
 }
diff --git a/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp b/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp
index a68674b221d38e..e2b5aa465e1d78 100644
--- a/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp
+++ b/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp
@@ -838,6 +838,9 @@ bool RISCVRegisterInfo::getRegAllocationHints(
         tryAddHint(MO, MI.getOperand(0), NeedGPRC);
       }
     }
+    if (MI.isFusible())
+      if (OpIdx == 1 || (OpIdx == 2 && MI.isCommutable()))
+        tryAddHint(MO, MI.getOperand(0), false);
   }
 
   for (MCPhysReg OrderReg : Order)
diff --git a/llvm/test/CodeGen/RISCV/pr76779.ll b/llvm/test/CodeGen/RISCV/pr76779.ll
new file mode 100644
index 00000000000000..cec132251e3c4c
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/pr76779.ll
@@ -0,0 +1,37 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+;RUN: llc < %s -mtriple=riscv64 -mattr=+f -target-abi=lp64f \
+;RUN:   | FileCheck %s --check-prefix=NOFUSION
+;RUN: llc < %s -mtriple=riscv64 -mattr=+f,+lui-addi-fusion \
+;RUN:   -target-abi=lp64f | FileCheck %s --check-prefix=FUSION
+;RUN: llc < %s -mtriple=riscv64 -mattr=+f,+lui-addi-fusion,+use-postra-scheduler \
+;RUN:   -target-abi=lp64f | FileCheck %s --check-prefixes=FUSION-POSTRA
+
+define void @foo(i32 noundef signext %0, i32 noundef signext %1) {
+; NOFUSION-LABEL: foo:
+; NOFUSION:       # %bb.0:
+; NOFUSION-NEXT:    lui a0, 3014
+; NOFUSION-NEXT:    addiw a2, a0, 334
+; NOFUSION-NEXT:    mv a0, a1
+; NOFUSION-NEXT:    mv a1, a2
+; NOFUSION-NEXT:    tail bar
+;
+; FUSION-LABEL: foo:
+; FUSION:       # %bb.0:
+; FUSION-NEXT:    lui a2, 3014
+; FUSION-NEXT:    addiw a2, a2, 334
+; FUSION-NEXT:    mv a0, a1
+; FUSION-NEXT:    mv a1, a2
+; FUSION-NEXT:    tail bar
+;
+; FUSION-POSTRA-LABEL: foo:
+; FUSION-POSTRA:       # %bb.0:
+; FUSION-POSTRA-NEXT:    lui a2, 3014
+; FUSION-POSTRA-NEXT:    addiw a2, a2, 334
+; FUSION-POSTRA-NEXT:    mv a0, a1
+; FUSION-POSTRA-NEXT:    mv a1, a2
+; FUSION-POSTRA-NEXT:    tail bar
+  tail call void @bar(i32 noundef signext %1, i32 noundef signext 12345678)
+  ret void
+}
+
+declare void @bar(i32 noundef signext, i32 noundef signext)