[RTG] Add memory type and operations

frontends/PyRTG/src/CMakeLists.txt

@@ -22,6 +22,7 @@ declare_mlir_python_sources(PyRTGSources
   pyrtg/index.py
   pyrtg/integers.py
   pyrtg/labels.py
+  pyrtg/memories.py
   pyrtg/rtg.py
   pyrtg/sequences.py
   pyrtg/sets.py

frontends/PyRTG/src/pyrtg/__init__.py

@@ -19,3 +19,4 @@
 from .arrays import Array
 from .contexts import CPUCore
 from .control_flow import If, Else, EndIf, For, Foreach
+from .memories import Memory, MemoryBlock

frontends/PyRTG/src/pyrtg/memories.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+from .core import Value
+from .circt import ir
+from .index import index
+from .rtg import rtg
+from .integers import Integer
+from .resources import Immediate
+
+from typing import Union
+
+
+class MemoryBlock(Value):
+
+  def __init__(self, value: ir.Value):
+    """
+    For library internal usage only.
+    """
+
+    self._value = value
+
+  def declare(size: int, base_address: int, address_width: int) -> MemoryBlock:
+    """
+    Declare a new memory block with the specified parameters.
+
+    Args:
+      size: The size of the memory block in bytes.
+      address_width: The width of the memory block addresses in bits.
+    """
+
+    return rtg.MemoryBlockDeclareOp(
+        size,
+        ir.IntegerAttr.get(ir.IntegerType.get_signless(address_width),
+                           base_address))
+
+  def _get_ssa_value(self) -> ir.Value:
+    return self._value
+
+  def get_type(self) -> ir.Type:
+    return self._value.type
+
+  def type(address_width: int) -> ir.Type:
+    return rtg.MemoryBlockType.get(address_width)
+
+
+class Memory(Value):
+
+  def __init__(self, value: ir.Value):
+    """
+    For library internal usage only.
+    """
+
+    self._value = value
+
+  def alloc(mem_block: MemoryBlock, size: Union[Integer, int],
+            align: Union[Integer, int]) -> Memory:
+    """
+    Allocate a new memory from a memory block with the specified parameters.
+
+    Args:
+      size: The size of the memory in bytes.
+      align: The alignment of the memory in bytes.
+    """
+
+    if isinstance(size, int):
+      size = index.ConstantOp(size)
+    if isinstance(align, int):
+      align = index.ConstantOp(align)
+    return rtg.MemoryAllocOp(mem_block, size, align)
+
+  def size(self) -> Integer:
+    """
+    Get the size of the memory in bytes.
+    """
+
+    return rtg.MemorySizeOp(self._value)
+
+  def base_address(self) -> Immediate:
+    """
+    Get the base address of the memory as an immediate matching the memories
+    address width.
+    """
+
+    return rtg.MemoryBaseAddressOp(self._value)
+
+  def _get_ssa_value(self) -> ir.Value:
+    return self._value
+
+  def get_type(self) -> ir.Type:
+    return self._value.type
+
+  def type(address_width: int) -> ir.Type:
+    return rtg.MemoryType.get(address_width)

frontends/PyRTG/src/pyrtg/support.py

@@ -43,6 +43,12 @@ def _FromCirctValue(value: ir.Value) -> Value:
   if isinstance(type, rtg.ImmediateType):
     from .resources import Immediate
     return Immediate(type.width, value)
+  if isinstance(type, rtg.MemoryType):
+    from .memories import Memory
+    return Memory(value)
+  if isinstance(type, rtg.MemoryBlockType):
+    from .memories import MemoryBlock
+    return MemoryBlock(value)
   assert False, "Unsupported value"
 
 

frontends/PyRTG/test/basic.py

@@ -2,7 +2,7 @@
 # RUN: %rtgtool% %s --seed=0 --output-format=elaborated | FileCheck %s --check-prefix=ELABORATED
 # RUN: %rtgtool% %s --seed=0 -o %t --output-format=asm && FileCheck %s --input-file=%t --check-prefix=ASM
 
-from pyrtg import test, sequence, target, entry, rtg, Label, Set, Integer, Bag, rtgtest, Immediate, IntegerRegister, Array, Bool
+from pyrtg import test, sequence, target, entry, rtg, Label, Set, Integer, Bag, rtgtest, Immediate, IntegerRegister, Array, Bool, MemoryBlock, Memory
 
 # MLIR-LABEL: rtg.target @Tgt0 : !rtg.dict<entry0: !rtg.set<index>>
 # MLIR-NEXT: [[C0:%.+]] = index.constant 0
@@ -39,6 +39,19 @@ def entry1():
     return Label.declare("l0")
 
 
+# MLIR-LABEL: rtg.target @Tgt2
+# MLIR-NEXT: [[V0:%.+]] = rtg.isa.memoryblock_declare 64, 0 : i32
+# MLIR-NEXT: rtg.yield [[V0]] : !rtg.isa.memoryblock<32>
+
+
+@target
+class Tgt2:
+
+  @entry
+  def mem_blk_0():
+    return MemoryBlock.declare(size=64, base_address=0, address_width=32)
+
+
 # MLIR-LABEL: rtg.target @Tgt4
 # MLIR-NEXT: [[IDX12:%.+]] = index.constant 12
 # MLIR-NEXT: [[IDX11:%.+]] = index.constant 11
@@ -344,6 +357,25 @@ def test4_integer_to_immediate():
                Immediate(12, Integer(2)))
 
 
