ui: Wattson: Optimize selection aggregators and track attribution

Migrate Wattson UI to directly use newly added MACROS and tables. This
reduces some work done by the UI and has more code and table reuse.

Bug: 354290668
Signed-off-by: Samuel Wu <wusamuel@google.com>
diff --git a/ui/src/plugins/org.kernel.Wattson/estimate_aggregator.ts b/ui/src/plugins/org.kernel.Wattson/estimate_aggregator.ts
index 2350529..9f200cd 100644
--- a/ui/src/plugins/org.kernel.Wattson/estimate_aggregator.ts
+++ b/ui/src/plugins/org.kernel.Wattson/estimate_aggregator.ts
@@ -65,17 +65,10 @@
     let query = `
       INCLUDE PERFETTO MODULE wattson.estimates;
 
-      CREATE OR REPLACE PERFETTO TABLE wattson_plugin_ui_selection_window AS
-      SELECT
-        ${area.start} as ts,
-        ${duration} as dur;
-
-      DROP TABLE IF EXISTS wattson_plugin_windowed_subsystems_estimate;
-      CREATE VIRTUAL TABLE wattson_plugin_windowed_subsystems_estimate
-      USING
-        SPAN_JOIN(wattson_plugin_ui_selection_window, _system_state_mw);
-
       CREATE PERFETTO VIEW ${this.id} AS
+      WITH window_stats AS (
+        SELECT * FROM _windowed_system_state_mw(${area.start}, ${duration})
+      )
     `;
 
     // Convert average power track to total energy in UI window, then divide by
@@ -87,9 +80,9 @@
       query += `
         SELECT
         '${estimateTrack}' as name,
-        ROUND(SUM(${estimateTrack}_mw * dur) / ${duration}, 3) as power_mw,
-        ROUND(SUM(${estimateTrack}_mw * dur) / 1000000000, 3) as energy_mws
-        FROM wattson_plugin_windowed_subsystems_estimate
+        ROUND(${estimateTrack}_mw, 3) as power_mw,
+        ROUND(${estimateTrack}_mw * ${duration} / 1000000000, 3) as energy_mws
+        FROM window_stats
       `;
     });
     query += `;`;
