ANDROID: KVM: arm64: s2mpu: Create SysMMU_SYNC driver

SysMMU_SYNC devices expose an interface to start a sync counter and
poll its SFR until the device signals that all memory transactions in
flight at the start have drained. This gives the hypervisor a reliable
indicator that S2MPU invalidation has fully completed and all new
transactions will use the new MPTs.

Add a new pKVM IOMMU driver that the host can use to register
SysMMU_SYNCs. Each device is expected to be a supplier to exactly one
S2MPU (parent), but multiple SYNCs can supply a single S2MPU.

To keep things simple, the SYNCs do not implement suspend/resume and are
assumed to follow the power transitions of their parent.

Following an invalidation, the S2MPU driver iterates over its children
and waits for each SYNC to signal that its transactions have drained.
The algorithm currently waits on each SYNC in turn. If latency proves to
be an issue, this could be optimized to initiate a SYNC on all powered
devices before starting to poll.

Bug: 190463801
Signed-off-by: David Brazdil <dbrazdil@google.com>
Change-Id: I45b832fd11d76b65987935c8548e2a214ee2fa2a

arch/arm64/include/asm/kvm_host.h

@@ -376,6 +376,7 @@ extern u64 kvm_nvhe_sym(hyp_cpu_logical_map)[NR_CPUS];
 
 enum pkvm_iommu_driver_id {
 	PKVM_IOMMU_DRIVER_S2MPU,
+	PKVM_IOMMU_DRIVER_SYSMMU_SYNC,
 	PKVM_IOMMU_NR_DRIVERS,
 };
 
@@ -391,6 +392,8 @@ int pkvm_iommu_suspend(struct device *dev);
 int pkvm_iommu_resume(struct device *dev);
 
 int pkvm_iommu_s2mpu_register(struct device *dev, phys_addr_t pa);
+int pkvm_iommu_sysmmu_sync_register(struct device *dev, phys_addr_t pa,
+				    struct device *parent);
 
 struct vcpu_reset_state {
 	unsigned long	pc;

arch/arm64/include/asm/kvm_s2mpu.h

@@ -12,6 +12,10 @@
 #include <asm/kvm_mmu.h>
 
 #define S2MPU_MMIO_SIZE				SZ_64K
+#define SYSMMU_SYNC_MMIO_SIZE			SZ_64K
+#define SYSMMU_SYNC_S2_OFFSET			SZ_32K
+#define SYSMMU_SYNC_S2_MMIO_SIZE		(SYSMMU_SYNC_MMIO_SIZE - \
+						 SYSMMU_SYNC_S2_OFFSET)
 
 #define NR_VIDS					8
 #define NR_CTX_IDS				8
@@ -128,6 +132,13 @@ static_assert(SMPT_GRAN <= PAGE_SIZE);
 #define SMPT_NUM_PAGES				(SMPT_SIZE / PAGE_SIZE)
 #define SMPT_ORDER				get_order(SMPT_SIZE)
 
+/* SysMMU_SYNC registers, relative to SYSMMU_SYNC_S2_OFFSET. */
+#define REG_NS_SYNC_CMD				0x0
+#define REG_NS_SYNC_COMP			0x4
+
+#define SYNC_CMD_SYNC				BIT(0)
+#define SYNC_COMP_COMPLETE			BIT(0)
+
 /*
  * Iterate over S2MPU gigabyte regions. Skip those that cannot be modified
  * (the MMIO registers are read only, with reset value MPT_PROT_NONE).

arch/arm64/kvm/hyp/include/nvhe/iommu.h

@@ -92,5 +92,6 @@ void pkvm_iommu_host_stage2_idmap(phys_addr_t start, phys_addr_t end,
 				  enum kvm_pgtable_prot prot);
 
 extern const struct pkvm_iommu_ops pkvm_s2mpu_ops;
+extern const struct pkvm_iommu_ops pkvm_sysmmu_sync_ops;
 
 #endif	/* __ARM64_KVM_NVHE_IOMMU_H__ */

arch/arm64/kvm/hyp/nvhe/iommu.c

@@ -68,6 +68,8 @@ static const struct pkvm_iommu_ops *get_driver_ops(enum pkvm_iommu_driver_id id)
 	switch (id) {
 	case PKVM_IOMMU_DRIVER_S2MPU:
 		return IS_ENABLED(CONFIG_KVM_S2MPU) ? &pkvm_s2mpu_ops : NULL;
+	case PKVM_IOMMU_DRIVER_SYSMMU_SYNC:
+		return IS_ENABLED(CONFIG_KVM_S2MPU) ? &pkvm_sysmmu_sync_ops : NULL;
 	default:
 		return NULL;
 	}

arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c

@@ -28,6 +28,9 @@
 	(CONTEXT_CFG_VALID_VID_CTX_VID(ctxid, vid) \
 	 | (((ctxid) < (nr_ctx)) ? CONTEXT_CFG_VALID_VID_CTX_VALID(ctxid) : 0))
 
