diff --git a/src/trace_processor/perfetto_sql/stdlib/android/memory/memory_breakdown.sql b/src/trace_processor/perfetto_sql/stdlib/android/memory/memory_breakdown.sql
index 3e6467e31c..481f89b04d 100644
--- a/src/trace_processor/perfetto_sql/stdlib/android/memory/memory_breakdown.sql
+++ b/src/trace_processor/perfetto_sql/stdlib/android/memory/memory_breakdown.sql
@@ -21,23 +21,21 @@ INCLUDE PERFETTO MODULE counters.intervals;
 -- Create a table containing intervals of memory counters values, adjusted to process lifetime.
 CREATE PERFETTO TABLE _memory_breakdown_mem_intervals_raw AS
 WITH
-  -- We deny tracks that have large swings in value
-  -- This can happen because of rss_stat accounting issue: see b/418231246 for details.
   mem_process_counter_tracks AS (
     SELECT
       id,
-      name,
+      iif(name = 'Heap size (KB)', 'mem.heap', name) AS name,
       upid
     FROM process_counter_track
     WHERE
-      name IN ('mem.rss.anon', 'mem.swap', 'mem.rss.file')
+      name IN ('mem.rss.anon', 'mem.swap', 'mem.rss.file', 'mem.rss.shmem', 'Heap size (KB)', 'mem.dmabuf_rss', 'mem.locked', 'GPU Memory')
   ),
   mem_counters AS (
     SELECT
       c.id,
       c.ts,
       c.track_id,
-      c.value
+      iif(name = 'mem.heap', cast_int!(c.value) * 1024, cast_int!(c.value)) AS value
     FROM counter AS c
     JOIN mem_process_counter_tracks AS t
       ON c.track_id = t.id
@@ -51,6 +49,8 @@ WITH
       delta_value
     FROM counter_leading_intervals!(mem_counters)
   ),
