[Memory Viz] Add Compile Context to Visualizer (#152862)

Summary: Adds PT2 info to visualizer. Also makes sure we have a case when compile context is not in pickle file.

Test Plan: {F1977637362}

Differential Revision: D74202811

Pull Request resolved: https://github.com/pytorch/pytorch/pull/152862
Approved by: https://github.com/aaronenyeshi

torch/utils/viz/MemoryViz.js

@@ -980,6 +980,8 @@ function process_alloc_data(snapshot, device, plot_segments, max_entries) {
       text = `${text}, Total memory used after allocation: ${formatSize(
         elem.max_allocated_mem,
       )}`;
+      const context = elem?.compile_context ?? 'None';
+      text = `${text}, Compile context: ${context}`;
       if (elem.stream !== null) {
         text = `${text}, stream ${elem.stream}`;
       }