+#define for_each_child(child, dev) \
+	list_for_each_entry((child), &(dev)->children, siblings)
+
 struct s2mpu_drv_data {
 	u32 version;
 	u32 context_cfg_valid_vid;
@@ -155,6 +158,13 @@ static void __set_control_regs(struct pkvm_iommu *dev)
 	writel_relaxed(ctrl0, dev->va + REG_NS_CTRL0);
 }
 
+/* Poll the given SFR until its value has all bits of a given mask set. */
+static void __wait_until(void __iomem *addr, u32 mask)
+{
+	while ((readl_relaxed(addr) & mask) != mask)
+		continue;
+}
+
 /* Poll the given SFR as long as its value has all bits of a given mask set. */
 static void __wait_while(void __iomem *addr, u32 mask)
 {
@@ -164,6 +174,17 @@ static void __wait_while(void __iomem *addr, u32 mask)
 
 static void __wait_for_invalidation_complete(struct pkvm_iommu *dev)
 {
+	struct pkvm_iommu *sync;
+
+	/*
+	 * Wait for transactions to drain if SysMMU_SYNCs were registered.
+	 * Assumes that they are in the same power domain as the S2MPU.
+	 */
+	for_each_child(sync, dev) {
+		writel_relaxed(SYNC_CMD_SYNC, sync->va + REG_NS_SYNC_CMD);
+		__wait_until(sync->va + REG_NS_SYNC_COMP, SYNC_COMP_COMPLETE);
+	}
+
 	/* Must not access SFRs while S2MPU is busy invalidating (v9 only). */
 	if (is_version(dev, S2MPU_VERSION_9)) {
 		__wait_while(dev->va + REG_NS_STATUS,
@@ -473,9 +494,29 @@ static int s2mpu_validate(struct pkvm_iommu *dev)
 	return 0;
 }
 
+static int s2mpu_validate_child(struct pkvm_iommu *dev, struct pkvm_iommu *child)
+{
+	if (child->ops != &pkvm_sysmmu_sync_ops)
+		return -EINVAL;
+
+	return 0;
+}
+
+static int sysmmu_sync_validate(struct pkvm_iommu *dev)
+{
+	if (dev->size != SYSMMU_SYNC_S2_MMIO_SIZE)
+		return -EINVAL;
+
+	if (!dev->parent || dev->parent->ops != &pkvm_s2mpu_ops)
+		return -EINVAL;
+
+	return 0;
+}
+
 const struct pkvm_iommu_ops pkvm_s2mpu_ops = (struct pkvm_iommu_ops){
 	.init = s2mpu_init,
 	.validate = s2mpu_validate,
+	.validate_child = s2mpu_validate_child,
 	.resume = s2mpu_resume,
 	.suspend = s2mpu_suspend,
 	.host_stage2_idmap_prepare = s2mpu_host_stage2_idmap_prepare,
@@ -483,3 +524,7 @@ const struct pkvm_iommu_ops pkvm_s2mpu_ops = (struct pkvm_iommu_ops){
 	.host_dabt_handler = s2mpu_host_dabt_handler,
 	.data_size = sizeof(struct s2mpu_drv_data),
 };
+
+const struct pkvm_iommu_ops pkvm_sysmmu_sync_ops = (struct pkvm_iommu_ops){
+	.validate = sysmmu_sync_validate,
+};

arch/arm64/kvm/iommu/s2mpu.c

@@ -84,3 +84,37 @@ int pkvm_iommu_s2mpu_register(struct device *dev, phys_addr_t addr)
 				   addr, S2MPU_MMIO_SIZE, NULL);
 }
 EXPORT_SYMBOL_GPL(pkvm_iommu_s2mpu_register);
+
+static int init_sysmmu_sync_driver(void)
+{
+	static DEFINE_MUTEX(lock);
+	static bool init_done;
+
+	int ret = 0;
+
+	mutex_lock(&lock);
+	if (!init_done) {
+		ret = pkvm_iommu_driver_init(PKVM_IOMMU_DRIVER_SYSMMU_SYNC, NULL, 0);
+		init_done = !ret;
+	}
+	mutex_unlock(&lock);
+	return ret;
+}
+
+int pkvm_iommu_sysmmu_sync_register(struct device *dev, phys_addr_t addr,
+				    struct device *parent)
+{
+	int ret;
+
+	if (!is_protected_kvm_enabled())
+		return -ENODEV;
+
+	ret = init_sysmmu_sync_driver();
+	if (ret)
+		return ret;
+
+	return pkvm_iommu_register(dev, PKVM_IOMMU_DRIVER_SYSMMU_SYNC,
+				   addr + SYSMMU_SYNC_S2_OFFSET,
+				   SYSMMU_SYNC_S2_MMIO_SIZE, parent);
+}
+EXPORT_SYMBOL_GPL(pkvm_iommu_sysmmu_sync_register);