[rocprofiler-compute] Improve iteration multiplexing code and documentation (#2080)

* Improve Iteration multiplexing * Improve iteration multiplexing documentation by adding usage note and listing caveats * Bugfixes for iteration mulitplexing * Use merge iteration multiplexing in analysis webui and db mode * Do not remove Dispatch_ID column in merge iteration multiplexing since it is needed for analysis of top dispatches based on duration * Bugfixes for analysis logic * Graceful handling of missing counters in case of iteration multiplexing * Improved warnings when metrics could not be calculated due to missing counter data * Fix the check to prevent showing table when a column is full of N/A * Improve detection of empty values when metric evaludation fails due to missing counter data * Bugfixes for profile logic * Fix kernel filtering during roofline benchmark phase * Update changelog for bugfixes * Remove unnecessary columns when merging dispatches for iteration multiplexing * bugfix * Better analysis warnings * fix to_std() in parser * Use median in merge iteration multiplex * Address review comments * Fix cmake formatting * fix None handling of parser util functions * Enable stochastic counter accuracy test * fix cmake formatting
2025-12-18 11:51:21 -05:00
@@ -60,6 +60,7 @@ from utils.roofline_calc import (
    SUPPORTED_DATATYPES,
 )
 from utils.utils import get_uuid, get_version
+import numpy as np


 class db_analysis(OmniAnalyze_Base):
@@ -77,22 +78,7 @@ class db_analysis(OmniAnalyze_Base):
            )
        self._roofline_ceilings_per_workload = self.calc_roofline_ceilings()
        self._pc_sampling_data_per_workload = self.calc_pc_sampling_data()
-        self._pmc_df_per_workload = {
-            workload_path: rocpd_data.process_rocpd_csv(
-                pd.read_csv(Path(workload_path) / "pmc_perf.csv")
-            )
-            for workload_path in self._runs.keys()
-        }
-        self._top_kernels_per_workload = {
-            workload_path: pmc_df.assign(
-                duration=pmc_df["End_Timestamp"] - pmc_df["Start_Timestamp"]
-            )
-            .sort_values(by="duration", ascending=False)
-            .drop_duplicates("Kernel_Name")["Kernel_Name"]
-            .to_list()
-            for workload_path, pmc_df in self._pmc_df_per_workload.items()
-        }
-        console_debug("Collected dispatch data")
+        self._pmc_df_per_workload = self.calc_pmc_df_data()
        self._pmc_df_per_workload = self.apply_pmc_filters()
        self._dispatch_data_per_workload = self.calc_dispatch_data()
        self._metrics_info_data_per_workload, self._values_data_per_workload = (
@@ -233,6 +219,34 @@ class db_analysis(OmniAnalyze_Base):
        console_debug("Completed writing database")
        console_warning(f"Created file: {db_name}")

+    def calc_pmc_df_data(self) -> dict[str, pd.DataFrame]:
+        pmc_df_per_workload: dict[str, pd.DataFrame] = {}
+        args = self.get_args()
+
+        for workload_path in self._runs.keys():
+            pmc_df = rocpd_data.process_rocpd_csv(
+                pd.read_csv(Path(workload_path) / "pmc_perf.csv")
+            )
+
+            # Create multi index df with collection level as pmc_perf
+            raw_pmc = pd.concat([pmc_df], keys=["pmc_perf"], axis=1, copy=False)
+
+            if args.spatial_multiplexing:
+                raw_pmc = self.spatial_multiplex_merge_counters(
+                    raw_pmc
+                )
+
+            if self._profiling_config.get("iteration_multiplexing") is not None:
+                raw_pmc = self.iteration_multiplex_merge_counters(
+                    raw_pmc,
+                    policy=self._profiling_config["iteration_multiplexing"],
+                )
+
+            pmc_df_per_workload[workload_path] = raw_pmc["pmc_perf"]
+
+        console_debug("Collected dispatch data")
+        return pmc_df_per_workload
+
    def calc_roofline_ceilings(self) -> dict[str, dict[str, Any]]:
        roofline_ceilings_per_workload: dict[str, dict[str, Any]] = {}

@@ -404,7 +418,7 @@ class db_analysis(OmniAnalyze_Base):
                value,
            )
        try:
