Enabling rocTX markers for migraphx-driver via roctx knob (#946)

Added features:
This enables wrapping each migraphx operator with rocTX markers.
It adds new knob trace to migraphx-driver binary.

Limitation:

rocTX standalone does not output a file, it needs to be used with rocprof. Example command line:

/opt/rocm/bin/rocprof -i ./in.txt --hip-trace --roctx-trace --flush-rate 10ms --timestamp on -d cagri_out --obj-tracking on /opt/rocm/bin/migraphx-driver trace ./resnet50-v2-7.onnx --onnx --gpu
Co-authored-by: Shucai Xiao <shucai@gmail.com>

src/driver/CMakeLists.txt

@@ -6,6 +6,7 @@ add_executable(driver
     resnet50.cpp
     inceptionv3.cpp
     alexnet.cpp
+    marker_roctx.cpp
 )
 set_target_properties(driver PROPERTIES OUTPUT_NAME migraphx-driver)
 # Copy driver for backwards compatibility

src/driver/main.cpp

@@ -3,6 +3,7 @@
 #include "verify.hpp"
 #include "perf.hpp"
 #include "models.hpp"
+#include "marker_roctx.hpp"
 
 #include <migraphx/tf.hpp>
 #include <migraphx/onnx.hpp>
@@ -483,6 +484,23 @@ struct perf : command<perf>
     }
 };
 
