x86/tdx: Enable direct iomap MMIO optimizations

Using the infrastructure added in a previous patch, enable direct
TDCALLs for MMIO accesses through iomap. Add functions for each
of the read and write operations and hook them in using the iomap_mmio
call vector when TDX is active.

This improves the TD virtio performance for disk workloads by about 4%

Signed-off-by: Andi Kleen <ak@linux.intel.com>

arch/x86/Kconfig

@@ -881,6 +881,7 @@ config INTEL_TDX_GUEST
 	select X86_MEM_ENCRYPT
 	select UNACCEPTED_MEMORY
 	select ARCH_HAS_RESTRICTED_VIRTIO_MEMORY_ACCESS
+	select IOMAP_IND_MMIO
 	help
 	  Support running as a guest under Intel TDX.  Without this support,
 	  the guest kernel can not boot or run under TDX.

arch/x86/kernel/tdx.c

@@ -6,6 +6,7 @@
 
 #include <linux/cpufeature.h>
 #include <linux/swiotlb.h>
+#include <linux/io.h>
 #include <asm/tdx.h>
 #include <asm/i8259.h>
 #include <asm/apic.h>
@@ -645,6 +646,97 @@ static int tdx_handle_mmio(struct pt_regs *regs, struct ve_info *ve)
 	return insn.length;
 }
 
+static unsigned long tdx_virt_mmio(int size, bool write, unsigned long vaddr,
+				   unsigned long *val)
+{
+	pte_t *pte;
+	int level;
+
+	pte = lookup_address(vaddr, &level);
+	if (!pte)
+		return -EIO;
+
+	return tdx_mmio(size, write,
+			(pte_pfn(*pte) << PAGE_SHIFT) +
+			(vaddr & ~page_level_mask(level)),
+			val);
+}
+
+static unsigned char tdx_mmio_readb(void __iomem *addr)
+{
+	unsigned long val;
+
+	if (tdx_virt_mmio(1, false, (unsigned long)addr, &val))
+		return 0xff;
+	return val;
+}
+
+static unsigned short tdx_mmio_readw(void __iomem *addr)
+{
+	unsigned long val;
+
+	if (tdx_virt_mmio(2, false, (unsigned long)addr, &val))
+		return 0xffff;
+	return val;
+}
+
+static unsigned int tdx_mmio_readl(void __iomem *addr)
+{
+	unsigned long val;
+
+	if (tdx_virt_mmio(4, false, (unsigned long)addr, &val))
+		return 0xffffffff;
+	return val;
+}
+
+unsigned long tdx_mmio_readq(void __iomem *addr)
+{
+	unsigned long val;
+
+	if (tdx_virt_mmio(8, false, (unsigned long)addr, &val))
+		return 0xffffffffffffffff;
+	return val;
+}
+
+static void tdx_mmio_writeb(unsigned char v, void __iomem *addr)
+{
+	unsigned long val = v;
+
+	tdx_virt_mmio(1, true, (unsigned long)addr, &val);
+}
+
+static void tdx_mmio_writew(unsigned short v, void __iomem *addr)
+{
+	unsigned long val = v;
+
+	tdx_virt_mmio(2, true, (unsigned long)addr, &val);
+}
+
+static void tdx_mmio_writel(unsigned int v, void __iomem *addr)
+{
+	unsigned long val = v;
+
+	tdx_virt_mmio(4, true, (unsigned long)addr, &val);
+}
+
+static void tdx_mmio_writeq(unsigned long v, void __iomem *addr)
+{
+	unsigned long val = v;
+
+	tdx_virt_mmio(8, true, (unsigned long)addr, &val);
+}
+
+static const struct iomap_mmio tdx_iomap_mmio = {
+	.ireadb  = tdx_mmio_readb,
+	.ireadw  = tdx_mmio_readw,
+	.ireadl  = tdx_mmio_readl,
+	.ireadq  = tdx_mmio_readq,
+	.iwriteb = tdx_mmio_writeb,
+	.iwritew = tdx_mmio_writew,
+	.iwritel = tdx_mmio_writel,
+	.iwriteq = tdx_mmio_writeq,
+};
+
 /*
  * Emulate I/O using hypercall.
  *
@@ -899,6 +991,8 @@ void __init tdx_early_init(void)
 	 */
 	random_enable_trust_cpu();
 
+	iomap_mmio = &tdx_iomap_mmio;
+
 	/*
 	 * Make sure there is a panic if something goes wrong,
 	 * just in case it's some kind of host attack.