[RISCV][MC] Support imm symbol in parseCSRSystemRegister (#112007)

Co-authored-by: Alex Richardson <[email protected]>

Author: zqb-all

llvm/lib/Target/RISCV/AsmParser/RISCVAsmParser.cpp

@@ -1878,6 +1878,25 @@ ParseStatus RISCVAsmParser::parseCSRSystemRegister(OperandVector &Operands) {
   SMLoc S = getLoc();
   const MCExpr *Res;
 
+  auto SysRegFromConstantInt = [this](const MCExpr *E, SMLoc S) {
+    if (auto *CE = dyn_cast<MCConstantExpr>(E)) {
+      int64_t Imm = CE->getValue();
+      if (isUInt<12>(Imm)) {
+        auto Range = RISCVSysReg::lookupSysRegByEncoding(Imm);
+        // Accept an immediate representing a named Sys Reg if it satisfies the
+        // the required features.
+        for (auto &Reg : Range) {
+          if (Reg.haveRequiredFeatures(STI->getFeatureBits()))
+            return RISCVOperand::createSysReg(Reg.Name, S, Imm);
+        }
+        // Accept an immediate representing an un-named Sys Reg if the range is
+        // valid, regardless of the required features.
+        return RISCVOperand::createSysReg("", S, Imm);
+      }
+    }
+    return std::unique_ptr<RISCVOperand>();
+  };
+
   switch (getLexer().getKind()) {
   default:
     return ParseStatus::NoMatch;
@@ -1891,24 +1910,9 @@ ParseStatus RISCVAsmParser::parseCSRSystemRegister(OperandVector &Operands) {
     if (getParser().parseExpression(Res))
       return ParseStatus::Failure;
 
-    auto *CE = dyn_cast<MCConstantExpr>(Res);
-    if (CE) {
-      int64_t Imm = CE->getValue();
-      if (isUInt<12>(Imm)) {
-        auto Range = RISCVSysReg::lookupSysRegByEncoding(Imm);
-        // Accept an immediate representing a named Sys Reg if it satisfies the
-        // the required features.
-        for (auto &Reg : Range) {
-          if (Reg.haveRequiredFeatures(STI->getFeatureBits())) {
-            Operands.push_back(RISCVOperand::createSysReg(Reg.Name, S, Imm));
-            return ParseStatus::Success;
-          }
-        }
-        // Accept an immediate representing an un-named Sys Reg if the range is
-        // valid, regardless of the required features.
-        Operands.push_back(RISCVOperand::createSysReg("", S, Imm));
-        return ParseStatus::Success;
-      }
+    if (auto SysOpnd = SysRegFromConstantInt(Res, S)) {
+      Operands.push_back(std::move(SysOpnd));
+      return ParseStatus::Success;
     }
 
     return generateImmOutOfRangeError(S, 0, (1 << 12) - 1);
@@ -1951,6 +1955,18 @@ ParseStatus RISCVAsmParser::parseCSRSystemRegister(OperandVector &Operands) {
       return ParseStatus::Success;
     }
 
+    // Accept a symbol name that evaluates to an absolute value.
+    MCSymbol *Sym = getContext().lookupSymbol(Identifier);
+    if (Sym && Sym->isVariable()) {
+      // Pass false for SetUsed, since redefining the value later does not
+      // affect this instruction.
+      if (auto SysOpnd = SysRegFromConstantInt(
+              Sym->getVariableValue(/*SetUsed=*/false), S)) {
+        Operands.push_back(std::move(SysOpnd));
+        return ParseStatus::Success;
+      }
+    }
+
     return generateImmOutOfRangeError(S, 0, (1 << 12) - 1,
                                       "operand must be a valid system register "
                                       "name or an integer in the range");

llvm/test/MC/RISCV/rv32i-invalid.s

@@ -93,6 +93,19 @@ csrrsi a0, mhpm12counter, a0 # CHECK: :[[@LINE]]:12: error: operand must be a va
 csrrwi a0, mhpmcounter32, 0 # CHECK: :[[@LINE]]:12: error: operand must be a valid system register name or an integer in the range [0, 4095]
 csrrsi a0, A, a0 # CHECK: :[[@LINE]]:12: error: operand must be a valid system register name or an integer in the range [0, 4095]
 
+## symbol in place of uimm12
+.set out_of_range, 4096
+csrr a0, out_of_range # CHECK: [[#@LINE]]:10: error: operand must be a valid system register name or an integer in the range [0, 4095]
+csrr a0, undef_symbol # CHECK: [[#@LINE]]:10: error: operand must be a valid system register name or an integer in the range [0, 4095]
+local_label:
+csrr a0, local_label # CHECK: [[#@LINE]]:10: error: operand must be a valid system register name or an integer in the range [0, 4095]
+.Lstart:
+.space 10
+.Lend:
+csrr a0, .Lstart-.Lend # CHECK: [[#@LINE]]:10: error: operand must be a valid system register name or an integer in the range [0, 4095]
+.set dot_set_sym_diff, .Lstart-.Lend
+csrr a0, dot_set_sym_diff # CHECK: [[#@LINE]]:10: error: operand must be a valid system register name or an integer in the range [0, 4095]
+
 ## simm13_lsb0
 beq t0, t1, %lo(1) # CHECK: :[[@LINE]]:13: error: immediate must be a multiple of 2 bytes in the range [-4096, 4094]
 bne t0, t1, %lo(a) # CHECK: :[[@LINE]]:13: error: immediate must be a multiple of 2 bytes in the range [-4096, 4094]

llvm/test/MC/RISCV/rv32i-valid.s

@@ -373,3 +373,32 @@ csrrsi t2, 0xfff, 31
 # CHECK-ASM-AND-OBJ: csrrci t1, sscratch, 5
 # CHECK-ASM: encoding: [0x73,0xf3,0x02,0x14]
 csrrci t1, 0x140, 5
+
+## Check that we can use an absolute symbol value as a CSR number
+# CHECK-ASM-AND-OBJ: csrrs a0, fflags, zero
+# CHECK-ASM: encoding: [0x73,0x25,0x10,0x00]
+.set fflags_abs_sym, 1
+csrr a0, fflags_abs_sym
+# CHECK-ASM-AND-OBJ: csrrs a0, frm, zero
+# CHECK-ASM: encoding: [0x73,0x25,0x20,0x00]
+csrr a0, (fflags_abs_sym+1)
+# CHECK-ASM-AND-OBJ: csrrs a0, frm, zero
+# CHECK-ASM: encoding: [0x73,0x25,0x20,0x00]
+.equ fplus_one_abs_sym, fflags_abs_sym + 1
+csrr a0, fplus_one_abs_sym
+
+## Check that redefining the value is allowed
+# CHECK-ASM-AND-OBJ: csrrs a0, fflags, zero
+# CHECK-ASM: encoding: [0x73,0x25,0x10,0x00]
+.set csr_index, 1
+csrr a0, csr_index
+# CHECK-ASM-AND-OBJ: csrrs a0, frm, zero
+# CHECK-ASM: encoding: [0x73,0x25,0x20,0x00]
+.set csr_index, 2
+csrr a0, csr_index
+
+## Check that select the CSR first.
+.set frm, 1
+# CHECK-ASM-AND-OBJ: csrrs a0, frm, zero
+# CHECK-ASM: encoding: [0x73,0x25,0x20,0x00]
+csrr a0, frm