8355003: Implement Ahead-of-Time Method Profiling

make/hotspot/lib/JvmFeatures.gmk

@@ -127,7 +127,8 @@ ifneq ($(call check-jvm-feature, cds), true)
   JVM_EXCLUDE_FILES += \
       classLoaderDataShared.cpp \
       classLoaderExt.cpp \
-      systemDictionaryShared.cpp
+      systemDictionaryShared.cpp \
+      trainingData.cpp
   JVM_EXCLUDE_PATTERNS += cds/
 endif
 

src/hotspot/share/cds/aotArtifactFinder.cpp

@@ -33,6 +33,7 @@
 #include "memory/metaspaceClosure.hpp"
 #include "oops/instanceKlass.hpp"
 #include "oops/objArrayKlass.hpp"
+#include "oops/trainingData.hpp"
 #include "utilities/resourceHash.hpp"
 
 // All the classes that should be included in the AOT cache (in at least the "allocated" state)
@@ -162,6 +163,8 @@ void AOTArtifactFinder::find_artifacts() {
   });
 
   end_scanning_for_oops();
+
+  TrainingData::cleanup_training_data();
 }
 
 void AOTArtifactFinder::start_scanning_for_oops() {

src/hotspot/share/cds/aotLinkedClassBulkLoader.cpp

@@ -32,10 +32,12 @@
 #include "classfile/systemDictionary.hpp"
 #include "classfile/systemDictionaryShared.hpp"
 #include "classfile/vmClasses.hpp"
+#include "compiler/compilationPolicy.hpp"
 #include "gc/shared/gcVMOperations.hpp"
 #include "memory/resourceArea.hpp"
 #include "oops/instanceKlass.hpp"
 #include "oops/klass.inline.hpp"
+#include "oops/trainingData.hpp"
 #include "runtime/handles.inline.hpp"
 #include "runtime/java.hpp"
 
@@ -48,6 +50,17 @@ void AOTLinkedClassBulkLoader::serialize(SerializeClosure* soc, bool is_static_a
   AOTLinkedClassTable::get(is_static_archive)->serialize(soc);
 }
 
+bool AOTLinkedClassBulkLoader::class_preloading_finished() {
+  if (!CDSConfig::is_using_aot_linked_classes()) {
+    return true;
+  } else {
+    // The ConstantPools of preloaded classes have references to other preloaded classes. We don't
+    // want any Java code (including JVMCI compiler) to use these classes until all of them
+    // are loaded.
+    return Atomic::load_acquire(&_all_completed);
+  }
+}
+
 void AOTLinkedClassBulkLoader::load_javabase_classes(JavaThread* current) {
   assert(CDSConfig::is_using_aot_linked_classes(), "sanity");
   load_classes_in_loader(current, AOTLinkedClassCategory::BOOT1, nullptr); // only java.base classes
@@ -70,8 +83,14 @@ void AOTLinkedClassBulkLoader::load_non_javabase_classes(JavaThread* current) {
   _platform_completed = true;
 
   load_classes_in_loader(current, AOTLinkedClassCategory::APP, SystemDictionary::java_system_loader());
+
+  if (AOTPrintTrainingInfo) {
+    tty->print_cr("==================== archived_training_data ** after all classes preloaded ====================");
+    TrainingData::print_archived_training_data_on(tty);
+  }
+
   _app_completed = true;
-  _all_completed = true;
+  Atomic::release_store(&_all_completed, true);
 }
 
 void AOTLinkedClassBulkLoader::load_classes_in_loader(JavaThread* current, AOTLinkedClassCategory class_category, oop class_loader_oop) {
@@ -394,3 +413,25 @@ bool AOTLinkedClassBulkLoader::is_pending_aot_linked_class(Klass* k) {
     return false;
   }
 }
+
+void AOTLinkedClassBulkLoader::replay_training_at_init(Array<InstanceKlass*>* classes, TRAPS) {
+  if (classes != nullptr) {
+    for (int i = 0; i < classes->length(); i++) {
+      InstanceKlass* ik = classes->at(i);
+      if (ik->has_aot_initialized_mirror() && ik->is_initialized() && !ik->has_init_deps_processed()) {
+        CompilationPolicy::replay_training_at_init(ik, CHECK);
+      }
+    }
+  }
+}
+
+void AOTLinkedClassBulkLoader::replay_training_at_init_for_preloaded_classes(TRAPS) {
+  if (CDSConfig::is_using_aot_linked_classes() && TrainingData::have_data()) {
+    // Only static archive can have training data.
+    AOTLinkedClassTable* table = AOTLinkedClassTable::for_static_archive();
+    replay_training_at_init(table->boot(),     CHECK);
+    replay_training_at_init(table->boot2(),    CHECK);
+    replay_training_at_init(table->platform(), CHECK);
+    replay_training_at_init(table->app(),      CHECK);
+  }
+}

src/hotspot/share/cds/aotLinkedClassBulkLoader.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2024, 2025, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -55,6 +55,7 @@ class AOTLinkedClassBulkLoader :  AllStatic {
                                 const char* category_name, Handle loader, TRAPS);
   static void load_hidden_class(ClassLoaderData* loader_data, InstanceKlass* ik, TRAPS);
   static void init_required_classes_for_loader(Handle class_loader, Array<InstanceKlass*>* classes, TRAPS);
+  static void replay_training_at_init(Array<InstanceKlass*>* classes, TRAPS) NOT_CDS_RETURN;
 public:
   static void serialize(SerializeClosure* soc, bool is_static_archive) NOT_CDS_RETURN;
 
@@ -63,6 +64,8 @@ class AOTLinkedClassBulkLoader :  AllStatic {
   static void finish_loading_javabase_classes(TRAPS) NOT_CDS_RETURN;
   static void exit_on_exception(JavaThread* current);
 
+  static void replay_training_at_init_for_preloaded_classes(TRAPS) NOT_CDS_RETURN;
+  static bool class_preloading_finished();
   static bool is_pending_aot_linked_class(Klass* k) NOT_CDS_RETURN_(false);
 };
 

src/hotspot/share/cds/archiveBuilder.cpp

@@ -52,9 +52,12 @@
 #include "memory/resourceArea.hpp"
 #include "oops/compressedKlass.inline.hpp"
 #include "oops/instanceKlass.hpp"
+#include "oops/methodCounters.hpp"
+#include "oops/methodData.hpp"
 #include "oops/objArrayKlass.hpp"
 #include "oops/objArrayOop.inline.hpp"
 #include "oops/oopHandle.inline.hpp"
+#include "oops/trainingData.hpp"
 #include "runtime/arguments.hpp"
 #include "runtime/fieldDescriptor.inline.hpp"
 #include "runtime/globals_extension.hpp"
@@ -129,13 +132,27 @@ class RelocateEmbeddedPointers : public BitMapClosure {
     size_t field_offset = size_t(bit_offset - _start_idx) * sizeof(address);
     address* ptr_loc = (address*)(_buffered_obj + field_offset);
 
-    address old_p = *ptr_loc;
+    address old_p_with_tags = *ptr_loc;
+    assert(old_p_with_tags != nullptr, "null ptrs shouldn't have been marked");
+
+    address old_p = MetaspaceClosure::strip_tags(old_p_with_tags);
+    uintx tags = MetaspaceClosure::decode_tags(old_p_with_tags);
     address new_p = _builder->get_buffered_addr(old_p);
 
-    log_trace(cds)("Ref: [" PTR_FORMAT "] -> " PTR_FORMAT " => " PTR_FORMAT,
-                   p2i(ptr_loc), p2i(old_p), p2i(new_p));
+    bool nulled;
+    if (new_p == nullptr) {
+      // old_p had a FollowMode of set_to_null
+      nulled = true;
+    } else {
+      new_p = MetaspaceClosure::add_tags(new_p, tags);
+      nulled = false;
+    }
+
+    log_trace(cds)("Ref: [" PTR_FORMAT "] -> " PTR_FORMAT " => " PTR_FORMAT " %zu",
+                   p2i(ptr_loc), p2i(old_p) + tags, p2i(new_p), tags);
 
     ArchivePtrMarker::set_and_mark_pointer(ptr_loc, new_p);
+    ArchiveBuilder::current()->count_relocated_pointer(tags != 0, nulled);
     return true; // keep iterating the bitmap
   }
 };
@@ -175,6 +192,9 @@ ArchiveBuilder::ArchiveBuilder() :
   _klasses = new (mtClassShared) GrowableArray<Klass*>(4 * K, mtClassShared);
   _symbols = new (mtClassShared) GrowableArray<Symbol*>(256 * K, mtClassShared);
   _entropy_seed = 0x12345678;
+  _relocated_ptr_info._num_ptrs = 0;
+  _relocated_ptr_info._num_tagged_ptrs = 0;
+  _relocated_ptr_info._num_nulled_ptrs = 0;
   assert(_current == nullptr, "must be");
   _current = this;
 }
