diff --git a/block/blk-cgroup.c b/block/blk-cgroup.c
index 3d34ac0..857898f 100644
--- a/block/blk-cgroup.c
+++ b/block/blk-cgroup.c
@@ -1732,7 +1732,7 @@ static void blkcg_maybe_throttle_blkg(struct blkcg_gq *blkg, bool use_memdelay)
 	delay_nsec = min_t(u64, delay_nsec, 250 * NSEC_PER_MSEC);
 
 	if (use_memdelay)
-		psi_memstall_enter(&pflags);
+		psi_memstall_enter(&pflags, TSK_MEMSTALL_WORKINGSET);
 
 	exp = ktime_add_ns(now, delay_nsec);
 	tok = io_schedule_prepare();
@@ -1744,7 +1744,7 @@ static void blkcg_maybe_throttle_blkg(struct blkcg_gq *blkg, bool use_memdelay)
 	io_schedule_finish(tok);
 
 	if (use_memdelay)
-		psi_memstall_leave(&pflags);
+		psi_memstall_leave(&pflags, TSK_MEMSTALL_WORKINGSET);
 }
 
 /**
diff --git a/block/blk-core.c b/block/blk-core.c
index d221322..ebbbe49 100644
--- a/block/blk-core.c
+++ b/block/blk-core.c
@@ -1185,12 +1185,12 @@ blk_qc_t submit_bio(struct bio *bio)
 	 * submission can be a significant part of overall IO time.
 	 */
 	if (workingset_read)
-		psi_memstall_enter(&pflags);
+		psi_memstall_enter(&pflags, TSK_MEMSTALL_WORKINGSET);
 
 	ret = generic_make_request(bio);
 
 	if (workingset_read)
-		psi_memstall_leave(&pflags);
+		psi_memstall_leave(&pflags, TSK_MEMSTALL_WORKINGSET);
 
 	return ret;
 }
diff --git a/include/linux/psi_types.h b/include/linux/psi_types.h
index 194ea78..8200623 100644
--- a/include/linux/psi_types.h
+++ b/include/linux/psi_types.h
@@ -174,9 +174,9 @@ struct psi_group {
 	u64 total[NR_PSI_AGGREGATORS][NR_PSI_STATES - 1];
 	unsigned long avg[NR_PSI_STATES - 1][3];
 
-	u64 total_mems[PSI_MEM_STATES - 1];
-	unsigned long avg_mems[PSI_MEM_STATES - 1][3];
-	u64 avg_total_mems[PSI_MEM_STATES - 1];
+	u64 total_mems[PSI_MEM_STATES];
+	unsigned long avg_mems[PSI_MEM_STATES][3];
+	u64 avg_total_mems[PSI_MEM_STATES];
 
 	/* Monitor work control */
 	atomic_t poll_scheduled;
diff --git a/mm/compaction.c b/mm/compaction.c
index 903aea9..62d1416 100644
--- a/mm/compaction.c
+++ b/mm/compaction.c
@@ -2664,9 +2664,9 @@ static int kcompactd(void *p)
 		wait_event_freezable(pgdat->kcompactd_wait,
 				kcompactd_work_requested(pgdat));
 
-		psi_memstall_enter(&pflags);
+		psi_memstall_enter(&pflags, TSK_MEMSTALL_KCOMPACTED);
 		kcompactd_do_work(pgdat);
-		psi_memstall_leave(&pflags);
+		psi_memstall_leave(&pflags, TSK_MEMSTALL_KCOMPACTED);
 	}
 
 	return 0;
diff --git a/mm/filemap.c b/mm/filemap.c
index 3d43769..cf08388 100755
--- a/mm/filemap.c
+++ b/mm/filemap.c
@@ -1151,7 +1151,7 @@ static inline int wait_on_page_bit_common(wait_queue_head_t *q,
 			delayacct_thrashing_start();
 			delayacct = true;
 		}
-		psi_memstall_enter(&pflags);
+		psi_memstall_enter(&pflags, TSK_MEMSTALL_WORKINGSET);
 		thrashing = true;
 	}
 
