Working prototype

Author: tcgu-amd

source/bin/rocprof-sys-sample/impl.cpp

@@ -343,6 +343,15 @@ parse_args(int argc, char** argv, std::vector<char*>& _env)
     %{INDENT}%- discard     : new data is ignored
     %{INDENT}%- ring_buffer : new data overwrites oldest data)";
 
+    const auto* _pre_attach_desc =
+        R"(Launch the application in pre-attach mode. This mode sets up and configures the profiling tools without initiating them. 
+           Applications running in pre-attach mode can later be profiled by executing rocprof-sys-sample with the --attach option. 
+           Please note that using pre-attach mode may negatively impact application performance. The extent of this impact varies based 
+           on the specific profiling configuration and the nature of the application.)";
+
+    const auto* _attach_desc =
+        R"(Attach to a running process launched in pre-attach mode. )";
+
     parser.set_use_color(true);
     parser.enable_help();
     parser.enable_version("rocprof-sys-sample", ROCPROFSYS_ARGPARSE_VERSION_INFO);
@@ -823,6 +832,21 @@ parse_args(int argc, char** argv, std::vector<char*>& _env)
             update_env(_env, "HSA_ENABLE_INTERRUPT", p.get<int>("hsa-interrupt"));
         });
 
+    parser.add_argument({ "--pre-attach" }, _pre_attach_desc)
+        .count(0)
+        .action([&](parser_t& p) { update_env(_env, "ROCPROFSYS_ATTACH", true); });
+
+    parser.add_argument({ "--attach" }, _attach_desc)
+        .count(1)
+        .dtype("pid")
+        .action([&](parser_t& p) {
+            auto _v = p.get<std::deque<std::string>>("attach");
+            if(!_v.empty())
+            {
+                *(get_attach_pid()) = std::stoi(_v.front());
+                _v.pop_front();
+            }
+        });
     parser.end_group();
 
     auto _inpv = std::vector<char*>{};
@@ -861,3 +885,21 @@ parse_args(int argc, char** argv, std::vector<char*>& _env)
 
     return _outv;
 }
+
+int*
+get_attach_pid()
+{
+    static int _v = -1;
+    return &_v;
+}
+
+int
+attach(int _pid)
+{
+    kill(_pid, 10);
+    stream(std::cout, color::info())
+        << "Entering into attach mode. Press any key to detach.\n";
+    std::cin.get();
+    kill(_pid, 10);
+    return 0;
+}

source/bin/rocprof-sys-sample/rocprof-sys-sample.cpp

@@ -53,6 +53,11 @@ main(int argc, char** argv)
             _argv.emplace_back(argv[i]);
     }
 
+    if(int _pid = *(get_attach_pid()); _pid > 0)
+    {
+        return attach(_pid);
+    }
+
     print_updated_environment(_env);
 
     if(!_argv.empty())

source/bin/rocprof-sys-sample/rocprof-sys-sample.hpp

@@ -59,3 +59,9 @@ remove_env(std::vector<char*>& _environ, std::string_view _env_var);
 
 std::vector<char*>
 parse_args(int argc, char** argv, std::vector<char*>& envp);
+
+int*
+get_attach_pid();
+
+int
+attach(int _pid);

source/lib/core/config.cpp

@@ -247,6 +247,8 @@ configure_settings(bool _init)
                                ROCPROFSYS_ROCM_VERSION_PATCH);
 #endif
 
+    is_pre_attach_mode() = tim::get_env("ROCPROFSYS_ATTACH", false);
+
     auto _config = *get_config_impl();
 
     // if using timemory, default to perfetto being off
@@ -258,6 +260,10 @@ configure_settings(bool _init)
     auto _rocprofsys_debug = _config->get<bool>("ROCPROFSYS_DEBUG");
     if(_rocprofsys_debug) tim::set_env("TIMEMORY_DEBUG_SETTINGS", "1", 0);
 
