Merge pull request #1871 from xlsynth:cdleary/2025-01-19-bytecode-cac…

…he-import-data PiperOrigin-RevId: 718415111
google · Jan 22, 2025 · 456ee42 · 456ee42
2 parents 2a74d76 + 38f4282
commit 456ee42
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 16 deletions.
diff --git a/xls/dslx/bytecode/bytecode_cache.cc b/xls/dslx/bytecode/bytecode_cache.cc
@@ -31,18 +31,15 @@
 
 namespace xls::dslx {
 
-BytecodeCache::BytecodeCache(ImportData* import_data)
-    : import_data_(import_data) {}
-
 absl::StatusOr<BytecodeFunction*> BytecodeCache::GetOrCreateBytecodeFunction(
-    const Function& f, const TypeInfo* type_info,
+    ImportData& import_data, const Function& f, const TypeInfo* type_info,
     const std::optional<ParametricEnv>& caller_bindings) {
   XLS_RET_CHECK(type_info != nullptr);
   Key key = std::make_tuple(&f, type_info, caller_bindings);
   if (!cache_.contains(key)) {
     XLS_ASSIGN_OR_RETURN(
         std::unique_ptr<BytecodeFunction> bf,
-        BytecodeEmitter::Emit(import_data_, type_info, f, caller_bindings));
+        BytecodeEmitter::Emit(&import_data, type_info, f, caller_bindings));
     cache_.emplace(key, std::move(bf));
   }
 

diff --git a/xls/dslx/bytecode/bytecode_cache.h b/xls/dslx/bytecode/bytecode_cache.h
@@ -31,17 +31,16 @@ namespace xls::dslx {
 
 class BytecodeCache : public BytecodeCacheInterface {
  public:
-  explicit BytecodeCache(ImportData* import_data);
+  BytecodeCache() = default;
 
   absl::StatusOr<BytecodeFunction*> GetOrCreateBytecodeFunction(
-      const Function& f, const TypeInfo* type_info,
+      ImportData& import_data, const Function& f, const TypeInfo* type_info,
       const std::optional<ParametricEnv>& caller_bindings) override;
 
  private:
   using Key = std::tuple<const Function*, const TypeInfo*,
                          std::optional<ParametricEnv>>;
 
-  ImportData* import_data_;
   absl::flat_hash_map<Key, std::unique_ptr<BytecodeFunction>> cache_;
 };
 

diff --git a/xls/dslx/bytecode/bytecode_cache_interface.h b/xls/dslx/bytecode/bytecode_cache_interface.h
@@ -24,6 +24,9 @@
 
 namespace xls::dslx {
 
+// Forward decl.
+class ImportData;
+
 // Defines the interface a type must provide in order to serve as a bytecode
 // cache. In practice, this type exists to avoid attaching too many concrete
 // dependencies onto ImportData, which is the primary cache owner.
@@ -35,7 +38,7 @@ class BytecodeCacheInterface {
   // constants are held inside the given TypeInfo - different instances of a
   // parametric function will have different TypeInfos associated with them.
   virtual absl::StatusOr<BytecodeFunction*> GetOrCreateBytecodeFunction(
-      const Function& f, const TypeInfo* type_info,
+      ImportData& import_data, const Function& f, const TypeInfo* type_info,
       const std::optional<ParametricEnv>& caller_bindings) = 0;
 };
 

diff --git a/xls/dslx/bytecode/bytecode_interpreter.cc b/xls/dslx/bytecode/bytecode_interpreter.cc
@@ -580,7 +580,7 @@ absl::StatusOr<BytecodeFunction*> BytecodeInterpreter::GetBytecodeFn(
                          import_data_->GetRootTypeInfo(f.owner()));
   }
 
-  return cache->GetOrCreateBytecodeFunction(f, callee_type_info,
+  return cache->GetOrCreateBytecodeFunction(*import_data(), f, callee_type_info,
                                             callee_bindings);
 }
 

diff --git a/xls/dslx/create_import_data.cc b/xls/dslx/create_import_data.cc
@@ -34,7 +34,7 @@ ImportData CreateImportData(
     WarningKindSet warnings, std::unique_ptr<VirtualizableFilesystem> vfs) {
   ImportData import_data(stdlib_path, additional_search_paths, warnings,
                          std::move(vfs));
-  import_data.SetBytecodeCache(std::make_unique<BytecodeCache>(&import_data));
+  import_data.SetBytecodeCache(std::make_unique<BytecodeCache>());
   return import_data;
 }
 
@@ -46,7 +46,7 @@ ImportData CreateImportDataForTest(
   absl::Span<const std::filesystem::path> additional_search_paths = {};
   ImportData import_data(xls::kDefaultDslxStdlibPath, additional_search_paths,
                          kDefaultWarningsSet, std::move(vfs));
-  import_data.SetBytecodeCache(std::make_unique<BytecodeCache>(&import_data));
+  import_data.SetBytecodeCache(std::make_unique<BytecodeCache>());
   return import_data;
 }
 
@@ -55,8 +55,7 @@ std::unique_ptr<ImportData> CreateImportDataPtrForTest() {
       new ImportData(xls::kDefaultDslxStdlibPath,
                      /*additional_search_paths=*/{}, kDefaultWarningsSet,
                      std::make_unique<RealFilesystem>()));
-  import_data->SetBytecodeCache(
-      std::make_unique<BytecodeCache>(import_data.get()));
+  import_data->SetBytecodeCache(std::make_unique<BytecodeCache>());
   return import_data;
 }
 

diff --git a/xls/dslx/run_routines/run_routines.cc b/xls/dslx/run_routines/run_routines.cc
@@ -132,7 +132,7 @@ void HandleError(TestResultData& result, const absl::Status& status,
 absl::Status RunDslxTestFunction(ImportData* import_data, TypeInfo* type_info,
                                  Module* module, TestFunction* tf,
                                  const BytecodeInterpreterOptions& options) {
-  auto cache = std::make_unique<BytecodeCache>(import_data);
+  auto cache = std::make_unique<BytecodeCache>();
   import_data->SetBytecodeCache(std::move(cache));
   XLS_ASSIGN_OR_RETURN(
       std::unique_ptr<BytecodeFunction> bf,
@@ -149,7 +149,7 @@ absl::Status RunDslxTestFunction(ImportData* import_data, TypeInfo* type_info,
 absl::Status RunDslxTestProc(ImportData* import_data, TypeInfo* type_info,
                              Module* module, TestProc* tp,
                              const BytecodeInterpreterOptions& options) {
-  auto cache = std::make_unique<BytecodeCache>(import_data);
+  auto cache = std::make_unique<BytecodeCache>();
   import_data->SetBytecodeCache(std::move(cache));
 
   XLS_ASSIGN_OR_RETURN(TypeInfo * ti,