diff --git a/zephyr/wrapper.c b/zephyr/wrapper.c
index 55f3f7e17210..25858f46b570 100644
--- a/zephyr/wrapper.c
+++ b/zephyr/wrapper.c
@@ -151,10 +151,12 @@ void *rmalloc(enum mem_zone zone, uint32_t flags, uint32_t caps, size_t bytes)
 {
 	if (zone_is_cached(zone))
 		return heap_alloc_aligned_cached(&sof_heap, 0, bytes);
-	else
-		return heap_alloc_aligned(&sof_heap, 8, bytes);
 
+#ifdef CONFIG_IMX
+	return heap_alloc_aligned(&sof_heap, 8, bytes);
+#else
 	return heap_alloc_aligned(&sof_heap_shared, 8, bytes);
+#endif
 }
 
 /* Use SOF_MEM_ZONE_BUFFER at the moment */