[DO NOT MERGE ANYWHERE] Revert "mali_pixel: Refactor SLC partition management"

Revert submission 2753879-gpu-slcv2-gs201

Reason for revert: Prebuild did not land before cutoff

Reverted changes: /q/submissionid:2753879-gpu-slcv2-gs201

Bug: 329447972
Change-Id: I59fcb56a434c21cc3bdfdacaf3c3978d548330ad

5 files changed, 486 insertions, 324 deletions

diff --git a/common/include/uapi/gpu/arm/midgard/platform/pixel/pixel_memory_group_manager.h b/common/include/uapi/gpu/arm/midgard/platform/pixel/pixel_memory_group_manager.h
index 893cdca..b575c79 100644
--- a/common/include/uapi/gpu/arm/midgard/platform/pixel/pixel_memory_group_manager.h
+++ b/common/include/uapi/gpu/arm/midgard/platform/pixel/pixel_memory_group_manager.h
@@ -7,4 +7,49 @@
 #ifndef _UAPI_PIXEL_MEMORY_GROUP_MANAGER_H_
 #define _UAPI_PIXEL_MEMORY_GROUP_MANAGER_H_
 
+/**
+ * enum pixel_mgm_group_id - Symbolic names for used memory groups
+ */
+enum pixel_mgm_group_id
+{
+	/* The Mali driver requires that allocations made on one of the groups
+	 * are not treated specially.
+	 */
+	MGM_RESERVED_GROUP_ID = 0,
+
+	/* Group for memory that should be cached in the system level cache. */
+	MGM_SLC_GROUP_ID = 1,
+
+	/* Group for memory explicitly allocated in SLC. */
+	MGM_SLC_EXPLICIT_GROUP_ID = 2,
+
+	/* Imported memory is handled by the allocator of the memory, and the Mali
+	 * DDK will request a group_id for such memory via mgm_get_import_memory_id().
+	 * We specify which group we want to use for this here.
+	 */
+	MGM_IMPORTED_MEMORY_GROUP_ID = (MEMORY_GROUP_MANAGER_NR_GROUPS - 1),
+};
+
+/**
+ * pixel_mgm_query_group_size - Query the current size of a memory group
+ *
+ * @mgm_dev:   The memory group manager through which the request is being made.
+ * @group_id:  Memory group to query.
+ *
+ * Returns the actual size of the memory group's active partition
+ */
+extern u64 pixel_mgm_query_group_size(struct memory_group_manager_device* mgm_dev,
+                                      enum pixel_mgm_group_id group_id);
+
+/**
+ * pixel_mgm_resize_group_to_fit - Resize a memory group to meet @demand, if possible
+ *
+ * @mgm_dev:   The memory group manager through which the request is being made.
+ * @group_id:  Memory group for which we will change the backing partition.
+ * @demand:    The demanded space from the memory group.
+ */
+extern void pixel_mgm_resize_group_to_fit(struct memory_group_manager_device* mgm_dev,
+                                          enum pixel_mgm_group_id group_id,
+                                          u64 demand);
+
 #endif /* _UAPI_PIXEL_MEMORY_GROUP_MANAGER_H_ */
diff --git a/mali_pixel/Kbuild b/mali_pixel/Kbuild
index 4967b34..4b519a9 100644
--- a/mali_pixel/Kbuild
+++ b/mali_pixel/Kbuild
@@ -37,7 +37,6 @@ endif
 ifeq ($(CONFIG_MALI_MEMORY_GROUP_MANAGER),m)
 	DEFINES += -DCONFIG_MALI_MEMORY_GROUP_MANAGER
 	mali_pixel-objs += memory_group_manager.o
-	mali_pixel-objs += pixel_slc.o
 endif
 ifeq ($(CONFIG_MALI_PRIORITY_CONTROL_MANAGER),m)
 	DEFINES += -DCONFIG_MALI_PRIORITY_CONTROL_MANAGER
diff --git a/mali_pixel/memory_group_manager.c b/mali_pixel/memory_group_manager.c
index 3015468..0c9a241 100644
--- a/mali_pixel/memory_group_manager.c
+++ b/mali_pixel/memory_group_manager.c
@@ -23,7 +23,7 @@
 
 #include <linux/memory_group_manager.h>
 
-#include "pixel_slc.h"
+#include <soc/google/pt.h>
 
 #include <uapi/gpu/arm/midgard/platform/pixel/pixel_memory_group_manager.h>
 
@@ -31,30 +31,29 @@
 #define ORDER_SMALL_PAGE 0
 #define ORDER_LARGE_PAGE const_ilog2(NUM_PAGES_IN_2MB_LARGE_PAGE)
 
-/**
- * enum mgm_group_id - Symbolic names for used memory groups
+/* Borr does not have "real" PBHA support. However, since we only use a 36-bit PA on the bus,
+ * AxADDR[39:36] is wired up to the GPU AxUSER[PBHA] field seen by the rest of the system.
+ * Those AxADDR bits come from [39:36] in the page descriptor.
+ *
+ * Odin and Turse have "real" PBHA support using a dedicated output signal and page descriptor field.
+ * The AxUSER[PBHA] field is driven by the GPU's PBHA signal, and AxADDR[39:36] is dropped.
+ * The page descriptor PBHA field is [62:59].
+ *
+ * We could write to both of these locations, as each SoC only reads from its respective PBHA
+ * location with the other being ignored or dropped.
+ *
+ * b/148988078 contains confirmation of the above description.
  */