-            return eval(
+            eval_result = eval(
                compile(value, "<string>", "eval"),
                {},  # no globals
                {
@@ -424,6 +438,28 @@ class db_analysis(OmniAnalyze_Base):
                    "to_sum": to_sum,
                },
            )
+
+            # eval_result can be None if expression has None explicitly specified
+            # Do not give warning for this case and simply return None
+            if eval_result is None or "None" in value:
+                return None
+
+            # Only return None for scalar NA values
+            # For vectors/Series, return as-is to preserve shape for downstream operations
+            # Note: pd.NA is not detected as scalar by np.isscalar()
+            is_scalar_na = (
+                eval_result is pd.NA
+                or (np.isscalar(eval_result) and pd.isna(eval_result))
+            )
+
+            if is_scalar_na:
+                console_warning(
+                    f"Could not evaluate expression for {name}: {value} - likely due to missing "
+                    "counter data."
+                )
+                return None
+            else:
+                return eval_result
        except Exception as e:
            console_warning(f"Failed to evaluate expression for {name}: {value} - {e}")
            return None
@@ -587,6 +623,14 @@ class db_analysis(OmniAnalyze_Base):
        pmc_df_per_workload = self._pmc_df_per_workload.copy()

        for workload_path, pmc_df in pmc_df_per_workload.items():
+            top_kernels = (
+                pmc_df.assign(
+                    duration=pmc_df["End_Timestamp"] - pmc_df["Start_Timestamp"]
+                )
+                .sort_values(by="duration", ascending=False)
+                .drop_duplicates("Kernel_Name")["Kernel_Name"]
+                .to_list()
+            )
            # Filter gpu_ids
            if self._runs[workload_path].filter_gpu_ids:
                pmc_df = pmc_df.loc[
@@ -598,7 +642,7 @@ class db_analysis(OmniAnalyze_Base):
            if self._runs[workload_path].filter_kernel_ids:
                pmc_df = pmc_df.loc[
                    pmc_df["Kernel_Name"].isin([
-                        self._top_kernels_per_workload[workload_path][id]
+                        top_kernels[id]
                        for id in self._runs[workload_path].filter_kernel_ids
                    ])
                ]
@@ -639,6 +683,14 @@ class db_analysis(OmniAnalyze_Base):
                "l2_cache_data": roofline_data_expressions.get("AI L2", ""),
                "hbm_cache_data": roofline_data_expressions.get("AI HBM", ""),
            }
+            top_kernels = (
+                pmc_df.assign(
+                    duration=pmc_df["End_Timestamp"] - pmc_df["Start_Timestamp"]
+                )
+                .sort_values(by="duration", ascending=False)
+                .drop_duplicates("Kernel_Name")["Kernel_Name"]
+                .to_list()
+            )

            roofline_df = pd.DataFrame([
                {
@@ -653,7 +705,7 @@ class db_analysis(OmniAnalyze_Base):
                        for metric_name in roofline_data_expressions
                    },
                }
-                for kernel_name in self._top_kernels_per_workload[workload_path][
+                for kernel_name in top_kernels[
                    : self.get_args().max_stat_num
                ]
            ])
@@ -354,6 +354,12 @@ class webui_analysis(OmniAnalyze_Base):
                self._runs[self.dest_dir].raw_pmc
            )

