diff --git a/frontends/PyRTG/src/CMakeLists.txt b/frontends/PyRTG/src/CMakeLists.txt index a27cc56394ea..990d83684a69 100644 --- a/frontends/PyRTG/src/CMakeLists.txt +++ b/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 diff --git a/frontends/PyRTG/src/pyrtg/__init__.py b/frontends/PyRTG/src/pyrtg/__init__.py index f18dcf836d6e..b03266a2cbd0 100644 --- a/frontends/PyRTG/src/pyrtg/__init__.py +++ b/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 diff --git a/frontends/PyRTG/src/pyrtg/memories.py b/frontends/PyRTG/src/pyrtg/memories.py new file mode 100644 index 000000000000..d5f632d088b3 --- /dev/null +++ b/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) diff --git a/frontends/PyRTG/src/pyrtg/support.py b/frontends/PyRTG/src/pyrtg/support.py index 473cda671d61..0748a29deb5c 100644 --- a/frontends/PyRTG/src/pyrtg/support.py +++ b/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" diff --git a/frontends/PyRTG/test/basic.py b/frontends/PyRTG/test/basic.py index 9abba3867e51..a6db53282b16 100644 --- a/frontends/PyRTG/test/basic.py +++ b/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> # 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 diff --git a/include/circt-c/Dialect/RTG.h b/include/circt-c/Dialect/RTG.h index db6c36a85491..e267e9647b07 100644 --- a/include/circt-c/Dialect/RTG.h +++ b/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); diff --git a/include/circt/Dialect/RTG/IR/RTGOps.td b/include/circt/Dialect/RTG/IR/RTGOps.td index 8089c2430ad9..50afae64283a 100644 --- a/include/circt/Dialect/RTG/IR/RTGOps.td +++ b/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($_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, +]> { + 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"; +} diff --git a/include/circt/Dialect/RTG/IR/RTGTypes.td b/include/circt/Dialect/RTG/IR/RTGTypes.td index af5837db1380..355461a699cc 100644 --- a/include/circt/Dialect/RTG/IR/RTGTypes.td +++ b/include/circt/Dialect/RTG/IR/RTGTypes.td @@ -175,6 +175,18 @@ class ImmediateOfWidth : 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 = [{ diff --git a/include/circt/Dialect/RTG/IR/RTGVisitors.h b/include/circt/Dialect/RTG/IR/RTGVisitors.h index 7881aff17501..5f3cadbfef6a 100644 --- a/include/circt/Dialect/RTG/IR/RTGVisitors.h +++ b/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 }; diff --git a/integration_test/Bindings/Python/dialects/rtg.py b/integration_test/Bindings/Python/dialects/rtg.py index d40f96d3fc47..986836653b44 100644 --- a/integration_test/Bindings/Python/dialects/rtg.py +++ b/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() diff --git a/lib/Bindings/Python/RTGModule.cpp b/lib/Bindings/Python/RTGModule.cpp index abaf683b8ba2..a805eb73807d 100644 --- a/lib/Bindings/Python/RTGModule.cpp +++ b/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 //===--------------------------------------------------------------------===// diff --git a/lib/Bindings/Python/support.py b/lib/Bindings/Python/support.py index c9abd86b2d5e..12e4c34d5ae1 100644 --- a/lib/Bindings/Python/support.py +++ b/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: diff --git a/lib/CAPI/Dialect/RTG.cpp b/lib/CAPI/Dialect/RTG.cpp index 3fc9cb6e6a66..1d1b9b25d865 100644 --- a/lib/CAPI/Dialect/RTG.cpp +++ b/lib/CAPI/Dialect/RTG.cpp @@ -144,6 +144,19 @@ uint32_t rtgImmediateTypeGetWidth(MlirType type) { return cast(unwrap(type)).getWidth(); } +// MemoryType +//===----------------------------------------------------------------------===// + +bool rtgTypeIsAMemory(MlirType type) { return isa(unwrap(type)); } + +MlirType rtgMemoryTypeGet(MlirContext ctxt, uint32_t addressWidth) { + return wrap(MemoryType::get(unwrap(ctxt), addressWidth)); +} + +uint32_t rtgMemoryTypeGetAddressWidth(MlirType type) { + return cast(unwrap(type)).getAddressWidth(); +} + // MemoryBlockType //===----------------------------------------------------------------------===// diff --git a/lib/Dialect/RTG/IR/RTGOps.cpp b/lib/Dialect/RTG/IR/RTGOps.cpp index 62c97d4d3e97..ff47b9894c79 100644 --- a/lib/Dialect/RTG/IR/RTGOps.cpp +++ b/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 loc, ValueRange operands, + DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + if (operands.empty()) + return failure(); + auto memTy = dyn_cast(operands[0].getType()); + if (!memTy) + return failure(); + inferredReturnTypes.push_back( + ImmediateType::get(context, memTy.getAddressWidth())); + return success(); +} + //===----------------------------------------------------------------------===// // TableGen generated logic. //===----------------------------------------------------------------------===// diff --git a/test/CAPI/rtg.c b/test/CAPI/rtg.c index f1bd2365463e..8e44149942eb 100644 --- a/test/CAPI/rtg.c +++ b/test/CAPI/rtg.c @@ -173,7 +173,7 @@ static void testImmediate(MlirContext ctx) { fprintf(stderr, "value=%llu\n", rtgImmediateAttrGetValue(immediateAttr)); } -static void testMemoryBlock(MlirContext ctx) { +static void testMemories(MlirContext ctx) { MlirType memoryBlockTy = rtgMemoryBlockTypeGet(ctx, 32); // CHECK: is_memoryblock @@ -184,6 +184,15 @@ static void testMemoryBlock(MlirContext ctx) { // CHECK: address_width=32 fprintf(stderr, "address_width=%u\n", rtgMemoryBlockTypeGetAddressWidth(memoryBlockTy)); + + MlirType memoryTy = rtgMemoryTypeGet(ctx, 32); + // CHECK: is_memory + fprintf(stderr, + rtgTypeIsAMemory(memoryTy) ? "is_memory\n" : "isnot_memory\n"); + // CHECK: addressWidth=32 + fprintf(stderr, "addressWidth=%u\n", rtgMemoryTypeGetAddressWidth(memoryTy)); + // CHECK: !rtg.isa.memory<32> + mlirTypeDump(memoryTy); } int main(int argc, char **argv) { @@ -202,7 +211,7 @@ int main(int argc, char **argv) { testLabelVisibilityAttr(ctx); testDefaultContextAttr(ctx); testImmediate(ctx); - testMemoryBlock(ctx); // Add the new test + testMemories(ctx); mlirContextDestroy(ctx); diff --git a/test/Dialect/RTG/IR/basic.mlir b/test/Dialect/RTG/IR/basic.mlir index 5723d44f9c2a..81c8476f3378 100644 --- a/test/Dialect/RTG/IR/basic.mlir +++ b/test/Dialect/RTG/IR/basic.mlir @@ -154,12 +154,21 @@ rtg.test @interleaveSequences(seq0 = %seq0: !rtg.randomized_sequence, seq1 = %se rtg.interleave_sequences %seq0, %seq1 batch 4 {rtg.some_attr} } -// CHECK-LABEL: @memoryBlocks : !rtg.dict> -rtg.target @memoryBlocks : !rtg.dict> { +// CHECK-LABEL: @memoryBlocks +rtg.target @memoryBlocks : !rtg.dict, mem_block: !rtg.isa.memoryblock<32>, mem_size: index> { // CHECK: rtg.isa.memoryblock_declare 8, 0 : i32 %0 = rtg.isa.memoryblock_declare 8, 0 : i32 + + // CHECK: [[IDX8:%.+]] = index.constant 8 + // CHECK: [[V1:%.+]] = rtg.isa.memory_alloc %0, [[IDX8]], [[IDX8]] : !rtg.isa.memoryblock<32> + // CHECK: [[V2:%.+]] = rtg.isa.memory_base_address [[V1]] : !rtg.isa.memory<32> + // CHECK: [[V3:%.+]] = rtg.isa.memory_size [[V1]] : !rtg.isa.memory<32> + %idx8 = index.constant 8 + %1 = rtg.isa.memory_alloc %0, %idx8, %idx8 : !rtg.isa.memoryblock<32> + %2 = rtg.isa.memory_base_address %1 : !rtg.isa.memory<32> + %3 = rtg.isa.memory_size %1 : !rtg.isa.memory<32> - rtg.yield %0 : !rtg.isa.memoryblock<32> + rtg.yield %2, %0, %3 : !rtg.isa.immediate<32>, !rtg.isa.memoryblock<32>, index } // CHECK-LABEL: rtg.test @arrays