[llvm] [AMDGPU][AsmParser] Allow `v_writelane_b32` to use SGPR and M0 as source operands at the same time (PR #78827)

[Shilei Tian via llvm-commits]

https://github.com/shiltian updated https://github.com/llvm/llvm-project/pull/78827

>From 6bcae09cd8759255516b48fa98e2bddc3ade3859 Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Mon, 22 Jan 2024 16:46:47 -0500
Subject: [PATCH] [AMDGPU][AsmParser] Allow `v_writelane_b32` to use SGPR and
 M0 as source operands at the same time

Currently the asm parser takes `v_writelane_b32 v1, s13, m0` as illegal
instruction for pre-gfx11 because it uses two constant buses while the hardware
can only allow one. However, based on the comment of `AMDGPUInstructionSelector::selectWritelane`,
it is allowed to have M0 as lane selector and a SGPR used as SRC0 because the
lane selector doesn't count as a use of constant bus. In fact, codegen can already
generate this form, but this inconsistency is not exposed because the validation
of constant bus limitation only happens when paring an assembly but we don't have
a test case when both SGPR and M0 used as source operands for the instruction.
---
 .../AMDGPU/AsmParser/AMDGPUAsmParser.cpp      | 21 +++++++++++++++++++
 llvm/test/MC/AMDGPU/writelane_m0.s            | 14 +++++++++++++
 2 files changed, 35 insertions(+)
 create mode 100644 llvm/test/MC/AMDGPU/writelane_m0.s

diff --git a/llvm/lib/Target/AMDGPU/AsmParser/AMDGPUAsmParser.cpp b/llvm/lib/Target/AMDGPU/AsmParser/AMDGPUAsmParser.cpp
index 66267c9255f419..b810442053ca02 100644
--- a/llvm/lib/Target/AMDGPU/AsmParser/AMDGPUAsmParser.cpp
+++ b/llvm/lib/Target/AMDGPU/AsmParser/AMDGPUAsmParser.cpp
@@ -3507,6 +3507,24 @@ bool AMDGPUAsmParser::usesConstantBus(const MCInst &Inst, unsigned OpIdx) {
   }
 }
 
+// Based on the comment for `AMDGPUInstructionSelector::selectWritelane`:
+// Writelane is special in that it can use SGPR and M0 (which would normally
+// count as using the constant bus twice - but in this case it is allowed since
+// the lane selector doesn't count as a use of the constant bus). However, it is
+// still required to abide by the 1 SGPR rule.
+static bool checkWriteLane(const MCInst &Inst) {
+  const unsigned Opcode = Inst.getOpcode();
+  if (Opcode != V_WRITELANE_B32_gfx6_gfx7 && Opcode != V_WRITELANE_B32_vi)
+    return false;
+  const MCOperand &LaneSelOp = Inst.getOperand(2);
+  if (!LaneSelOp.isReg())
+    return false;
+  auto LaneSelReg = mc2PseudoReg(LaneSelOp.getReg());
+  if (LaneSelReg == M0 || LaneSelReg == M0_gfxpre11)
+    return true;
+  return false;
+}
+
 bool AMDGPUAsmParser::validateConstantBusLimitations(
     const MCInst &Inst, const OperandVector &Operands) {
   const unsigned Opcode = Inst.getOpcode();
@@ -3522,6 +3540,9 @@ bool AMDGPUAsmParser::validateConstantBusLimitations(
       !isVOPD(Opcode))
     return true;
 
+  if (checkWriteLane(Inst))
+    return true;
+
   // Check special imm operands (used by madmk, etc)
   if (AMDGPU::hasNamedOperand(Opcode, AMDGPU::OpName::imm)) {
     ++NumLiterals;
diff --git a/llvm/test/MC/AMDGPU/writelane_m0.s b/llvm/test/MC/AMDGPU/writelane_m0.s
new file mode 100644
index 00000000000000..f98b114d19b3a8
--- /dev/null
+++ b/llvm/test/MC/AMDGPU/writelane_m0.s
@@ -0,0 +1,14 @@
+// RUN: llvm-mc --triple=amdgcn --mcpu=gfx904 -show-encoding %s | FileCheck %s -check-prefix=GFX9
+// RUN: llvm-mc --triple=amdgcn --mcpu=gfx940 -show-encoding %s | FileCheck %s -check-prefix=GFX9
+// RUN: llvm-mc --triple=amdgcn --mcpu=gfx1010 -show-encoding %s | FileCheck %s -check-prefix=GFX10
+// RUN: llvm-mc --triple=amdgcn --mcpu=gfx1030 -show-encoding %s | FileCheck %s -check-prefix=GFX10
+// RUN: llvm-mc --triple=amdgcn --mcpu=gfx1100 -show-encoding %s | FileCheck %s -check-prefix=GFX11
+// RUN: llvm-mc --triple=amdgcn --mcpu=gfx1200 -show-encoding %s | FileCheck %s -check-prefix=GFX12
+
+.text
+  v_writelane_b32 v1, s13, m0
+
+// GFX9: v_writelane_b32 v1, s13, m0 ; encoding: [0x01,0x00,0x8a,0xd2,0x0d,0xf8,0x00,0x00]
+// GFX10: v_writelane_b32 v1, s13, m0 ; encoding: [0x01,0x00,0x61,0xd7,0x0d,0xf8,0x00,0x00]
+// GFX11: v_writelane_b32 v1, s13, m0 ; encoding: [0x01,0x00,0x61,0xd7,0x0d,0xfa,0x00,0x00]
+// GFX12: v_writelane_b32 v1, s13, m0 ; encoding: [0x01,0x00,0x61,0xd7,0x0d,0xfa,0x00,0x00]