#include <linux/swiotlb.h>
 #include <linux/bootmem.h>
 #include <linux/dma-mapping.h>
+#include <linux/mem_encrypt.h>
 
 #include <asm/iommu.h>
 #include <asm/swiotlb.h>
 #include <asm/dma.h>
 #include <asm/xen/swiotlb-xen.h>
 #include <asm/iommu_table.h>
+
 int swiotlb __read_mostly;
 
 void *x86_swiotlb_alloc_coherent(struct device *hwdev, size_t size,
                  pci_swiotlb_late_init);
 
 /*
- * if 4GB or more detected (and iommu=off not set) return 1
- * and set swiotlb to 1.
+ * If 4GB or more detected (and iommu=off not set) or if SME is active
+ * then set swiotlb to 1 and return 1.
  */
 int __init pci_swiotlb_detect_4gb(void)
 {
        if (!no_iommu && max_possible_pfn > MAX_DMA32_PFN)
                swiotlb = 1;
 #endif
+
+       /*
+        * If SME is active then swiotlb will be set to 1 so that bounce
+        * buffers are allocated and used for devices that do not support
+        * the addressing range required for the encryption mask.
+        */
+       if (sme_active())
+               swiotlb = 1;
+
        return swiotlb;
 }
 IOMMU_INIT(pci_swiotlb_detect_4gb,
 
 #include <linux/linkage.h>
 #include <linux/init.h>
 #include <linux/mm.h>
+#include <linux/dma-mapping.h>
+#include <linux/swiotlb.h>
 
 #include <asm/tlbflush.h>
 #include <asm/fixmap.h>
 #include <asm/setup.h>
 #include <asm/bootparam.h>
+#include <asm/set_memory.h>
 
 /*
  * Since SME related variables are set early in the boot process they must
                protection_map[i] = pgprot_encrypted(protection_map[i]);
 }
 
+/* Architecture __weak replacement functions */
+void __init mem_encrypt_init(void)
+{
+       if (!sme_me_mask)
+               return;
+
+       /* Call into SWIOTLB to update the SWIOTLB DMA buffers */
+       swiotlb_update_mem_attributes();
+}
+
+void swiotlb_set_mem_attributes(void *vaddr, unsigned long size)
+{
+       WARN(PAGE_ALIGN(size) != size,
+            "size is not page-aligned (%#lx)\n", size);
+
+       /* Make the SWIOTLB buffer area decrypted */
+       set_memory_decrypted((unsigned long)vaddr, size >> PAGE_SHIFT);
+}
+
 void __init sme_encrypt_kernel(void)
 {
 }
 
 #include <linux/highmem.h>
 #include <linux/gfp.h>
 #include <linux/scatterlist.h>
+#include <linux/mem_encrypt.h>
 
 #include <asm/io.h>
 #include <asm/dma.h>
        return size ? size : (IO_TLB_DEFAULT_SIZE);
 }
 
+void __weak swiotlb_set_mem_attributes(void *vaddr, unsigned long size) { }
+
+/* For swiotlb, clear memory encryption mask from dma addresses */
+static dma_addr_t swiotlb_phys_to_dma(struct device *hwdev,
+                                     phys_addr_t address)
+{
+       return __sme_clr(phys_to_dma(hwdev, address));
+}
+
 /* Note that this doesn't work with highmem page */
 static dma_addr_t swiotlb_virt_to_bus(struct device *hwdev,
                                      volatile void *address)
               bytes >> 20, vstart, vend - 1);
 }
 
