[RTG] Add MemoryAllocation pass to pipeline (#8397)

frontends/PyRTG/src/rtgtool/rtgtool.py

@@ -84,6 +84,10 @@ def parse_args() -> argparse.Namespace:
       default="-",
       help=
       "Output Path (directory if 'split-output=True' or a filepath otherwise)")
+  parser.add_argument("--memories-as-immediates",
+                      type=bool,
+                      default=True,
+                      help="Lower memories as immediates")
 
   args = parser.parse_args()
 
@@ -147,7 +151,8 @@ def compile(mlir_module: ir.Module, args: argparse.Namespace) -> None:
       verbose_pass_execution=args.verbose_pass_executions,
       output_format=to_output_format(args.output_format),
       output_path=args.output_path,
-      split_output=False)
+      split_output=False,
+      memories_as_immediates=args.memories_as_immediates)
   rtgtool_support.populate_randomizer_pipeline(pm, options)
   pm.run(mlir_module.operation)
 

frontends/PyRTG/test/test_memories.py

@@ -0,0 +1,28 @@
+# RUN: %rtgtool% %s --seed=0 --output-format=asm | FileCheck %s
+
+from pyrtg import test, MemoryBlock, Memory, target, entry, rtgtest, IntegerRegister
+
+
+@target
+class Target:
+
+  @entry
+  def mem_blk():
+    return MemoryBlock.declare(0, 31, 32)
+
+
+# CHECK-LABEL: Begin of test0
+# CHECK-EMPTY:
+# CHECK-NEXT: la t0, 0
+# CHECK-NEXT: la t1, 8
+# CHECK-EMPTY:
+# CHECK-NEXT: End of test0
+
+
+@test(("mem_blk", MemoryBlock.type(32)))
+def test0(mem_blk):
+  mem0 = Memory.alloc(mem_blk, 8, 4)
+  mem1 = Memory.alloc(mem_blk, 8, 4)
+
+  rtgtest.LA(IntegerRegister.t0(), mem0)
+  rtgtest.LA(IntegerRegister.t1(), mem1)

include/circt-c/RtgTool.h

@@ -71,6 +71,10 @@ circtRtgToolOptionsSetSplitOutput(CirctRtgToolOptions options, bool enable);
 MLIR_CAPI_EXPORTED void
 circtRtgToolOptionsSetOutputPath(CirctRtgToolOptions options, const char *path);
 
+MLIR_CAPI_EXPORTED void
+circtRtgToolOptionsSetMemoriesAsImmediates(CirctRtgToolOptions options,
+                                           bool enable);
+
 //===----------------------------------------------------------------------===//
 // Pipeline Population API.
 //===----------------------------------------------------------------------===//

include/circt/Tools/rtgtool/RtgToolOptions.h

@@ -85,6 +85,12 @@ public:
   }
   std::string getOutputPath() const { return outputPath; }
 
+  RtgToolOptions &setMemoriesAsImmediates(bool enable) {
+    memoriesAsImmediates = enable;
+    return *this;
+  }
+  bool getMemoriesAsImmediates() const { return memoriesAsImmediates; }
+
 private:
   OutputFormat outputFormat = OutputFormat::ElaboratedMLIR;
   unsigned seed;
@@ -94,6 +100,7 @@ private:
   std::string unsupportedInstructionsFile;
   bool splitOutput = false;
   std::string outputPath;
+  bool memoriesAsImmediates = true;
 };
 
 /// Populates the passes necessary to lower IR with RTG randomization operations

lib/Bindings/Python/RTGToolModule.cpp

@@ -25,7 +25,8 @@ public:
                    bool verifyPasses, bool verbosePassExecution,
                    const std::vector<std::string> &unsupportedInstructions,
                    const std::string &unsupportedInstructionsFile,
