diff --git a/arch/arm64/include/asm/kvm_host.h b/arch/arm64/include/asm/kvm_host.h
index e158b25f6df7..ac63d40cb6a7 100644
--- a/arch/arm64/include/asm/kvm_host.h
+++ b/arch/arm64/include/asm/kvm_host.h
@@ -1668,7 +1668,7 @@ int __pkvm_topup_hyp_alloc_mgt_gfp(unsigned long id, unsigned long nr_pages,
 struct kvm_iommu_driver {
 	int (*init_driver)(void);
 	void (*remove_driver)(void);
-	pkvm_handle_t (*get_iommu_id)(struct device *dev);
+	pkvm_handle_t (*get_iommu_id_by_of)(struct device_node *np);
 	ANDROID_KABI_RESERVE(1);
 	ANDROID_KABI_RESERVE(2);
 	ANDROID_KABI_RESERVE(3);
@@ -1685,6 +1685,7 @@ int kvm_iommu_init_hyp(struct kvm_iommu_ops *hyp_ops,
 		       struct kvm_hyp_memcache *atomic_mc);
 int kvm_iommu_init_driver(void);
 void kvm_iommu_remove_driver(void);
+pkvm_handle_t kvm_get_iommu_id_by_of(struct device_node *np);
 
 int pkvm_iommu_suspend(struct device *dev);
 int pkvm_iommu_resume(struct device *dev);
diff --git a/arch/arm64/kvm/iommu.c b/arch/arm64/kvm/iommu.c
index 1686e04cc342..31b0d4982bf8 100644
--- a/arch/arm64/kvm/iommu.c
+++ b/arch/arm64/kvm/iommu.c
@@ -40,7 +40,7 @@ EXPORT_SYMBOL(kvm_iommu_init_hyp);
 
 int kvm_iommu_init_driver(void)
 {
-	if (!smp_load_acquire(&iommu_driver) || !iommu_driver->get_iommu_id) {
+	if (!smp_load_acquire(&iommu_driver) || !iommu_driver->get_iommu_id_by_of) {
 		kvm_err("pKVM enabled without an IOMMU driver, do not run confidential workloads in virtual machines\n");
 		return 0;
 	}
@@ -62,9 +62,18 @@ void kvm_iommu_remove_driver(void)
 		iommu_driver->remove_driver();
 }
 
+
+pkvm_handle_t kvm_get_iommu_id_by_of(struct device_node *np)
+{
+	if (!iommu_driver)
+		return 0;
+
+	return iommu_driver->get_iommu_id_by_of(np);
+}
+
 static pkvm_handle_t kvm_get_iommu_id(struct device *dev)
 {
-	return iommu_driver->get_iommu_id(dev);
+	return kvm_get_iommu_id_by_of(dev_of_node(dev));
 }
 
 int pkvm_iommu_suspend(struct device *dev)
diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-kvm.c b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-kvm.c
index 1995f70fead5..cc7169bd913c 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-kvm.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-kvm.c
@@ -1201,10 +1201,23 @@ static pkvm_handle_t kvm_arm_smmu_v3_id(struct device *dev)
 	return host_smmu->id;
 }
 
+static pkvm_handle_t kvm_arm_v3_id_by_of(struct device_node *np)
+{
+	struct device *dev;
+
+	dev = driver_find_device_by_of_node(&kvm_arm_smmu_driver.driver, np);
+	if (!dev)
+		return 0;
+
+	put_device(dev);
+
+	return kvm_arm_smmu_v3_id(dev);
+}
+
 struct kvm_iommu_driver kvm_smmu_v3_ops = {
 	.init_driver = kvm_arm_smmu_v3_init_drv,
 	.remove_driver = kvm_arm_smmu_v3_remove_drv,
-	.get_iommu_id = kvm_arm_smmu_v3_id,
+	.get_iommu_id_by_of = kvm_arm_v3_id_by_of,
 };
 
 static int kvm_arm_smmu_v3_register(void)