+  -- We deny tracks that have large swings in value
+  -- This can happen because of rss_stat accounting issue: see b/418231246 for details.
   denied_tracks AS (
     SELECT DISTINCT
       track_id
@@ -68,6 +68,7 @@ WITH
       p.start_ts,
       p.end_ts,
       t.name AS track_name,
+      i.track_id,
       i.value
     FROM mem_intervals AS i
     JOIN mem_process_counter_tracks AS t
@@ -86,6 +87,7 @@ SELECT
   min(raw_end_ts, coalesce(end_ts, raw_end_ts)) - max(ts, coalesce(start_ts, ts)) AS dur,
   upid,
   track_name,
+  track_id,
   value
 FROM mem_intervals_with_process_lifetime
 -- Only keep rows where the clipping resulted in a positive duration.
@@ -97,46 +99,69 @@ WHERE
 -- Create a table containing intervals of OOM adjustment scores.
 -- This table will be used as the right side of a span join.
 CREATE PERFETTO TABLE _memory_breakdown_oom_intervals_prepared AS
+WITH
+  process_mem_track_ids AS (
+    SELECT
+      track_id,
+      upid
+    FROM _memory_breakdown_mem_intervals_raw
+    GROUP BY
+      track_id,
+      upid
+  )
 SELECT
-  ts,
-  dur,
-  upid,
-  bucket
-FROM android_oom_adj_intervals
+  t.track_id,
+  o.ts,
+  o.dur,
+  o.bucket
+FROM android_oom_adj_intervals AS o
+JOIN process_mem_track_ids AS t
+  USING (upid)
 WHERE
-  dur > 0;
+  o.dur > 0;
 
--- Create a virtual table that joins memory counter intervals with OOM
--- adjustment score intervals.
 CREATE VIRTUAL TABLE _memory_breakdown_mem_oom_span_join USING SPAN_LEFT_JOIN (
-  _memory_breakdown_mem_intervals_raw PARTITIONED upid,
-  _memory_breakdown_oom_intervals_prepared PARTITIONED upid
+  _memory_breakdown_mem_intervals_raw PARTITIONED track_id,
+  _memory_breakdown_oom_intervals_prepared PARTITIONED track_id
 );
 
 -- Create a table containing memory counter intervals with OOM buckets.
 CREATE PERFETTO TABLE _memory_breakdown_mem_with_buckets AS
 WITH
   -- Get the baseline values for RSS anon and swap from the zygote process.
+  zygote_upid AS (
+    SELECT
+      upid
+    FROM process
+    -- TODO: improve zygote process detection
+    WHERE
+      name IN ('zygote', 'zygote64', 'webview_zygote')
+  ),
+  zygote_tracks AS (
+    SELECT
+      iif(t.name = 'Heap size (KB)', 'mem.heap', t.name) AS track_name,
+      t.id AS track_id
+    FROM process_counter_track AS t
+    JOIN zygote_upid AS z
+      USING (upid)
+    WHERE
+      t.name IN ('mem.rss.anon', 'mem.swap', 'mem.rss.file', 'Heap size (KB)')
+  ),
   zygote_baseline AS (
     SELECT
       max(CASE WHEN track_name = 'mem.rss.anon' THEN avg_val END) AS rss_anon_base,
       max(CASE WHEN track_name = 'mem.swap' THEN avg_val END) AS swap_base,
-      max(CASE WHEN track_name = 'mem.rss.file' THEN avg_val END) AS rss_file_base
+      max(CASE WHEN track_name = 'mem.rss.file' THEN avg_val END) AS rss_file_base,
+      max(CASE WHEN track_name = 'mem.heap' THEN avg_val END) AS heap_base
     FROM (
       SELECT
-        t.name AS track_name,
-        avg(c.value) AS avg_val
+        z.track_name,
+        avg(iif(z.track_name = 'mem.heap', cast_int!(c.value) * 1024, cast_int!(c.value))) AS avg_val
       FROM counter AS c
-      JOIN process_counter_track AS t
-        ON c.track_id = t.id
-      JOIN process AS p
-        USING (upid)
-      WHERE
-        -- TODO: improve zygote process detection
-        p.name IN ('zygote', 'zygote64', 'webview_zygote')
-        AND t.name IN ('mem.rss.anon', 'mem.swap', 'mem.rss.file')
+      JOIN zygote_tracks AS z
+        USING (track_id)
       GROUP BY
-        t.name
+        z.track_name
     )
   ),
   mem_with_zygote_baseline AS (
@@ -149,6 +174,8 @@ WITH
         THEN b.swap_base
         WHEN track_name = 'mem.rss.file'
         THEN b.rss_file_base
+        WHEN track_name = 'mem.heap'
+        THEN b.heap_base
         ELSE 0
       END AS zygote_baseline_value
     FROM _memory_breakdown_mem_oom_span_join AS s
@@ -159,17 +186,17 @@ SELECT
   ts,
   dur,
   track_name,
-  app.name AS process_name,
+  p.name AS process_name,
   upid,
   pid,
   coalesce(bucket, 'unknown') AS bucket,
   CASE
-    WHEN app.upid IS NOT NULL
+    WHEN NOT p.upid IS NULL AND NOT p.name IN ('zygote', 'zygote64', 'webview_zygote')
     THEN max(0, cast_int!(value) - cast_int!(IFNULL(zygote_baseline_value, 0)))
     ELSE cast_int!(value)
   END AS zygote_adjusted_value
 FROM mem_with_zygote_baseline
-LEFT JOIN process AS app
+LEFT JOIN process AS p
   USING (upid)
 WHERE
   dur > 0;
diff --git a/ui/src/plugins/com.android.MemoryViz/index.ts b/ui/src/plugins/com.android.MemoryViz/index.ts
index c98b4a0073..27d13986b3 100644
--- a/ui/src/plugins/com.android.MemoryViz/index.ts
+++ b/ui/src/plugins/com.android.MemoryViz/index.ts
@@ -34,16 +34,31 @@ export default class MemoryViz implements PerfettoPlugin {
       name: 'Memory: Visualize (over selection)',
       callback: async () => {
         const window = await getTimeSpanOfSelectionOrVisibleWindow(ctx);
-        const rssAnonSwapTrack = await this.createRssAnonSwapTrack(ctx, window);
-        ctx.defaultWorkspace.pinnedTracksNode.addChildLast(rssAnonSwapTrack);
 
-        const rssFileTrack = await this.createBreakdownTrack(
-          ctx,
-          window,
-          'mem.rss.file',
-          'RSS File',
+        const tracks = [
+          ['mem.rss.file', 'RSS File'],
+          ['mem.rss.shmem', 'RSS Shmem'],
+          ['mem.dmabuf_rss', 'DMA buffer RSS'],
+          ['mem.heap', 'Heap Size'],
+          ['mem.locked', 'Locked Memory'],
+          ['GPU Memory', 'GPU Memory'],
+        ];
+
+        const trackPromises: Promise<TrackNode | undefined>[] = [
+          this.createRssAnonSwapTrack(ctx, window),
+        ];
+        trackPromises.push(
+          ...tracks.map(([trackName, displayName]) =>
+            this.createBreakdownTrack(ctx, window, trackName, displayName),
+          ),
         );
-        ctx.defaultWorkspace.pinnedTracksNode.addChildLast(rssFileTrack);
+        const createdTracks = await Promise.all(trackPromises);
+
+        for (const track of createdTracks) {
+          if (track) {
+            ctx.defaultWorkspace.pinnedTracksNode.addChildLast(track);
+          }
+        }
       },
     });
   }
@@ -55,6 +70,7 @@ export default class MemoryViz implements PerfettoPlugin {
     name: string,
     description: string,
     removable = true,
+    sortOrder?: number,
   ): Promise<TrackNode> {
     const track = await createQueryCounterTrack({
       trace: ctx,
@@ -78,13 +94,14 @@ export default class MemoryViz implements PerfettoPlugin {
       uri,
       name,
       removable,
+      sortOrder,
     });
   }
 
   private async createRssAnonSwapTrack(
     ctx: Trace,
     window: TimeSpan,
-  ): Promise<TrackNode> {
+  ): Promise<TrackNode | undefined> {
     const uri = `${MemoryViz.id}.rss_anon_swap.${uuidv4()}`;
     const sqlSource = this.getSqlSource(window, [
       `track_name IN ('mem.rss.anon', 'mem.swap')`,
@@ -111,8 +128,12 @@ export default class MemoryViz implements PerfettoPlugin {
       'Swap',
     );
 
-    rootNode.addChildLast(rssAnonNode);
-    rootNode.addChildLast(swapNode);
+    if (rssAnonNode) {
+      rootNode.addChildLast(rssAnonNode);
+    }
+    if (swapNode) {
+      rootNode.addChildLast(swapNode);
+    }
 
     return rootNode;
   }
@@ -122,7 +143,7 @@ export default class MemoryViz implements PerfettoPlugin {
     window: TimeSpan,
     trackName: string,
     name: string,
-  ): Promise<TrackNode> {
+  ): Promise<TrackNode | undefined> {
     const uri = `${MemoryViz.id}.${trackName}.${uuidv4()}`;
     const sqlSource = this.getSqlSource(window, [
       `track_name = '${trackName}'`,
@@ -151,7 +172,9 @@ export default class MemoryViz implements PerfettoPlugin {
         trackName,
         bucket,
       );
-      breakdownNode.addChildLast(bucketNode);
+      if (bucketNode) {
+        breakdownNode.addChildLast(bucketNode);
+      }
     }
 
     return breakdownNode;
@@ -162,25 +185,29 @@ export default class MemoryViz implements PerfettoPlugin {
     window: TimeSpan,
     trackName: string,
     bucket: string,
-  ): Promise<TrackNode> {
+  ): Promise<TrackNode | undefined> {
     const processes = await ctx.engine.query(`
-      SELECT upid, pid, process_name, MAX(IIF(iss.interval_ends_at_ts = FALSE, m.zygote_adjusted_value, 0)) as max_value FROM interval_self_intersect!((
-        SELECT
-          id,
-          MAX(ts, ${window.start}) as ts,
-          MIN(ts + dur, ${window.end}) - MAX(ts, ${window.start}) as dur
-        FROM _memory_breakdown_mem_with_buckets
-        WHERE bucket = '${bucket}' AND track_name = '${trackName}' AND ts < ${
-          window.end
-        } AND ts + dur > ${window.start}
-      )) iss
-      JOIN _memory_breakdown_mem_with_buckets m USING(id)
+      SELECT
+        upid,
+        pid,
+        process_name,
+        MAX(zygote_adjusted_value) AS max_value
+      FROM _memory_breakdown_mem_with_buckets
+      WHERE
+        bucket = '${bucket}' AND
+        track_name = '${trackName}' AND
+        ts < ${window.end} AND
+        ts + dur > ${window.start}
       GROUP BY upid, pid, process_name
       HAVING max_value > 0
       ORDER BY max_value DESC
-      `);
+    `);
 
     const numProcesses = processes.numRows();
+    if (numProcesses === 0) {
+      return undefined;
+    }
+
     const plural = numProcesses === 1 ? '' : 'es';
     const name = `${bucket} (${numProcesses} process${plural})`;
 
@@ -189,6 +216,11 @@ export default class MemoryViz implements PerfettoPlugin {
       `bucket = '${bucket}'`,
       `track_name = '${trackName}'`,
     ]);
+    const peak = await ctx.engine.query(
+      `SELECT MAX(value) AS peak FROM (${sqlSource})`,
+    );
+    const peakValue = peak.iter({}).get('peak') as number;
+
     const bucketNode = await this.createTrack(
       ctx,
       uri,
@@ -198,6 +230,7 @@ export default class MemoryViz implements PerfettoPlugin {
         bucket
       }' OOM bucket (0 otherwise).`,
       true,
+      -peakValue, // Sort buckets in descending order of peak memory usage
     );
 
     for (