+        if self._profiling_config.get("iteration_multiplexing") is not None:
+            self._runs[self.dest_dir].raw_pmc = self.iteration_multiplex_merge_counters(
+                self._runs[self.dest_dir].raw_pmc,
+                policy=self._profiling_config["iteration_multiplexing"],
+            )
+
        file_io.create_df_kernel_top_stats(
            df_in=self._runs[self.dest_dir].raw_pmc,
            raw_data_dir=self.dest_dir,
@@ -184,7 +184,8 @@ class Roofline:
        df_list = df_pmc["Kernel_Name"].tolist()

        for idx in range(len(df_list)):
-            if df_list[idx].split("(")[0] not in args.kernel:
+            # If there is no any kernel match, drop the row
+            if not any([kernel in df_list[idx] for kernel in args.kernel]):
                df_filtered.drop(index=idx, inplace=True)

        # Verify that final filtered kernel df matches the kernel list requested
@@ -118,7 +118,7 @@ PC_SAMPLING_NOT_ISSUE_PREFIX = "ROCPROFILER_PC_SAMPLING_INSTRUCTION_NOT_ISSUED_R
 # ------------------------------------------------------------------------------


-def to_min(*args: Any) -> Union[float, None]:
+def to_min(*args: Any) -> float:
    if len(args) == 1 and isinstance(args[0], pd.Series):
        return args[0].min()
    elif min(args) is None:
@@ -127,7 +127,7 @@ def to_min(*args: Any) -> Union[float, None]:
        return min(args)


-def to_max(*args: Any) -> Union[float, np.ndarray, None]:
+def to_max(*args: Any) -> Union[float, np.ndarray]:
    if len(args) == 1 and isinstance(args[0], pd.Series):
        return args[0].max()
    elif len(args) == 2 and (
@@ -142,9 +142,11 @@ def to_max(*args: Any) -> Union[float, np.ndarray, None]:

 def to_avg(
    a: Union[pd.Series, np.ndarray, list, int, float, str, np.number, None],
-) -> Union[float, np.floating, None]:
+) -> Union[float, np.floating]:
    if a is None:
        return np.nan
+    if np.isscalar(a) and pd.isna(a):
+        return np.nan
    elif isinstance(a, pd.Series):
        if a.empty:
            return np.nan
@@ -173,9 +175,9 @@ def to_avg(
        raise Exception(f"to_avg: unsupported type: {type(a)}")


-def to_median(a: Union[pd.Series, None]) -> Union[float, None]:
+def to_median(a: Union[pd.Series, None]) -> float:
    if a is None:
-        return None
+        return np.nan
    elif isinstance(a, pd.Series):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", category=RuntimeWarning)
@@ -186,6 +188,9 @@ def to_median(a: Union[pd.Series, None]) -> Union[float, None]:

 def to_std(a: pd.Series) -> float:
    if isinstance(a, pd.Series):
+        # Define std as 0.0 if there is only one element
+        if len(a) <= 1:
+            return 0.0
        return a.std()
    else:
        raise Exception("to_std: unsupported type.")
@@ -193,20 +198,23 @@ def to_std(a: pd.Series) -> float:

 def to_int(
    a: Union[int, float, str, np.integer, pd.Series, None],
-) -> Union[int, pd.Series, None]:
+) -> Union[int, float, pd.Series]:
    if a is None:
-        return None
+        return np.nan
+    if np.isscalar(a) and pd.isna(a):
+        return np.nan
    elif isinstance(a, (int, float, np.integer)):
        return int(a)
    elif isinstance(a, pd.Series):
-        return a.astype(int)
+        # "Int64" handles null values
+        return a.astype("Int64")
    elif isinstance(a, str):
        return int(a)
    else:
        raise Exception("to_int: unsupported type.")


-def to_sum(a: Union[pd.Series, None]) -> Union[float, None]:
+def to_sum(a: Union[pd.Series, None]) -> float:
    if a is None:
        return np.nan
    elif np.isnan(a).all():
@@ -226,9 +234,9 @@ def to_round(a: Union[pd.Series, float], b: int) -> Union[pd.Series, float]:
        return round(a, b)


-def to_quantile(a: Union[pd.Series, None], b: float) -> Union[float, None]:
+def to_quantile(a: Union[pd.Series, None], b: float) -> float:
    if a is None:
-        return None
+        return np.nan
    elif isinstance(a, pd.Series):
        return a.quantile(b)
    else:
@@ -346,7 +354,26 @@ class MetricEvaluator:
                local_expr_context,
            )

