SWDEV-542550 [rocpd] Convert to perfetto does not display scratch_memory correctly

source/lib/output/generateRocpd.cpp

@@ -1167,8 +1167,16 @@ write_rocpd(
             }
         };
 
+    struct agent_and_size
+    {
+        uint64_t agent_abs_index = {};
+        uint64_t size            = {0};
+    };
+    auto address_to_agent_and_size = std::unordered_map<rocprofiler_address_t, agent_and_size>{};
+
     auto insert_memory_alloc_data =
-        [&conn, &tool_metadata, &string_entries, node_id, this_pid](const auto& _gen) {
+        [&conn, &tool_metadata, &string_entries, node_id, this_pid, &address_to_agent_and_size](
+            const auto& _gen) {
             for(auto pitr : _gen)
             {
                 auto _deferred = sql::deferred_transaction{conn};
@@ -1188,17 +1196,60 @@ write_rocpd(
                     ROCP_FATAL_IF(_level != "REAL" && _level != "VIRTUAL" && _level != "SCRATCH")
                         << "erroneous db level: " << _level;
 
-                    auto _node_id = std::optional<uint64_t>{};
-                    if(_type == "ALLOC")
-                    {
-                        _node_id = tool_metadata.get_agent(itr.agent_id)->node_id;
-                    }
-
                     auto _stream_id       = get_stream_id(extract_stream_field(itr));
                     auto _queue_id        = get_queue_id(extract_queue_field(itr));
                     auto _address         = extract_address_field(itr);
                     auto _allocation_size = extract_allocation_size_field(itr);
 
+                    // memory allocation counter track
+                    struct free_memory_information
+                    {
+                        rocprofiler_timestamp_t start_timestamp = 0;
+                        rocprofiler_timestamp_t end_timestamp   = 0;
+                        rocprofiler_address_t   address         = {.handle = 0};
+                    };
+
+                    auto _node_id      = std::optional<uint64_t>{};
+                    auto free_mem_info = std::vector<free_memory_information>{};
+                    if(_type == "ALLOC")
+                    {
+                        _node_id = tool_metadata.get_agent(itr.agent_id)->node_id;
+                        address_to_agent_and_size.emplace(
+                            rocprofiler_address_t{.handle = _address.handle},
+                            agent_and_size{_node_id.value(), _allocation_size});
+                    }
+                    else if(_type == "FREE")
+                    {
+                        // Store free memory operations in seperate vector to pair with agent
+                        // and allocation size in following loop
+                        free_mem_info.push_back(free_memory_information{
+                            itr.start_timestamp, itr.end_timestamp, _address});
+                        if(address_to_agent_and_size.count(_address) == 0)
+                        {
+                            if(_address.handle == 0)
+                            {
+                                // Freeing null pointers is expected behavior and is occurs in HSA
+                                // functions like hipStreamDestroy
+                                ROCP_INFO << "null pointer freed due to HSA operation";
+                            }
+                            else
+                            {
+                                // Following should not occur
+                                ROCP_INFO << "Unpaired free operation occurred";
+                            }
+                        }
+                        else
+                        {
+                            auto [agent_abs_index, size] = address_to_agent_and_size[_address];
+                            _node_id                     = agent_abs_index;
+                            _allocation_size             = 0;
+                        }
+                    }
+                    else
+                    {
+                        ROCP_CI_LOG(WARNING) << "unhandled memory allocation type " << _type;
+                    }
+
                     auto evt_id = create_event(
                         conn,
                         {

source/lib/python/rocpd/libpyrocpd.cpp

@@ -460,6 +460,9 @@ PYBIND11_MODULE(libpyrocpd, pyrocpd)
 
                         auto memory_copies = rocpd::sql_generator<rocpd::types::memory_copies>{
                             conn, select_guid_nid_pid("memory_copies")};
+
+                        auto scratch_memory = rocpd::sql_generator<rocpd::types::scratch_memory>{
+                            conn, select_guid_nid_pid("scratch_memory")};
 
                         auto counters = rocpd::sql_generator<rocpd::types::counter>{
                             conn, select_guid_nid_pid("counters_collection")};
@@ -496,6 +499,7 @@ PYBIND11_MODULE(libpyrocpd, pyrocpd)
                                                       samples,
                                                       kernels,
                                                       memory_copies,
+                                                      scratch_memory,
                                                       memory_allocations,
                                                       counters);
                     }

source/lib/python/rocpd/source/perfetto.cpp

@@ -174,6 +174,7 @@ write_perfetto(
     const tool::generator<types::sample>&            sample_gen,
     const tool::generator<types::kernel_dispatch>&   kernel_dispatch_gen,
     const tool::generator<types::memory_copies>&     memory_copy_gen,
+    const tool::generator<types::scratch_memory>&    scratch_memory_gen,
     const tool::generator<types::memory_allocation>& memory_allocation_gen,
     const tool::generator<types::counter>&           counter_collection_gen)
 {
@@ -687,13 +688,15 @@ write_perfetto(
             rocprofiler_timestamp_t start_timestamp = 0;
             rocprofiler_timestamp_t end_timestamp   = 0;
             rocprofiler_address_t   address         = {.handle = 0};
+            rocprofiler_queue_id_t  queue           = {.handle = 0};
         };
 
         struct memory_information
         {
-            uint64_t              alloc_size  = {0};
-            rocprofiler_address_t address     = {.handle = 0};
-            bool                  is_alloc_op = {false};
+            uint64_t               alloc_size  = {0};
+            rocprofiler_address_t  address     = {.handle = 0};
+            rocprofiler_queue_id_t queue       = {.handle = 0};
+            bool                   is_alloc_op = {false};
         };
 
         struct agent_and_size
@@ -708,7 +711,8 @@ write_perfetto(
             std::numeric_limits<uint64_t>::max(), std::numeric_limits<uint64_t>::min()};
         auto address_to_agent_and_size =
             std::unordered_map<rocprofiler_address_t, agent_and_size>{};
-        auto free_mem_info = std::vector<free_memory_information>{};
+        auto queue_to_agent_and_size = std::unordered_map<rocprofiler_queue_id_t, agent_and_size>{};
+        auto free_mem_info           = std::vector<free_memory_information>{};
 
         // Load memory allocation endpoints
         for(auto ditr : memory_allocation_gen)
@@ -719,24 +723,47 @@ write_perfetto(
                 {
                     LOG_IF(FATAL, itr.agent_name.empty())
                         << "Missing agent id for memory allocation trace";
-                    mem_alloc_endpoints[itr.agent_abs_index].emplace(
-                        itr.start,
-                        memory_information{
-                            itr.size, rocprofiler_address_t{.handle = itr.address}, true});
-                    mem_alloc_endpoints[itr.agent_abs_index].emplace(
-                        itr.end,
-                        memory_information{
-                            itr.size, rocprofiler_address_t{.handle = itr.address}, true});
-                    address_to_agent_and_size.emplace(
-                        rocprofiler_address_t{.handle = itr.address},
-                        agent_and_size{itr.agent_abs_index, itr.size});
+
+                    if(itr.level == "REAL")
+                    {
+                        mem_alloc_endpoints[itr.agent_abs_index].emplace(
+                            itr.start,
+                            memory_information{itr.size,
+                                               rocprofiler_address_t{.handle = itr.address},
+                                               rocprofiler_queue_id_t{.handle = itr.queue_id},
+                                               true});
+                        mem_alloc_endpoints[itr.agent_abs_index].emplace(
+                            itr.end,
+                            memory_information{itr.size,
+                                               rocprofiler_address_t{.handle = itr.address},
+                                               rocprofiler_queue_id_t{.handle = itr.queue_id},
+                                               true});
+
+                        address_to_agent_and_size.emplace(
+                            rocprofiler_address_t{.handle = itr.address},
+                            agent_and_size{itr.agent_abs_index, itr.size});
+                    }
+                    // Scratch memory operations operations are indexed by queue id as agent
+                    // id is not available
+                    else if(itr.level == "SCRATCH")
+                    {
+                        queue_to_agent_and_size.emplace(
+                            rocprofiler_queue_id_t{.handle = itr.queue_id},
+                            agent_and_size{itr.agent_abs_index, itr.size});
+                    }
                 }
                 else if(itr.type == "FREE")
                 {
                     // Store free memory operations in seperate vector to pair with agent
                     // and allocation size in following loop
-                    free_mem_info.push_back(free_memory_information{
-                        itr.start, itr.end, rocprofiler_address_t{.handle = itr.address}});
+                    if(itr.level == "REAL")
+                    {
+                        free_mem_info.push_back(free_memory_information{
+                            itr.start,
+                            itr.end,
+                            rocprofiler_address_t{.handle = itr.address},
+                            rocprofiler_queue_id_t{.handle = itr.queue_id}});
+                    }
                 }
                 else
                 {
@@ -765,9 +792,9 @@ write_perfetto(
             }
             auto [agent_abs_index, size] = address_to_agent_and_size[itr.address];
             mem_alloc_endpoints[agent_abs_index].emplace(
-                itr.start_timestamp, memory_information{size, itr.address, false});
+                itr.start_timestamp, memory_information{size, itr.address, itr.queue, false});
             mem_alloc_endpoints[agent_abs_index].emplace(
-                itr.end_timestamp, memory_information{size, itr.address, false});
+                itr.end_timestamp, memory_information{size, itr.address, itr.queue, false});
         }
         // Create running sum of allocated memory
         for(auto& [_, endpoint_map] : mem_alloc_endpoints)
@@ -817,10 +844,10 @@ write_perfetto(
         {
             mem_alloc_endpoints[abs_index].emplace(
                 mem_alloc_extremes.first - extremes_endpoint_buffer,
-                memory_information{0, {0}, false});
+                memory_information{0, {0}, {0}, false});
             mem_alloc_endpoints[abs_index].emplace(
                 mem_alloc_extremes.second + extremes_endpoint_buffer,
-                memory_information{0, {0}, false});
+                memory_information{0, {0}, {0}, false});
 
             auto _track_name = std::stringstream{};
 
@@ -853,9 +880,110 @@ write_perfetto(
                               mem_alloc_tracks.at(alloc_itr.first),
                               itr.first,
                               itr.second.alloc_size / bytes_multiplier);
+                tracing_session->FlushBlocking();
+            }
+        }
+
+        // scratch memory counter track
+        auto scratch_mem_endpoints =
+            std::unordered_map<uint64_t, std::map<rocprofiler_timestamp_t, uint64_t>>{};
+        auto scratch_mem_extremes = std::pair<uint64_t, uint64_t>{
+            std::numeric_limits<uint64_t>::max(), std::numeric_limits<uint64_t>::min()};
+
+        // Load scratch memory usage endpoints
+        for(auto ditr : scratch_memory_gen)
+            for(auto itr : scratch_memory_gen.get(ditr))
+            {
+                auto agent_abs_index = itr.agent_abs_index;
+                if(itr.operation == "FREE")
+                {
+                    auto [agent_index, size] =
+                        queue_to_agent_and_size[rocprofiler_queue_id_t{.handle = itr.queue_id}];
+                    agent_abs_index = agent_index;
+                }
+
+                // Track start and end timestamps for this scratch memory record
+                scratch_mem_endpoints[agent_abs_index].emplace(itr.start, 0);
+                scratch_mem_endpoints[agent_abs_index].emplace(itr.end, 0);
+
+                // Update overall time range
+                scratch_mem_extremes =
+                    std::make_pair(std::min(scratch_mem_extremes.first, itr.start),
+                                   std::max(scratch_mem_extremes.second, itr.end));
+            }
+
+        // Load values at each endpoint
+        for(auto ditr : scratch_memory_gen)
+            for(auto itr : scratch_memory_gen.get(ditr))
+            {
+                auto agent_abs_index = itr.agent_abs_index;
+                if(itr.operation == "FREE")
+                {
+                    auto [agent_index, size] =
+                        queue_to_agent_and_size[rocprofiler_queue_id_t{.handle = itr.queue_id}];
+                    agent_abs_index = agent_index;
+                }
+                // For each timestamp in the range of this record
+                auto begin = scratch_mem_endpoints.at(agent_abs_index).lower_bound(itr.start);
+                auto end   = scratch_mem_endpoints.at(agent_abs_index).upper_bound(itr.end);
+
+                for(auto mitr = begin; mitr != end; ++mitr)
+                {
+                    // Add scratch memory size to the counter value at this timestamp
+                    if(itr.operation == "ALLOC")
+                        mitr->second = itr.size;
+                    else if(itr.operation == "FREE")
+                        mitr->second = 0;  // For all free events current allocation drops to 0.
+                }
+            }
+
+        // Create counter tracks for visualization
+        auto scratch_mem_tracks = std::unordered_map<uint64_t, ::perfetto::CounterTrack>{};
+        auto scratch_mem_names  = std::vector<std::string>{};
+        scratch_mem_names.reserve(scratch_mem_endpoints.size());
+
+        for(auto& [abs_index, ts_map] : scratch_mem_endpoints)
+        {
+            // Add buffer timestamps for better visualization
+            if(!ts_map.empty())
+            {
+                scratch_mem_endpoints[abs_index].emplace(
+                    scratch_mem_extremes.first - extremes_endpoint_buffer, 0);
+                scratch_mem_endpoints[abs_index].emplace(
+                    scratch_mem_extremes.second + extremes_endpoint_buffer, 0);
+
+                auto       _track_name      = std::stringstream{};
+                const auto _agent           = agent_data.at(abs_index).first;
+                auto       agent_index_info = agent_data.at(abs_index).second;
+
+                _track_name << "SCRATCH MEMORY on " << agent_index_info.label << " ["
+                            << agent_index_info.index << "] (" << agent_index_info.type << ")";
+
+                constexpr auto _unit = ::perfetto::CounterTrack::Unit::UNIT_SIZE_BYTES;
+                auto&          _name = scratch_mem_names.emplace_back(_track_name.str());
+                scratch_mem_tracks.emplace(abs_index,
+                                           ::perfetto::CounterTrack{_name.c_str(), this_pid_track}
+                                               .set_unit(_unit)
+                                               .set_unit_multiplier(bytes_multiplier)
+                                               .set_is_incremental(false));
+            }
+        }
+
+        // Write counter values to perfetto trace
+        for(auto& mitr : scratch_mem_endpoints)
+        {
+            if(scratch_mem_tracks.count(mitr.first) > 0)
+            {
+                for(auto itr : mitr.second)
+                {
+                    TRACE_COUNTER(sdk::perfetto_category<sdk::category::scratch_memory>::name,
+                                  scratch_mem_tracks.at(mitr.first),
+                                  itr.first,
+                                  itr.second / bytes_multiplier);
+                    tracing_session->FlushBlocking();
+                }
             }
         }
-        tracing_session->FlushBlocking();
     }
 
     // Create counter tracks per agent

source/lib/python/rocpd/source/perfetto.hpp

@@ -63,6 +63,7 @@ write_perfetto(
     const tool::generator<types::sample>&            sample_gen,
     const tool::generator<types::kernel_dispatch>&   kernel_dispatch_gen,
     const tool::generator<types::memory_copies>&     memory_copy_gen,
+    const tool::generator<types::scratch_memory>&    scratch_memory_gen,
     const tool::generator<types::memory_allocation>& memory_allocation_gen,
     const tool::generator<types::counter>&           counter_collection_gen);
 }  // namespace output