+    ROCPROFSYS_CONFIG_SETTING(bool, "ROCPROFSYS_ATTACH",
+                              "Determines if rocprofsys is launched in pre-attach mode.",
+                              false, "attach");
+
     ROCPROFSYS_CONFIG_SETTING(
         std::string, "ROCPROFSYS_MODE",
         "Data collection mode. Used to set default values for ROCPROFSYS_USE_* options. "
@@ -1184,18 +1190,73 @@ get_signal_handler()
     return _v;
 }
 
+std::atomic<signal_handler_t>&
+get_attach_signal_handler()
+{
+    static auto _v = std::atomic<signal_handler_t>{ nullptr };
+    return _v;
+}
+
 void
 rocprofsys_exit_action(int nsig)
 {
     tim::signals::block_signals(get_sampling_signals(),
                                 tim::signals::sigmask_scope::process);
     ROCPROFSYS_BASIC_PRINT("Finalizing after signal %i :: %s\n", nsig,
                            signal_settings::str(static_cast<sys_signal>(nsig)).c_str());
-    auto _handler = get_signal_handler().load();
-    if(_handler) (*_handler)();
+
+    // Tim: Handles the case where State is never active i.e. attaching never occured.
+    if(get_state() == State::Active)
+    {
+        auto _handler = get_signal_handler().load();
+        if(_handler) (*_handler)();
+    }
     kill(process::get_id(), nsig);
 }
 
+// Tim: This handles signals for triggering attach/detach. It prevents the process from
+// being killed at the end.
+void
+rocprofsys_attach_detach_action(int, siginfo_t*, void*)
+{
+    if(!config::is_pre_attach_mode())
+    {
+        ROCPROFSYS_BASIC_PRINT(R"(
+    EEEEEEEEEEEEEEEEEEEEEERRRRRRRRRRRRRRRRR   RRRRRRRRRRRRRRRRR        OOOOOOOOO     RRRRRRRRRRRRRRRRR   
+    E::::::::::::::::::::ER::::::::::::::::R  R::::::::::::::::R     OO:::::::::OO   R::::::::::::::::R  
+    E::::::::::::::::::::ER::::::RRRRRR:::::R R::::::RRRRRR:::::R  OO:::::::::::::OO R::::::RRRRRR:::::R 
+    EE::::::EEEEEEEEE::::ERR:::::R     R:::::RRR:::::R     R:::::RO:::::::OOO:::::::ORR:::::R     R:::::R
+    E:::::E       EEEEEE  R::::R     R:::::R  R::::R     R:::::RO::::::O   O::::::O  R::::R     R:::::R
+    E:::::E               R::::R     R:::::R  R::::R     R:::::RO:::::O     O:::::O  R::::R     R:::::R
+    E::::::EEEEEEEEEE     R::::RRRRRR:::::R   R::::RRRRRR:::::R O:::::O     O:::::O  R::::RRRRRR:::::R 
+    E:::::::::::::::E     R:::::::::::::RR    R:::::::::::::RR  O:::::O     O:::::O  R:::::::::::::RR  
+    E:::::::::::::::E     R::::RRRRRR:::::R   R::::RRRRRR:::::R O:::::O     O:::::O  R::::RRRRRR:::::R 
+    E::::::EEEEEEEEEE     R::::R     R:::::R  R::::R     R:::::RO:::::O     O:::::O  R::::R     R:::::R
+    E:::::E               R::::R     R:::::R  R::::R     R:::::RO:::::O     O:::::O  R::::R     R:::::R
+    E:::::E       EEEEEE  R::::R     R:::::R  R::::R     R:::::RO::::::O   O::::::O  R::::R     R:::::R
+    EE::::::EEEEEEEE:::::ERR:::::R     R:::::RRR:::::R     R:::::RO:::::::OOO:::::::ORR:::::R     R:::::R
+    E::::::::::::::::::::ER::::::R     R:::::RR::::::R     R:::::R OO:::::::::::::OO R::::::R     R:::::R
+    E::::::::::::::::::::ER::::::R     R:::::RR::::::R     R:::::R   OO:::::::::OO   R::::::R     R:::::R
+    EEEEEEEEEEEEEEEEEEEEEERRRRRRRR     RRRRRRRRRRRRRRR     RRRRRRR     OOOOOOOOO     RRRRRRRR     RRRRRR
+                                    
+    Trying to attach to an uninitialized process. Only applications launched by rocprof-sys in 
+    pre-attach mode can be attached to. 
+
+    To launch an application in pre-attach mode, either 
+    
+        - set env var `ROCPROFSYS_ATTACH = true` or,
+        - launch the process with `rocprof-sys-sample --pre-attach <OPTIONS> -- <APP> <ARGS>`
+
+    Ignoring attach attempt and continuing...
+        )");
+        return;
+    }
+    tim::signals::block_signals(get_sampling_signals(),
+                                tim::signals::sigmask_scope::process);
+    auto _handler = get_attach_signal_handler().load();
+    if(_handler) (*_handler)();
+}
+
 void
 rocprofsys_trampoline_handler(int _v)
 {
@@ -1234,6 +1295,28 @@ set_signal_handler(signal_handler_t _func)
     return get_signal_handler().load();
 }
 