+struct roctx : command<roctx>
+{
+    compiler c;
+    void parse(argument_parser& ap) { c.parse(ap); }
+
+    void run()
+    {
+        std::cout << "Compiling ... " << std::endl;
+        auto p = c.compile();
+        std::cout << "Allocating params ... " << std::endl;
+        auto m = c.params(p);
+        std::cout << "rocTX:\tLoading rocTX library..." << std::endl;
+        auto rtx = create_marker_roctx();
+        p.mark(m, std::move(rtx));
+    }
+};
+
 struct op : command<op>
 {
     bool show_ops = false;

src/driver/marker_roctx.cpp

+#include "marker_roctx.hpp"
+
+#include <migraphx/dynamic_loader.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/instruction_ref.hpp>
+
+namespace migraphx {
+namespace driver {
+inline namespace MIGRAPHX_INLINE_NS {
+
+class marker_roctx
+{
+    std::function<void(const char*)> sym_roctx_mark;
+    std::function<uint64_t(const char*)> sym_roctx_range_start;
+    std::function<void(uint64_t)> sym_roctx_range_stop;
+
+    std::function<int(const char*)> sym_roctx_range_push;
+    std::function<int()> sym_roctx_range_pop;
+
+    uint64_t range_id;
+
+    public:
+    marker_roctx()
+    {
+        dynamic_loader lib    = migraphx::dynamic_loader{"libroctx64.so"};
+        sym_roctx_mark        = lib.get_function<void(const char*)>("roctxMarkA");
+        sym_roctx_range_start = lib.get_function<uint64_t(const char*)>("roctxRangeStartA");
+        sym_roctx_range_stop  = lib.get_function<void(uint64_t)>("roctxRangeStop");
+
+        sym_roctx_range_push = lib.get_function<int(const char*)>("roctxRangePushA");
+        sym_roctx_range_pop  = lib.get_function<int()>("roctxRangePop");
+
+        sym_roctx_mark("rocTX marker created.");
+    }
+
+    void mark_start(instruction_ref ins_ref)
+    {
+        std::string text = "Marker start: " + ins_ref->name();
+        sym_roctx_range_push(text.c_str());
+    }
+    void mark_stop(instruction_ref) { sym_roctx_range_pop(); }
+    void mark_start(const program&) { range_id = sym_roctx_range_start("0"); }
+    void mark_stop(const program&) { sym_roctx_range_stop(range_id); }
+};
+
+marker create_marker_roctx() { return marker_roctx(); }
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace driver
+} // namespace migraphx

src/driver/marker_roctx.hpp

+#ifndef MIGRAPHX_GUARD_RTGLIB_MARKER_ROCTX_HPP
+#define MIGRAPHX_GUARD_RTGLIB_MARKER_ROCTX_HPP
+
+#include <migraphx/marker.hpp>
+
+namespace migraphx {
+namespace driver {
+inline namespace MIGRAPHX_INLINE_NS {
+
+marker create_marker_roctx();
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace driver
+} // namespace migraphx
+
+#endif

src/dynamic_loader.cpp

@@ -45,6 +45,7 @@ dynamic_loader::dynamic_loader(const std::vector<char>& buffer)
 
 std::shared_ptr<void> dynamic_loader::get_symbol(const std::string& name) const
 {
+    dlerror();
     void* symbol = dlsym(impl->handle.get(), name.c_str());
     if(symbol == nullptr)
         MIGRAPHX_THROW("Symbol not found: " + name);

src/include/migraphx/marker.hpp

+#ifndef MIGRAPHX_GUARD_MARKER_HPP
+#define MIGRAPHX_GUARD_MARKER_HPP
+
+#include <cassert>
+#include <string>
+#include <functional>
+#include <memory>
+#include <type_traits>
+#include <utility>
+#include <migraphx/config.hpp>
+#include <migraphx/instruction_ref.hpp>
+#include <migraphx/program.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+#ifdef DOXYGEN
+
+/// Marker is an interface to general marking functions, such as rocTX markers.
+
+#else
+
+/*
+ * Type-erased interface for:
+ *
+ * struct marker
+ * {
+ *      void mark_start(instruction_ref ins_ref) ;
+ *      void mark_start(const program& prog) ;
+ *      void mark_stop(instruction_ref ins) ;
+ *      void mark_stop(const program& prog) ;
+ * };
+ *
+ */
+
+struct marker
+{
+    // Constructors
+    marker() = default;
+
+    template <typename PrivateDetailTypeErasedT>
+    marker(PrivateDetailTypeErasedT value)
+        : private_detail_te_handle_mem_var(
+              std::make_shared<private_detail_te_handle_type<
+                  typename std::remove_reference<PrivateDetailTypeErasedT>::type>>(
+                  std::forward<PrivateDetailTypeErasedT>(value)))
+    {
+    }
+
+    // Assignment
+    template <typename PrivateDetailTypeErasedT>
+    marker& operator=(PrivateDetailTypeErasedT value)
+    {
+        using std::swap;
+        auto* derived = this->any_cast<PrivateDetailTypeErasedT>();
+        if(derived and private_detail_te_handle_mem_var.unique())
+        {
+            *derived = std::forward<PrivateDetailTypeErasedT>(value);
+        }
+        else
+        {
+            marker rhs(value);
+            swap(private_detail_te_handle_mem_var, rhs.private_detail_te_handle_mem_var);
+        }
+        return *this;
+    }
+
+    // Cast
+    template <typename PrivateDetailTypeErasedT>
+    PrivateDetailTypeErasedT* any_cast()
+    {
+        return this->type_id() == typeid(PrivateDetailTypeErasedT)
+                   ? std::addressof(static_cast<private_detail_te_handle_type<
+                                        typename std::remove_cv<PrivateDetailTypeErasedT>::type>&>(
+                                        private_detail_te_get_handle())
+                                        .private_detail_te_value)
+                   : nullptr;
+    }
+
+    template <typename PrivateDetailTypeErasedT>
+    const typename std::remove_cv<PrivateDetailTypeErasedT>::type* any_cast() const
+    {
+        return this->type_id() == typeid(PrivateDetailTypeErasedT)
+                   ? std::addressof(static_cast<const private_detail_te_handle_type<
+                                        typename std::remove_cv<PrivateDetailTypeErasedT>::type>&>(
+                                        private_detail_te_get_handle())
+                                        .private_detail_te_value)
+                   : nullptr;
+    }
+
+    const std::type_info& type_id() const
+    {
+        if(private_detail_te_handle_empty())
+            return typeid(std::nullptr_t);
+        else
+            return private_detail_te_get_handle().type();
+    }
+
+    void mark_start(instruction_ref ins_ref)
+    {
+        assert((*this).private_detail_te_handle_mem_var);
+        (*this).private_detail_te_get_handle().mark_start(ins_ref);
+    }
+
+    void mark_start(const program& prog)
+    {
+        assert((*this).private_detail_te_handle_mem_var);
+        (*this).private_detail_te_get_handle().mark_start(prog);
+    }
+
+    void mark_stop(instruction_ref ins)
+    {
+        assert((*this).private_detail_te_handle_mem_var);
+        (*this).private_detail_te_get_handle().mark_stop(ins);
+    }
+
+    void mark_stop(const program& prog)
+    {
+        assert((*this).private_detail_te_handle_mem_var);
+        (*this).private_detail_te_get_handle().mark_stop(prog);
+    }
+
+    friend bool is_shared(const marker& private_detail_x, const marker& private_detail_y)
+    {
+        return private_detail_x.private_detail_te_handle_mem_var ==
+               private_detail_y.private_detail_te_handle_mem_var;
+    }
+
+    private:
+    struct private_detail_te_handle_base_type
+    {
+        virtual ~private_detail_te_handle_base_type() {}
+        virtual std::shared_ptr<private_detail_te_handle_base_type> clone() const = 0;
+        virtual const std::type_info& type() const                                = 0;
+
+        virtual void mark_start(instruction_ref ins_ref) = 0;
+        virtual void mark_start(const program& prog)     = 0;
+        virtual void mark_stop(instruction_ref ins)      = 0;
+        virtual void mark_stop(const program& prog)      = 0;
+    };
+
+    template <typename PrivateDetailTypeErasedT>
+    struct private_detail_te_handle_type : private_detail_te_handle_base_type
+    {
+        template <typename PrivateDetailTypeErasedU = PrivateDetailTypeErasedT>
+        private_detail_te_handle_type(
+            PrivateDetailTypeErasedT value,
+            typename std::enable_if<std::is_reference<PrivateDetailTypeErasedU>::value>::type* =
+                nullptr)
+            : private_detail_te_value(value)
+        {
+        }
+
+        template <typename PrivateDetailTypeErasedU = PrivateDetailTypeErasedT>
+        private_detail_te_handle_type(
+            PrivateDetailTypeErasedT value,
+            typename std::enable_if<!std::is_reference<PrivateDetailTypeErasedU>::value,
+                                    int>::type* = nullptr) noexcept
+            : private_detail_te_value(std::move(value))
+        {
+        }
+
+        std::shared_ptr<private_detail_te_handle_base_type> clone() const override
+        {
+            return std::make_shared<private_detail_te_handle_type>(private_detail_te_value);
+        }
+
+        const std::type_info& type() const override { return typeid(private_detail_te_value); }
+
+        void mark_start(instruction_ref ins_ref) override
+        {
+
+            private_detail_te_value.mark_start(ins_ref);
+        }
+
+        void mark_start(const program& prog) override { private_detail_te_value.mark_start(prog); }
+
+        void mark_stop(instruction_ref ins) override { private_detail_te_value.mark_stop(ins); }
+
+        void mark_stop(const program& prog) override { private_detail_te_value.mark_stop(prog); }
+
+        PrivateDetailTypeErasedT private_detail_te_value;
+    };
+
+    template <typename PrivateDetailTypeErasedT>
+    struct private_detail_te_handle_type<std::reference_wrapper<PrivateDetailTypeErasedT>>
+        : private_detail_te_handle_type<PrivateDetailTypeErasedT&>
+    {
+        private_detail_te_handle_type(std::reference_wrapper<PrivateDetailTypeErasedT> ref)
+            : private_detail_te_handle_type<PrivateDetailTypeErasedT&>(ref.get())
+        {
+        }
+    };
+
+    bool private_detail_te_handle_empty() const
+    {
+        return private_detail_te_handle_mem_var == nullptr;
+    }
+
+    const private_detail_te_handle_base_type& private_detail_te_get_handle() const
+    {
+        assert(private_detail_te_handle_mem_var != nullptr);
+        return *private_detail_te_handle_mem_var;
+    }
+
+    private_detail_te_handle_base_type& private_detail_te_get_handle()
+    {
+        assert(private_detail_te_handle_mem_var != nullptr);
+        if(!private_detail_te_handle_mem_var.unique())
+            private_detail_te_handle_mem_var = private_detail_te_handle_mem_var->clone();
+        return *private_detail_te_handle_mem_var;
+    }
+
+    std::shared_ptr<private_detail_te_handle_base_type> private_detail_te_handle_mem_var;
+};
+
+template <typename ValueType>
+inline const ValueType* any_cast(const marker* x)
+{
+    return x->any_cast<ValueType>();
+}
+
+template <typename ValueType>
+inline ValueType* any_cast(marker* x)
+{
+    return x->any_cast<ValueType>();
+}
+
+template <typename ValueType>
+inline ValueType& any_cast(marker& x)
+{
+    auto* y = x.any_cast<typename std::remove_reference<ValueType>::type>();
+    if(y == nullptr)
+        throw std::bad_cast();
+    return *y;
+}
+
+template <typename ValueType>
+inline const ValueType& any_cast(const marker& x)
+{
+    const auto* y = x.any_cast<typename std::remove_reference<ValueType>::type>();
+    if(y == nullptr)
+        throw std::bad_cast();
+    return *y;
+}
+
+#endif
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/program.hpp

@@ -23,6 +23,8 @@ MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_TRACE_EVAL)
 
 struct program_impl;
 
+struct marker;
+
 /**
  * @brief Stores the instruction stream
  */
@@ -67,6 +69,8 @@ struct program
 
     void perf_report(std::ostream& os, std::size_t n, parameter_map params) const;
 
+    void mark(const parameter_map& params, marker&& m);
+
     value to_value() const;
     void from_value(const value& v);
 

src/program.cpp

@@ -13,6 +13,7 @@
 #include <migraphx/algorithm.hpp>
 #include <migraphx/output_iterator.hpp>
 #include <migraphx/make_op.hpp>
+#include <migraphx/marker.hpp>
 #include <iostream>
 #include <sstream>
 #include <algorithm>
@@ -507,6 +508,24 @@ std::string perf_group(const operation& op)
     return op.name();
 }
 
+void program::mark(const parameter_map& params, marker&& m)
+{
+    auto& ctx = this->impl->ctx;
+    // Run once by itself
+    eval(params);
+    ctx.finish();
+    // Start marking
+    m.mark_start(*this);
+    generic_eval(*this, ctx, params, always([&](auto ins, auto f) {
+        argument result;
+        m.mark_start(ins);
+        result = f();
+        m.mark_stop(ins);
+        return result;
+    }));
+    m.mark_stop(*this);
+}
+
 void program::perf_report(std::ostream& os, std::size_t n, parameter_map params) const
 {
     auto& ctx = this->impl->ctx;

test/marker.cpp

+#include <migraphx/program.hpp>
+#include <migraphx/ref/target.hpp>
+#include <migraphx/ranges.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/marker.hpp>
+#include <migraphx/instruction.hpp>
+
+#include "test.hpp"
+
+struct mock_marker
+{
+    std::shared_ptr<std::stringstream> ss = std::make_shared<std::stringstream>();
+
+    void mark_start(migraphx::instruction_ref ins_ref)
+    {
+        std::string text = "Mock marker instruction start:" + ins_ref->name();
+        (*ss) << text;
+    }
+    void mark_stop(migraphx::instruction_ref)
+    {
+        std::string text = "Mock marker instruction stop.";
+        (*ss) << text;
+    }
+    void mark_start(const migraphx::program&)
+    {
+        std::string text = "Mock marker program start.";
+        (*ss) << text;
+    }
+    void mark_stop(const migraphx::program&)
+    {
+        std::string text = "Mock marker program stop.";
+        (*ss) << text;
+    }
+};
+
+TEST_CASE(marker)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+
+    auto one = mm->add_literal(1);
+    auto two = mm->add_literal(2);
+    mm->add_instruction(migraphx::make_op("add"), one, two);
+    p.compile(migraphx::ref::target{});
+
+    mock_marker temp_marker;
+    p.mark({}, temp_marker);
+
+    std::string output = temp_marker.ss->str();
+    EXPECT(migraphx::contains(output, "Mock marker instruction start:@literal"));
+    EXPECT(migraphx::contains(output, "Mock marker instruction start:ref::op"));
+    EXPECT(migraphx::contains(output, "Mock marker instruction stop."));
+    EXPECT(migraphx::contains(output, "Mock marker program start."));
+    EXPECT(migraphx::contains(output, "Mock marker program stop."));
+}
+
+int main(int argc, const char* argv[]) { test::run(argc, argv); }

tools/include/marker.hpp

+#ifndef MIGRAPHX_GUARD_MARKER_HPP
+#define MIGRAPHX_GUARD_MARKER_HPP
+
+#include <cassert>
+#include <string>
+#include <functional>
+#include <memory>
+#include <type_traits>
+#include <utility>
+#include <migraphx/config.hpp>
+#include <migraphx/instruction_ref.hpp>
+#include <migraphx/program.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+#ifdef DOXYGEN
+
+/// Marker is an interface to general marking functions, such as rocTX markers.
+
+#else
+
+<%
+interface('marker',
+           virtual('mark_start', ins_ref = 'instruction_ref', returns = 'void'),
+           virtual('mark_start', prog = 'const program&', returns = 'void'),
+           virtual('mark_stop', ins = 'instruction_ref', returns = 'void'),
+           virtual('mark_stop', prog = 'const program&', returns = 'void')
+        ) %>
+#endif
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif