Store AotPackageEntrypointsProto within JIT wrappers.

This change modifies BlockBaseJitWrapper, FunctionBaseJitWrapper, and ProcBaseJitWrapper to store the AotPackageEntrypointsProto used during their creation from AOT data. A new accessor method is added to retrieve the stored proto.

PiperOrigin-RevId: 860316111

Author: allight

xls/jit/block_base_jit_wrapper.h

@@ -157,11 +157,19 @@ class BaseBlockJitWrapper {
     XLS_ASSIGN_OR_RETURN(
         auto jit, BlockJit::CreateFromAot(proto.entrypoint(0),
                                           proto.data_layout(), entrypoint));
-    return std::unique_ptr<RealType>(new RealType(std ::move(jit)));
+
+    auto res = std::unique_ptr<RealType>(new RealType(std ::move(jit)));
+    res->aot_entrypoints_proto_ = std::move(proto);
+    return res;
+  }
+
+  const AotPackageEntrypointsProto& aot_entrypoints_proto() const {
+    return aot_entrypoints_proto_;
   }
 
  private:
   std::unique_ptr<BlockJit> jit_;
+  AotPackageEntrypointsProto aot_entrypoints_proto_;
 };
 
 }  // namespace xls

xls/jit/function_base_jit_wrapper.h

@@ -74,11 +74,18 @@ class BaseFunctionJitWrapper {
     XLS_ASSIGN_OR_RETURN(auto jit, FunctionJit::CreateFromAot(
                                        proto.entrypoint(0), proto.data_layout(),
                                        unpacked_entrypoint, packed_entrypoint));
-    return std::unique_ptr<RealType>(
+
+    auto res = std::unique_ptr<RealType>(
         new RealType(std::move(jit),
                      MatchesImplicitToken(entrypoint_proto.function_metadata()
                                               .function_interface()
                                               .parameters())));
+    res->aot_entrypoints_proto_ = std::move(proto);
+    return res;
+  }
+
+  const AotPackageEntrypointsProto& aot_entrypoints_proto() const {
+    return aot_entrypoints_proto_;
   }
 
   // Matches the parameter signature for an "implicit token/activation taking"
@@ -137,6 +144,7 @@ class BaseFunctionJitWrapper {
 
   std::unique_ptr<FunctionJit> jit_;
   const bool needs_fake_token_;
+  AotPackageEntrypointsProto aot_entrypoints_proto_;
 };
 
 }  // namespace xls

xls/jit/proc_base_jit_wrapper.h

@@ -98,8 +98,14 @@ class BaseProcJitWrapper {
     XLS_ASSIGN_OR_RETURN(auto* man, aot->GetJitChannelQueueManager());
     JitRuntime& runtime = man->runtime();
 
-    return std::unique_ptr<RealType>(
+    auto result = std::unique_ptr<RealType>(
         new RealType(std::move(package), proc, std::move(aot), runtime));
+    result->aot_entrypoints_proto_ = std::move(proto);
+    return result;
+  }
+
+  const AotPackageEntrypointsProto& aot_entrypoints_proto() const {
+    return aot_entrypoints_proto_;
   }
 
   Package* package() const { return package_.get(); }
@@ -203,6 +209,8 @@ class BaseProcJitWrapper {
   Proc* proc_;
   std::unique_ptr<ProcRuntime> runtime_;
   JitRuntime& jit_runtime_;
+  // Filled in after initialization. The entrypoints we wrap.
+  AotPackageEntrypointsProto aot_entrypoints_proto_;
 
  private:
   std::tuple<std::unique_ptr<Package>, std::unique_ptr<ProcRuntime>>