+# MLIR-LABEL: rtg.test @test6_memories
+# MLIR-NEXT: [[REG:%.+]] = rtg.fixed_reg #rtgtest.t0 : !rtgtest.ireg
+# MLIR-NEXT: [[IDX8:%.+]] = index.constant 8
+# MLIR-NEXT: [[IDX4:%.+]] = index.constant 4
+# MLIR-NEXT: [[MEM:%.+]] = rtg.isa.memory_alloc %mem_blk, [[IDX8]], [[IDX4]] : !rtg.isa.memoryblock<32>
+# MLIR-NEXT: [[SIZE:%.+]] = rtg.isa.memory_size [[MEM]] : !rtg.isa.memory<32>
+# MLIR-NEXT: [[IMM:%.+]] = rtg.isa.int_to_immediate [[SIZE]] : !rtg.isa.immediate<32>
+# MLIR-NEXT: rtgtest.rv32i.auipc [[REG]], [[IMM]] : !rtg.isa.immediate<32>
+# MLIR-NEXT: [[BASE:%.+]] = rtg.isa.memory_base_address [[MEM]] : !rtg.isa.memory<32>
+# MLIR-NEXT: rtgtest.rv32i.auipc [[REG]], [[BASE]] : !rtg.isa.immediate<32>
+
+
+@test(("mem_blk", MemoryBlock.type(32)))
+def test6_memories(mem_blk):
+  mem = Memory.alloc(mem_blk, size=8, align=4)
+  rtgtest.AUIPC(IntegerRegister.t0(), Immediate(32, mem.size()))
+  rtgtest.AUIPC(IntegerRegister.t0(), mem.base_address())
+
+
 # MLIR-LABEL: rtg.test @test7_bools
 # MLIR: index.bool.constant false
 # MLIR: index.bool.constant true

include/circt-c/Dialect/RTG.h

