[RTGTest] Add some registers

This adds some registers to later test passes such as register alloction, assembly emission, correct behavior of the elaboration pass in presence of registers, etc.

Co-authored-by: Andrew Lenharth <[email protected]>

include/circt/Dialect/RTGTest/IR/CMakeLists.txt

@@ -4,3 +4,7 @@ set(LLVM_TARGET_DEFINITIONS RTGTest.td)
 
 add_circt_dialect_doc(RTGTest rtgtest)
 
+set(LLVM_TARGET_DEFINITIONS RTGTestAttributes.td)
+mlir_tablegen(RTGTestEnums.h.inc -gen-enum-decls)
+mlir_tablegen(RTGTestEnums.cpp.inc -gen-enum-defs)
+add_public_tablegen_target(CIRCTRTGTestEnumsIncGen)

include/circt/Dialect/RTGTest/IR/RTGTest.td

@@ -20,6 +20,7 @@ include "mlir/IR/OpAsmInterface.td"
 include "circt/Dialect/RTG/IR/RTGInterfaces.td"
 
 include "circt/Dialect/RTGTest/IR/RTGTestDialect.td"
+include "circt/Dialect/RTGTest/IR/RTGTestAttributes.td"
 include "circt/Dialect/RTGTest/IR/RTGTestTypes.td"
 include "circt/Dialect/RTGTest/IR/RTGTestOps.td"
 

include/circt/Dialect/RTGTest/IR/RTGTestAttributes.td

@@ -0,0 +1,60 @@
+//===- RTGTestAttributes.td - RTGTest attributes -----------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This describes the RTGTest attributes.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef CIRCT_DIALECT_RTGTEST_IR_RTGTESTATTRIBUTES_TD
+#define CIRCT_DIALECT_RTGTEST_IR_RTGTESTATTRIBUTES_TD
+
+include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/EnumAttr.td"
+
+// Flat allocation of unique IDs to all registers. The actual ID value does not
+// matter.
+def RegisterAttr : I32EnumAttr<
+  "Registers", "Unique IDs for all RTGTest registers", [
+    I32EnumAttrCase<"zero",  0>,
+    I32EnumAttrCase<"ra",  1>,
+    I32EnumAttrCase<"sp",  2>,
+    I32EnumAttrCase<"gp",  3>,
+    I32EnumAttrCase<"tp",  4>,
+    I32EnumAttrCase<"t0",  5>,
+    I32EnumAttrCase<"t1",  6>,
+    I32EnumAttrCase<"t2",  7>,
+    I32EnumAttrCase<"s0",  8>,
+    I32EnumAttrCase<"s1",  9>,
+    I32EnumAttrCase<"a0",  10>,
+    I32EnumAttrCase<"a1",  11>,
+    I32EnumAttrCase<"a2",  12>,
+    I32EnumAttrCase<"a3",  13>,
+    I32EnumAttrCase<"a4",  14>,
+    I32EnumAttrCase<"a5",  15>,
+    I32EnumAttrCase<"a6",  16>,
+    I32EnumAttrCase<"a7",  17>,
+    I32EnumAttrCase<"s2",  18>,
+    I32EnumAttrCase<"s3",  19>,
+    I32EnumAttrCase<"s4",  20>,
+    I32EnumAttrCase<"s5",  21>,
+    I32EnumAttrCase<"s6",  22>,
+    I32EnumAttrCase<"s7",  23>,
+    I32EnumAttrCase<"s8",  24>,
+    I32EnumAttrCase<"s9",  25>,
+    I32EnumAttrCase<"s10",  26>,
+    I32EnumAttrCase<"s11",  27>,
+    I32EnumAttrCase<"t3",  28>,
+    I32EnumAttrCase<"t4",  29>,
+    I32EnumAttrCase<"t5",  30>,
+    I32EnumAttrCase<"t6",  31>,
+    I32EnumAttrCase<"Virtual",  32>
+  ]> {
+  let cppNamespace = "::circt::rtgtest";
+}
+
+#endif // CIRCT_DIALECT_RTGTEST_IR_RTGTESTATTRIBUTES_TD

include/circt/Dialect/RTGTest/IR/RTGTestDialect.h

@@ -21,4 +21,7 @@
 // Pull in the Dialect definition.
 #include "circt/Dialect/RTGTest/IR/RTGTestDialect.h.inc"
 
+// Pull in all enum type definitions and utility function declarations.
+#include "circt/Dialect/RTGTest/IR/RTGTestEnums.h.inc"
+
 #endif // CIRCT_DIALECT_RTGTEST_IR_RTGTESTDIALECT_H

