Block-Jit with block instantiations

Adds support for block-instantiations to the block-jit. This can be used to run blocks generated using multi-proc codegen or other stitchers. There is no current support for Fifos. PiperOrigin-RevId: 649131040
google · Jul 3, 2024 · d67ce25 · d67ce25
1 parent 71a3959
commit d67ce25
Show file tree

Hide file tree

Showing 7 changed files with 254 additions and 68 deletions.
diff --git a/xls/interpreter/block_evaluator_test_base.cc b/xls/interpreter/block_evaluator_test_base.cc
@@ -116,11 +116,6 @@ TEST_P(BlockEvaluatorTest, OutputOnlyBlock) {
 }
 
 TEST_P(BlockEvaluatorTest, StatelessInstantiatedPassthrough) {
-  // TODO(rigge): add instantiation support to block jit and remove this guard.
-  if (!SupportsHierarchicalBlocks()) {
-    GTEST_SKIP();
-    return;
-  }
   auto package = CreatePackage();
   Block* inner;
   {
@@ -147,11 +142,6 @@ TEST_P(BlockEvaluatorTest, StatelessInstantiatedPassthrough) {
 }
 
 TEST_P(BlockEvaluatorTest, PipelinedHierarchicalRotate) {
-  // TODO(rigge): add instantiation support to block jit and remove this guard.
-  if (!SupportsHierarchicalBlocks()) {
-    GTEST_SKIP();
-    return;
-  }
   auto package = CreatePackage();
   auto rot = [&](int64_t n, int64_t r) -> absl::StatusOr<Block*> {
     BlockBuilder b(absl::StrCat(TestName(), "_rot", n, "_", r), package.get());
@@ -304,7 +294,7 @@ MATCHER(NoTryPopValue, "") {
 
 TEST_P(BlockEvaluatorTest, SingleElementFifoInstantiationNoBypassWorks) {
   // TODO(rigge): add instantiation support to block jit and remove this guard.
-  if (!SupportsHierarchicalBlocks()) {
+  if (!SupportsFifos()) {
     GTEST_SKIP();
     return;
   }
@@ -377,7 +367,7 @@ TEST_P(BlockEvaluatorTest, SingleElementFifoInstantiationNoBypassWorks) {
 
 TEST_P(BlockEvaluatorTest, SingleElementFifoInstantiationWithBypassWorks) {
   // TODO(rigge): add instantiation support to block jit and remove this guard.
-  if (!SupportsHierarchicalBlocks()) {
+  if (!SupportsFifos()) {
     GTEST_SKIP();
     return;
   }
@@ -450,7 +440,7 @@ TEST_P(BlockEvaluatorTest, SingleElementFifoInstantiationWithBypassWorks) {
 
 TEST_P(BlockEvaluatorTest, FifoInstantiationNoBypassWorks) {
   // TODO(rigge): add instantiation support to block jit and remove this guard.
-  if (!SupportsHierarchicalBlocks()) {
+  if (!SupportsFifos()) {
     GTEST_SKIP();
     return;
   }
@@ -549,7 +539,7 @@ TEST_P(BlockEvaluatorTest, FifoInstantiationNoBypassWorks) {
 
 TEST_P(BlockEvaluatorTest, FifoInstantiationWithBypassWorks) {
   // TODO(rigge): add instantiation support to block jit and remove this guard.
-  if (!SupportsHierarchicalBlocks()) {
+  if (!SupportsFifos()) {
     GTEST_SKIP();
     return;
   }

diff --git a/xls/interpreter/block_evaluator_test_base.h b/xls/interpreter/block_evaluator_test_base.h
@@ -23,7 +23,7 @@ namespace xls {
 
 struct BlockEvaluatorTestParam {
   const BlockEvaluator* evaluator;
-  bool supports_hierarhical_blocks;
+  bool supports_fifos;
 };
 
 class BlockEvaluatorTest
@@ -32,9 +32,7 @@ class BlockEvaluatorTest
  public:
   const BlockEvaluator& evaluator() { return *GetParam().evaluator; }
 
-  bool SupportsHierarchicalBlocks() {
-    return GetParam().supports_hierarhical_blocks;
-  }
+  bool SupportsFifos() { return GetParam().supports_fifos; }
 };
 }  // namespace xls
 

diff --git a/xls/interpreter/block_interpreter_test.cc b/xls/interpreter/block_interpreter_test.cc
@@ -25,7 +25,8 @@ namespace {
 INSTANTIATE_TEST_SUITE_P(BlockInterpreterTest, BlockEvaluatorTest,
                          testing::Values(BlockEvaluatorTestParam{
                              .evaluator = &kInterpreterBlockEvaluator,
-                             .supports_hierarhical_blocks = true}),
+                             .supports_fifos = true,
+                         }),
                          [](const auto& v) -> std::string {
                            return std::string(v.param.evaluator->name());
                          });

diff --git a/xls/jit/BUILD b/xls/jit/BUILD
@@ -701,11 +701,15 @@ cc_library(
         ":jit_callbacks",
         ":jit_runtime",
         ":orc_jit",
+        "//xls/codegen:block_inlining_pass",
+        "//xls/codegen:codegen_pass",
         "//xls/codegen:module_signature_cc_proto",
         "//xls/common/status:ret_check",
         "//xls/common/status:status_macros",
         "//xls/interpreter:block_evaluator",
         "//xls/ir",
+        "//xls/ir:block_elaboration",
+        "//xls/ir:clone_package",
         "//xls/ir:elaboration",
         "//xls/ir:events",
         "//xls/ir:register",
@@ -717,6 +721,7 @@ cc_library(
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/types:span",
     ],
@@ -732,10 +737,12 @@ cc_test(
         "//xls/interpreter:block_evaluator_test_base",
         "//xls/ir",
         "//xls/ir:bits",
+        "//xls/ir:foreign_function_data_cc_proto",
         "//xls/ir:function_builder",
         "//xls/ir:ir_test_base",
         "//xls/ir:value",
         "//xls/ir:value_view",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest",
     ],