Add debug symbols for parsed and compiler pass replaced instructions

CHANGELOG.md

@@ -16,6 +16,7 @@ Full documentation for MIGraphX is available at
 * Added Cubic resize jit kernel (#4652).
 * Added JIT compiler for `fill` operation (#4666).
 * Added JIT compiler for `multinomial` operation (#4721).
+* Added debug symbols for MIGraphX instructions such that parsed and compiled instructions can be tracked back to their ONNX origin node (#4626)
 
 ### Changed
 

src/api/api.cpp

@@ -178,6 +178,11 @@ static void set_limit_loop_iterations(onnx_options& options, int64_t value)
     options.limit_max_iterations = value;
 }
 
+static void set_use_debug_symbols(onnx_options& options, bool value)
+{
+    options.use_debug_symbols = value;
+}
+
 static void set_nhwc(tf_options& options, bool is_nhwc) { options.is_nhwc = is_nhwc; }
 
 static void set_default_dim_value(tf_options& options, size_t value) { options.batch_size = value; }
@@ -1991,6 +1996,17 @@ migraphx_onnx_options_set_external_data_path(migraphx_onnx_options_t onnx_option
     return api_error_result;
 }
 
+extern "C" migraphx_status
+migraphx_onnx_options_set_use_debug_symbols(migraphx_onnx_options_t onnx_options, bool value)
+{
+    auto api_error_result = migraphx::try_([&] {
+        if(onnx_options == nullptr)
+            MIGRAPHX_THROW(migraphx_status_bad_param, "Bad parameter onnx_options: Null pointer");
+        migraphx::set_use_debug_symbols((onnx_options->object), (value));
+    });
+    return api_error_result;
+}
+
 extern "C" migraphx_status migraphx_file_options_destroy(migraphx_file_options_t file_options)
 {
     auto api_error_result = migraphx::try_([&] { destroy((file_options)); });

src/api/include/migraphx/migraphx.h

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2025 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -535,6 +535,9 @@ MIGRAPHX_C_EXPORT migraphx_status migraphx_onnx_options_set_limit_loop_iteration
 MIGRAPHX_C_EXPORT migraphx_status migraphx_onnx_options_set_external_data_path(
     migraphx_onnx_options_t onnx_options, const char* external_data_path);
 
+MIGRAPHX_C_EXPORT migraphx_status
+migraphx_onnx_options_set_use_debug_symbols(migraphx_onnx_options_t onnx_options, bool value);
+
 MIGRAPHX_C_EXPORT migraphx_status
 migraphx_file_options_destroy(migraphx_file_options_t file_options);
 

src/api/include/migraphx/migraphx.hpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2025 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -1366,6 +1366,12 @@ struct onnx_options : MIGRAPHX_HANDLE_BASE(onnx_options)
              this->get_handle_ptr(),
              external_data_path.c_str());
     }
+
+    /// Enable debug symbols from ONNX node names
+    void set_use_debug_symbols(bool value = true)
+    {
+        call(&migraphx_onnx_options_set_use_debug_symbols, this->get_handle_ptr(), value);
+    }
 };
 
 /// Parse an onnx file into a migraphx program

src/api/migraphx.py

@@ -1,7 +1,7 @@
 #####################################################################################
 # The MIT License (MIT)
 #
-# Copyright (c) 2015-2025 Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -370,6 +370,11 @@ def onnx_options(h):
         api.params(external_data_path='const char*'),
         invoke='migraphx::set_external_data_path($@)',
     )
+    h.method(
+        'set_use_debug_symbols',
+        api.params(value='bool'),
+        invoke='migraphx::set_use_debug_symbols($@)',
+    )
 
 
 @auto_handle()

src/driver/main.cpp

@@ -191,6 +191,7 @@ struct loader
     bool replace_literals       = false;
     bool brief                  = false;
     bool verbose                = false;
+    bool use_debug_symbols      = false;
     std::string output_type;
     std::string output;
     std::string default_dyn_dim;
@@ -223,6 +224,11 @@ struct loader
            ap.help("Skip unknown operators when parsing and continue to parse."),
            ap.set_value(true));
         ap(is_nhwc, {"--nchw"}, ap.help("Treat tensorflow format as nchw"), ap.set_value(false));