+/*
+ * Early SWIOTLB allocation may be too early to allow an architecture to
+ * perform the desired operations.  This function allows the architecture to
+ * call SWIOTLB when the operations are possible.  It needs to be called
+ * before the SWIOTLB memory is used.
+ */
+void __init swiotlb_update_mem_attributes(void)
+{
+       void *vaddr;
+       unsigned long bytes;
+
+       if (no_iotlb_memory || late_alloc)
+               return;
+
+       vaddr = phys_to_virt(io_tlb_start);
+       bytes = PAGE_ALIGN(io_tlb_nslabs << IO_TLB_SHIFT);
+       swiotlb_set_mem_attributes(vaddr, bytes);
+       memset(vaddr, 0, bytes);
+
+       vaddr = phys_to_virt(io_tlb_overflow_buffer);
+       bytes = PAGE_ALIGN(io_tlb_overflow);
+       swiotlb_set_mem_attributes(vaddr, bytes);
+       memset(vaddr, 0, bytes);
+}
+
 int __init swiotlb_init_with_tbl(char *tlb, unsigned long nslabs, int verbose)
 {
        void *v_overflow_buffer;
        io_tlb_start = virt_to_phys(tlb);
        io_tlb_end = io_tlb_start + bytes;
 
+       swiotlb_set_mem_attributes(tlb, bytes);
        memset(tlb, 0, bytes);
 
        /*
        if (!v_overflow_buffer)
                goto cleanup2;
 
+       swiotlb_set_mem_attributes(v_overflow_buffer, io_tlb_overflow);
+       memset(v_overflow_buffer, 0, io_tlb_overflow);
        io_tlb_overflow_buffer = virt_to_phys(v_overflow_buffer);
 
        /*
                return SWIOTLB_MAP_ERROR;
        }
 
-       start_dma_addr = phys_to_dma(hwdev, io_tlb_start);
+       start_dma_addr = swiotlb_phys_to_dma(hwdev, io_tlb_start);
        return swiotlb_tbl_map_single(hwdev, start_dma_addr, phys, size,
                                      dir, attrs);
 }
                        goto err_warn;
 
                ret = phys_to_virt(paddr);
-               dev_addr = phys_to_dma(hwdev, paddr);
+               dev_addr = swiotlb_phys_to_dma(hwdev, paddr);
 
                /* Confirm address can be DMA'd by device */
                if (dev_addr + size - 1 > dma_mask) {
        map = map_single(dev, phys, size, dir, attrs);
        if (map == SWIOTLB_MAP_ERROR) {
                swiotlb_full(dev, size, dir, 1);
-               return phys_to_dma(dev, io_tlb_overflow_buffer);
+               return swiotlb_phys_to_dma(dev, io_tlb_overflow_buffer);
        }
 
-       dev_addr = phys_to_dma(dev, map);
+       dev_addr = swiotlb_phys_to_dma(dev, map);
 
        /* Ensure that the address returned is DMA'ble */
        if (dma_capable(dev, dev_addr, size))
        attrs |= DMA_ATTR_SKIP_CPU_SYNC;
        swiotlb_tbl_unmap_single(dev, map, size, dir, attrs);
 
-       return phys_to_dma(dev, io_tlb_overflow_buffer);
+       return swiotlb_phys_to_dma(dev, io_tlb_overflow_buffer);
 }
 EXPORT_SYMBOL_GPL(swiotlb_map_page);
 
                                sg_dma_len(sgl) = 0;
                                return 0;
                        }
-                       sg->dma_address = phys_to_dma(hwdev, map);
+                       sg->dma_address = swiotlb_phys_to_dma(hwdev, map);
                } else
                        sg->dma_address = dev_addr;
                sg_dma_len(sg) = sg->length;
 int
 swiotlb_dma_mapping_error(struct device *hwdev, dma_addr_t dma_addr)
 {
-       return (dma_addr == phys_to_dma(hwdev, io_tlb_overflow_buffer));
+       return (dma_addr == swiotlb_phys_to_dma(hwdev, io_tlb_overflow_buffer));
 }
 EXPORT_SYMBOL(swiotlb_dma_mapping_error);
 
 int
 swiotlb_dma_supported(struct device *hwdev, u64 mask)
 {
-       return phys_to_dma(hwdev, io_tlb_end - 1) <= mask;
+       return swiotlb_phys_to_dma(hwdev, io_tlb_end - 1) <= mask;
 }
 EXPORT_SYMBOL(swiotlb_dma_supported);