+signal_handler_t
+set_attach_signal_handler(signal_handler_t _func)
+{
+    if(_func)
+    {
+        auto _handler = get_attach_signal_handler().load(std::memory_order_relaxed);
+        if(get_attach_signal_handler().compare_exchange_strong(_handler, _func,
+                                                               std::memory_order_relaxed))
+        {
+            return _handler;
+        }
+        else
+        {
+            _handler = get_attach_signal_handler().load(std::memory_order_seq_cst);
+            get_attach_signal_handler().store(_func);
+            return _handler;
+        }
+    }
+
+    return get_attach_signal_handler().load();
+}
+
 void
 configure_signal_handler(const std::shared_ptr<settings>& _config)
 {
@@ -1266,6 +1349,15 @@ configure_signal_handler(const std::shared_ptr<settings>& _config)
         _action.sa_handler = rocprofsys_trampoline_handler;
         sigaction(_dyninst_trampoline_signal, &_action, nullptr);
     }
+
+    // Set up custom signals handlers for detaching.
+    // Avoid using timeory's signal handler because it kills the process.
+    int              DETACH_SIG = 10;  // SIGUSR1
+    struct sigaction sa;
+    memset(&sa, 0, sizeof(sa));
+    sa.sa_sigaction = rocprofsys_attach_detach_action;
+    sa.sa_flags     = SA_SIGINFO;
+    sigaction(DETACH_SIG, &sa, nullptr);
 }
 
 bool
@@ -1722,6 +1814,13 @@ is_attached()
     return _v;
 }
 
+bool&
+is_pre_attach_mode()
+{
+    static bool _v = false;
+    return _v;
+}
+
 bool&
 is_binary_rewrite()
 {

source/lib/core/config.hpp

@@ -48,6 +48,8 @@ using signal_handler_t = void (*)(void);
 // if arg is non-null, returns replaced signal handler
 signal_handler_t set_signal_handler(signal_handler_t);
 
+signal_handler_t set_attach_signal_handler(signal_handler_t);
+
 bool
 settings_are_configured() ROCPROFSYS_HOT;
 
@@ -162,6 +164,9 @@ get_mode();
 bool&
 is_attached();
 
+bool&
+is_pre_attach_mode();
+
 bool&
 is_binary_rewrite();
 

source/lib/rocprof-sys-dl/dl.cpp

@@ -1522,11 +1522,17 @@ extern "C"
             }
         }
 
+        auto _attach = get_env("ROCPROFSYS_ATTACH", false);
+
+        // Tim: Disable tooling initalization in attach mode.
+        if(_attach) setenv("ROCPROFSYS_INIT_TOOLING", "false", 1);
         auto _mode = get_env("ROCPROFSYS_MODE", get_default_mode());
         rocprofsys_init(_mode.c_str(),
                         dl::get_instrumented() == dl::InstrumentMode::BinaryRewrite,
                         argv[0]);
 
+        // Tim: Enable tooling initialization during runtime.
+        if(_attach) setenv("ROCPROFSYS_INIT_TOOLING", "true", 1);
         int ret = (*::rocprofsys::dl::main_real)(argc, argv, envp);
 
         rocprofsys_pop_trace(basename(argv[0]));