-                   bool splitOutput, const std::string &outputPath)
+                   bool splitOutput, const std::string &outputPath,
+                   bool memoriesAsImmediates)
       : options(circtRtgToolOptionsCreateDefault(seed)) {
     setOutputFormat(outputFormat);
     setVerifyPasses(verifyPasses);
@@ -34,6 +35,7 @@ public:
     setUnsupportedInstructionsFile(unsupportedInstructionsFile);
     setSplitOutput(splitOutput);
     setOutputPath(outputPath);
+    setMemoriesAsImmediates(memoriesAsImmediates);
   }
   ~PyRtgToolOptions() { circtRtgToolOptionsDestroy(options); }
 
@@ -79,6 +81,10 @@ public:
     circtRtgToolOptionsSetOutputPath(options, path.c_str());
   }
 
+  void setMemoriesAsImmediates(bool enable) {
+    circtRtgToolOptionsSetMemoriesAsImmediates(options, enable);
+  }
+
 private:
   CirctRtgToolOptions options;
 };
@@ -96,14 +102,15 @@ void circt::python::populateDialectRTGToolSubmodule(nb::module_ &m) {
   nb::class_<PyRtgToolOptions>(m, "Options")
       .def(nb::init<unsigned, CirctRtgToolOutputFormat, bool, bool,
                     const std::vector<std::string> &, const std::string &, bool,
-                    const std::string &>(),
+                    const std::string &, bool>(),
            nb::arg("seed"),
            nb::arg("output_format") = CIRCT_RTGTOOL_OUTPUT_FORMAT_ASM,
            nb::arg("verify_passes") = true,
            nb::arg("verbose_pass_execution") = false,
            nb::arg("unsupported_instructions") = std::vector<const char *>(),
            nb::arg("unsupported_instructions_file") = "",
-           nb::arg("split_output") = false, nb::arg("output_path") = "")
+           nb::arg("split_output") = false, nb::arg("output_path") = "",
+           nb::arg("memories_as_immediates") = true)
       .def("set_output_format", &PyRtgToolOptions::setOutputFormat,
            "Specify the output format of the tool", nb::arg("format"))
       .def("set_seed", &PyRtgToolOptions::setSeed,
@@ -131,11 +138,15 @@ void circt::python::populateDialectRTGToolSubmodule(nb::module_ &m) {
            nb::arg("filename"))
       .def("set_split_output", &PyRtgToolOptions::setSplitOutput,
            "Determines whether each test should be emitted to a separate file.",
-           nb::arg("filename"))
+           nb::arg("enable"))
       .def("output_path", &PyRtgToolOptions::setOutputPath,
            "The path of a file to be emitted to or a directory if "
            "'split_output' is enabled.",
-           nb::arg("filename"));
+           nb::arg("filename"))
+      .def("set_memories_as_immediates",
+           &PyRtgToolOptions::setMemoriesAsImmediates,
+           "Determines whether memories are lowered to immediates or labels.",
+           nb::arg("enable"));
 
   m.def("populate_randomizer_pipeline",
         [](MlirPassManager pm, const PyRtgToolOptions &options) {

lib/CAPI/RtgTool/RtgTool.cpp

@@ -93,6 +93,11 @@ void circtRtgToolOptionsSetOutputPath(CirctRtgToolOptions options,
   unwrap(options)->setOutputPath(std::string(path));
 }
 
+void circtRtgToolOptionsSetMemoriesAsImmediates(CirctRtgToolOptions options,
+                                                bool enable) {
+  unwrap(options)->setMemoriesAsImmediates(enable);
+}
+
 //===----------------------------------------------------------------------===//
 // Pipeline Population API.
 //===----------------------------------------------------------------------===//

lib/Tools/rtgtool/RtgToolOptions.cpp

@@ -45,6 +45,11 @@ void rtg::populateRandomizerPipeline(mlir::PassManager &pm,
     pm.addPass(rtg::createElaborationPass(passOptions));
   }
   pm.addPass(rtg::createInlineSequencesPass());
+  {
+    MemoryAllocationPassOptions passOptions;
+    passOptions.useImmediates = options.getMemoriesAsImmediates();
+    pm.addNestedPass<rtg::TestOp>(rtg::createMemoryAllocationPass());
+  }
   pm.addPass(rtg::createLowerUniqueLabelsPass());
   pm.addNestedPass<rtg::TestOp>(rtg::createLinearScanRegisterAllocationPass());
   {