-enum mgm_group_id
-{
-	/**
-	 * @MGM_RESERVED_GROUP_ID: The Mali driver requires that allocations made on one of the
-	 *                         groups are not treated specially.
-	 */
-	MGM_RESERVED_GROUP_ID = 0,
+#if IS_ENABLED(CONFIG_SOC_GS101)
+#define PBHA_BIT_POS  (36)
+#else
+#define PBHA_BIT_POS  (59)
+#endif
+#define PBHA_BIT_MASK (0xf)
 
-	/**
-	 * @MGM_SLC_GROUP_ID: Group for memory that should be cached in the system level cache.
-	 */
-	MGM_SLC_GROUP_ID = 1,
+#define MGM_PBHA_DEFAULT 0
 
-	/**
-	 * @MGM_IMPORTED_MEMORY_GROUP_ID: Imported memory is handled by the allocator of the memory,
-	 *                                and the Mali DDK will request a group_id for such memory
-	 *                                via mgm_get_import_memory_id(). We specify which group we
-	 *                                want to use for this here.
-	 */
-	MGM_IMPORTED_MEMORY_GROUP_ID = (MEMORY_GROUP_MANAGER_NR_GROUPS - 1),
-};
+#define MGM_SENTINEL_PT_SIZE U64_MAX
 
 #define INVALID_GROUP_ID(group_id) \
 	WARN_ON((group_id) >= MEMORY_GROUP_MANAGER_NR_GROUPS)
