[PATCH 3/5] ttm: add initial memcg integration.

[Dave Airlie <airlied@gmail.com>]

From: Dave Airlie <airlied@redhat.com>

Doing proper integration of TTM system memory allocations with
memcg is a difficult ask, primarily due to difficulties around
accounting for evictions properly.

However there are systems where userspace will be allocating
objects in system memory and they won't be prone to migrating
or evicting and we should start with at least accounting those.

This adds a memcg group to ttm bo and tt objects. If this object
allocates pages directly (not from a TTM pool, so cached non-dma
pages), then they will be accounted for using __GFP_ACCOUNT,
and added to the memcg GPU statistic.

Only operations which set the account_op flag in the ttm operation
context will have this accounting happen. This patch disables the
flag around object evictions, but any operation that could populate
a TTM tt object in process context should set the account_op flag.

This doesn't account for uncached pages due to the page pool,
and it doesn't account for dma allocated pages. However for a lot
of use cases this should be a sufficient first step. Adding uncached
page support would mean hooking into memcg for the uncached pool
interactions, and I need to consider dma paths further.

Signed-off-by: Dave Airlie <airlied@redhat.com>
---
 drivers/gpu/drm/ttm/ttm_bo.c    |  4 ++++
 drivers/gpu/drm/ttm/ttm_bo_vm.c |  3 ++-
 drivers/gpu/drm/ttm/ttm_pool.c  | 19 ++++++++++++++++++-
 drivers/gpu/drm/ttm/ttm_tt.c    |  1 +
 include/drm/ttm/ttm_bo.h        |  8 ++++++++
 include/drm/ttm/ttm_tt.h        |  6 +++++-
 6 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/drivers/gpu/drm/ttm/ttm_bo.c b/drivers/gpu/drm/ttm/ttm_bo.c
index 95b86003c50d..631984fd459b 100644
--- a/drivers/gpu/drm/ttm/ttm_bo.c
+++ b/drivers/gpu/drm/ttm/ttm_bo.c
@@ -744,8 +744,12 @@ static int ttm_bo_alloc_resource(struct ttm_buffer_object *bo,
 				continue;
 			}
 
+			/* we don't want to account evictions at this point */
+			bool old_ctx_account = ctx->account_op;
+			ctx->account_op = false;
 			ret = ttm_bo_evict_alloc(bdev, man, place, bo, ctx,
 						 ticket, res, limit_pool);
+			ctx->account_op = old_ctx_account;
 			dmem_cgroup_pool_state_put(limit_pool);
 			if (ret == -EBUSY)
 				continue;
diff --git a/drivers/gpu/drm/ttm/ttm_bo_vm.c b/drivers/gpu/drm/ttm/ttm_bo_vm.c
index a194db83421d..163039cf40a5 100644
--- a/drivers/gpu/drm/ttm/ttm_bo_vm.c
+++ b/drivers/gpu/drm/ttm/ttm_bo_vm.c
@@ -220,7 +220,8 @@ vm_fault_t ttm_bo_vm_fault_reserved(struct vm_fault *vmf,
 		struct ttm_operation_ctx ctx = {
 			.interruptible = true,
 			.no_wait_gpu = false,
-			.force_alloc = true
+			.force_alloc = true,
+			.account_op = true,
 		};
 
 		ttm = bo->ttm;
diff --git a/drivers/gpu/drm/ttm/ttm_pool.c b/drivers/gpu/drm/ttm/ttm_pool.c
index 83b10706ba89..934bbf2c0ff6 100644
--- a/drivers/gpu/drm/ttm/ttm_pool.c
+++ b/drivers/gpu/drm/ttm/ttm_pool.c
@@ -672,6 +672,9 @@ static void ttm_pool_free_range(struct ttm_pool *pool, struct ttm_tt *tt,
 				tt->dma_address + i : NULL;
 
 			nr = ttm_pool_unmap_and_free(pool, p, dma_addr, caching);
+			if (tt->page_flags & TTM_TT_FLAG_ACCOUNTED) {
+				mod_memcg_state(tt->memcg, MEMCG_GPU, -nr);
+			}
 		}
 	}
 }