@@ -1210,7 +1210,7 @@ static inline int wait_on_page_bit_common(wait_queue_head_t *q,
 	if (thrashing) {
 		if (delayacct)
 			delayacct_thrashing_end();
-		psi_memstall_leave(&pflags);
+		psi_memstall_leave(&pflags, TSK_MEMSTALL_WORKINGSET);
 	}
 
 	/*
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 94e9a1c..fab06b7 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -2589,9 +2589,9 @@ void mem_cgroup_handle_over_high(void)
 	 * schedule_timeout_killable sets TASK_KILLABLE). This means we don't
 	 * need to account for any ill-begotten jiffies to pay them off later.
 	 */
-	psi_memstall_enter(&pflags);
+	psi_memstall_enter(&pflags, TSK_MEMSTALL_CGROUP);
 	schedule_timeout_killable(penalty_jiffies);
-	psi_memstall_leave(&pflags);
+	psi_memstall_leave(&pflags, TSK_MEMSTALL_CGROUP);
 
 out:
 	css_put(&memcg->css);
diff --git a/mm/page_alloc.c b/mm/page_alloc.c
index ef077b8..52d86c2 100644
--- a/mm/page_alloc.c
+++ b/mm/page_alloc.c
@@ -3988,7 +3988,7 @@ __alloc_pages_direct_compact(gfp_t gfp_mask, unsigned int order,
 	if (!order)
 		return NULL;
 
-	psi_memstall_enter(&pflags);
+	psi_memstall_enter(&pflags, TSK_MEMSTALL_DCOMPACT);
 	delayacct_compact_start();
 	noreclaim_flag = memalloc_noreclaim_save();
 
@@ -3996,7 +3996,7 @@ __alloc_pages_direct_compact(gfp_t gfp_mask, unsigned int order,
 								prio, &page);
 
 	memalloc_noreclaim_restore(noreclaim_flag);
-	psi_memstall_leave(&pflags);
+	psi_memstall_leave(&pflags, TSK_MEMSTALL_DCOMPACT);
 	delayacct_compact_end();
 
 	/*
@@ -4212,7 +4212,7 @@ __perform_reclaim(gfp_t gfp_mask, unsigned int order,
 
 	/* We now go into synchronous reclaim */
 	cpuset_memory_pressure_bump();
-	psi_memstall_enter(&pflags);
+	psi_memstall_enter(&pflags, TSK_MEMSTALL_DRECLAIM);
 	fs_reclaim_acquire(gfp_mask);
 	noreclaim_flag = memalloc_noreclaim_save();
 
@@ -4221,7 +4221,7 @@ __perform_reclaim(gfp_t gfp_mask, unsigned int order,
 
 	memalloc_noreclaim_restore(noreclaim_flag);
 	fs_reclaim_release(gfp_mask);
-	psi_memstall_leave(&pflags);
+	psi_memstall_leave(&pflags, TSK_MEMSTALL_DRECLAIM);
 
 	cond_resched();
 
diff --git a/mm/vmscan.c b/mm/vmscan.c
index 075da44..c2038b4 100644
--- a/mm/vmscan.c
+++ b/mm/vmscan.c
@@ -3416,13 +3416,13 @@ unsigned long try_to_free_mem_cgroup_pages(struct mem_cgroup *memcg,
 
 	trace_mm_vmscan_memcg_reclaim_begin(0, sc.gfp_mask);
 
-	psi_memstall_enter(&pflags);
+	psi_memstall_enter(&pflags, TSK_MEMSTALL_CGROUP);
 	noreclaim_flag = memalloc_noreclaim_save();
 
 	nr_reclaimed = do_try_to_free_pages(zonelist, &sc);
 
 	memalloc_noreclaim_restore(noreclaim_flag);
-	psi_memstall_leave(&pflags);
+	psi_memstall_leave(&pflags, TSK_MEMSTALL_CGROUP);
 
 	trace_mm_vmscan_memcg_reclaim_end(nr_reclaimed);
 	set_task_reclaim_state(current, NULL);
@@ -3794,7 +3794,7 @@ static int balance_pgdat(pg_data_t *pgdat, int order, int classzone_idx)
 	};
 
 	set_task_reclaim_state(current, &sc.reclaim_state);
-	psi_memstall_enter(&pflags);
+	psi_memstall_enter(&pflags, TSK_MEMSTALL_KSWAPD);
 	__fs_reclaim_acquire();
 
 	count_vm_event(PAGEOUTRUN);
@@ -3973,7 +3973,7 @@ static int balance_pgdat(pg_data_t *pgdat, int order, int classzone_idx)
 
 	snapshot_refaults(NULL, pgdat);
 	__fs_reclaim_release();
-	psi_memstall_leave(&pflags);
+	psi_memstall_leave(&pflags, TSK_MEMSTALL_KSWAPD);
 	set_task_reclaim_state(current, NULL);
 
 	/*