include/circt/Dialect/RTGTest/IR/RTGTestOps.h

@@ -13,6 +13,7 @@
 #ifndef CIRCT_DIALECT_RTGTEST_IR_RTGTESTOPS_H
 #define CIRCT_DIALECT_RTGTEST_IR_RTGTESTOPS_H
 
+#include "circt/Dialect/RTG/IR/RTGISAAssemblyOpInterfaces.h"
 #include "circt/Dialect/RTG/IR/RTGOpInterfaces.h"
 #include "circt/Dialect/RTG/IR/RTGOps.h"
 #include "circt/Dialect/RTGTest/IR/RTGTestDialect.h"

include/circt/Dialect/RTGTest/IR/RTGTestOps.td

@@ -11,7 +11,9 @@
 //===----------------------------------------------------------------------===//
 
 include "mlir/IR/CommonAttrConstraints.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "circt/Dialect/RTG/IR/RTGInterfaces.td"
+include "circt/Dialect/RTG/IR/RTGISAAssemblyInterfaces.td"
 
 // Base class for the operation in this dialect.
 class RTGTestOp<string mnemonic, list<Trait> traits = []> :
@@ -33,3 +35,21 @@ def CPUDeclOp : RTGTestOp<"cpu_decl", [
 
   let assemblyFormat = "$id attr-dict";
 }
+
+def RegisterOp : RTGTestOp<"reg", [
+  DeclareOpInterfaceMethods<RegisterOpInterface>,
+  DeclareOpInterfaceMethods<InferTypeOpInterface>,
+]> {
+  let summary = "returns a value representing a register";
+  let description = [{
+    This operation creates a value representing the register given as the 'reg'
+    attribute. A register can be a concrete register or a virtual register.
+    Virtual registers will be assigned a concrete register when running register
+    allocation.
+  }];
+
+  let arguments = (ins RegisterAttr:$reg);
+  let results = (outs RegisterType:$result);
+
+  let assemblyFormat = "$reg attr-dict";
+}

include/circt/Dialect/RTGTest/IR/RTGTestTypes.td

@@ -32,4 +32,13 @@ def CPUType : RTGTestTypeDef<"CPU", [ContextResourceTypeInterface]> {
   let assemblyFormat = "";
 }
 
+def IntegerRegisterType : TypeDef<RTGTestDialect, "IntegerRegister", []> {
+  let summary = "represents an integer register";
+
+  let mnemonic = "ireg";
+  let assemblyFormat = "";
+}
+
+def RegisterType : AnyTypeOf<[IntegerRegisterType]>;
+
 #endif // CIRCT_DIALECT_RTGTEST_IR_RTGTESTTYPES_TD

integration_test/Bindings/Python/dialects/rtg.py

@@ -20,8 +20,8 @@
     target = rtg.TargetOp('target_name', TypeAttr.get(dictTy))
     targetBlock = Block.create_at_start(target.bodyRegion, [])
     with InsertionPoint(targetBlock):
-      cpu0 = rtgtest.CPUDeclOp(cpuTy, 0)
-      cpu1 = rtgtest.CPUDeclOp(cpuTy, 1)
+      cpu0 = rtgtest.CPUDeclOp(0)
+      cpu1 = rtgtest.CPUDeclOp(1)
       rtg.YieldOp([cpu0, cpu1])
 
     test = rtg.TestOp('test_name', TypeAttr.get(dictTy))

lib/Dialect/RTGTest/IR/CMakeLists.txt

@@ -8,8 +8,10 @@ add_circt_dialect_library(CIRCTRTGTestDialect
 
   DEPENDS
   MLIRRTGTestIncGen
+  CIRCTRTGTestEnumsIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
+  MLIRInferTypeOpInterface
   CIRCTRTGDialect
 )

lib/Dialect/RTGTest/IR/RTGTestDialect.cpp

@@ -33,4 +33,6 @@ void RTGTestDialect::initialize() {
       >();
 }
 
+#include "circt/Dialect/RTGTest/IR/RTGTestEnums.cpp.inc"
+
 #include "circt/Dialect/RTGTest/IR/RTGTestDialect.cpp.inc"

lib/Dialect/RTGTest/IR/RTGTestOps.cpp

@@ -23,6 +23,53 @@ using namespace rtgtest;
 
 size_t CPUDeclOp::getIdentifier(size_t idx) { return getId().getZExtValue(); }
 