@@ -742,9 +745,17 @@ static int __ttm_pool_alloc(struct ttm_pool *pool, struct ttm_tt *tt,
 		 * that behaviour.
 		 */
 		if (!p) {
+			gfp_t account_flags = gfp_flags;
 			page_caching = ttm_cached;
 			allow_pools = false;
-			p = ttm_pool_alloc_page(pool, gfp_flags, order);
+
+			if (!pool->use_dma_alloc && tt->memcg && ctx->account_op) {
+				account_flags |= __GFP_ACCOUNT;
+				tt->page_flags |= TTM_TT_FLAG_ACCOUNTED;
+			} else {
+				tt->page_flags &= ~TTM_TT_FLAG_ACCOUNTED;
+			}
+			p = ttm_pool_alloc_page(pool, account_flags, order);
 		}
 		/* If that fails, lower the order if possible and retry. */
 		if (!p) {
@@ -757,6 +768,12 @@ static int __ttm_pool_alloc(struct ttm_pool *pool, struct ttm_tt *tt,
 			r = -ENOMEM;
 			goto error_free_all;
 		}
+
+		/* only deal with non-pool pages for memcg for now */
+		if (tt->page_flags & TTM_TT_FLAG_ACCOUNTED) {
+			mod_memcg_state(tt->memcg, MEMCG_GPU, (1 << order));
+		}
+
 		r = ttm_pool_page_allocated(pool, order, p, page_caching, alloc,
 					    restore);
 		if (r)
diff --git a/drivers/gpu/drm/ttm/ttm_tt.c b/drivers/gpu/drm/ttm/ttm_tt.c
index df0aa6c4b8b8..cf9852560d9b 100644
--- a/drivers/gpu/drm/ttm/ttm_tt.c
+++ b/drivers/gpu/drm/ttm/ttm_tt.c
@@ -161,6 +161,7 @@ static void ttm_tt_init_fields(struct ttm_tt *ttm,
 	ttm->caching = caching;
 	ttm->restore = NULL;
 	ttm->backup = NULL;
+	ttm->memcg = bo->memcg;
 }
 
 int ttm_tt_init(struct ttm_tt *ttm, struct ttm_buffer_object *bo,
diff --git a/include/drm/ttm/ttm_bo.h b/include/drm/ttm/ttm_bo.h
index 903cd1030110..03e30a056427 100644
--- a/include/drm/ttm/ttm_bo.h
+++ b/include/drm/ttm/ttm_bo.h
@@ -135,6 +135,12 @@ struct ttm_buffer_object {
 	 * reservation lock.
 	 */
 	struct sg_table *sg;
+
+	/**
+	 * @memcg: memory cgroup to charge this to if it ends up using system memory.
+	 * NULL means don't charge.
+	 */
+	struct mem_cgroup *memcg;
 };
 
 #define TTM_BO_MAP_IOMEM_MASK 0x80
@@ -174,6 +180,7 @@ struct ttm_bo_kmap_obj {
  * BOs share the same reservation object.
  * @force_alloc: Don't check the memory account during suspend or CPU page
  * faults. Should only be used by TTM internally.
+ * @account_op: use __GFP_ACCOUNT for this.
  * @resv: Reservation object to allow reserved evictions with.
  * @bytes_moved: Statistics on how many bytes have been moved.
  *
@@ -186,6 +193,7 @@ struct ttm_operation_ctx {
 	bool gfp_retry_mayfail;
 	bool allow_res_evict;
 	bool force_alloc;
+	bool account_op;
 	struct dma_resv *resv;
 	uint64_t bytes_moved;
 };
diff --git a/include/drm/ttm/ttm_tt.h b/include/drm/ttm/ttm_tt.h
index 13cf47f3322f..23e5bcbaec3a 100644
--- a/include/drm/ttm/ttm_tt.h
+++ b/include/drm/ttm/ttm_tt.h
@@ -101,8 +101,9 @@ struct ttm_tt {
 #define TTM_TT_FLAG_EXTERNAL_MAPPABLE	BIT(3)
 #define TTM_TT_FLAG_DECRYPTED		BIT(4)
 #define TTM_TT_FLAG_BACKED_UP	        BIT(5)
+#define TTM_TT_FLAG_ACCOUNTED		BIT(6)
 
-#define TTM_TT_FLAG_PRIV_POPULATED	BIT(6)
+#define TTM_TT_FLAG_PRIV_POPULATED	BIT(7)
 	uint32_t page_flags;
 	/** @num_pages: Number of pages in the page array. */
 	uint32_t num_pages;
@@ -126,6 +127,9 @@ struct ttm_tt {
 	enum ttm_caching caching;
 	/** @restore: Partial restoration from backup state. TTM private */
 	struct ttm_pool_tt_restore *restore;
+
+	/** @memcg: Memory cgroup to account this to */
+	struct mem_cgroup *memcg;
 };
 
 /**
-- 
2.49.0