Spatial multiplexing: part2, analysis (#542)

Co-authored-by: Yang Wang <ywang103@amd.com> [ROCm/rocprofiler-compute commit: 808a9a79ed]
2025-02-13 16:47:25 -05:00
parent a36e280903
commit 952b675890
7 changed files with 202 additions and 17 deletions
@@ -1091,3 +1091,125 @@ def set_locale_encoding():
            exit=False,
        )
        console_error(error)
+
+
+def reverse_multi_index_df_pmc(final_df):
+    """
+    Util function to decompose multi-index dataframe.
+    """
+    # Check if the columns have more than one level
+    if len(final_df.columns.levels) < 2:
+        raise ValueError("Input DataFrame does not have a multi-index column.")
+
+    # Extract the first level of the MultiIndex columns (the file names)
+    coll_levels = final_df.columns.get_level_values(0).unique().tolist()
+
+    # Initialize the list of DataFrames
+    dfs = []
+
+    # Loop through each 'coll_level' and rebuild the DataFrames
+    for level in coll_levels:
+        # Select columns that belong to the current 'coll_level'
+        columns_for_level = final_df.xs(level, axis=1, level=0)
+
+        # Append the DataFrame for this level
+        dfs.append(columns_for_level)
+
+    # Return the list of DataFrames and the column levels
+    return dfs, coll_levels
+
+
+def merge_counters_spatial_multiplex(df_multi_index):
+    """
+    For spatial multiplexing, this merges counter values for the same kernel that runs on different devices. For time stamp, start time stamp will use median while for end time stamp, it will be equal to the summation between median start stamp and median delta time.
+    """
+    non_counter_column_index = [
+        "Dispatch_ID",
+        "GPU_ID",
+        "Queue_ID",
+        "PID",
+        "TID",
+        "Grid_Size",
+        "Workgroup_Size",
+        "LDS_Per_Workgroup",
+        "Scratch_Per_Workitem",
+        "Arch_VGPR",
+        "Accum_VGPR",
+        "SGPR",
+        "Wave_Size",
+        "Kernel_Name",
+        "Start_Timestamp",
+        "End_Timestamp",
+        "Correlation_ID",
+        "Kernel_ID",
+        "Node",
+    ]
+
+    expired_column_index = [
+        "Node",
+        "PID",
+        "TID",
+        "Queue_ID",
+    ]
+
+    result_dfs = []
+
+    # TODO: will need optimize to avoid this convertion to single index format and do merge directly on multi-index dataframe
+    dfs, coll_levels = reverse_multi_index_df_pmc(df_multi_index)
+
+    for df in dfs:
+        kernel_name_column_name = "Kernel_Name"
+        if not "Kernel_Name" in df and "Name" in df:
+            kernel_name_column_name = "Name"
+
+        # Find the values in Kernel_Name that occur more than once
+        kernel_single_occurances = df[kernel_name_column_name].value_counts().index
+
+        # Define a list to store the merged rows
+        result_data = []
+
+        for kernel_name in kernel_single_occurances:
+            # Get all rows for the current kernel_name
+            group = df[df[kernel_name_column_name] == kernel_name]
+
+            # Create a dictionary to store the merged row for the current group
+            merged_row = {}
+
+            # Process non-counter columns
+            for col in [
+                col for col in non_counter_column_index if col not in expired_column_index
+            ]:
+                if col == "Start_Timestamp":
+                    # For Start_Timestamp, take the median
+                    merged_row[col] = group["Start_Timestamp"].median()
+                elif col == "End_Timestamp":
+                    # For End_Timestamp, calculate the median delta time
+                    delta_time = group["End_Timestamp"] - group["Start_Timestamp"]
+                    median_delta_time = delta_time.median()
+                    merged_row[col] = merged_row["Start_Timestamp"] + median_delta_time
+                else:
+                    # For other non-counter columns, take the first occurrence (0th row)
+                    merged_row[col] = group.iloc[0][col]
+
+            # Process counter columns (assumed to be all columns not in non_counter_column_index)
+            counter_columns = [
+                col for col in group.columns if col not in non_counter_column_index
+            ]
+            for counter_col in counter_columns:
+                # for counter columns, take the first non-none (or non-nan) value
+                current_valid_counter_group = group[group[counter_col].notna()]
+                first_valid_value = (
+                    current_valid_counter_group.iloc[0][counter_col]
+                    if len(current_valid_counter_group) > 0
+                    else None
+                )
+                merged_row[counter_col] = first_valid_value
+
+            # Append the merged row to the result list
+            result_data.append(merged_row)
+
+        # Create a new DataFrame from the merged rows
+        result_dfs.append(pd.DataFrame(result_data))
+
+    final_df = pd.concat(result_dfs, keys=coll_levels, axis=1, copy=False)
+    return final_df