diff --git a/tritonparse/ir_analysis.py b/tritonparse/ir_analysis.py index 20addd3..3f5a183 100644 --- a/tritonparse/ir_analysis.py +++ b/tritonparse/ir_analysis.py @@ -40,7 +40,7 @@ def process_amd_gcn_bufferops( ) -> dict[str, int]: ir_content = load_ir_contents(key, file_content, file_path) # TODO: Add atomics - io_keys = ["global_load_", "global_store_", "buffer_load_", "buffer_store_"] + io_keys = ["global_load", "global_store", "buffer_load", "buffer_store"] return process_amd_bufferop(ir_content, io_keys) @@ -64,9 +64,12 @@ def _generate_ir_analysis(entry: str): gcn_bufferops_info = process_amd_gcn_bufferops( amdgcn_key, file_content, file_path ) + io_counts = {} # NDJSON format requires a newline at the end of each line if ttgir_bufferops_info: - ir_analysis["amd_ttgir_bufferops_count"] = ttgir_bufferops_info + io_counts["amd_ttgir_bufferops_count"] = ttgir_bufferops_info if gcn_bufferops_info: - ir_analysis["amd_gcn_bufferops_count"] = gcn_bufferops_info - return {"ir_analysis": ir_analysis} + io_counts["amd_gcn_bufferops_count"] = gcn_bufferops_info + if io_counts: + ir_analysis["io_counts"] = io_counts + return ir_analysis diff --git a/tritonparse/trace_processor.py b/tritonparse/trace_processor.py index b321e6b..904e052 100644 --- a/tritonparse/trace_processor.py +++ b/tritonparse/trace_processor.py @@ -299,8 +299,13 @@ def parse_single_file( ) if compilation_event: - ir_analysis_event = _generate_ir_analysis(compilation_event) - if ir_analysis_event: + ir_analysis = _generate_ir_analysis(compilation_event) + if ir_analysis: + ir_analysis_event = { + "event_type": "ir_analysis", + "hash": _kernel_hash, + "ir_analysis": ir_analysis, + } all_output_lines[output_file].append( json.dumps(ir_analysis_event, separators=(",", ":")) + "\n" ) diff --git a/website/src/App.tsx b/website/src/App.tsx index c6d8d2c..142812f 100644 --- a/website/src/App.tsx +++ b/website/src/App.tsx @@ -12,6 +12,7 @@ import CodeView from "./pages/CodeView"; import FileDiffView from "./pages/FileDiffView"; import SingleCodeViewer from "./components/SingleCodeViewer"; import KernelOverview from "./pages/KernelOverview"; +import IRAnalysis from "./pages/IRAnalysis"; import DataSourceSelector from "./components/DataSourceSelector"; import WelcomeScreen from "./components/WelcomeScreen"; import ExternalLink from "./components/ExternalLink"; @@ -409,7 +410,7 @@ function App() { ); } else { - // Show either overview, IR code, or file diff based on active tab + // Show either overview, IR code, IR analysis, or file diff based on active tab if (activeTab === "overview") { return ( ); } + if (activeTab === "ir_analysis") { + return ( + + ); + } if (activeTab === "comparison") { return ( )} - + {dataLoaded && kernels.length > 0 && ( + + )} diff --git a/website/src/pages/IRAnalysis.tsx b/website/src/pages/IRAnalysis.tsx new file mode 100644 index 0000000..39d97e5 --- /dev/null +++ b/website/src/pages/IRAnalysis.tsx @@ -0,0 +1,98 @@ +import React from "react"; +import { ProcessedKernel } from "../utils/dataLoader"; + +interface IRAnalysisProps { + kernels: ProcessedKernel[]; + selectedKernel: number; +} + +const IRAnalysis: React.FC = ({ kernels, selectedKernel }) => { + if (kernels.length === 0) { + return ( +
+
No kernel data available
+
+ ); + } + + const kernel = kernels[selectedKernel]; + if (kernel.ir_analysis === null) { + return ( +
+
No IR Analysis available
+
+ ); + } + + const io_counts = kernel.ir_analysis?.io_counts; + const ttgir_info = io_counts?.["amd_ttgir_bufferops_count"]; + const amdgcn_info = io_counts?.["amd_gcn_bufferops_count"]; + const getCount = (info: Record | undefined, key: string): string => { return info?.[key]?.toString() ?? "N/A"; }; + + return ( +
+

Triton Kernel IR Analysis

+ +
+

+ Kernel: {kernel.name} +

+ + {io_counts && (ttgir_info || amdgcn_info) && ( + <> +

+ AMD BufferOps Information +

+ +
+
+ {ttgir_info && ( + <> +
+ Tiled Buffer Load Count + {getCount(ttgir_info, "tt.load_count")} +
+
+ Tiled Buffer Store Count + {getCount(ttgir_info, "tt.store_count")} +
+
+ Tiled Global Load Count + {getCount(ttgir_info, "amdgpu.buffer_load_count")} +
+
+ Tiled Global Store Count + {getCount(ttgir_info, "amdgpu.buffer_store_count")} +
+ + )} + {amdgcn_info && ( + <> +
+ AMDGCN Global Load Instruction Count + {getCount(amdgcn_info, "global_load_count")} +
+
+ AMDGCN Global Store Instruction Count + {getCount(amdgcn_info, "global_store_count")} +
+
+ AMDGCN Buffer Load Instruction Count + {getCount(amdgcn_info, "buffer_load_count")} +
+
+ AMDGCN Buffer Store Instruction Count + {getCount(amdgcn_info, "buffer_store_count")} +
+ + )} +
+
+ + )} +
+
+ ); +}; + +export default IRAnalysis; diff --git a/website/src/utils/dataLoader.ts b/website/src/utils/dataLoader.ts index 924f990..c064e49 100644 --- a/website/src/utils/dataLoader.ts +++ b/website/src/utils/dataLoader.ts @@ -167,6 +167,11 @@ export interface CompilationMetadata { [key: string]: any; // Allow additional unknown fields } +export interface IRAnalysisData { + // Mapping from IR stage -> count> + io_counts?: Record>; +} + /** * Extracted argument information */ @@ -224,6 +229,7 @@ export interface LogEntry { launch_index_map?: LaunchRange[]; diffs?: LaunchDiffData; sames?: LaunchSamesData; + ir_analysis?: IRAnalysisData; // Stored IR Analysis information. } /** @@ -239,6 +245,7 @@ export interface ProcessedKernel { pythonSourceInfo?: PythonSourceCodeInfo; // Python source code information metadata?: KernelMetadata; // Compilation metadata launchDiff?: LogEntry; // Aggregated launch event differences + ir_analysis?: IRAnalysisData; // Stored IR Analysis information. } /** @@ -503,6 +510,15 @@ export function processKernelData(logEntries: LogEntry[]): ProcessedKernel[] { console.warn(`Could not find matching kernel for launch_diff hash: ${hash}`); } } + if (entry.event_type === "ir_analysis") { + const hash = entry.hash; + if (hash && kernelsByHash.has(hash)) { + const kernel = kernelsByHash.get(hash)!; + kernel.ir_analysis = entry.ir_analysis!; // Attach the ir_analysis + } else { + console.warn(`Could not find matching kernel for ir_analysis hash: ${hash}`); + } + } } const finalKernels = Array.from(kernelsByHash.values());