+        ap(use_debug_symbols,
+           {"--debug-symbols"},
+           ap.help(
+               "Parse ONNX node names into MIGX instructions and propagate them as debug symbols."),
+           ap.set_value(true));
         ap(trim, {"--trim", "-t"}, ap.help("Trim instructions from the end"));
         ap(trim_size, {"--trim-size", "-s"}, ap.help("Number of instructions in the trim model"));
         ap(param_dims,
@@ -410,6 +416,7 @@ struct loader
         }
         options.skip_unknown_operators = skip_unknown_operators;
         options.print_program_on_error = true;
+        options.use_debug_symbols      = use_debug_symbols;
         options.map_input_dims         = map_input_dims;
         options.map_dyn_input_dims     = map_dyn_input_dims;
         options.dim_params             = map_dim_params;

src/include/migraphx/instruction.hpp

@@ -31,8 +31,9 @@
 #include <migraphx/operation.hpp>
 #include <migraphx/erase.hpp>
 #include <migraphx/config.hpp>
-#include <string>
+#include <set>
 #include <unordered_set>
+#include <string>
 #include <utility>
 
 namespace migraphx {
@@ -95,6 +96,14 @@ struct MIGRAPHX_EXPORT instruction
     /// Where this instruction is used as an input to another instruction
     const std::vector<instruction_ref>& outputs() const;
 
+    const std::set<std::string>& get_debug_symbols() const;
+
+    /// Avoid using directly because module will not track number of debug symbols
+    void add_debug_symbols(const std::set<std::string>& symbols);
+
+    /// Avoid using directly because module will not track number of debug symbols
+    void remove_debug_symbols();
+
     MIGRAPHX_EXPORT friend bool operator==(const instruction& x, const instruction& y);
 
     MIGRAPHX_EXPORT friend bool operator!=(const instruction& x, const instruction& y);
@@ -188,6 +197,7 @@ struct MIGRAPHX_EXPORT instruction
     std::vector<instruction_ref> output;
     std::vector<instruction_ref> arguments;
     std::vector<module_ref> module_args;
+    std::set<std::string> debug_symbols;
     literal lit;
     bool normalized       = false;
     std::size_t target_id = 0;

src/include/migraphx/module.hpp

@@ -24,9 +24,6 @@
 #ifndef MIGRAPHX_GUARD_MIGRAPHLIB_MODULE_HPP
 #define MIGRAPHX_GUARD_MIGRAPHLIB_MODULE_HPP
 
-#include <list>
-#include <unordered_set>
-#include <unordered_map>
 #include <migraphx/operation.hpp>
 #include <migraphx/literal.hpp>
 #include <migraphx/builtin.hpp>
@@ -38,6 +35,10 @@
 #include <migraphx/config.hpp>
 #include <algorithm>
 #include <iostream>
+#include <list>
+#include <set>
+#include <unordered_set>
+#include <unordered_map>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -81,6 +82,13 @@ struct MIGRAPHX_EXPORT module
     bool bypass() const;
     void set_bypass(bool b = true);
 
+    /// If any instructions in this module have debug symbols
+    bool has_debug_symbols() const;
+    /// Merge given symbols with instruction's symbols
+    void add_debug_symbols(instruction_ref ins, const std::set<std::string>& symbols);
+    /// Clear all debug symbols from instruction
+    void remove_debug_symbols(instruction_ref ins);
+
     template <class... Ts, MIGRAPHX_REQUIRES(std::is_same<Ts, instruction_ref>{}...)>
     instruction_ref add_instruction(operation op, Ts... args)
     {
@@ -122,6 +130,19 @@ struct MIGRAPHX_EXPORT module
 
     instruction_ref replace_instruction(instruction_ref ins, instruction_ref rep);
 
+    struct instruction_replacement
+    {
+        instruction_ref ins;
+        operation op;
+        std::vector<instruction_ref> args;
+        std::vector<module_ref> module_args;
+    };
+
+    /// Replaces an array of instructions within the same function to properly handle debug symbols
+    /// propagation. Returns vector of instruction_ref to replaced instructions.
+    std::vector<instruction_ref>
+    batch_replace_instruction(const std::vector<instruction_replacement>& replacers);
+
     instruction_ref remove_instruction(instruction_ref ins);
     instruction_ref remove_instructions(instruction_ref first, instruction_ref last);
 

src/include/migraphx/onnx.hpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2025 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -58,6 +58,8 @@ struct onnx_options
     int64_t limit_max_iterations = std::numeric_limits<uint16_t>::max();
     /// Use dynamic output for operators when available
     bool use_dyn_output = false;
+    /// Parse in ONNX node names as debug symbols
+    bool use_debug_symbols = false;
     /// Path to use for the external data if it is stored at different location compared to onnx
     /// file
     std::string external_data_path = "";

src/include/migraphx/stringutils.hpp

@@ -69,9 +69,12 @@ inline bool ends_with(const std::string& value, const std::string& suffix)
     else
         return std::equal(suffix.rbegin(), suffix.rend(), value.rbegin());
 }
-
+// NOLINTBEGIN(performance-unnecessary-value-param)
 template <class Strings>
-inline std::string join_strings(Strings strings, const std::string& delim)
+inline std::string
+join_strings(Strings strings,
+             const std::string& delim) 
+// NOLINTEND(performance-unnecessary-value-param)
 {
     auto it = strings.begin();
     if(it == strings.end())

src/instruction.cpp

@@ -28,6 +28,7 @@
 #include <migraphx/ranges.hpp>
 #include <migraphx/output_iterator.hpp>
 #include <migraphx/iterator.hpp>
+#include <migraphx/stringutils.hpp>
 #include <migraphx/iterator_for.hpp>
 #include <bitset>
 #include <queue>
@@ -189,6 +190,15 @@ const std::vector<module_ref>& instruction::module_inputs() const { return modul
 
 const std::vector<instruction_ref>& instruction::outputs() const { return output; }
 
+const std::set<std::string>& instruction::get_debug_symbols() const { return debug_symbols; }
+
+void instruction::add_debug_symbols(const std::set<std::string>& symbols)
+{
+    debug_symbols.insert(symbols.begin(), symbols.end());
+}
+
+void instruction::remove_debug_symbols() { debug_symbols.clear(); }
+
 bool operator==(const instruction& x, const instruction& y)
 {
     if(not std::equal(x.arguments.begin(),
@@ -274,7 +284,7 @@ void instruction::replace(operation o,
                           std::vector<module_ref> mdl_args)
 {
     lit = literal{};
-    op = std::move(o);
+    op  = std::move(o);
     replace(r);
     replace(std::move(args), std::move(mdl_args));
 }
@@ -442,6 +452,12 @@ void instruction::print(std::ostream& os,
     // print tid
     if(ins->target_id != 0)
         os << ", target_id=" << ins->target_id;
+
+    // print debug symbols if they exist
+    if(not ins->debug_symbols.empty())
+    {
+        os << " # " << join_strings(ins->debug_symbols, ", ");
+    }
 }
 
 static void debug_name(std::ostream& os, const instruction& ins)
@@ -472,7 +488,14 @@ void instruction::debug_print() const
     }
     if(not this->inputs().empty())
         std::cout << ")";
-    std::cout << " -> " << this->get_shape() << std::endl;
+    std::cout << " -> " << this->get_shape();
+
+    // print debug symbols if they exist
+    if(not debug_symbols.empty())
+    {
+        std::cout << " # " << join_strings(debug_symbols, ", ");
+    }
+    std::cout << std::endl;
 }
 
 std::vector<instruction_ref> instruction::get_output_alias(instruction_ref ins, bool shallow)
@@ -578,6 +601,8 @@ static auto track_visits(instruction_ref start, instruction_ref end, F f)
     std::size_t n           = std::distance(start, end);
     if(n < small)
     {
+        // Stop condition is ins distance to end > N or
+        // same instruction already visited.
         std::bitset<small> visited;
         auto stop = [&](auto ins) {
             auto i = std::distance(ins, end);
@@ -592,6 +617,9 @@ static auto track_visits(instruction_ref start, instruction_ref end, F f)
     }
     else
     {
+        // Make a hashmap of instructions between start and end.
+        // Stop condition is instruction not in the hashmap or
+        // same instruction already visited.
         auto instructions     = range(start, std::next(end));
         auto instruction_refs = iterator_for(instructions);
         std::unordered_set<instruction_ref> in_range(instruction_refs.begin(),