-            if eval_result is None or np.isnan(eval_result).any():
+            # Only return "N/A" for scalar NA values
+            # For vectors/Series, return as-is to preserve shape for
+            # downstream operations
+            # Note: None and pd.NA are not detected as scalar by np.isscalar()
+            if (
+                eval_result is None
+                or eval_result is pd.NA
+                or (np.isscalar(eval_result) and pd.isna(eval_result))
+            ):
+                # Do not give warning if None is explicitly specified in expression
+                if "None" not in expr:
+                    console_warning(
+                        f"Could not evaluate expression '{expr}' - likely "
+                        "due to missing counter data."
+                    )
+                else:
+                    console_debug(
+                        f"Expression '{expr}' evaluated to None - likely "
+                        "explicitly specified."
+                    )
                return "N/A"
            else:
                return eval_result
@@ -360,18 +387,18 @@ class MetricEvaluator:
                return "N/A"

        except AttributeError as attribute_error:
-            if str(attribute_error) == "'NoneType' object has no attribute 'get'":
-                console_warning(
-                    f"Failed to evaluate expression '{expr}': {attribute_error}."
-                )
-                return "N/A"
-            else:
-                console_error("analysis", str(attribute_error))
-                return "N/A"
+            console_warning(
+                f"Failed to evaluate expression '{expr}': {attribute_error}."
+            )
+            return "N/A"

        except pd.errors.IntCastingNaNError as exception:
-            console_warning(f"Missing data: {exception}. Using empty value.")
-            return ""
+            console_warning(f"Failed to evaluate expression '{expr}': {exception}.")
+            return "N/A"
+
+        except ValueError as value_error:
+            console_warning(f"Failed to evaluate expression '{expr}': {value_error}.")
+            return "N/A"


 def build_eval_string(equation: str, coll_level: str, config: dict) -> str:
@@ -929,9 +956,12 @@ def calc_builtin_vars(
            # Pass sys_vars so that $num_xcd and other system variables are available
            temporary_evaluator = MetricEvaluator(raw_pmc_df, sys_vars, {})
            calculation_result = temporary_evaluator.eval_expression(eval_string)
+            # Convert "N/A" string to np.nan to maintain numeric type for calculations
+            if np.isscalar(calculation_result) and calculation_result == "N/A":
+                calculation_result = np.nan
            builtin_vars_collection[f"ammolite__{variable_key}"] = calculation_result
        except (TypeError, NameError, KeyError, AttributeError):
-            builtin_vars_collection[f"ammolite__{variable_key}"] = None
+            builtin_vars_collection[f"ammolite__{variable_key}"] = np.nan

    # Second pass: calculate remaining variables that depend on per-XCD values
    for variable_key, variable_value in BUILD_IN_VARS.items():
@@ -946,9 +976,12 @@ def calc_builtin_vars(
            combined_vars = {**sys_vars, **builtin_vars_collection}
            temporary_evaluator = MetricEvaluator(raw_pmc_df, combined_vars, {})
            calculation_result = temporary_evaluator.eval_expression(eval_string)
+            # Convert "N/A" string to np.nan to maintain numeric type for calculations
+            if np.isscalar(calculation_result) and calculation_result == "N/A":
+                calculation_result = np.nan
            builtin_vars_collection[f"ammolite__{variable_key}"] = calculation_result
        except (TypeError, NameError, KeyError, AttributeError):
-            builtin_vars_collection[f"ammolite__{variable_key}"] = None
+            builtin_vars_collection[f"ammolite__{variable_key}"] = np.nan

    return builtin_vars_collection

@@ -1653,9 +1686,7 @@ def load_pc_sampling_data(
    csv_kernel_trace_file_path = Path(dir_path) / f"{file_prefix}_kernel_trace.csv"

    if not csv_kernel_trace_file_path.exists():
-        console_error(
-            f"PC sampling: can not read {csv_kernel_trace_file_path}", exit=False
-        )
+        console_warning(f"PC sampling: can not read {csv_kernel_trace_file_path}")
        return pd.DataFrame()

    if stochastic_path.exists():
@@ -1722,7 +1753,7 @@ def load_pc_sampling_data(

    elif len(workload.filter_kernel_ids) == 1:
        if not json_file_path.exists():
-            console_error(f"PC sampling: can not read {json_file_path}", exit=False)
+            console_warning(f"PC sampling: can not read {json_file_path}")
            return pd.DataFrame()
        else:
            # NB:
@@ -382,7 +382,7 @@ def format_table_output(

    # Check if any column in df is empty
    is_empty_columns_exist = any(
-        df.replace("", None).iloc[:, col_idx].isnull().all()
+        df.replace(["", "N/A"], None).iloc[:, col_idx].isnull().all()
        for col_idx in range(len(df.columns))
    )

@@ -1384,10 +1384,6 @@ def merge_counters_iteration_multiplex(
        "Kernel_ID",
    ]

-    expired_column_index = [
-        "Dispatch_ID",
-    ]
-
    result_dfs: list[pd.DataFrame] = []

    # TODO: will need to optimize to avoid this conversion to single index format
@@ -1419,30 +1415,32 @@ def merge_counters_iteration_multiplex(

        pd.set_option("display.max_columns", None)

+        # Reset Dispatch_ID
+        dispatch_id_counter = 0
+
        for name, group in unique_occurences:
            # Create a dictionary to store the merged row for the current group
            merged_row: dict[str, Any] = {}

            # Process non-counter columns
-            for col in [
-                col
-                for col in non_counter_column_index
-                if col not in expired_column_index
-            ]:
+            for col in non_counter_column_index:
                if 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
-                    merged_row["Median_Time"] = median_delta_time
-                    merged_row["Mean_Time"] = delta_time.mean()
+                    delta_time = group[col] - group["Start_Timestamp"]
+                    merged_row[col] = group["Start_Timestamp"] + delta_time.median()
+                if col == "Dispatch_ID":
+                    # Assign new Dispatch_ID
+                    merged_row[col] = dispatch_id_counter
+                    dispatch_id_counter += 1
                elif pd.api.types.is_numeric_dtype(group[col]):
                    # For other non-counter numeric columns, take the median value
                    merged_row[col] = group[col].median()
                    if pd.api.types.is_integer_dtype(group[col]):
                        merged_row[col] = merged_row[col].astype(int)
                else:
-                    # For other non-counter columns, take the first occurrence (0th row)
+                    # For other non-counter non-numeric columns,
+                    # take the first occurrence (0th row)
+                    # Only Kernel_Name should be non-numeric here
                    merged_row[col] = group.iloc[0][col]

            # Process counter columns (assumed to be all columns not in
@@ -1451,16 +1449,19 @@ def merge_counters_iteration_multiplex(
                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
-
-            merged_row["Count"] = group["Dispatch_ID"].nunique()
+                # For counter columns, calculate median only across non-NaN values
+                # Preserve original data type
+                valid_values = group[counter_col].dropna()
+                if not valid_values.empty:
+                    median_value = valid_values.median()
+                    # Preserve original data type - check if all
+                    # non-null values are integers
+                    if (valid_values == valid_values.astype(int)).all():
+                        merged_row[counter_col] = int(median_value)
+                    else:
+                        merged_row[counter_col] = median_value
+                else:
+                    merged_row[counter_col] = None

            # Append the merged row to the result list
            result_data.append(merged_row)
@@ -1543,9 +1544,8 @@ def merge_counters_spatial_multiplex(df_multi_index: pd.DataFrame) -> pd.DataFra
                    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
+                    delta_time = group[col] - group["Start_Timestamp"]
+                    merged_row[col] = group["Start_Timestamp"] + delta_time.median()
                else:
                    # For other non-counter columns, take the first occurrence (0th row)
                    merged_row[col] = group.iloc[0][col]