+//===----------------------------------------------------------------------===//
+// RegisterOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult RegisterOp::inferReturnTypes(
+    MLIRContext *context, std::optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, OpaqueProperties properties,
+    mlir::RegionRange regions, SmallVectorImpl<Type> &inferredReturnTypes) {
+  Registers reg = properties.as<Properties *>()->getReg().getValue();
+  if (static_cast<uint32_t>(reg) >= 32 && reg != Registers::Virtual) {
+    if (location)
+      return mlir::emitError(*location) << "unsupported register";
+
+    return failure();
+  }
+
+  inferredReturnTypes.push_back(IntegerRegisterType::get(context));
+  return success();
+}
+
+unsigned RegisterOp::getClassIndex() { return static_cast<uint32_t>(getReg()); }
+
+APInt RegisterOp::getClassIndexBinary() { return APInt(5, getClassIndex()); }
+
+std::string RegisterOp::getRegisterAssembly() {
+  return stringifyRegisters(getReg()).str();
+}
+
+llvm::BitVector RegisterOp::getAllowedRegs() {
+  llvm::BitVector retval(getMaxEnumValForRegisters(), false);
+  if (getReg() == Registers::Virtual)
+    return retval.set(0, 32);
+  return retval.set(static_cast<uint32_t>(getReg()));
+}
+
+unsigned RegisterOp::getFixedReg() {
+  if (getReg() == Registers::Virtual)
+    return ~0U;
+  return static_cast<uint32_t>(getReg());
+}
+
+void RegisterOp::setFixedReg(unsigned reg) {
+  auto sym = symbolizeRegisters(reg);
+  assert(sym.has_value() && sym.value() != Registers::Virtual);
+  setReg(sym.value());
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen generated logic.
 //===----------------------------------------------------------------------===//

test/Dialect/RTGTest/IR/basic.mlir

@@ -7,3 +7,75 @@ rtg.target @cpus : !rtg.dict<cpu: !rtgtest.cpu> {
   %0 = rtgtest.cpu_decl 0
   rtg.yield %0 : !rtgtest.cpu
 }
+
+// CHECK-LABEL: rtg.test @registers
+// CHECK-SAME: !rtgtest.ireg
+rtg.test @registers : !rtg.dict<reg: !rtgtest.ireg> {
+^bb0(%reg: !rtgtest.ireg):
+  // CHECK: rtgtest.reg zero
+  // CHECK: rtgtest.reg ra
+  // CHECK: rtgtest.reg sp
+  // CHECK: rtgtest.reg gp
+  // CHECK: rtgtest.reg tp
+  // CHECK: rtgtest.reg t0
+  // CHECK: rtgtest.reg t1
+  // CHECK: rtgtest.reg t2
+  // CHECK: rtgtest.reg s0
+  // CHECK: rtgtest.reg s1
+  // CHECK: rtgtest.reg a0
+  // CHECK: rtgtest.reg a1
+  // CHECK: rtgtest.reg a2
+  // CHECK: rtgtest.reg a3
+  // CHECK: rtgtest.reg a4
+  // CHECK: rtgtest.reg a5
+  // CHECK: rtgtest.reg a6
+  // CHECK: rtgtest.reg a7
+  // CHECK: rtgtest.reg s2
+  // CHECK: rtgtest.reg s3
+  // CHECK: rtgtest.reg s4
+  // CHECK: rtgtest.reg s5
+  // CHECK: rtgtest.reg s6
+  // CHECK: rtgtest.reg s7
+  // CHECK: rtgtest.reg s8
+  // CHECK: rtgtest.reg s9
+  // CHECK: rtgtest.reg s10
+  // CHECK: rtgtest.reg s11
+  // CHECK: rtgtest.reg t3
+  // CHECK: rtgtest.reg t4
+  // CHECK: rtgtest.reg t5
+  // CHECK: rtgtest.reg t6
+  // CHECK: rtgtest.reg Virtual
+  %1 = rtgtest.reg zero
+  %2 = rtgtest.reg ra
+  %3 = rtgtest.reg sp
+  %4 = rtgtest.reg gp
+  %5 = rtgtest.reg tp
+  %6 = rtgtest.reg t0
+  %7 = rtgtest.reg t1
+  %8 = rtgtest.reg t2
+  %9 = rtgtest.reg s0
+  %10 = rtgtest.reg s1
+  %11 = rtgtest.reg a0
+  %12 = rtgtest.reg a1
+  %13 = rtgtest.reg a2
+  %14 = rtgtest.reg a3
+  %15 = rtgtest.reg a4
+  %16 = rtgtest.reg a5
+  %17 = rtgtest.reg a6
+  %18 = rtgtest.reg a7
+  %19 = rtgtest.reg s2
+  %20 = rtgtest.reg s3
+  %21 = rtgtest.reg s4
+  %22 = rtgtest.reg s5
+  %23 = rtgtest.reg s6
+  %24 = rtgtest.reg s7
+  %25 = rtgtest.reg s8
+  %26 = rtgtest.reg s9
+  %27 = rtgtest.reg s10
+  %28 = rtgtest.reg s11
+  %29 = rtgtest.reg t3
+  %30 = rtgtest.reg t4
+  %31 = rtgtest.reg t5
+  %32 = rtgtest.reg t6
+  %33 = rtgtest.reg Virtual
+}

unittests/Dialect/CMakeLists.txt

@@ -3,4 +3,5 @@ add_subdirectory(FIRRTL)
 add_subdirectory(ESI)
 add_subdirectory(HW)
 add_subdirectory(OM)
+add_subdirectory(RTGTest)
 add_subdirectory(SMT)

unittests/Dialect/RTGTest/CMakeLists.txt

@@ -0,0 +1,8 @@
+add_circt_unittest(CIRCTRTGTestTests
+  RegisterTest.cpp
+)
+
+target_link_libraries(CIRCTRTGTestTests
+  PRIVATE
+  CIRCTRTGTestDialect
+)

unittests/Dialect/RTGTest/RegisterTest.cpp

@@ -0,0 +1,59 @@
+//===- RegisterTest.cpp - RTGTest register unit tests ---------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "circt/Dialect/RTGTest/IR/RTGTestOps.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+using namespace circt;
+using namespace rtgtest;
+
+namespace {
+
+TEST(RegisterInterfaceTest, IntegerRegisters) {
+  MLIRContext context;
+  context.loadDialect<RTGTestDialect>();
+  Location loc(UnknownLoc::get(&context));
+
+  SmallVector<std::tuple<Registers, std::string, unsigned>> regs{
+      {Registers::zero, "zero", 0}, {Registers::ra, "ra", 1},
+      {Registers::sp, "sp", 2},     {Registers::gp, "gp", 3},
+      {Registers::tp, "tp", 4},     {Registers::t0, "t0", 5},
+      {Registers::t1, "t1", 6},     {Registers::t2, "t2", 7},
+      {Registers::s0, "s0", 8},     {Registers::s1, "s1", 9},
+      {Registers::a0, "a0", 10},    {Registers::a1, "a1", 11},
+      {Registers::a2, "a2", 12},    {Registers::a3, "a3", 13},
+      {Registers::a4, "a4", 14},    {Registers::a5, "a5", 15},
+      {Registers::a6, "a6", 16},    {Registers::a7, "a7", 17},
+      {Registers::s2, "s2", 18},    {Registers::s3, "s3", 19},
+      {Registers::s4, "s4", 20},    {Registers::s5, "s5", 21},
+      {Registers::s6, "s6", 22},    {Registers::s7, "s7", 23},
+      {Registers::s8, "s8", 24},    {Registers::s9, "s9", 25},
+      {Registers::s10, "s10", 26},  {Registers::s11, "s11", 27},
+      {Registers::t3, "t3", 28},    {Registers::t4, "t4", 29},
+      {Registers::t5, "t5", 30},    {Registers::t6, "t6", 31}};
+
+  auto moduleOp = ModuleOp::create(loc);
+  OpBuilder builder = OpBuilder::atBlockBegin(moduleOp.getBody());
+  auto regOp = builder.create<RegisterOp>(loc, Registers::Virtual);
+  ASSERT_EQ(regOp.getAllowedRegs(),
+            llvm::BitVector(getMaxEnumValForRegisters(), true));
+  ASSERT_EQ(regOp.getFixedReg(), ~0U);
+
+  for (auto [reg, str, idx] : regs) {
+    regOp.setFixedReg(idx);
+    ASSERT_EQ(regOp.getClassIndex(), idx);
+    ASSERT_EQ(regOp.getClassIndexBinary(), APInt(5, idx));
+    ASSERT_EQ(regOp.getRegisterAssembly(), str);
+    ASSERT_EQ(regOp.getAllowedRegs(),
+              llvm::BitVector(getMaxEnumValForRegisters(), false).set(idx));
+    ASSERT_EQ(regOp.getFixedReg(), idx);
+  }
+}
+
+} // namespace