diff --git a/drivers/gpu/host1x/dev.c b/drivers/gpu/host1x/dev.c
index 1f916b579e95..03db71173f5d 100644
--- a/drivers/gpu/host1x/dev.c
+++ b/drivers/gpu/host1x/dev.c
@@ -218,19 +218,24 @@ static int host1x_probe(struct platform_device *pdev)
 		return err;
 	}
 
-	if (iommu_present(&platform_bus_type)) {
+	host->group = iommu_group_get(&pdev->dev);
+	if (host->group) {
 		struct iommu_domain_geometry *geometry;
 		unsigned long order;
 
 		host->domain = iommu_domain_alloc(&platform_bus_type);
-		if (!host->domain)
-			return -ENOMEM;
+		if (!host->domain) {
+			err = -ENOMEM;
+			goto put_group;
+		}
 
-		err = iommu_attach_device(host->domain, &pdev->dev);
+		err = iommu_attach_group(host->domain, host->group);
 		if (err) {
 			if (err == -ENODEV) {
 				iommu_domain_free(host->domain);
 				host->domain = NULL;
+				iommu_group_put(host->group);
+				host->group = NULL;
 				goto skip_iommu;
 			}
 
@@ -296,13 +301,15 @@ fail_unprepare_disable:
 fail_free_channels:
 	host1x_channel_list_free(&host->channel_list);
 fail_detach_device:
-	if (host->domain) {
+	if (host->group && host->domain) {
 		put_iova_domain(&host->iova);
-		iommu_detach_device(host->domain, &pdev->dev);
+		iommu_detach_group(host->domain, host->group);
 	}
 fail_free_domain:
 	if (host->domain)
 		iommu_domain_free(host->domain);
+put_group:
+	iommu_group_put(host->group);
 
 	return err;
 }
@@ -319,8 +326,9 @@ static int host1x_remove(struct platform_device *pdev)
 
 	if (host->domain) {
 		put_iova_domain(&host->iova);
-		iommu_detach_device(host->domain, &pdev->dev);
+		iommu_detach_group(host->domain, host->group);
 		iommu_domain_free(host->domain);
+		iommu_group_put(host->group);
 	}
 
 	return 0;
diff --git a/drivers/gpu/host1x/dev.h b/drivers/gpu/host1x/dev.h
index 502769726480..43e9fabb43a1 100644
--- a/drivers/gpu/host1x/dev.h
+++ b/drivers/gpu/host1x/dev.h
@@ -117,6 +117,7 @@ struct host1x {
 	struct clk *clk;
 	struct reset_control *rst;
 
+	struct iommu_group *group;
 	struct iommu_domain *domain;
 	struct iova_domain iova;
 	dma_addr_t iova_end;