tests/pytest-packages/pytest_utils/perfetto_reader.py

@@ -288,7 +288,11 @@ def extract_tp_data(self, **kwargs):
                 counter_track.id as slice_id,
                 counter.track_id,
                 counter_track.name as track_name,
-                'counter_collection' as category,
+                CASE
+                    WHEN counter_track.name LIKE '%SCRATCH MEMORY%' THEN 'scratch_memory'
+                    WHEN counter_track.name LIKE '%ALLOCATE BYTES%' THEN 'memory_allocation'
+                    ELSE 'counter_collection'
+                END as category,
                 0 as depth,
                 0 as stack_id,
                 0 as parent_stack_id,
@@ -297,7 +301,9 @@ def extract_tp_data(self, **kwargs):
                 counter_track.name as name
             FROM counter_track
             JOIN counter ON counter.track_id = counter_track.id
-            WHERE counter_track.name LIKE 'AGENT%'
+            WHERE (counter_track.name LIKE 'AGENT%'
+                   OR counter_track.name LIKE '%SCRATCH MEMORY%'
+                   OR counter_track.name LIKE '%ALLOCATE BYTES%')
             AND counter.value > 0
             GROUP BY counter.track_id"""
         )
@@ -328,7 +334,7 @@ def extract_tp_data(self, **kwargs):
                     "tp_index": counter_df["tp_index"],
                     "slice_id": counter_df["slice_id"],
                     "track_id": counter_df["track_id"],
-                    "category": "counter_collection",
+                    "category": counter_df["category"],
                     "depth": 0,
                     "stack_id": 0,
                     "parent_stack_id": 0,

tests/rocprofv3/CMakeLists.txt

@@ -46,3 +46,4 @@ add_subdirectory(minimum-bytes)
 add_subdirectory(conversion-script)
 add_subdirectory(python-bindings)
 add_subdirectory(rocpd)
+add_subdirectory(rocpd-scratch)