diff --git a/drivers/rknpu/rknpu_drv.c b/drivers/rknpu/rknpu_drv.c
index eea2e50caa02..055e20177401 100644
--- a/drivers/rknpu/rknpu_drv.c
+++ b/drivers/rknpu/rknpu_drv.c
@@ -566,17 +566,20 @@ static int rknpu_release(struct inode *inode, struct file *file)
 			"Fd close free rknpu_obj: %#llx, rknpu_obj->dma_addr: %#llx\n",
 			(__u64)(uintptr_t)entry, (__u64)entry->dma_addr);
 
-		vunmap(entry->kv_addr);
-		entry->kv_addr = NULL;
-
-		if (!entry->owner) {
-			dma_buf_put(entry->dmabuf);
-		} else {
-			dma_buf_unmap_attachment(entry->attachment, entry->sgt,
-						 DMA_BIDIRECTIONAL);
-			dma_buf_detach(entry->dmabuf, entry->attachment);
+		if (entry->kv_addr) {
+			struct iosys_map map =
+				IOSYS_MAP_INIT_VADDR(entry->kv_addr);
+			dma_buf_vunmap(entry->dmabuf, &map);
+			entry->kv_addr = NULL;
 		}
 
+		dma_buf_unmap_attachment(entry->attachment, entry->sgt,
+					 DMA_BIDIRECTIONAL);
+		dma_buf_detach(entry->dmabuf, entry->attachment);
+
+		if (!entry->owner)
+			dma_buf_put(entry->dmabuf);
+
 		list_del(&entry->head);
 		kfree(entry);
 	}
