[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 created https://github.com/llvm/llvm-project/pull/78827

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.


>From 70306f7439a34861ab82ec08c64f2f97351a52a1 Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Fri, 19 Jan 2024 21:26:38 -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      | 28 +++++++++++++++++++
 llvm/test/MC/AMDGPU/writelane_m0.s            | 10 +++++++
 2 files changed, 38 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 bd68054589b112..491ff9cb655cf9 100644
--- a/llvm/lib/Target/AMDGPU/AsmParser/AMDGPUAsmParser.cpp
+++ b/llvm/lib/Target/AMDGPU/AsmParser/AMDGPUAsmParser.cpp
@@ -3524,6 +3524,34 @@ bool AMDGPUAsmParser::validateConstantBusLimitations(
       !isVOPD(Opcode))
     return true;
 
+  // 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.
+  switch (Opcode) {
+  default:
+    break;
+  case V_WRITELANE_B32_e64_gfx11:
+  case V_WRITELANE_B32_e64_gfx12:
+  case V_WRITELANE_B32_gfx10:
+  case V_WRITELANE_B32_gfx6_gfx7:
+  case V_WRITELANE_B32_vi: {
+    const MCOperand &LaneSelOp = Inst.getOperand(2);
+    if (LaneSelOp.isReg()) {
+      auto LaneSelReg = mc2PseudoReg(LaneSelOp.getReg());
+      switch (LaneSelReg) {
+      default:
+        break;
+      case M0:
+      case M0_gfxpre11:
+      case M0_gfx11plus:
+        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..4b5c265719f8f8
--- /dev/null
+++ b/llvm/test/MC/AMDGPU/writelane_m0.s
@@ -0,0 +1,10 @@
+// RUN: llvm-mc --triple=amdgcn --mcpu=gfx904 %s | FileCheck %s
+// RUN: llvm-mc --triple=amdgcn --mcpu=gfx940 %s | FileCheck %s
+// RUN: llvm-mc --triple=amdgcn --mcpu=gfx1010 %s | FileCheck %s
+// RUN: llvm-mc --triple=amdgcn --mcpu=gfx1030 %s | FileCheck %s
+// RUN: llvm-mc --triple=amdgcn --mcpu=gfx1100 %s | FileCheck %s
+
+.text
+  v_writelane_b32 v1, s13, m0
+
+// CHECK: v_writelane_b32 v1, s13, m0