[BE] [TritonParse] Add a web frontend for IR Analysis

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

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"
                 )

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() {
         </div>
       );
     } 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 (
           <KernelOverview
@@ -420,6 +421,14 @@ function App() {
           />
         );
       }
+      if (activeTab === "ir_analysis") {
+        return (
+          <IRAnalysis
+            kernels={kernels}
+            selectedKernel={selectedKernel}
+          />
+        );
+      }
       if (activeTab === "comparison") {
         return (
           <CodeView
@@ -550,7 +559,6 @@ function App() {
                   </button>
                 </>
               )}
-
               <button
                 className={`px-3 py-2 text-sm font-medium rounded-md ${activeTab === "file_diff" ? "bg-blue-700 text-white shadow-md" : "bg-blue-100 text-blue-700 hover:bg-blue-200"
                   }`}
@@ -561,6 +569,24 @@ function App() {
               >
                 File Diff
               </button>
+              {dataLoaded && kernels.length > 0 && (
+              <button
+                    className={`px-3 py-2 text-sm font-medium rounded-md ${activeTab === "ir_analysis" ? "bg-blue-700 text-white shadow-md" : "bg-blue-100 text-blue-700 hover:bg-blue-200"
+                      }`}
+                    onClick={() => {
+                      if (sess.preview?.active) sess.clearPreview();
+                      setActiveTab("ir_analysis");
+
+                      if (loadedUrl) {
+                        const newUrl = new URL(window.location.href);
+                        newUrl.searchParams.set("view", "ir_analysis");
+                        window.history.replaceState({}, "", newUrl.toString());
+                      }
+                    }}
+                  >
+                    IR Analysis (Beta)
+                  </button>
+              )}
             </div>
           </div>
         </div>

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<IRAnalysisProps> = ({ kernels, selectedKernel }) => {
+  if (kernels.length === 0) {
+    return (
+      <div className="flex items-center justify-center h-screen">
+        <div className="text-gray-800">No kernel data available</div>
+      </div>
+    );
+  }
+
+  const kernel = kernels[selectedKernel];
+  if (kernel.ir_analysis === null) {
+    return (
+      <div className="flex items-center justify-center h-screen">
+        <div className="text-gray-800">No IR Analysis available</div>
+      </div>
+    );
+  }
+
+  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<string, number> | undefined, key: string): string => { return info?.[key]?.toString() ?? "N/A"; };
+
+  return (
+    <div className="p-6">
+      <h1 className="text-2xl font-bold text-gray-800 mb-6">Triton Kernel IR Analysis</h1>
+
+      <div className="bg-white rounded-lg p-4 mb-4 shadow-sm border border-gray-200">
+        <h2 className="text-xl font-semibold mb-4 text-gray-800">
+          Kernel: {kernel.name}
+        </h2>
+
+        {io_counts && (ttgir_info || amdgcn_info) && (
+          <>
+            <h3 className="text-lg font-medium mb-3 text-gray-800">
+              AMD BufferOps Information
+            </h3>
+
+            <div className="bg-gray-50 p-4 rounded-md border border-gray-200">
+              <div className="grid grid-cols-[repeat(auto-fit,_minmax(180px,_1fr))] gap-3">
+                {ttgir_info && (
+                  <>
+                    <div className="flex flex-col">
+                      <span className="text-sm font-medium text-gray-500">Tiled Buffer Load Count</span>
+                      <span className="font-mono text-sm break-words">{getCount(ttgir_info, "tt.load_count")}</span>
+                    </div>
+                    <div className="flex flex-col">
+                      <span className="text-sm font-medium text-gray-500">Tiled Buffer Store Count</span>
+                      <span className="font-mono text-sm break-words">{getCount(ttgir_info, "tt.store_count")}</span>
+                    </div>
+                    <div className="flex flex-col">
+                      <span className="text-sm font-medium text-gray-500">Tiled Global Load Count</span>
+                      <span className="font-mono text-sm break-words">{getCount(ttgir_info, "amdgpu.buffer_load_count")}</span>
+                    </div>
+                    <div className="flex flex-col">
+                      <span className="text-sm font-medium text-gray-500">Tiled Global Store Count</span>
+                      <span className="font-mono text-sm break-words">{getCount(ttgir_info, "amdgpu.buffer_store_count")}</span>
+                    </div>
+                  </>
+                )}
+                {amdgcn_info && (
+                  <>
+                    <div className="flex flex-col">
+                      <span className="text-sm font-medium text-gray-500">AMDGCN Global Load Instruction Count</span>
+                      <span className="font-mono text-sm break-words">{getCount(amdgcn_info, "global_load_count")}</span>
+                    </div>
+                    <div className="flex flex-col">
+                      <span className="text-sm font-medium text-gray-500">AMDGCN Global Store Instruction Count</span>
+                      <span className="font-mono text-sm break-words">{getCount(amdgcn_info, "global_store_count")}</span>
+                    </div>
+                    <div className="flex flex-col">
+                      <span className="text-sm font-medium text-gray-500">AMDGCN Buffer Load Instruction Count</span>
+                      <span className="font-mono text-sm break-words">{getCount(amdgcn_info, "buffer_load_count")}</span>
+                    </div>
+                    <div className="flex flex-col">
+                      <span className="text-sm font-medium text-gray-500">AMDGCN Buffer Store Instruction Count</span>
+                      <span className="font-mono text-sm break-words">{getCount(amdgcn_info, "buffer_store_count")}</span>
+                    </div>
+                  </>
+                )}
+              </div>
+            </div>
+          </>
+        )}
+      </div>
+    </div>
+  );
+};
+
+export default IRAnalysis;

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 -> <IO type -> count>
+    io_counts?: Record<string, Record<string, number>>;
+}
+
 /**
  * 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());