@@ -89,6 +89,16 @@ MLIR_CAPI_EXPORTED MlirType rtgImmediateTypeGet(MlirContext ctx,
 /// Returns the width of the RTG immediate type.
 MLIR_CAPI_EXPORTED uint32_t rtgImmediateTypeGetWidth(MlirType type);
 
+/// If the type is an RTG memory.
+MLIR_CAPI_EXPORTED bool rtgTypeIsAMemory(MlirType type);
+
+/// Creates an RTG memory type in the context.
+MLIR_CAPI_EXPORTED MlirType rtgMemoryTypeGet(MlirContext ctx,
+                                             uint32_t addressWidth);
+
+/// Returns the address with of an RTG memory type.
+MLIR_CAPI_EXPORTED uint32_t rtgMemoryTypeGetAddressWidth(MlirType type);
+
 /// If the type is an RTG memory block.
 MLIR_CAPI_EXPORTED bool rtgTypeIsAMemoryBlock(MlirType type);
 

include/circt/Dialect/RTG/IR/RTGOps.td

@@ -704,3 +704,55 @@ def MemoryBlockDeclareOp : RTGISAOp<"memoryblock_declare", [
 
   let assemblyFormat = "$size `,` $baseAddress attr-dict";
 }
+
+//===- ISA Memory Handling Operations -------------------------------------===//
+
+def MemoryAllocOp : RTGISAOp<"memory_alloc", [
+  TypesMatchWith<"memory must have the same address width as the memory block",
+                 "memoryBlock", "result",
+                 "MemoryType::get($_ctxt, " #
+                 "cast<MemoryBlockType>($_self).getAddressWidth())">,
+]> {
+  let summary = "allocate a memory with the provided properties";
+  let description = [{
+    This operation declares a memory to be allocated with the provided
+    properties. It is only allowed to declare new memories in the `rtg.target`
+    operations and must be passed as argument to the `rtg.test`.
+  }];
+
+  let arguments = (ins MemoryBlockType:$memoryBlock,
+                       Index:$size,
+                       Index:$alignment);
+
+  let results = (outs MemoryType:$result);
+
+  let assemblyFormat = [{
+    $memoryBlock `,` $size `,` $alignment
+    `:` qualified(type($memoryBlock)) attr-dict
+  }];
+}
+
+def MemoryBaseAddressOp : RTGISAOp<"memory_base_address", [
+  Pure,
+  DeclareOpInterfaceMethods<InferTypeOpInterface>,
+]> {
+  let summary = "get the memory base address as an immediate";
+  let description = [{
+    This operation returns the base address of the given memory. The bit-width
+    of the returned immediate must match the address width of the given memory.
+  }];
+
+  let arguments = (ins MemoryType:$memory);
+  let results = (outs ImmediateType:$result);
+
+  let assemblyFormat = "$memory `:` qualified(type($memory)) attr-dict";
+}
+
+def MemorySizeOp : RTGISAOp<"memory_size", [Pure]> {
+  let summary = "get the size of the memory in bytes";
+
+  let arguments = (ins MemoryType:$memory);
+  let results = (outs Index:$result);
+
+  let assemblyFormat = "$memory `:` qualified(type($memory)) attr-dict";
+}

include/circt/Dialect/RTG/IR/RTGTypes.td

@@ -175,6 +175,18 @@ class ImmediateOfWidth<int width> : Type<
   "a " # width # "-bit immediate">,
   BuildableType<"::circt::rtg::ImmediateType::get($_builder.getContext(), " # width # ")">;
 
+def MemoryType : RTGISATypeDef<"Memory"> {
+  let summary = "handle to a memory";
+  let description = [{
+    This type is used to represent memory resources that are allocated from
+    memory blocks and can be accessed and manipulated by payload dialect
+    operations.
+  }];
+
+  let parameters = (ins "uint32_t":$addressWidth);
+  let assemblyFormat = "`<` $addressWidth `>`";
+}
+
 def MemoryBlockType : RTGISATypeDef<"MemoryBlock"> {
   let summary = "handle to a memory block";
   let description = [{

include/circt/Dialect/RTG/IR/RTGVisitors.h

@@ -57,6 +57,8 @@ class RTGOpVisitor {
             ArrayCreateOp, ArrayExtractOp, ArrayInjectOp, ArraySizeOp,
             // Immediates
             IntToImmediateOp,
+            // Memories
+            MemoryAllocOp, MemoryBaseAddressOp, MemorySizeOp,
             // Memory Blocks
             MemoryBlockDeclareOp>([&](auto expr) -> ResultType {
           return thisCast->visitOp(expr, args...);
@@ -124,6 +126,9 @@ class RTGOpVisitor {
   HANDLE(VirtualRegisterOp, Unhandled);
   HANDLE(IntToImmediateOp, Unhandled);
   HANDLE(MemoryBlockDeclareOp, Unhandled);
+  HANDLE(MemoryAllocOp, Unhandled);
+  HANDLE(MemoryBaseAddressOp, Unhandled);
+  HANDLE(MemorySizeOp, Unhandled);
 #undef HANDLE
 };
 

integration_test/Bindings/Python/dialects/rtg.py

@@ -225,6 +225,12 @@
   # CHECK: !rtg.isa.memoryblock<32>
   print(memoryblock_type)
 
+  memoryTy = rtg.MemoryType.get(32)
+  # CHECK: address_width=32
+  print(f'address_width={memoryTy.address_width}')
+  # CHECK: !rtg.isa.memory<32>
+  print(memoryTy)
+
 with Context() as ctx, Location.unknown():
   circt.register_dialects(ctx)
   indexTy = IndexType.get()

lib/Bindings/Python/RTGModule.cpp

@@ -144,6 +144,17 @@ void circt::python::populateDialectRTGSubmodule(nb::module_ &m) {
         return rtgMemoryBlockTypeGetAddressWidth(self);
       });
 
+  mlir_type_subclass(m, "MemoryType", rtgTypeIsAMemory)
+      .def_classmethod(
+          "get",
+          [](nb::object cls, uint32_t addressWidth, MlirContext ctxt) {
+            return cls(rtgMemoryTypeGet(ctxt, addressWidth));
+          },
+          nb::arg("self"), nb::arg("address_width"), nb::arg("ctxt") = nullptr)
+      .def_property_readonly("address_width", [](MlirType self) {
+        return rtgMemoryTypeGetAddressWidth(self);
+      });
+
   //===--------------------------------------------------------------------===//
   // Attributes
   //===--------------------------------------------------------------------===//

lib/Bindings/Python/support.py

@@ -161,6 +161,14 @@ def type_to_pytype(t) -> ir.Type:
     return rtg.ArrayType(t)
   except ValueError:
     pass
+  try:
+    return rtg.MemoryType(t)
+  except ValueError:
+    pass
+  try:
+    return rtg.MemoryBlockType(t)
+  except ValueError:
+    pass
   try:
     return rtgtest.IntegerRegisterType(t)
   except ValueError:

lib/CAPI/Dialect/RTG.cpp

@@ -144,6 +144,19 @@ uint32_t rtgImmediateTypeGetWidth(MlirType type) {
   return cast<ImmediateType>(unwrap(type)).getWidth();
 }
 
+// MemoryType
+//===----------------------------------------------------------------------===//
+
+bool rtgTypeIsAMemory(MlirType type) { return isa<MemoryType>(unwrap(type)); }
+
+MlirType rtgMemoryTypeGet(MlirContext ctxt, uint32_t addressWidth) {
+  return wrap(MemoryType::get(unwrap(ctxt), addressWidth));
+}
+
+uint32_t rtgMemoryTypeGetAddressWidth(MlirType type) {
+  return cast<MemoryType>(unwrap(type)).getAddressWidth();
+}
+
 // MemoryBlockType
 //===----------------------------------------------------------------------===//
 

lib/Dialect/RTG/IR/RTGOps.cpp

@@ -591,6 +591,24 @@ void ArrayCreateOp::print(OpAsmPrinter &p) {
   p.printOptionalAttrDict((*this)->getAttrs(), {});
 }
 
+//===----------------------------------------------------------------------===//
+// MemoryBaseAddressOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult MemoryBaseAddressOp::inferReturnTypes(
+    MLIRContext *context, std::optional<Location> loc, ValueRange operands,
+    DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  if (operands.empty())
+    return failure();
+  auto memTy = dyn_cast<MemoryType>(operands[0].getType());
+  if (!memTy)
+    return failure();
+  inferredReturnTypes.push_back(
+      ImmediateType::get(context, memTy.getAddressWidth()));
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen generated logic.
 //===----------------------------------------------------------------------===//