And specify the memory to be 4K-aligned when running in SNP host, and
to be decrypted when running in an SEV guestso that the VMM can access
the memory successfully.

Signed-off-by: Suravee Suthikulpanit <suravee.suthikulpanit@amd.com>
---
 drivers/iommu/amd/amd_iommu_types.h |  2 +-
 drivers/iommu/amd/init.c            | 17 +++++++++++------
 drivers/iommu/amd/iommu.c           |  5 +++--
 3 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/drivers/iommu/amd/amd_iommu_types.h b/drivers/iommu/amd/amd_iommu_types.h
index 653955ab120d..e671e9220a21 100644
--- a/drivers/iommu/amd/amd_iommu_types.h
+++ b/drivers/iommu/amd/amd_iommu_types.h
@@ -798,7 +798,7 @@ struct amd_iommu {
 #endif
 
 	u32 flags;
-	volatile u64 *cmd_sem;
+	struct amd_iommu_mem cmd_sem_mem;
 	atomic64_t cmd_sem_val;
 
 #ifdef CONFIG_AMD_IOMMU_DEBUGFS
diff --git a/drivers/iommu/amd/init.c b/drivers/iommu/amd/init.c
index 77147dc3b79f..51861874656e 100644
--- a/drivers/iommu/amd/init.c
+++ b/drivers/iommu/amd/init.c
@@ -383,7 +383,7 @@ static void iommu_set_exclusion_range(struct amd_iommu *iommu)
 
 static void iommu_set_cwwb_range(struct amd_iommu *iommu)
 {
-	u64 start = iommu_virt_to_phys((void *)iommu->cmd_sem);
+	u64 start = amd_iommu_mem_to_phys(&iommu->cmd_sem_mem);
 	u64 entry = start & PM_ADDR_MASK;
 
 	if (!check_feature(FEATURE_SNP))
@@ -963,15 +963,20 @@ static int iommu_init_ga_log(struct amd_iommu *iommu)
 
 static int __init alloc_cwwb_sem(struct amd_iommu *iommu)
 {
-	iommu->cmd_sem = iommu_alloc_4k_pages(iommu, GFP_KERNEL, 1);
+	struct amd_iommu_mem *mem = &iommu->cmd_sem_mem;
 
-	return iommu->cmd_sem ? 0 : -ENOMEM;
+	mem->modes = ALLOC_MODE_4K | ALLOC_MODE_GUEST_MEM_DECRYPT;
+	mem->order = get_order(1);
+	mem->buf = amd_iommu_get_zeroed_mem(GFP_KERNEL, mem);
+	if (!mem->buf)
+		return -ENOMEM;
+
+	return 0;
 }
 
 static void __init free_cwwb_sem(struct amd_iommu *iommu)
 {
-	if (iommu->cmd_sem)
-		iommu_free_page((void *)iommu->cmd_sem);
+	amd_iommu_free_mem(&iommu->cmd_sem_mem);
 }
 
 static void iommu_enable_xt(struct amd_iommu *iommu)
@@ -3841,7 +3846,7 @@ int amd_iommu_snp_disable(void)
 		if (ret)
 			return ret;
 
-		ret = iommu_make_shared((void *)iommu->cmd_sem, PAGE_SIZE);
+		ret = iommu_make_shared(iommu->cmd_sem_mem.buf, PAGE_SIZE);
 		if (ret)
 			return ret;
 	}
diff --git a/drivers/iommu/amd/iommu.c b/drivers/iommu/amd/iommu.c
index 2b18134f1eb5..bd29d26c8d44 100644
--- a/drivers/iommu/amd/iommu.c
+++ b/drivers/iommu/amd/iommu.c
@@ -1056,8 +1056,9 @@ irqreturn_t amd_iommu_int_handler(int irq, void *data)
 static int wait_on_sem(struct amd_iommu *iommu, u64 data)
 {
 	int i = 0;
+	u64 *cmd_sem = (u64 *)iommu->cmd_sem_mem.buf;
 
-	while (*iommu->cmd_sem != data && i < LOOP_TIMEOUT) {
+	while (*cmd_sem != data && i < LOOP_TIMEOUT) {
 		udelay(1);
 		i += 1;
 	}
@@ -1092,7 +1093,7 @@ static void build_completion_wait(struct iommu_cmd *cmd,
 				  struct amd_iommu *iommu,
 				  u64 data)
 {
-	u64 paddr = iommu_virt_to_phys((void *)iommu->cmd_sem);
+	u64 paddr = amd_iommu_mem_to_phys(&iommu->cmd_sem_mem);
 
 	memset(cmd, 0, sizeof(*cmd));
 	cmd->data[0] = lower_32_bits(paddr) | CMD_COMPL_WAIT_STORE_MASK;
-- 
2.34.1