diff --git a/drivers/iommu/rockchip-iommu.c b/drivers/iommu/rockchip-iommu.c
index d0316996279e..ed39fbafd496 100644
--- a/drivers/iommu/rockchip-iommu.c
+++ b/drivers/iommu/rockchip-iommu.c
@@ -1108,6 +1108,14 @@ static int rk_iommu_add_device(struct device *dev)
 
 	data->defer_attach = false;
 
+	/* set max segment size for dev, needed for single chunk map */
+	if (!dev->dma_parms)
+		dev->dma_parms = kzalloc(sizeof(*dev->dma_parms), GFP_KERNEL);
+	if (!dev->dma_parms)
+		return -ENOMEM;
+
+	dma_set_max_seg_size(dev, DMA_BIT_MASK(32));
+
 	return 0;
 }
 
@@ -1118,6 +1126,9 @@ static void rk_iommu_remove_device(struct device *dev)
 
 	iommu = rk_iommu_from_dev(dev);
 
+	kfree(dev->dma_parms);
+	dev->dma_parms = NULL;
+
 	device_link_del(data->link);
 	iommu_device_unlink(&iommu->iommu, dev);
 	iommu_group_remove_device(dev);