[LoongArch] Prevent R0/R1 allocation for rj operand of {G}CSRXCHG

The {G}CSRXCHG instruction must not use R0 or R1 as the rj operand, as
encoding rj as 0 or 1 will be interpreted as {G}CSRRD OR {G}CSRWR,
respectively, rather than {G}CSRXCHG.

This patch introduces a new register class `GPRNoR0R1` and updates the
{G}CSRXCHG instruction definition to use it for the rj operand, ensuring
the register allocator avoids assigning R0 or R1.

Fixes llvm#140842

Author: heiher

llvm/lib/Target/LoongArch/Disassembler/LoongArchDisassembler.cpp

@@ -62,6 +62,14 @@ static DecodeStatus DecodeGPRRegisterClass(MCInst &Inst, uint64_t RegNo,
   return MCDisassembler::Success;
 }
 
+static DecodeStatus
+DecodeGPRNoR0R1RegisterClass(MCInst &Inst, uint64_t RegNo, uint64_t Address,
+                             const MCDisassembler *Decoder) {
+  if (RegNo <= 1)
+    return MCDisassembler::Fail;
+  return DecodeGPRRegisterClass(Inst, RegNo, Address, Decoder);
+}
+
 static DecodeStatus DecodeFPR32RegisterClass(MCInst &Inst, uint64_t RegNo,
                                              uint64_t Address,
                                              const MCDisassembler *Decoder) {

llvm/lib/Target/LoongArch/LoongArchInstrInfo.td

@@ -2431,7 +2431,7 @@ let hasSideEffects = 1, Constraints = "$rd = $dst" in {
 def CSRWR : FmtCSR<0x04000020, (outs GPR:$dst),
                    (ins GPR:$rd, uimm14:$csr_num), "$rd, $csr_num">;
 def CSRXCHG : FmtCSRXCHG<0x04000000, (outs GPR:$dst),
-                         (ins GPR:$rd, GPR:$rj, uimm14:$csr_num),
+                         (ins GPR:$rd, GPRNoR0R1:$rj, uimm14:$csr_num),
                          "$rd, $rj, $csr_num">;
 } // hasSideEffects = 1, Constraints = "$rd = $dst"
 
@@ -2478,8 +2478,8 @@ def IDLE : MISC_I15<0x06488000>;
 def : Pat<(loongarch_csrrd uimm14:$imm14), (CSRRD uimm14:$imm14)>;
 def : Pat<(loongarch_csrwr GPR:$rd, uimm14:$imm14),
           (CSRWR GPR:$rd, uimm14:$imm14)>;
-def : Pat<(loongarch_csrxchg GPR:$rd, GPR:$rj, uimm14:$imm14),
-          (CSRXCHG GPR:$rd, GPR:$rj, uimm14:$imm14)>;
+def : Pat<(loongarch_csrxchg GPR:$rd, GPRNoR0R1:$rj, uimm14:$imm14),
+          (CSRXCHG GPR:$rd, GPRNoR0R1:$rj, uimm14:$imm14)>;
 
 def : Pat<(loongarch_iocsrrd_b GPR:$rj), (IOCSRRD_B GPR:$rj)>;
 def : Pat<(loongarch_iocsrrd_h GPR:$rj), (IOCSRRD_H GPR:$rj)>;

llvm/lib/Target/LoongArch/LoongArchLVZInstrInfo.td

@@ -23,7 +23,7 @@ let Constraints = "$rd = $dst" in {
 def GCSRWR : FmtCSR<0x05000020, (outs GPR:$dst),
                     (ins GPR:$rd, uimm14:$csr_num), "$rd, $csr_num">;
 def GCSRXCHG : FmtCSRXCHG<0x05000000, (outs GPR:$dst),
-                          (ins GPR:$rd, GPR:$rj, uimm14:$csr_num),
+                          (ins GPR:$rd, GPRNoR0R1:$rj, uimm14:$csr_num),
                           "$rd, $rj, $csr_num">;
 } // Constraints = "$rd = $dst"
 

llvm/lib/Target/LoongArch/LoongArchRegisterInfo.td

@@ -127,6 +127,11 @@ def GPRT : GPRRegisterClass<(add // a0...a7, t0...t8
 // prediction.
 def GPRJR : GPRRegisterClass<(sub GPR, R1)>;
 
+// Don't use R0 or R1 for the rj operand of {G}CSRXCHG, because when rj is
+// encoded as 0 or 1, the instruction is interpreted as {G}CSRRD or {G}CSRWR,
+// respectively, rather than {G}CSRXCHG.
+def GPRNoR0R1 : GPRRegisterClass<(sub GPR, R0, R1)>;
+
 // Floating point registers
 
 let RegAltNameIndices = [RegAliasName] in {

llvm/test/CodeGen/LoongArch/csrxchg-intrinsic.ll

@@ -0,0 +1,24 @@
+; RUN: llc --mtriple=loongarch32 --mattr=+f --verify-machineinstrs < %s | FileCheck %s
+; RUN: llc --mtriple=loongarch64 --mattr=+f --verify-machineinstrs < %s | FileCheck %s
+
+declare i32 @llvm.loongarch.csrxchg.w(i32, i32, i32 immarg)
+
+;; Check that the rj operand of csrxchg is not R0.
+define void @csrxchg_w_rj_not_r0(i32 signext %a) {
+; CHECK-NOT:    csrxchg ${{[a-z]*}}, $r0, 0
+; CHECK-NOT:    csrxchg ${{[a-z]*}}, $zero, 0
+entry:
+  %0 = tail call i32 @llvm.loongarch.csrxchg.w(i32 %a, i32 0, i32 0)
+  ret void
+}
+
+;; Check that the rj operand of csrxchg is not R1.
+define i32 @csrxchg_w_rj_not_r1() {
+; CHECK-NOT:    csrxchg ${{[a-z]*}}, $r1, 0
+; CHECK-NOT:    csrxchg ${{[a-z]*}}, $ra, 0
+entry:
+  %0 = tail call i32 asm "", "=r,r,i,{r4},{r5},{r6},{r7},{r8},{r9},{r10},{r11},{r12},{r13},{r14},{r15},{r16},{r17},{r18},{r19},{r20},{r23},{r24},{r25},{r26},{r27},{r28},{r29},{r30},{r31},0"(i32 4, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0)
+  %1 = tail call i32 @llvm.loongarch.csrxchg.w(i32 %0, i32 4, i32 0)
+  %2 = tail call i32 asm "", "=r,r,i,{r4},{r5},{r6},{r7},{r8},{r9},{r10},{r11},{r12},{r13},{r14},{r15},{r16},{r17},{r18},{r19},{r20},{r23},{r24},{r25},{r26},{r27},{r28},{r29},{r30},{r31},0"(i32 4, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 %1)
+  ret i32 %2
+}

llvm/test/MC/LoongArch/Basic/Privilege/invalid.s

@@ -3,9 +3,9 @@
 
 ## csrxchg: rj != 0,1
 csrxchg $a0, $zero, 0
-# ERR: :[[#@LINE-1]]:15: error: must not be $r0 or $r1
+# ERR: :[[#@LINE-1]]:15: error: invalid operand for instruction
 csrxchg $a0, $ra, 0
-# ERR: :[[#@LINE-1]]:15: error: must not be $r0 or $r1
+# ERR: :[[#@LINE-1]]:15: error: invalid operand for instruction
 
 ## LoongArch64 mnemonics
 iocsrrd.d $a0, $a1

llvm/test/MC/LoongArch/lvz/lvz-err.s

@@ -19,10 +19,10 @@ gcsrxchg $a0, $a1, -1
 # CHECK: :[[#@LINE-1]]:20: error: immediate must be an integer in the range [0, 16383]
 
 gcsrxchg $a0, $ra, 1
-# CHECK: :[[#@LINE-1]]:16: error: must not be $r0 or $r1
+# CHECK: :[[#@LINE-1]]:16: error: invalid operand for instruction
 
 gcsrxchg $a0, $zero, 1
-# CHECK: :[[#@LINE-1]]:16: error: must not be $r0 or $r1
+# CHECK: :[[#@LINE-1]]:16: error: invalid operand for instruction
 
 hvcl 32768
 # CHECK: :[[#@LINE-1]]:6: error: immediate must be an integer in the range [0, 32767]