get_df_xla_perf raises KeyError for FP8 and s8 dtypes in XLA kernel operands

[brieflynn]

FP8 variants f8e4m3fnuz, f8e5m2fnuz or s8 cause KeyError

This PR addresses a KeyError in get_df_xla_perf when processing XLA kernel operands containing FP8 or int8 dtypes (e.g., f8e4m3fnuz, f8e5m2fnuz, s8). The current dtype_to_bytes lacks those types and uses direct indexing, causing the KeyError and breaking analysis flows in JAX profiler traces. This change adds these missing dtypes to the dtype_to_bytes mapping.

Problem

dtype_to_bytes in get_df_xla_perf of TraceLens/TreePerf/tree_perf.py does not include:

f8e4m3fnuz
f8e5m2fnuz
s8

Direct indexing dtype_to_bytes[dtype] crashes with KeyError when these appear in operands here:

for index, row in df_xla_events.iterrows():

            kernel_details = row.get("kernel_details")[0]
            operands = kernel_details.get("operands")

            total_input_bytes = 0
            for operand in operands:
                dtype, shape, layout = parse_dtype_shape_layout(operand)
                if shape and dtype:
                    # crashes here
                    total_input_bytes = (
                        total_input_bytes + np.prod(shape) * dtype_to_bytes[dtype]
                    )

Changes

dtype_to_bytes = {
    "f32": 4,
    "bf16": 2,
    "s32": 4,
    "fp16": 2,
    "u32": 4,
    "f16": 2,
    "u64": 8,
    # Added
    "s8": 1,
    "f8e4m3fnuz": 1,
    "f8e5m2fnuz": 1,
}

Note to AMDers:
This is a public repository. Please do not upload any confidential or customer data. Make sure all such data has been anonymized or removed before making this PR. If you need to attach any private files or links, please insert a Internal OneDrive Link or a Jira Ticket Link instead.