diff --git a/drivers/rknpu/rknpu_mem.c b/drivers/rknpu/rknpu_mem.c
index 70f363d7d21a..84ec2a723cf6 100644
--- a/drivers/rknpu/rknpu_mem.c
+++ b/drivers/rknpu/rknpu_mem.c
@@ -77,18 +77,14 @@ int rknpu_mem_create_ioctl(struct rknpu_device *rknpu_dev, struct file *file,
 	int ret = -EINVAL;
 	struct dma_buf_attachment *attachment;
 	struct sg_table *table;
-	struct scatterlist *sgl;
-	dma_addr_t phys;
 	struct dma_buf *dmabuf;
-	struct page **pages;
-	struct page *page;
 	struct rknpu_mem_object *rknpu_obj = NULL;
 	struct rknpu_session *session = NULL;
-	int i, fd;
-	unsigned int length, page_count;
+	int fd;
 	unsigned int in_size = _IOC_SIZE(cmd);
 	unsigned int k_size = sizeof(struct rknpu_mem_create);
 	char *k_data = (char *)&args;
+	struct iosys_map map = { 0 };
 
 	if (unlikely(copy_from_user(&args, (struct rknpu_mem_create *)data,
 				    in_size))) {
@@ -100,13 +96,6 @@ int rknpu_mem_create_ioctl(struct rknpu_device *rknpu_dev, struct file *file,
 	if (k_size > in_size)
 		memset(k_data + in_size, 0, k_size - in_size);
 
-	if (args.flags & RKNPU_MEM_NON_CONTIGUOUS) {
-		LOG_ERROR("%s: malloc iommu memory unsupported in current!\n",
-			  __func__);
-		ret = -EINVAL;
-		return ret;
-	}
-
 	rknpu_obj = kzalloc(sizeof(*rknpu_obj), GFP_KERNEL);
 	if (!rknpu_obj)
 		return -ENOMEM;
@@ -154,45 +143,23 @@ int rknpu_mem_create_ioctl(struct rknpu_device *rknpu_dev, struct file *file,
 
 	table = dma_buf_map_attachment(attachment, DMA_BIDIRECTIONAL);
 	if (IS_ERR(table)) {
-		LOG_ERROR("dma_buf_attach failed\n");
+		LOG_ERROR("dma_buf_map_attachment failed\n");
 		dma_buf_detach(dmabuf, attachment);
 		ret = PTR_ERR(table);
 		goto err_free_dma_buf;
 	}
 
-	for_each_sgtable_sg(table, sgl, i) {
-		phys = sg_dma_address(sgl);
-		page = sg_page(sgl);
-		length = sg_dma_len(sgl);
-		LOG_DEBUG("%s, %d, phys: %pad, length: %u\n", __func__,
-			  __LINE__, &phys, length);
-	}
-
 	if (args.flags & RKNPU_MEM_KERNEL_MAPPING) {
-		page_count = length >> PAGE_SHIFT;
-		pages = vmalloc(page_count * sizeof(struct page));
-		if (!pages) {
-			LOG_ERROR("alloc pages failed\n");
-			ret = -ENOMEM;
+		ret = dma_buf_vmap(dmabuf, &map);
+		if (ret) {
+			LOG_ERROR("dma_buf_vmap failed\n");
 			goto err_detach_dma_buf;
 		}
-
-		for (i = 0; i < page_count; i++)
-			pages[i] = &page[i];
-
-		rknpu_obj->kv_addr =
-			vmap(pages, page_count, VM_MAP, PAGE_KERNEL);
-		if (!rknpu_obj->kv_addr) {
-			LOG_ERROR("vmap pages addr failed\n");
-			ret = -ENOMEM;
-			goto err_free_pages;
-		}
-		vfree(pages);
-		pages = NULL;
+		rknpu_obj->kv_addr = map.vaddr;
 	}
 
 	rknpu_obj->size = PAGE_ALIGN(args.size);
-	rknpu_obj->dma_addr = phys;
+	rknpu_obj->dma_addr = sg_dma_address(table->sgl);
 	rknpu_obj->sgt = table;
 	rknpu_obj->attachment = attachment;
 
@@ -228,13 +195,9 @@ int rknpu_mem_create_ioctl(struct rknpu_device *rknpu_dev, struct file *file,
 	return 0;
 
 err_unmap_kv_addr:
-	vunmap(rknpu_obj->kv_addr);
+	dma_buf_vunmap(rknpu_obj->dmabuf, &map);
 	rknpu_obj->kv_addr = NULL;
 
-err_free_pages:
-	vfree(pages);
-	pages = NULL;
-
 err_detach_dma_buf:
 	dma_buf_unmap_attachment(attachment, table, DMA_BIDIRECTIONAL);
 	dma_buf_detach(dmabuf, attachment);
@@ -295,18 +258,20 @@ int rknpu_mem_destroy_ioctl(struct rknpu_device *rknpu_dev, struct file *file,
 	spin_unlock(&rknpu_dev->lock);
 
 	if (rknpu_obj == entry) {
-		vunmap(rknpu_obj->kv_addr);
-		rknpu_obj->kv_addr = NULL;
-
-		if (!rknpu_obj->owner) {
-			dma_buf_put(rknpu_obj->dmabuf);
-		} else {
-			dma_buf_unmap_attachment(rknpu_obj->attachment,
-						 rknpu_obj->sgt,
-						 DMA_BIDIRECTIONAL);
-			dma_buf_detach(rknpu_obj->dmabuf,
-				       rknpu_obj->attachment);
+		if (rknpu_obj->kv_addr) {
+			struct iosys_map map =
+				IOSYS_MAP_INIT_VADDR(rknpu_obj->kv_addr);
+			dma_buf_vunmap(rknpu_obj->dmabuf, &map);
+			rknpu_obj->kv_addr = NULL;
 		}
+
+		dma_buf_unmap_attachment(rknpu_obj->attachment, rknpu_obj->sgt,
+					 DMA_BIDIRECTIONAL);
+		dma_buf_detach(rknpu_obj->dmabuf, rknpu_obj->attachment);
+
+		if (!rknpu_obj->owner)
+			dma_buf_put(rknpu_obj->dmabuf);
+
 		kfree(rknpu_obj);
 	}