diff --git a/ui/src/plugins/org.kernel.Wattson/package_aggregator.ts b/ui/src/plugins/org.kernel.Wattson/package_aggregator.ts
index bb62d08..bab20b1 100644
--- a/ui/src/plugins/org.kernel.Wattson/package_aggregator.ts
+++ b/ui/src/plugins/org.kernel.Wattson/package_aggregator.ts
@@ -39,47 +39,25 @@
       prepareData: async (engine: Engine) => {
         await engine.query(`drop view if exists ${this.id};`);
 
-        const duration = area.end - area.start;
-
         // Prerequisite tables are already generated by Wattson thread aggregation,
         // which is run prior to execution of this module
         await engine.query(`
-          INCLUDE PERFETTO MODULE wattson.estimates;
-          INCLUDE PERFETTO MODULE wattson.tasks.idle_transitions_attribution;
-
-          -- Group idle attribution by package
-          CREATE OR REPLACE PERFETTO TABLE
-          wattson_plugin_per_package_idle_attribution AS
-          SELECT
-            SUM(idle_cost_mws) as idle_cost_mws,
-            uid
-          FROM wattson_plugin_idle_attribution
-          JOIN thread USING(utid)
-          JOIN process USING(upid)
-          GROUP BY uid;
-
-          -- Grouped by UID and made CPU agnostic
           CREATE PERFETTO VIEW ${this.id} AS
           WITH base AS (
             SELECT
-              ROUND(SUM(total_pws) / ${duration}, 3) as active_mw,
-              ROUND(SUM(total_pws) / 1000000000, 3) as active_mws,
-              ROUND(COALESCE(idle_cost_mws, 0), 3) as idle_cost_mws,
-              ROUND(
-                COALESCE(idle_cost_mws, 0) + SUM(total_pws) / 1000000000,
-                3
-              ) as total_mws,
-              wattson_plugin_unioned_per_cpu_total.uid AS uid,
-              package_name
-            FROM wattson_plugin_unioned_per_cpu_total
-            LEFT JOIN wattson_plugin_per_package_idle_attribution
-              ON wattson_plugin_unioned_per_cpu_total.uid IS
-              wattson_plugin_per_package_idle_attribution.uid
-            GROUP BY wattson_plugin_unioned_per_cpu_total.uid, package_name
+              ROUND(SUM(estimated_mw), 3) as active_mw,
+              ROUND(SUM(estimated_mws), 3) as active_mws,
+              ROUND(SUM(idle_transitions_mws), 3) as idle_cost_mws,
+              ROUND(SUM(total_mws), 3) as total_mws,
+              package_name,
+              uid
+            FROM wattson_plugin_thread_summary
+            GROUP BY uid, package_name
           )
-          select *,
+          SELECT
+            *,
             total_mws / (SUM(total_mws) OVER()) AS percent_of_total_energy
-            from base;
+          FROM base;
         `);
 
         return {
diff --git a/ui/src/plugins/org.kernel.Wattson/process_aggregator.ts b/ui/src/plugins/org.kernel.Wattson/process_aggregator.ts
index e0159ec..0cf674e 100644
--- a/ui/src/plugins/org.kernel.Wattson/process_aggregator.ts
+++ b/ui/src/plugins/org.kernel.Wattson/process_aggregator.ts
@@ -38,50 +38,27 @@
     return {
       prepareData: async (engine: Engine) => {
         await engine.query(`drop view if exists ${this.id};`);
-        const duration = area.end - area.start;
 
         // Prerequisite tables are already generated by Wattson thread aggregation,
         // which is run prior to execution of this module
         await engine.query(`
-          INCLUDE PERFETTO MODULE wattson.estimates;
-          INCLUDE PERFETTO MODULE wattson.tasks.idle_transitions_attribution;
-
-          -- Group idle attribution by process
-          CREATE OR REPLACE PERFETTO TABLE
-          wattson_plugin_per_process_idle_attribution AS
-          SELECT
-            SUM(idle_cost_mws) as idle_cost_mws,
-            upid
-          FROM wattson_plugin_idle_attribution
-          GROUP BY upid;
-
-          -- Grouped by UPID and made CPU agnostic
           CREATE PERFETTO VIEW ${this.id} AS
           WITH base AS (
             SELECT
-              ROUND(SUM(total_pws) / ${duration}, 3) as active_mw,
-              ROUND(SUM(total_pws) / 1000000000, 3) as active_mws,
-              ROUND(COALESCE(idle_cost_mws, 0), 3) as idle_cost_mws,
-              ROUND(
-                COALESCE(idle_cost_mws, 0) + SUM(total_pws) / 1000000000,
-                3
-              ) as total_mws,
+              ROUND(SUM(estimated_mw), 3) as active_mw,
+              ROUND(SUM(estimated_mws), 3) as active_mws,
+              ROUND(SUM(idle_transitions_mws), 3) as idle_cost_mws,
+              ROUND(SUM(total_mws), 3) as total_mws,
               pid,
-              process_name
-            FROM wattson_plugin_unioned_per_cpu_total
-            LEFT JOIN wattson_plugin_per_process_idle_attribution USING (upid)
+              process_name,
+              upid
+            FROM wattson_plugin_thread_summary
             GROUP BY upid
-          ),
-          secondary AS (
-            SELECT
-              pid,
-              total_mws / (SUM(total_mws) OVER()) AS percent_of_total_energy
-            FROM base
-            GROUP BY pid
           )
-          select *
-            from base INNER JOIN secondary
-            USING (pid);
+          SELECT
+            *,
+            total_mws / (SUM(total_mws) OVER()) AS percent_of_total_energy
+          FROM base;
         `);
 
         return {
diff --git a/ui/src/plugins/org.kernel.Wattson/thread_aggregator.ts b/ui/src/plugins/org.kernel.Wattson/thread_aggregator.ts
index f011187..9e13fd7 100644
--- a/ui/src/plugins/org.kernel.Wattson/thread_aggregator.ts
+++ b/ui/src/plugins/org.kernel.Wattson/thread_aggregator.ts
@@ -62,86 +62,54 @@
         const whereClause = `WHERE ${filters.join(' OR ')}`;
 
         await engine.query(`
-          INCLUDE PERFETTO MODULE wattson.tasks.attribution;
-          INCLUDE PERFETTO MODULE wattson.tasks.idle_transitions_attribution;
-          INCLUDE PERFETTO MODULE wattson.ui.continuous_estimates;
-
+          INCLUDE PERFETTO MODULE wattson.aggregation;
           CREATE OR REPLACE PERFETTO TABLE wattson_plugin_ui_selection_window AS
           SELECT
             ${area.start} as ts,
-            ${duration} as dur;
+            ${duration} as dur,
+            0 as period_id;
+
+          -- Prefilter tasks table to be smaller
+          CREATE OR REPLACE PERFETTO TABLE _wattson_ui_selected_tasks AS
+          SELECT *
+          FROM _estimates_w_tasks_attribution
+          ${whereClause}
+          AND ts + dur >= ${area.start}
+          AND ts < ${area.end};
 
           -- Processes filtered by CPU within the UI defined time window
           DROP TABLE IF EXISTS wattson_plugin_windowed_summary;
           CREATE VIRTUAL TABLE wattson_plugin_windowed_summary
           USING SPAN_JOIN(
             wattson_plugin_ui_selection_window,
-            _estimates_w_tasks_attribution
+            _wattson_ui_selected_tasks
           );
 
-          -- Only get idle attribution in user defined window and filter by selected
-          -- CPUs
-          CREATE OR REPLACE PERFETTO TABLE wattson_plugin_idle_attribution AS
-          SELECT
-            idle_cost_mws,
-            utid,
-            upid
-          FROM _filter_idle_attribution(${area.start}, ${duration})
-          ${whereClause};
+          -- Materialize the thread-level summary once.
+          CREATE OR REPLACE PERFETTO TABLE wattson_plugin_thread_summary AS
+          SELECT *
+          FROM _wattson_threads_aggregation!(
+            wattson_plugin_windowed_summary,
+            wattson_plugin_ui_selection_window
+          );
 
-          -- Group idle attribution by thread
-          CREATE OR REPLACE PERFETTO TABLE wattson_plugin_per_thread_idle_cost AS
-          SELECT
-            SUM(idle_cost_mws) as idle_cost_mws,
-            utid
-          FROM wattson_plugin_idle_attribution
-          GROUP BY utid;
-
-          CREATE OR REPLACE PERFETTO TABLE wattson_plugin_unioned_per_cpu_total AS
-          SELECT
-            SUM(estimated_mw * dur) AS total_pws,
-            SUM(dur) AS dur,
-            tid,
-            pid,
-            uid,
-            utid,
-            upid,
-            thread_name,
-            process_name,
-            package_name
-          FROM wattson_plugin_windowed_summary
-          ${whereClause}
-          GROUP BY utid;
-
-          -- Grouped again by UTID, but this time to make it CPU agnostic
           CREATE PERFETTO VIEW ${this.id} AS
           WITH base AS (
             SELECT
-              ROUND(SUM(total_pws) / ${duration}, 3) as active_mw,
-              ROUND(SUM(total_pws) / 1000000000, 3) as active_mws,
-              ROUND(COALESCE(idle_cost_mws, 0), 3) as idle_cost_mws,
-              ROUND(
-                COALESCE(idle_cost_mws, 0) + SUM(total_pws) / 1000000000,
-                3
-              ) as total_mws,
+              ROUND(estimated_mw, 3) as active_mw,
+              ROUND(estimated_mws, 3) as active_mws,
+              ROUND(idle_transitions_mws, 3) as idle_cost_mws,
+              ROUND(total_mws, 3) as total_mws,
               thread_name,
               utid,
               tid,
               pid
-            FROM wattson_plugin_unioned_per_cpu_total
-            LEFT JOIN wattson_plugin_per_thread_idle_cost USING (utid)
-            GROUP BY utid
-          ),
-          secondary AS (
-            SELECT
-              utid,
-              total_mws / (SUM(total_mws) OVER()) AS percent_of_total_energy
-            FROM base
-            GROUP BY utid
+            FROM wattson_plugin_thread_summary
           )
-          select *
-            from base INNER JOIN secondary
-            USING (utid);
+          SELECT
+            *,
+            total_mws / (SUM(total_mws) OVER()) AS percent_of_total_energy
+          FROM base;
         `);
 
         return {
diff --git a/ui/src/plugins/org.kernel.Wattson/wattson_thread_utils.ts b/ui/src/plugins/org.kernel.Wattson/wattson_thread_utils.ts
index 1fb44bf..02410f8 100644
--- a/ui/src/plugins/org.kernel.Wattson/wattson_thread_utils.ts
+++ b/ui/src/plugins/org.kernel.Wattson/wattson_thread_utils.ts
@@ -45,23 +45,11 @@
   `);
 
   const sqlSource = `
-    WITH _filtered_thread AS (
-      SELECT ts, dur, cpu, _auto_id
-      FROM _all_tasks_flattened_slices
+    WITH _per_thread_estimate AS (
+      SELECT ts, dur, estimated_mw
+      FROM _estimates_w_tasks_attribution
       WHERE utid = ${utid}
     ),
-    _per_thread_estimate AS (
-      SELECT ii.ts, ii.dur, uw.estimated_mw
-      FROM _interval_intersect!(
-        (
-          _ii_subquery!(_unioned_wattson_estimates_mw),
-          _ii_subquery!(_filtered_thread)
-        ),
-        (cpu)
-      ) AS ii
-      JOIN _unioned_wattson_estimates_mw AS uw ON uw._auto_id = id_0
-      JOIN _filtered_thread AS s ON s._auto_id = id_1
-    ),
     -- Need to fill in gaps where thread isn't running with 0mW entry
     gapless AS (
       SELECT ts, dur, estimated_mw FROM _per_thread_estimate