diff --git a/drivers/gpu/msm/adreno_drawctxt.c b/drivers/gpu/msm/adreno_drawctxt.c
index e31aaffb9f1a..cb2bc2af14df 100644
--- a/drivers/gpu/msm/adreno_drawctxt.c
+++ b/drivers/gpu/msm/adreno_drawctxt.c
@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: GPL-2.0-only
 /*
  * Copyright (c) 2002,2007-2020, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2024 Qualcomm Innovation Center, Inc. All rights reserved.
  */
 
 #include <linux/debugfs.h>
@@ -601,8 +602,6 @@ int adreno_drawctxt_switch(struct adreno_device *adreno_dev,
 	if (drawctxt != NULL && kgsl_context_detached(&drawctxt->base))
 		return -ENOENT;
 
-	trace_adreno_drawctxt_switch(rb, drawctxt);
-
 	/* Get a refcount to the new instance */
 	if (drawctxt) {
 		if (!_kgsl_context_get(&drawctxt->base))
@@ -616,7 +615,7 @@ int adreno_drawctxt_switch(struct adreno_device *adreno_dev,
 
 	ret = adreno_iommu_set_pt_ctx(rb, new_pt, drawctxt);
 	if (ret)
-		return ret;
+		goto err;
 
 	if (rb->drawctxt_active) {
 		/* Wait for the timestamp to expire */
@@ -626,7 +625,12 @@ int adreno_drawctxt_switch(struct adreno_device *adreno_dev,
 			kgsl_context_put(&rb->drawctxt_active->base);
 		}
 	}
+	trace_adreno_drawctxt_switch(rb, drawctxt);
 
 	rb->drawctxt_active = drawctxt;
 	return 0;
+err:
+	if (drawctxt)
+		kgsl_context_put(&drawctxt->base);
+	return ret;
 }
diff --git a/drivers/gpu/msm/adreno_iommu.c b/drivers/gpu/msm/adreno_iommu.c
index be88bafb12da..1ab5d6acea4a 100644
--- a/drivers/gpu/msm/adreno_iommu.c
+++ b/drivers/gpu/msm/adreno_iommu.c
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: GPL-2.0-only
 /*
  * Copyright (c) 2002,2007-2020, The Linux Foundation. All rights reserved.
- * Copyright (c) 2022 Qualcomm Innovation Center, Inc. All rights reserved.
+ * Copyright (c) 2022,2024 Qualcomm Innovation Center, Inc. All rights reserved.
  */
 
 #include <linux/slab.h>
@@ -353,63 +353,6 @@ static unsigned int __add_curr_ctxt_cmds(struct adreno_ringbuffer *rb,
 	return cmds - cmds_orig;
 }
 
-/**
- * _set_ctxt_gpu() - Add commands to set the current context in memstore
- * @rb: The ringbuffer in which commands to set memstore are added
- * @drawctxt: The context whose id is being set in memstore
- */
-static int _set_ctxt_gpu(struct adreno_ringbuffer *rb,
-			struct adreno_context *drawctxt)
-{
-	unsigned int link[15], *cmds;
-	int result;
-
-	cmds = &link[0];
-	cmds += __add_curr_ctxt_cmds(rb, cmds, drawctxt);
-	result = adreno_ringbuffer_issue_internal_cmds(rb, 0, link,
-			(unsigned int)(cmds - link));
-	return result;
-}
-
-/**
- * _set_pagetable_gpu() - Use GPU to switch the pagetable
- * @rb: The rb in which commands to switch pagetable are to be
- *    submitted
- * @new_pt: The pagetable to switch to
- */
-static int _set_pagetable_gpu(struct adreno_ringbuffer *rb,
-			struct kgsl_pagetable *new_pt)
-{
-	struct adreno_device *adreno_dev = ADRENO_RB_DEVICE(rb);
-	unsigned int *link = NULL, count;
-	int result;
-
-	link = kmalloc(PAGE_SIZE, GFP_KERNEL);
-	if (link == NULL)
-		return -ENOMEM;
-
-	/* If we are in a fault the MMU will be reset soon */
-	if (test_bit(ADRENO_DEVICE_FAULT, &adreno_dev->priv)) {
-		kfree(link);
-		return 0;
-	}
-
-	count = adreno_iommu_set_pt_generate_cmds(rb, link, new_pt);
-
-	WARN(count > (PAGE_SIZE / sizeof(unsigned int)),
-		"Temp command buffer overflow\n");
-
-	/*
-	 * This returns the per context timestamp but we need to
-	 * use the global timestamp for iommu clock disablement
-	 */
-	result = adreno_ringbuffer_issue_internal_cmds(rb,
-			KGSL_CMD_FLAGS_PMODE, link, count);
-
-	kfree(link);
-	return result;
-}
-
 /**
  * adreno_iommu_init() - Adreno iommu init
  * @adreno_dev: Adreno device
@@ -443,7 +386,6 @@ void adreno_iommu_init(struct adreno_device *adreno_dev)
 
 /**
  * adreno_iommu_set_pt_ctx() - Change the pagetable of the current RB
- * @device: Pointer to device to which the rb belongs
  * @rb: The RB pointer on which pagetable is to be changed
  * @new_pt: The new pt the device will change to
  * @drawctxt: The context whose pagetable the ringbuffer is switching to,
@@ -458,21 +400,34 @@ int adreno_iommu_set_pt_ctx(struct adreno_ringbuffer *rb,
 	struct adreno_device *adreno_dev = ADRENO_RB_DEVICE(rb);
 	struct kgsl_device *device = KGSL_DEVICE(adreno_dev);
 	struct kgsl_pagetable *cur_pt = device->mmu.defaultpagetable;
+	unsigned int *cmds = NULL, count = 0;
 	int result = 0;
 
+	cmds = kmalloc(PAGE_SIZE, GFP_KERNEL);
+	if (cmds == NULL)
+		return -ENOMEM;
+
 	/* Switch the page table if a MMU is attached */
 	if (kgsl_mmu_get_mmutype(device) != KGSL_MMU_TYPE_NONE) {
 		if (rb->drawctxt_active)
 			cur_pt = rb->drawctxt_active->base.proc_priv->pagetable;
 
-		/* Pagetable switch */
+		/* Add commands for pagetable switch */
 		if (new_pt != cur_pt)
-			result = _set_pagetable_gpu(rb, new_pt);
+			count += adreno_iommu_set_pt_generate_cmds(rb, cmds, new_pt);
 
-		if (result)
-			return result;
 	}
 
-	/* Context switch */
-	return _set_ctxt_gpu(rb, drawctxt);
+	/* Add commands to set the current context in memstore */
+	count += __add_curr_ctxt_cmds(rb, cmds + count, drawctxt);
+
+	WARN(count > (PAGE_SIZE / sizeof(unsigned int)),
+			"Temp command buffer overflow\n");
+
+	result = adreno_ringbuffer_issue_internal_cmds(rb, KGSL_CMD_FLAGS_PMODE,
+			cmds, count);
+
+	kfree(cmds);
+	return result;
+
 }