@@ -432,6 +452,11 @@ bool ArchiveBuilder::gather_one_source_obj(MetaspaceClosure::Ref* ref, bool read
   }
 #endif
 
+  if (ref->msotype() == MetaspaceObj::MethodDataType) {
+    MethodData* md = (MethodData*)ref->obj();
+    md->clean_method_data(false /* always_clean */);
+  }
+
   assert(p->read_only() == src_info.read_only(), "must be");
 
   if (created && src_info.should_copy()) {
@@ -529,8 +554,11 @@ ArchiveBuilder::FollowMode ArchiveBuilder::get_follow_mode(MetaspaceClosure::Ref
     // Don't dump existing shared metadata again.
     return point_to_it;
   } else if (ref->msotype() == MetaspaceObj::MethodDataType ||
-             ref->msotype() == MetaspaceObj::MethodCountersType) {
-    return set_to_null;
+             ref->msotype() == MetaspaceObj::MethodCountersType ||
+             ref->msotype() == MetaspaceObj::KlassTrainingDataType ||
+             ref->msotype() == MetaspaceObj::MethodTrainingDataType ||
+             ref->msotype() == MetaspaceObj::CompileTrainingDataType) {
+    return (TrainingData::need_data() || TrainingData::assembling_data()) ? make_a_copy : set_to_null;
   } else {
     if (ref->msotype() == MetaspaceObj::ClassType) {
       Klass* klass = (Klass*)ref->obj();
@@ -739,6 +767,10 @@ void ArchiveBuilder::relocate_metaspaceobj_embedded_pointers() {
   log_info(cds)("Relocating embedded pointers in core regions ... ");
   relocate_embedded_pointers(&_rw_src_objs);
   relocate_embedded_pointers(&_ro_src_objs);
+  log_info(cds)("Relocating %zu pointers, %zu tagged, %zu nulled",
+                _relocated_ptr_info._num_ptrs,
+                _relocated_ptr_info._num_tagged_ptrs,
+                _relocated_ptr_info._num_nulled_ptrs);
 }
 
 #define ADD_COUNT(x) \
@@ -940,6 +972,28 @@ void ArchiveBuilder::make_klasses_shareable() {
   DynamicArchive::make_array_klasses_shareable();
 }
 
+void ArchiveBuilder::make_training_data_shareable() {
+  auto clean_td = [&] (address& src_obj,  SourceObjInfo& info) {
+    if (!is_in_buffer_space(info.buffered_addr())) {
+      return;
+    }
+
+    if (info.msotype() == MetaspaceObj::KlassTrainingDataType ||
+        info.msotype() == MetaspaceObj::MethodTrainingDataType ||
+        info.msotype() == MetaspaceObj::CompileTrainingDataType) {
+      TrainingData* buffered_td = (TrainingData*)info.buffered_addr();
+      buffered_td->remove_unshareable_info();
+    } else if (info.msotype() == MetaspaceObj::MethodDataType) {
+      MethodData* buffered_mdo = (MethodData*)info.buffered_addr();
+      buffered_mdo->remove_unshareable_info();
+    } else if (info.msotype() == MetaspaceObj::MethodCountersType) {
+      MethodCounters* buffered_mc = (MethodCounters*)info.buffered_addr();
+      buffered_mc->remove_unshareable_info();
+    }
+  };
+  _src_obj_table.iterate_all(clean_td);
+}
+
 void ArchiveBuilder::serialize_dynamic_archivable_items(SerializeClosure* soc) {
   SymbolTable::serialize_shared_table_header(soc, false);
   SystemDictionaryShared::serialize_dictionary_headers(soc, false);
@@ -1564,6 +1618,12 @@ void ArchiveBuilder::write_region(FileMapInfo* mapinfo, int region_idx, DumpRegi
   mapinfo->write_region(region_idx, dump_region->base(), dump_region->used(), read_only, allow_exec);
 }
 
+void ArchiveBuilder::count_relocated_pointer(bool tagged, bool nulled) {
+  _relocated_ptr_info._num_ptrs ++;
+  _relocated_ptr_info._num_tagged_ptrs += tagged ? 1 : 0;
+  _relocated_ptr_info._num_nulled_ptrs += nulled ? 1 : 0;
+}
+
 void ArchiveBuilder::print_region_stats(FileMapInfo *mapinfo, ArchiveHeapInfo* heap_info) {
   // Print statistics of all the regions
   const size_t bitmap_used = mapinfo->region_at(MetaspaceShared::bm)->used();

src/hotspot/share/cds/archiveBuilder.hpp

@@ -237,6 +237,11 @@ class ArchiveBuilder : public StackObj {
   // statistics
   DumpAllocStats _alloc_stats;
   size_t _total_heap_region_size;
+  struct {
+    size_t _num_ptrs;
+    size_t _num_tagged_ptrs;
+    size_t _num_nulled_ptrs;
+  } _relocated_ptr_info;
 
   void print_region_stats(FileMapInfo *map_info, ArchiveHeapInfo* heap_info);
   void print_bitmap_region_stats(size_t size, size_t total_size);
@@ -257,6 +262,8 @@ class ArchiveBuilder : public StackObj {
     ~OtherROAllocMark();
   };
 
+  void count_relocated_pointer(bool tagged, bool nulled);
+
 private:
   FollowMode get_follow_mode(MetaspaceClosure::Ref *ref);
 
@@ -411,6 +418,7 @@ class ArchiveBuilder : public StackObj {
   void relocate_metaspaceobj_embedded_pointers();
   void record_regenerated_object(address orig_src_obj, address regen_src_obj);
   void make_klasses_shareable();
+  void make_training_data_shareable();
   void relocate_to_requested();
   void write_archive(FileMapInfo* mapinfo, ArchiveHeapInfo* heap_info);
   void write_region(FileMapInfo* mapinfo, int region_idx, DumpRegion* dump_region,
@@ -433,7 +441,8 @@ class ArchiveBuilder : public StackObj {
 
   address get_buffered_addr(address src_addr) const;
   template <typename T> T get_buffered_addr(T src_addr) const {
-    return (T)get_buffered_addr((address)src_addr);
+    CDS_ONLY(return (T)get_buffered_addr((address)src_addr);)
+    NOT_CDS(return nullptr;)
   }
 
   address get_source_addr(address buffered_addr) const;
@@ -446,7 +455,8 @@ class ArchiveBuilder : public StackObj {
   GrowableArray<Symbol*>* symbols() const { return _symbols; }
 
   static bool is_active() {
-    return (_current != nullptr);
+    CDS_ONLY(return (_current != nullptr));
+    NOT_CDS(return false;)
   }
 
   static ArchiveBuilder* current() {

src/hotspot/share/cds/cdsConfig.cpp

@@ -508,6 +508,8 @@ bool CDSConfig::check_vm_args_consistency(bool patch_mod_javabase, bool mode_fla
     FLAG_SET_ERGO_IF_DEFAULT(AOTClassLinking, true);
   }
 
+  setup_compiler_args();
+
   if (AOTClassLinking) {
     // If AOTClassLinking is specified, enable all AOT optimizations by default.
     FLAG_SET_ERGO_IF_DEFAULT(AOTInvokeDynamicLinking, true);
@@ -583,6 +585,28 @@ bool CDSConfig::check_vm_args_consistency(bool patch_mod_javabase, bool mode_fla
   return true;
 }
 
+void CDSConfig::setup_compiler_args() {
+  // AOT profiles are supported only in the JEP 483 workflow.
+  bool can_dump_profiles = AOTClassLinking && new_aot_flags_used();
+
+  if (is_dumping_preimage_static_archive() && can_dump_profiles) {
+    // JEP 483 workflow -- training
+    FLAG_SET_ERGO_IF_DEFAULT(AOTRecordTraining, true);
+    FLAG_SET_ERGO(AOTReplayTraining, false);
+  } else if (is_dumping_final_static_archive() && can_dump_profiles) {
+    // JEP 483 workflow -- assembly
+    FLAG_SET_ERGO(AOTRecordTraining, false); // This will be updated inside MetaspaceShared::preload_and_dump()
+    FLAG_SET_ERGO_IF_DEFAULT(AOTReplayTraining, true);
+  } else if (is_using_archive() && new_aot_flags_used()) {
+    // JEP 483 workflow -- production
+    FLAG_SET_ERGO(AOTRecordTraining, false);
+    FLAG_SET_ERGO_IF_DEFAULT(AOTReplayTraining, true);
+  } else {
+    FLAG_SET_ERGO(AOTReplayTraining, false);
+    FLAG_SET_ERGO(AOTRecordTraining, false);
+  }
+}
+
 void CDSConfig::prepare_for_dumping() {
   assert(CDSConfig::is_dumping_archive(), "sanity");
 

src/hotspot/share/cds/cdsConfig.hpp

@@ -67,6 +67,7 @@ class CDSConfig : public AllStatic {
   static void check_aotmode_auto_or_on();
   static void check_aotmode_record();
   static void check_aotmode_create();
+  static void setup_compiler_args();
   static void check_unsupported_dumping_module_options();
 
   // Called after Arguments::apply_ergo() has started

src/hotspot/share/cds/cds_globals.hpp

@@ -127,6 +127,23 @@
   product(bool, AOTCacheParallelRelocation, true, DIAGNOSTIC,               \
           "Use parallel relocation code to speed up startup.")              \
                                                                             \
+  /* flags to control training and deployment modes  */                     \
+                                                                            \
+  product(bool, AOTRecordTraining, false, DIAGNOSTIC,                       \
+          "Request output of training data for improved deployment.")       \
+                                                                            \
+  product(bool, AOTReplayTraining, false, DIAGNOSTIC,                       \
+          "Read training data, if available, for use in this execution")    \
+                                                                            \
+  product(bool, AOTPrintTrainingInfo, false, DIAGNOSTIC,                    \
+          "Print additional information about training")                    \
+                                                                            \
+  product(bool, AOTVerifyTrainingData, trueInDebug, DIAGNOSTIC,             \
+          "Verify archived training data")                                  \
+                                                                            \
+  product(bool, AOTCompileEagerly, false, DIAGNOSTIC,                       \
+          "Compile methods as soon as possible")                            \
+                                                                            \
 // end of CDS_FLAGS
 
 DECLARE_FLAGS(CDS_FLAGS)