@@ -82,6 +81,13 @@ static inline vm_fault_t vmf_insert_pfn_prot(struct vm_area_struct *vma,
  * @lp_size:  The number of allocated large(2MB) pages
  * @insert_pfn: The number of calls to map pages for CPU access.
  * @update_gpu_pte: The number of calls to update GPU page table entries.
+ * @ptid: The active partition ID for this group
+ * @pbha: The PBHA bits assigned to this group,
+ * @base_pt: The base partition ID available to this group.
+ * @pt_num: The number of partitions available to this group.
+ * @active_pt_idx: The relative index for the partition backing the group.
+ *                 Different from the absolute ptid.
+ * @state: The lifecycle state of the partition associated with this group
  * This structure allows page allocation information to be displayed via
  * debugfs. Display is organized per group with small and large sized pages.
  */
@@ -92,6 +98,30 @@ struct mgm_group {
 	atomic_t insert_pfn;
 	atomic_t update_gpu_pte;
 #endif
+
+	ptid_t ptid;
+	ptpbha_t pbha;
+
+	u32 base_pt;
+	u32 pt_num;
+	u32 active_pt_idx;
+	enum {
+		MGM_GROUP_STATE_NEW = 0,
+		MGM_GROUP_STATE_ENABLED = 10,
+		MGM_GROUP_STATE_DISABLED_NOT_FREED = 20,
+		MGM_GROUP_STATE_DISABLED = 30,
+	} state;
+};
+
+/**
+ * struct partition_stats - Structure for tracking sizing of a partition
+ *
+ * @capacity: The total capacity of each partition
+ * @size: The current size of each partition
+ */
+struct partition_stats {
+	u64 capacity;
+	atomic64_t size;
 };
 
 /**
@@ -100,22 +130,26 @@ struct mgm_group {
  * @groups: To keep track of the number of allocated pages of all groups
  * @ngroups: Number of groups actually used
  * @npartitions: Number of partitions used by all groups combined
+ * @pt_stats: The sizing info for each partition
  * @dev: device attached
+ * @pt_handle: Link to SLC partition data
  * @kobj: &sruct kobject used for linking to pixel_stats_sysfs node
  * @mgm_debugfs_root: debugfs root directory of memory group manager
- * @slc_data: To track GPU SLC partitions.
  *
  * This structure allows page allocation information to be displayed via
  * debugfs. Display is organized per group with small and large sized pages.
  */
 struct mgm_groups {
 	struct mgm_group groups[MEMORY_GROUP_MANAGER_NR_GROUPS];
+	size_t ngroups;
+	size_t npartitions;
+	struct partition_stats *pt_stats;
 	struct device *dev;
+	struct pt_handle *pt_handle;
 	struct kobject kobj;
 #ifdef CONFIG_MALI_MEMORY_GROUP_MANAGER_DEBUG_FS
 	struct dentry *mgm_debugfs_root;
 #endif
-	struct slc_data slc_data;
 };
 
 /*
@@ -124,6 +158,13 @@ struct mgm_groups {
 
 #ifdef CONFIG_MALI_MEMORY_GROUP_MANAGER_DEBUG_FS
 
+static int mgm_debugfs_state_get(void *data, u64 *val)
+{
+	struct mgm_group *group = data;
+	*val = (u64)group->state;
+	return 0;
+}
+
 static int mgm_debugfs_size_get(void *data, u64 *val)
 {
 	struct mgm_group *group = data;
@@ -152,6 +193,8 @@ static int mgm_debugfs_update_gpu_pte_get(void *data, u64 *val)
 	return 0;
 }
 
+DEFINE_SIMPLE_ATTRIBUTE(fops_mgm_state, mgm_debugfs_state_get,
+	NULL, "%llu\n");
 DEFINE_SIMPLE_ATTRIBUTE(fops_mgm_size, mgm_debugfs_size_get,
 	NULL, "%llu\n");
 DEFINE_SIMPLE_ATTRIBUTE(fops_mgm_lp_size, mgm_debugfs_lp_size_get,
@@ -175,6 +218,7 @@ static struct {
 	const char *name;
 	const struct file_operations *fops;
 } attribs[] = {
+	{ "state", &fops_mgm_state},
 	{ "size", &fops_mgm_size},
 	{ "lp_size", &fops_mgm_lp_size},
 	{ "insert_pfn", &fops_mgm_insert_pfn},
@@ -348,6 +392,20 @@ static void mgm_sysfs_term(struct mgm_groups *data)
 
 #endif /* CONFIG_MALI_PIXEL_STATS */
 
+static int group_pt_id(struct mgm_groups *data, enum pixel_mgm_group_id group_id, int pt_index)
+{
+	struct mgm_group *group = &data->groups[group_id];
+	if (WARN_ON_ONCE(pt_index >= group->pt_num))
+		return 0;
+
+	return group->base_pt + pt_index;
+}
+
+static int group_active_pt_id(struct mgm_groups *data, enum pixel_mgm_group_id group_id)
+{
+	return group_pt_id(data, group_id, data->groups[group_id].active_pt_idx);
+}
+
 static atomic64_t total_gpu_pages = ATOMIC64_INIT(0);
 
 static atomic_t* get_size_counter(struct memory_group_manager_device* mgm_dev, unsigned int group_id, unsigned int order)
@@ -404,6 +462,185 @@ static void update_size(struct memory_group_manager_device *mgm_dev, unsigned in
 		pr_warn("total_gpu_pages %lld\n", atomic64_read(&total_gpu_pages));
 }
 
+static void pt_size_invalidate(struct mgm_groups* data, int pt_idx)
+{
+	/* Set the size to a known sentinel value so that we can later detect an update */
+	atomic64_set(&data->pt_stats[pt_idx].size, MGM_SENTINEL_PT_SIZE);
+}
+
+static void pt_size_init(struct mgm_groups* data, int pt_idx, size_t size)
+{
+	/* The resize callback may have already been executed, which would have set
+	 * the correct size. Only update the size if this has not happened.
+	 * We can tell that no resize took place if the size is still a sentinel.
+	 */
+	atomic64_cmpxchg(&data->pt_stats[pt_idx].size, MGM_SENTINEL_PT_SIZE, size);
+}
+
+static void validate_ptid(struct mgm_groups* data, enum pixel_mgm_group_id group_id, int ptid)
+{
+	if (ptid == -EINVAL)
+		dev_err(data->dev, "Failed to get partition for group: %d\n", group_id);
+	else
+		dev_info(data->dev, "pt_client_mutate returned ptid=%d for group=%d", ptid, group_id);
+}
+
+static void update_group(struct mgm_groups* data,
+                         enum pixel_mgm_group_id group_id,
+                         int ptid,
+                         int relative_pt_idx)
+{
+	int const abs_pt_idx = group_pt_id(data, group_id, relative_pt_idx);
+	int const pbha = pt_pbha(data->dev->of_node, abs_pt_idx);
+
+	if (pbha == PT_PBHA_INVALID)
+		dev_err(data->dev, "Failed to get PBHA for group: %d\n", group_id);
+	else
+		dev_info(data->dev, "pt_pbha returned PBHA=%d for group=%d", pbha, group_id);
+
+	data->groups[group_id].ptid = ptid;
+	data->groups[group_id].pbha = pbha;
+	data->groups[group_id].state = MGM_GROUP_STATE_ENABLED;
+	data->groups[group_id].active_pt_idx = relative_pt_idx;
+}
+
+static void disable_partition(struct mgm_groups* data, enum pixel_mgm_group_id group_id)
+{
+	int const active_idx = group_active_pt_id(data, group_id);
+
+	/* Skip if not already enabled */
+	if (data->groups[group_id].state != MGM_GROUP_STATE_ENABLED)
+		return;
+
+	pt_client_disable_no_free(data->pt_handle, active_idx);
+	data->groups[group_id].state = MGM_GROUP_STATE_DISABLED_NOT_FREED;
+
+	pt_size_invalidate(data, active_idx);
+	pt_size_init(data, active_idx, 0);
+}
+
+static void enable_partition(struct mgm_groups* data, enum pixel_mgm_group_id group_id)
+{
+	int ptid;
+	size_t size = 0;
+	int const active_idx = group_active_pt_id(data, group_id);
+
+	/* Skip if already enabled */
+	if (data->groups[group_id].state == MGM_GROUP_STATE_ENABLED)
+		return;
+
+	pt_size_invalidate(data, active_idx);
+
+	ptid = pt_client_enable_size(data->pt_handle, active_idx, &size);
+
+	validate_ptid(data, group_id, ptid);
+
+	update_group(data, group_id, ptid, data->groups[group_id].active_pt_idx);
+
+	pt_size_init(data, active_idx, size);
+}
+
+static void set_group_partition(struct mgm_groups* data,
+                                enum pixel_mgm_group_id group_id,
+                                int new_pt_index)
+{
+	int ptid;
+	size_t size = 0;
+	int const active_idx = group_active_pt_id(data, group_id);
+	int const new_idx = group_pt_id(data, group_id, new_pt_index);
+
+	/* Early out if no changes are needed */
+	if (new_idx == active_idx)
+		return;
+
+	pt_size_invalidate(data, new_idx);
+
+	ptid = pt_client_mutate_size(data->pt_handle, active_idx, new_idx, &size);
+
+	validate_ptid(data, group_id, ptid);
+
+	update_group(data, group_id, ptid, new_pt_index);
+
+	pt_size_init(data, new_idx, size);
+	/* Reset old partition size */
+	atomic64_set(&data->pt_stats[active_idx].size, data->pt_stats[active_idx].capacity);
+}
+
+u64 pixel_mgm_query_group_size(struct memory_group_manager_device* mgm_dev,
+                               enum pixel_mgm_group_id group_id)
+{
+	struct mgm_groups *data;
+	struct mgm_group *group;
+	u64 size = 0;
+
+	/* Early out if the group doesn't exist */
+	if (INVALID_GROUP_ID(group_id))
+		goto done;
+
+	data = mgm_dev->data;
+	group = &data->groups[group_id];
+
+	/* Early out if the group has no partitions */
+	if (group->pt_num == 0)
+		goto done;
+
+	size = atomic64_read(&data->pt_stats[group_active_pt_id(data, group_id)].size);
+
+done:
+	return size;
+}
+EXPORT_SYMBOL(pixel_mgm_query_group_size);
+
+void pixel_mgm_resize_group_to_fit(struct memory_group_manager_device* mgm_dev,
+                                  enum pixel_mgm_group_id group_id,
+                                  u64 demand)
+{
+	struct mgm_groups *data;
+	struct mgm_group *group;
+	s64 diff, cur_size, min_diff = S64_MAX;
+	int pt_idx;
+
+	/* Early out if the group doesn't exist */
+	if (INVALID_GROUP_ID(group_id))
+		goto done;
+
+	data = mgm_dev->data;
+	group = &data->groups[group_id];
+
+	/* Early out if the group has no partitions */
+	if (group->pt_num == 0)
+		goto done;
+
+	/* We can disable the partition if there's no demand */
+	if (demand == 0)
+	{
+		disable_partition(data, group_id);
+		goto done;
+	}
+
+	/* Calculate best partition to use, by finding the nearest capacity */
+	for (pt_idx = 0; pt_idx < group->pt_num; ++pt_idx)
+	{
+		cur_size = data->pt_stats[group_pt_id(data, group_id, pt_idx)].capacity;
+		diff = abs(demand - cur_size);
+
+		if (diff > min_diff)
+			break;
+
+		min_diff = diff;
+	}
+
+	/* Ensure the partition is enabled before trying to mutate it */
+	enable_partition(data, group_id);
+	set_group_partition(data, group_id, pt_idx - 1);
+
+done:
+	dev_dbg(data->dev, "%s: resized memory_group_%d for demand: %lldB", __func__, group_id, demand);
+
+	return;
+}
+EXPORT_SYMBOL(pixel_mgm_resize_group_to_fit);
+
 static struct page *mgm_alloc_page(
 	struct memory_group_manager_device *mgm_dev, unsigned int group_id,
 	gfp_t gfp_mask, unsigned int order)
@@ -418,12 +655,35 @@ static struct page *mgm_alloc_page(
 	if (INVALID_GROUP_ID(group_id))
 		return NULL;
 
+	if (WARN_ON_ONCE((group_id != MGM_RESERVED_GROUP_ID) &&
+			 (group_active_pt_id(data, group_id) >= data->npartitions)))
+		return NULL;
+
 	/* We don't expect to be allocting pages into the group used for
 	 * external or imported memory
 	 */
 	if (WARN_ON(group_id == MGM_IMPORTED_MEMORY_GROUP_ID))
 		return NULL;
 
+	/* If we are allocating a page in this group for the first time then
+	 *  ensure that we have enabled the relevant partitions for it.
+	 */
+	if (group_id != MGM_RESERVED_GROUP_ID) {
+		switch (data->groups[group_id].state) {
+		case MGM_GROUP_STATE_NEW:
+			enable_partition(data, group_id);
+			break;
+		case MGM_GROUP_STATE_ENABLED:
+		case MGM_GROUP_STATE_DISABLED_NOT_FREED:
+		case MGM_GROUP_STATE_DISABLED:
+			/* Everything should already be set up*/
+			break;
+		default:
+			dev_err(data->dev, "Group %u in invalid state %d\n",
+				group_id, data->groups[group_id].state);
+		}
+	}
+
 	p = alloc_pages(gfp_mask, order);
 
 	if (p) {
@@ -482,7 +742,7 @@ static u64 mgm_update_gpu_pte(
 	int const mmu_level, u64 pte)
 {
 	struct mgm_groups *const data = mgm_dev->data;
-	u64 const old_pte = pte;
+	unsigned int pbha;
 
 	dev_dbg(data->dev,
 		"%s(mgm_dev=%p, group_id=%u, mmu_level=%d, pte=0x%llx)\n",
@@ -491,22 +751,40 @@ static u64 mgm_update_gpu_pte(
 	if (INVALID_GROUP_ID(group_id))
 		return pte;
 
+	/* Clear any bits set in the PBHA range */
+	if (pte & ((u64)PBHA_BIT_MASK << PBHA_BIT_POS)) {
+		dev_warn(data->dev,
+			"%s: updating pte with bits already set in PBHA range",
+			__func__);
+		pte &= ~((u64)PBHA_BIT_MASK << PBHA_BIT_POS);
+	}
+
 	switch (group_id) {
 	case MGM_RESERVED_GROUP_ID:
 	case MGM_IMPORTED_MEMORY_GROUP_ID:
 		/* The reserved group doesn't set PBHA bits */
-		pte = slc_wipe_pbha(pte);
+		/* TODO: Determine what to do with imported memory */
 		break;
-	case MGM_SLC_GROUP_ID:
-		/* Map requests for SLC memory groups to SLC */
-		pte = slc_set_pbha(&data->slc_data, pte);
 	default:
-		break;
+		/* All other groups will have PBHA bits */
+		if (data->groups[group_id].state > MGM_GROUP_STATE_NEW) {
+			u64 old_pte = pte;
+			pbha = data->groups[group_id].pbha;
+
+			pte |= ((u64)pbha & PBHA_BIT_MASK) << PBHA_BIT_POS;
+
+			dev_dbg(data->dev,
+				"%s: group_id=%u pbha=%d "
+				"pte=0x%llx -> 0x%llx\n",
+				__func__, group_id, pbha, old_pte, pte);
+
+		} else {
+			dev_err(data->dev,
+				"Tried to get PBHA of uninitialized group=%d",
+				group_id);
+		}
 	}
 
-	dev_dbg(data->dev, "%s: group_id=%u pte=0x%llx -> 0x%llx\n",
-		__func__, group_id, old_pte, pte);
-
 #ifdef CONFIG_MALI_MEMORY_GROUP_MANAGER_DEBUG_FS
 	atomic_inc(&data->groups[group_id].update_gpu_pte);
 #endif
@@ -517,10 +795,27 @@ static u64 mgm_update_gpu_pte(
 static u64 mgm_pte_to_original_pte(struct memory_group_manager_device *mgm_dev, unsigned int group_id,
 				int mmu_level, u64 pte)
 {
+	struct mgm_groups *const data = mgm_dev->data;
+	u64 old_pte;
+
 	if (INVALID_GROUP_ID(group_id))
 		return pte;
 
-	return slc_wipe_pbha(pte);
+	switch (group_id) {
+	case MGM_RESERVED_GROUP_ID:
+	case MGM_IMPORTED_MEMORY_GROUP_ID:
+		/* The reserved group doesn't set PBHA bits */
+		/* TODO: Determine what to do with imported memory */
+		break;
+	default:
+		/* All other groups will have PBHA bits, so clear them */
+		old_pte = pte;
+		pte &= ~((u64)PBHA_BIT_MASK << PBHA_BIT_POS);
+		dev_dbg(data->dev, "%s: group_id=%u pte=0x%llx -> 0x%llx\n", __func__, group_id,
+			old_pte, pte);
+	}
+
+	return pte;
 }
 
 static vm_fault_t mgm_vmf_insert_pfn_prot(
@@ -552,12 +847,49 @@ static vm_fault_t mgm_vmf_insert_pfn_prot(
 	return fault;
 }
 
+static void mgm_resize_callback(void *data, int id, size_t size_allocated)
+{
+	struct mgm_groups *const mgm_data = (struct mgm_groups *)data;
+	dev_dbg(mgm_data->dev, "Resize callback called, size_allocated: %zu\n", size_allocated);
+	/* Update the partition size for the group */
+	atomic64_set(&mgm_data->pt_stats[id].size, size_allocated);
+}
+
 static int mgm_initialize_data(struct mgm_groups *mgm_data)
 {
 	int i, ret;
 
-	if ((ret = slc_init_data(&mgm_data->slc_data, mgm_data->dev)))
+	/* +1 to include the required default group */
+	const int ngroups = of_property_count_strings(mgm_data->dev->of_node, "groups") + 1;
+	if (WARN_ON(ngroups < 0) ||
+	    WARN_ON(ngroups > MEMORY_GROUP_MANAGER_NR_GROUPS)) {
+		mgm_data->ngroups = 0;
+	} else {
+		mgm_data->ngroups = ngroups;
+	}
+	mgm_data->npartitions = of_property_count_strings(mgm_data->dev->of_node, "pt_id");
+
+	mgm_data->pt_stats = kzalloc(mgm_data->npartitions * sizeof(struct partition_stats), GFP_KERNEL);
+	if (mgm_data->pt_stats == NULL) {
+		dev_err(mgm_data->dev, "failed to allocate space for pt_stats");
+		ret = -ENOMEM;
 		goto out_err;
+	}
+
+	for (i = 0; i < mgm_data->npartitions; i++) {
+		struct partition_stats* stats;
+		u32 capacity_kb;
+		ret = of_property_read_u32_index(mgm_data->dev->of_node, "pt_size", i, &capacity_kb);
+		if (ret) {
+			dev_err(mgm_data->dev, "failed to read pt_size[%d]", i);
+			continue;
+		}
+
+		stats = &mgm_data->pt_stats[i];
+		// Convert from KB to bytes
+		stats->capacity = (u64)capacity_kb << 10;
+		atomic64_set(&stats->size, stats->capacity);
+	}
 
 	for (i = 0; i < MEMORY_GROUP_MANAGER_NR_GROUPS; i++) {
 		atomic_set(&mgm_data->groups[i].size, 0);
@@ -566,8 +898,50 @@ static int mgm_initialize_data(struct mgm_groups *mgm_data)
 		atomic_set(&mgm_data->groups[i].insert_pfn, 0);
 		atomic_set(&mgm_data->groups[i].update_gpu_pte, 0);
 #endif
+
+		mgm_data->groups[i].pbha = MGM_PBHA_DEFAULT;
+		mgm_data->groups[i].base_pt = 0;
+		mgm_data->groups[i].pt_num = 0;
+		mgm_data->groups[i].active_pt_idx = 0;
+		mgm_data->groups[i].state = MGM_GROUP_STATE_NEW;
+	}
+
+	/* Discover the partitions belonging to each memory group, skipping the reserved group */
+	for (i = 1; i < mgm_data->ngroups; i++) {
+		/* Device tree has no description for the reserved group */
+		int const dt_idx = i - 1;
+
+		int err = of_property_read_u32_index(
+		    mgm_data->dev->of_node, "group_base_pt", dt_idx, &mgm_data->groups[i].base_pt);
+		if (err) {
+			dev_warn(mgm_data->dev, "failed to read base pt index for group %d", i);
+			continue;
+		}
+
+		err = of_property_read_u32_index(
+		    mgm_data->dev->of_node, "group_pt_num", dt_idx, &mgm_data->groups[i].pt_num);
+		if (err)
+			dev_warn(mgm_data->dev, "failed to read pt number for group %d", i);
 	}
 
+	/*
+	 * Initialize SLC partitions. We don't enable partitions until
+	 * we actually allocate memory to the corresponding memory
+	 * group
+	 */
+	mgm_data->pt_handle =
+	    pt_client_register(mgm_data->dev->of_node, (void*)mgm_data, &mgm_resize_callback);
+
+	if (IS_ERR(mgm_data->pt_handle)) {
+		ret = PTR_ERR(mgm_data->pt_handle);
+		dev_err(mgm_data->dev, "pt_client_register returned %d\n", ret);
+		goto out_err;
+	}
+
+	/* We don't use PBHA bits for the reserved memory group, and so
+	 * it is effectively already initialized.
+	 */
+	mgm_data->groups[MGM_RESERVED_GROUP_ID].state = MGM_GROUP_STATE_ENABLED;
 
 	if ((ret = mgm_debugfs_init(mgm_data)))
 		goto out_err;
@@ -575,9 +949,20 @@ static int mgm_initialize_data(struct mgm_groups *mgm_data)
 	if ((ret = mgm_sysfs_init(mgm_data)))
 		goto out_err;
 
+#ifdef CONFIG_MALI_PIXEL_GPU_SLC
+	/* We enable the SLC partition by default to support dynamic SLC caching.
+	 * Enabling will initialize the partition, by querying the pbha and assigning a ptid.
+	 * We then immediately disable the partition, effectively resizing the group to zero,
+	 * whilst still retaining other properties such as pbha.
+	 */
+	enable_partition(mgm_data, MGM_SLC_GROUP_ID);
+	disable_partition(mgm_data, MGM_SLC_GROUP_ID);
+#endif
+
 	return ret;
 
 out_err:
+	kfree(mgm_data->pt_stats);
 	return ret;
 }
 
@@ -598,9 +983,29 @@ static void mgm_term_data(struct mgm_groups *data)
 			dev_warn(data->dev,
 				"%zu 9 order pages in group(%d) leaked\n",
 				(size_t)atomic_read(&group->lp_size), i);
+
+		/* Disable partition indices and free the partition */
+		switch (group->state) {
+
+		case MGM_GROUP_STATE_NEW:
+		case MGM_GROUP_STATE_DISABLED:
+			/* Nothing to do */
+			break;
+
+		case MGM_GROUP_STATE_ENABLED:
+			pt_client_disable(data->pt_handle, group_active_pt_id(data, i));
+			break;
+		case MGM_GROUP_STATE_DISABLED_NOT_FREED:
+			pt_client_free(data->pt_handle, group_active_pt_id(data, i));
+			break;
+
+		default:
+			dev_err(data->dev, "Group %d in invalid state %d\n",
+				i, group->state);
+		}
 	}
 
-	slc_term_data(&data->slc_data);
+	pt_client_unregister(data->pt_handle);
 
 	mgm_debugfs_term(data);
 	mgm_sysfs_term(data);
diff --git a/mali_pixel/pixel_slc.c b/mali_pixel/pixel_slc.c
deleted file mode 100644
index 62c6908..0000000
--- a/mali_pixel/pixel_slc.c
+++ /dev/null
@@ -1,215 +0,0 @@
-// SPDX-License-Identifier: GPL-2.0
-/*
- * Copyright 2024 Google LLC.
- *
- * Author: Jack Diver <diverj@google.com>
- */
-
-#include <linux/atomic.h>
-#include <linux/of.h>
-#include <linux/platform_device.h>
-#include <linux/slab.h>
-#include <linux/dev_printk.h>
-/* Pixel integration includes */
-#include "pixel_slc.h"
-
-
-/**
- * DOC: PBHA
- *
- * Borr does not have "real" PBHA support. However, since we only use a 36-bit PA on the bus,
- * AxADDR[39:36] is wired up to the GPU AxUSER[PBHA] field seen by the rest of the system.
- * Those AxADDR bits come from [39:36] in the page descriptor.
- *
- * Odin and Turse have "real" PBHA support using a dedicated output signal and page descriptor field.
- * The AxUSER[PBHA] field is driven by the GPU's PBHA signal, and AxADDR[39:36] is dropped.
- * The page descriptor PBHA field is [62:59].
- *
- * We could write to both of these locations, as each SoC only reads from its respective PBHA
- * location with the other being ignored or dropped.
- *
- * b/148988078 contains confirmation of the above description.
- */
-#if IS_ENABLED(CONFIG_SOC_GS101)
-#define PBHA_BIT_POS  (36)
-#else
-#define PBHA_BIT_POS  (59)
-#endif
-#define PBHA_BIT_MASK (0xf)
-
-
-/**
- * slc_wipe_pbha - Clear any set PBHA bits from the pte.
- *
- * @pte: The pte to strip of PBHA.
- *
- * Return: The PTE with all PBHA stripped.
- */
-u64 slc_wipe_pbha(u64 pte)
-{
-	return pte & ~((u64)PBHA_BIT_MASK << PBHA_BIT_POS);
-}
-
-/**
- * slc_set_pbha - Apply the PBHA to @pte.
- *
- * @data: The &struct slc_data tracking partition information.
- * @pte:  The pte to modify.
- *
- * Return: On success, returns a modified PTE. On failure the original PTE is returned.
- */
-u64 slc_set_pbha(struct slc_data const *data, u64 pte)
-{
-	/* Clear any bits set in the PBHA range */
-	pte = slc_wipe_pbha(pte);
-
-	/* Apply the PBHA for the given virtual partition */
-	return pte | (((u64)data->partition.pbha) & PBHA_BIT_MASK) << PBHA_BIT_POS;
-}
-
-/**
- * enable_partition - Enable @pt
- *
- * @data: The &struct slc_data tracking partition information.
- * @pt:   The &struct slc_partition representing the partition to enable.
- */
-static void enable_partition(struct slc_data *data, struct slc_partition *pt)
-{
-	/* Skip if already enabled */
-	if (pt->enabled)
-		return;
-
-	(void)pt_client_enable(data->pt_handle, pt->index);
-	pt->enabled = true;
-
-	dev_dbg(data->dev, "enabled partition %d", pt->index);
-}
-
-/**
- * disable_partition - Disable @pt
- *
- * @data: The &struct slc_data tracking partition information.
- * @pt:   The &struct slc_partition representing the partition to disable.
- */
-static void disable_partition(struct slc_data *data, struct slc_partition *pt)
-{
-	/* Skip if not enabled */
-	if (!pt->enabled)
-		return;
-
-	pt_client_disable_no_free(data->pt_handle, pt->index);
-	pt->enabled = false;
-
-	dev_dbg(data->dev, "disabled partition %d", pt->index);
-}
-
-/**
- * init_partition - Register and initialize a partition with the SLC driver.
- *
- * @data:  The &struct slc_data tracking partition information.
- * @pt:    The &struct slc_partition to store the configured partition information.
- * @index: The index of the partition, relative to the DT node.
- *
- * Returns EINVAL on error, otherwise 0.
- */
-static int init_partition(struct slc_data *data, struct slc_partition *pt, u32 index)
-{
-	ptid_t ptid;
-	ptpbha_t pbha;
-	int err = -EINVAL;
-
-	ptid = pt_client_enable(data->pt_handle, index);
-	if (ptid == PT_PTID_INVALID) {
-		dev_err(data->dev, "failed to enable pt: %d\n", index);
-		goto err_exit;
-	}
-
-	pbha = pt_pbha(data->dev->of_node, index);
-	if (pbha == PT_PBHA_INVALID) {
-		dev_err(data->dev, "failed to get PBHA for pt: %d\n", index);
-		goto err_exit;
-	}
-
-	/* This retains the allocated ptid */
-	pt_client_disable_no_free(data->pt_handle, index);
-
-	/* Success */
-	err = 0;
-
-	*pt = (struct slc_partition) {
-		.index = index,
-		.ptid = ptid,
-		.pbha = pbha,
-		.enabled = false,
-	};
-
-err_exit:
-	return err;
-}
-
-
-/**
- * term_partition - Disable and free a partition, unregistering it.
- *
- * @data:  The &struct slc_data tracking partition information.
- * @pt:    The &struct slc_partition to terminate.
- *
- * Returns EINVAL on error, otherwise 0.
- */
-static void term_partition(struct slc_data *data, struct slc_partition *pt)
-{
-	disable_partition(data, pt);
-	pt_client_free(data->pt_handle, pt->index);
-}
-
-/**
- * slc_init_data - Read all SLC partition information, init the partitions, and track within @data.
- *
- * @data:  The &struct slc_data tracking partition information.
- * @dev:   The platform device associated with the parent node.
- *
- * Return: On success, returns 0. On failure an error code is returned.
- */
-int slc_init_data(struct slc_data *data, struct device* dev)
-{
-	int ret = -EINVAL;
-
-	if (data == NULL || dev == NULL)
-		goto err_exit;
-
-	/* Inherit the platform device */
-	data->dev = dev;
-
-	/* Register our node with the SLC driver.
-	 * This detects our partitions defined within the DT.
-	 */
-	data->pt_handle = pt_client_register(data->dev->of_node, NULL, NULL);
-	if (IS_ERR(data->pt_handle)) {
-		ret = PTR_ERR(data->pt_handle);
-		dev_err(data->dev, "pt_client_register failed with: %d\n", ret);
-		goto err_exit;
-	}
-
-	if ((ret = init_partition(data, &data->partition, 0)))
-		goto pt_init_err_exit;
-
-	return 0;
-
-pt_init_err_exit:
-	pt_client_unregister(data->pt_handle);
-
-err_exit:
-	return ret;
-}
-
-/**
- * slc_term_data - Tear down SLC partitions and free tracking data.
- *
- * @data:  The &struct slc_data tracking partition information.
- */
-void slc_term_data(struct slc_data *data)
-{
-	term_partition(data, &data->partition);
-
-	pt_client_unregister(data->pt_handle);
-}
diff --git a/mali_pixel/pixel_slc.h b/mali_pixel/pixel_slc.h
deleted file mode 100644
index 40b5ad7..0000000
--- a/mali_pixel/pixel_slc.h
+++ /dev/null
@@ -1,72 +0,0 @@
-// SPDX-License-Identifier: GPL-2.0
-/*
- * Copyright 2024 Google LLC.
- *
- * Author: Jack Diver <diverj@google.com>
- */
-#ifndef _PIXEL_SLC_H_
-#define _PIXEL_SLC_H_
-
-#include <soc/google/pt.h>
-
-/**
- * DOC: SLC partition management
- *
- * Key definitions:
- * + Partition index - The unique index of a partition, relative to the dt node that owns it.
- *                     This index is used when communicating with the underlying SLC driver.
- * + ptid - This is the HW level ID associated with an enabled partition. These id's are allocated
- *          at partition enable time. The GPU driver will never directly use the ptid, but will
- *          track it.
- *          External analysis of the caching behavior (e.g. hit and eviction counters), are
- *          associated with a ptid, not a physical partition index.
- *          This driver attempts to hold on to any allocated ptids until driver termination to make
- *          profiling of caching performance easier.
- * + PBHA - Acronym: Page Based Hardware Attributes. Every physical partition has a PBHA value
- *          associated with it. We insert these attributes into PTEs so that transactions with a
- *          page carry the PBHA within their high bits.
- *          Transactions with PBHA bits set are intercepted by the SLC, where the corresponding
- *          partition and it's caching behavior (Read/write alloc etc.) are looked up and applied to
- *          the transaction.
- */
-
-/**
- * struct slc_partition - Structure for tracking partition state.
- */
-struct slc_partition {
-	/** @index: The active partition ID for this virtual partition */
-	u32 index;
-
-	/** @ptid: The active partition ID for this virtual partition */
-	ptid_t ptid;
-
-	/** @pbha: The page based HW attributes for this partition */
-	ptpbha_t pbha;
-
-	/** @enabled: Is the partition currently enabled */
-	bool enabled;
-};
-
-/**
- * struct slc_data - Structure for tracking SLC context.
- */
-struct slc_data {
-	/** @pt_handle: Link to ACPM SLC partition data */
-	struct pt_handle *pt_handle;
-
-	/** @partition: Information specific to an individual SLC partition */
-	struct slc_partition partition;
-
-	/** @dev: Inherited pointer to device attached */
-	struct device *dev;
-};
-
-int slc_init_data(struct slc_data *data, struct device* dev);
-
-void slc_term_data(struct slc_data *data);
-
-u64 slc_set_pbha(struct slc_data const *data, u64 pte);
-
-u64 slc_wipe_pbha(u64 pte);
-
-#endif /* _PIXEL_SLC_H_ */