diff --git a/bdk/mem/smmu.c b/bdk/mem/smmu.c
index dff576d9..f8dcf414 100644
--- a/bdk/mem/smmu.c
+++ b/bdk/mem/smmu.c
@@ -156,6 +156,7 @@ void smmu_disable()
 	smmu_enable_payload[SMMU_PAYLOAD_EN_SHIFT] = SMMU_PAYLOAD_EN_SET;
 
 	smmu_enabled = false;
+	smmu_heap    = (void *)SMMU_HEAP_ADDR;
 }
 
 void smmu_reset_heap()
@@ -163,7 +164,7 @@ void smmu_reset_heap()
 	smmu_heap = (void *)SMMU_HEAP_ADDR;
 }
 
-void *smmu_init_domain(u32 dev_base, u32 asid)
+void *smmu_domain_init(u32 dev_base, u32 asid)
 {
 	void *ptb = _smmu_pdir_alloc();
 
@@ -178,7 +179,7 @@ void *smmu_init_domain(u32 dev_base, u32 asid)
 	return ptb;
 }
 
-void smmu_deinit_domain(u32 dev_base, u32 asid)
+void smmu_domain_deinit(u32 dev_base, u32 asid)
 {
 	MC(MC_SMMU_PTB_ASID) = asid;
 	MC(MC_SMMU_PTB_DATA) = 0;
diff --git a/bdk/mem/smmu.h b/bdk/mem/smmu.h
index 9952457b..07a283f1 100644
--- a/bdk/mem/smmu.h
+++ b/bdk/mem/smmu.h
@@ -15,11 +15,10 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-#include <utils/types.h>
-
 #ifndef _SMMU_H_
 #define _SMMU_H_
 
+#include <utils/types.h>
 #include <assert.h>
 
 #define MC_SMMU_AVPC_ASID            0x23C
@@ -68,8 +67,8 @@ void  smmu_init();
 void  smmu_enable();
 void  smmu_disable();
 void  smmu_reset_heap();
-void *smmu_init_domain(u32 dev_base, u32 asid);
-void  smmu_deinit_domain(u32 dev_base, u32 asid);
+void *smmu_domain_init(u32 dev_base, u32 asid);
+void  smmu_domain_deinit(u32 dev_base, u32 asid);
 void  smmu_domain_bypass(u32 dev_base, bool bypass);
 void  smmu_map(void *ptb, u32 iova, u64 iopa, u32 pages, u32 attr);
 void  smmu_map_huge(void *ptb, u32 iova, u64 iopa, u32 regions, u32 attr);
diff --git a/bdk/sec/tsec.c b/bdk/sec/tsec.c
index 63059576..c7553f96 100644
--- a/bdk/sec/tsec.c
+++ b/bdk/sec/tsec.c
@@ -146,7 +146,7 @@ int tsec_query(void *tsec_keys, tsec_ctxt_t *tsec_ctxt)
 	if (type == TSEC_FW_TYPE_EMU)
 	{
 		// Init SMMU translation for TSEC.
-		ptb = smmu_init_domain(MC_SMMU_TSEC_ASID, 1);
+		ptb = smmu_domain_init(MC_SMMU_TSEC_ASID, 1);
 		smmu_init();
 
 		// Enable SMMU.
@@ -230,7 +230,7 @@ int tsec_query(void *tsec_keys, tsec_ctxt_t *tsec_ctxt)
 		if (kidx != 8)
 		{
 			res = -6;
-			smmu_deinit_domain(MC_SMMU_TSEC_ASID, 1);
+			smmu_domain_deinit(MC_SMMU_TSEC_ASID, 1);
 
 			goto out_free;
 		}
@@ -241,7 +241,7 @@ int tsec_query(void *tsec_keys, tsec_ctxt_t *tsec_ctxt)
 		memcpy(tsec_keys, &key, 0x20);
 		memcpy(tsec_ctxt->pkg1, iram, 0x30000);
 
-		smmu_deinit_domain(MC_SMMU_TSEC_ASID, 1);
+		smmu_domain_deinit(MC_SMMU_TSEC_ASID, 1);
 
 		// for (int i = 0; i < kidx; i++)
 		// 	gfx_printf("key %08X\n", key[i]);