1 files changed, 17 insertions, 3 deletions

diff --git a/drivers/gpu/drm/nouveau/nvif/mmu.c b/drivers/gpu/drm/nouveau/nvif/mmu.c
index 15d0dcbf7ab4..ae08a1ca8044 100644
--- a/drivers/gpu/drm/nouveau/nvif/mmu.c
+++ b/drivers/gpu/drm/nouveau/nvif/mmu.c
@@ -36,6 +36,12 @@ nvif_mmu_fini(struct nvif_mmu *mmu)
 int
 nvif_mmu_init(struct nvif_object *parent, s32 oclass, struct nvif_mmu *mmu)
 {
+	static const struct nvif_mclass mems[] = {
+		{ NVIF_CLASS_MEM_GF100, -1 },
+		{ NVIF_CLASS_MEM_NV50 , -1 },
+		{ NVIF_CLASS_MEM_NV04 , -1 },
+		{}
+	};
 	struct nvif_mmu_v0 args;
 	int ret, i;
 
@@ -54,12 +60,20 @@ nvif_mmu_init(struct nvif_object *parent, s32 oclass, struct nvif_mmu *mmu)
 	mmu->type_nr = args.type_nr;
 	mmu->kind_nr = args.kind_nr;
 
-	mmu->heap = kmalloc(sizeof(*mmu->heap) * mmu->heap_nr, GFP_KERNEL);
-	mmu->type = kmalloc(sizeof(*mmu->type) * mmu->type_nr, GFP_KERNEL);
+	ret = nvif_mclass(&mmu->object, mems);
+	if (ret < 0)
+		goto done;
+	mmu->mem = mems[ret].oclass;
+
+	mmu->heap = kmalloc_array(mmu->heap_nr, sizeof(*mmu->heap),
+				  GFP_KERNEL);
+	mmu->type = kmalloc_array(mmu->type_nr, sizeof(*mmu->type),
+				  GFP_KERNEL);
 	if (ret = -ENOMEM, !mmu->heap || !mmu->type)
 		goto done;
 
-	mmu->kind = kmalloc(sizeof(*mmu->kind) * mmu->kind_nr, GFP_KERNEL);
+	mmu->kind = kmalloc_array(mmu->kind_nr, sizeof(*mmu->kind),
+				  GFP_KERNEL);
 	if (!mmu->kind && mmu->kind_nr)
 		goto done;