[PATCH 14/15] iommu, cxl: Configure IOMMU for CXL.cache

[Ben Cheatham <Benjamin.Cheatham@amd.com>]

Some IOMMU implementations require additional set up for enabling
ATS requests past enabling the base PCI ATS support. Create a callback
in the IOMMU core to be used by the CXL driver during device set up
that, when called, configures the IOMMU to handle ATS requests with the
CXL source bit set. For CXL.cache devices that don't support ATS, the
expectation is for the device to be attached to an identity domain, or
the iommu be disabled.

Update the AMD IOMMU driver with an implementation of the CXL ATS
callback and defaulting CXL.cache devices without ATS support to an
identity domain.

Signed-off-by: Ben Cheatham <Benjamin.Cheatham@amd.com>
---

I have a note next to the definition of FEATURE_CXLMEMATTR that it's
missing in the AMD IOMMU spec. The spec I'm referring to is the "AMD I/O
Virtualization Technology (IOMMU) Specification, Rev. 3.11, April 2026"
spec available on AMD's website.

I *think* the bit is actually supposed to be 0 based on the figures for
the register it's located in, but I had to guess since it's missing from
the register field table and isn't mentioned elsewhere AFAIK. I figured
it's fine while this is in RFC, but I'll update with the actual bit
when/if this comes out of RFC. Thanks!

---
 drivers/cxl/cache.c                 |  1 +
 drivers/cxl/core/cache.c            | 43 ++++++++++++++++++++++
 drivers/iommu/amd/amd_iommu_types.h |  4 +++
 drivers/iommu/amd/init.c            | 14 ++++++++
 drivers/iommu/amd/iommu.c           | 56 ++++++++++++++++++++++++++++-
 drivers/iommu/iommu.c               | 28 +++++++++++++++
 include/cxl/cxl.h                   | 22 ++++++++++++
 include/linux/iommu.h               |  8 +++++
 8 files changed, 175 insertions(+), 1 deletion(-)

diff --git a/drivers/cxl/cache.c b/drivers/cxl/cache.c
index b4574a76ac0b..f73bf96d231e 100644
--- a/drivers/cxl/cache.c
+++ b/drivers/cxl/cache.c
@@ -3,6 +3,7 @@
 
 #include <linux/pci.h>
 #include <cxl/cxl.h>
+#include <linux/iommu.h>
 
 #include "cxlcache.h"
 
diff --git a/drivers/cxl/core/cache.c b/drivers/cxl/core/cache.c
index 429c997b2ada..44323dfdb759 100644
--- a/drivers/cxl/core/cache.c
+++ b/drivers/cxl/core/cache.c
@@ -2,6 +2,7 @@
 /* Copyright (C) Advanced Micro Device, Inc. */
 
 #include <linux/iopoll.h>
+#include <linux/iommu.h>
 #include <linux/pci.h>
 #include <cxlcache.h>
 
@@ -806,3 +807,45 @@ void cxl_destroy_snoop_filters(void)
 
 	xa_destroy(&snoop_filters);
 }
+
+/**
+ * cxl_cache_configure_iommu() - Configure a device's IOMMU for CXL.cache
+ * @cxlds: struct cxl_dev_state of a cxl_cachedev that has been through
+ *	   CXL.cache probe
+ *
+ * Fails if the underlying PCI device supports ATS and IOMMU can't be
+ * configured, or if the device doesn't support ATS and is not attached to an
+ * identity IOMMU domain.
+ */
+int cxl_cache_configure_iommu(struct cxl_dev_state *cxlds)
+{
+	struct device *dev = cxlds->dev;
+	struct iommu_domain *domain;
+	int rc;
+
+	lockdep_assert_held(&dev->mutex);
+
+	if (!dev->iommu || !dev->iommu->iommu_dev)
+		return 0;
+
+	if (!device_iommu_capable(dev, IOMMU_CAP_PCI_ATS_SUPPORTED)) {
+		domain = iommu_get_domain_for_dev(dev);
+		if (!domain || domain->type != IOMMU_DOMAIN_IDENTITY)
+			return -EINVAL;
+
+		return 0;
+	}
+
+	rc = iommu_enable_cxl_ats(cxlds->dev);
+	if (rc == -EOPNOTSUPP) {
+		dev_warn(cxlds->dev,
+			"IOMMU doesn't support enabling CXL ATS requests; CXL.cache may not function properly.");
+		rc = 0;
+	} else if (rc) {
+		dev_err(cxlds->dev, "Failed to enable CXL ATS requests: %d\n",
+			rc);
+	}
+
+	return rc;
+}
+EXPORT_SYMBOL_NS_GPL(cxl_cache_configure_iommu, "CXL");
diff --git a/drivers/iommu/amd/amd_iommu_types.h b/drivers/iommu/amd/amd_iommu_types.h
index 3dbe20023456..696d9879ce9b 100644
--- a/drivers/iommu/amd/amd_iommu_types.h
+++ b/drivers/iommu/amd/amd_iommu_types.h
@@ -107,6 +107,7 @@
 
 
 /* Extended Feature 2 Bits */
+#define FEATURE_CXLMEMATTR	BIT_ULL(0) /* WARNING: This bit isn't in the spec as of 09/26 */
 #define FEATURE_SEVSNPIO_SUP	BIT_ULL(1)
 #define FEATURE_GCR3TRPMODE	BIT_ULL(3)
 #define FEATURE_SNPAVICSUP	GENMASK_ULL(7, 5)
@@ -189,6 +190,7 @@
 #define CONTROL_EPH_EN		45
 #define CONTROL_XT_EN		50
 #define CONTROL_INTCAPXT_EN	51
+#define CONTROL_CXLMEMATTR_EN	57
 #define CONTROL_GCR3TRPMODE	58
 #define CONTROL_IRTCACHEDIS	59
 #define CONTROL_SNPAVIC_EN	61
@@ -355,6 +357,7 @@
  */
 #define DTE_FLAG_V	BIT_ULL(0)
 #define DTE_FLAG_TV	BIT_ULL(1)
+#define DTE_CXLMEM_MASK	GENMASK_ULL(6, 4)
 #define DTE_FLAG_HAD	(3ULL << 7)
 #define DTE_MODE_MASK	GENMASK_ULL(11, 9)
 #define DTE_HOST_TRP	GENMASK_ULL(51, 12)
@@ -831,6 +834,7 @@ struct iommu_dev_data {
 	u8 ppr          :1;		  /* Enable device PPR support */
 	bool use_vapic;			  /* Enable device to use vapic mode */
 	bool defer_attach;
+	u8 cxl_memattr;			  /* CXLMemAttr DTE setting */
 
 	struct ratelimit_state rs;        /* Ratelimit IOPF messages */
 };
diff --git a/drivers/iommu/amd/init.c b/drivers/iommu/amd/init.c
index 40726dfef273..e11f98db826e 100644
--- a/drivers/iommu/amd/init.c
+++ b/drivers/iommu/amd/init.c
@@ -1125,6 +1125,14 @@ static void iommu_enable_gt(struct amd_iommu *iommu)
 		iommu_feature_enable(iommu, CONTROL_GCR3TRPMODE);
 }
 
+static void iommu_enable_cxlmemattr(struct amd_iommu *iommu)
+{
+	if (!check_feature2(FEATURE_CXLMEMATTR) || !amd_iommu_iotlb_sup)
+		return;
+
+	iommu_feature_enable(iommu, CONTROL_CXLMEMATTR_EN);
+}
+
 /* sets a specific bit in the device table entry. */
 static void set_dte_bit(struct dev_table_entry *dte, u8 bit)
 {
@@ -2245,6 +2253,12 @@ static int __init iommu_init_pci(struct amd_iommu *iommu)
 			return ret;
 	}
 
+	/* 
+	 * Set the CXLMemAttr control bit and do the Device Table Entry parts
+	 * of set up when a CXL.cache device shows up.
+	 */
+	iommu_enable_cxlmemattr(iommu);
+
 	ret = iommu_device_register(&iommu->iommu, &amd_iommu_ops, NULL);
 	if (ret || amd_iommu_pgtable == PD_MODE_NONE) {
 		/*
diff --git a/drivers/iommu/amd/iommu.c b/drivers/iommu/amd/iommu.c
index 4dc306a4b5c6..78d34b907450 100644
--- a/drivers/iommu/amd/iommu.c
+++ b/drivers/iommu/amd/iommu.c
@@ -41,6 +41,7 @@
 #include <asm/dma.h>
 #include <uapi/linux/iommufd.h>
 #include <linux/generic_pt/iommu.h>
+#include <cxl/cxl.h>
 
 #include "amd_iommu.h"
 #include "iommufd.h"
@@ -2199,6 +2200,13 @@ static void set_dte_passthrough(struct iommu_dev_data *dev_data,
 
 }
 
+static void set_dte_cxl_mem_attr(struct iommu_dev_data *dev_data,
+				 struct dev_table_entry *new)
+{
+	new->data[0] &= ~DTE_CXLMEM_MASK;
+	new->data[0] |= FIELD_PREP(DTE_CXLMEM_MASK, dev_data->cxl_memattr);
+}
+
 static void set_dte_entry(struct amd_iommu *iommu,
 			  struct iommu_dev_data *dev_data,
 			  phys_addr_t top_paddr, unsigned int top_level)
@@ -2217,11 +2225,14 @@ static void set_dte_entry(struct amd_iommu *iommu,
 	else if (domain->domain.type == IOMMU_DOMAIN_IDENTITY)
 		set_dte_passthrough(dev_data, domain, &new);
 	else if ((domain->domain.type & __IOMMU_DOMAIN_PAGING) &&
-		 domain->pd_mode == PD_MODE_V1)
+		   domain->pd_mode == PD_MODE_V1)
 		set_dte_v1(dev_data, domain, domain->id, top_paddr, top_level, &new);
 	else
 		WARN_ON(true);
 
+	if (amd_iommu_iotlb_sup && check_feature2(FEATURE_CXLMEMATTR))
+		set_dte_cxl_mem_attr(dev_data, &new);
+
 	amd_iommu_update_dte(iommu, dev_data, &new);
 
 	/*
@@ -3190,6 +3201,14 @@ static int amd_iommu_def_domain_type(struct device *dev)
 		return IOMMU_DOMAIN_IDENTITY;
 	}
 
+	/*
+	 * CXL.cache devices that have no ATS capability need a passthrough
+	 * domain to function correctly.
+	 */
+	if (dev_is_pci(dev) && !pci_ats_supported(to_pci_dev(dev)) &&
+	    cxl_cache_supported(to_pci_dev(dev)))
+		return IOMMU_DOMAIN_IDENTITY;
+
 	return 0;
 }
 
@@ -3199,6 +3218,40 @@ static bool amd_iommu_enforce_cache_coherency(struct iommu_domain *domain)
 	return true;
 }
 
+static int amd_iommu_enable_cxl_ats(struct device *dev)
+{
+	struct iommu_dev_data *dev_data = dev_iommu_priv_get(dev);
+	int ret = 0;
+
+	if (!dev_data || !amd_iommu_iotlb_sup)
+		return -EINVAL;
+
+	if (!check_feature2(FEATURE_CXLMEMATTR))
+		return -ENXIO;
+
+	mutex_lock(&dev_data->mutex);
+
+	/* Already set up */
+	if (dev_data->cxl_memattr != 0)
+		goto out;
+
+	if (!dev_data->ats_enabled) {
+		ret = -EINVAL;
+		goto out;
+	}
+
+	/*
+	 * Give device the choice of whether to use CXL.io or CXL.cache and
+	 * whether accesses are snooped
+	 */
+	dev_data->cxl_memattr = 1;
+	dev_update_dte(dev_data, true);
+
+out:
+	mutex_unlock(&dev_data->mutex);
+	return ret;
+}
+
 const struct iommu_ops amd_iommu_ops = {
 	.capable = amd_iommu_capable,
 	.hw_info = amd_iommufd_hw_info,
@@ -3216,6 +3269,7 @@ const struct iommu_ops amd_iommu_ops = {
 	.page_response = amd_iommu_page_response,
 	.get_viommu_size = amd_iommufd_get_viommu_size,
 	.viommu_init = amd_iommufd_viommu_init,
+	.enable_cxl_ats = amd_iommu_enable_cxl_ats,
 };
 
 #ifdef CONFIG_IRQ_REMAP
diff --git a/drivers/iommu/iommu.c b/drivers/iommu/iommu.c
index cd1bca7ede9a..5d82b5d09a3f 100644
--- a/drivers/iommu/iommu.c
+++ b/drivers/iommu/iommu.c
@@ -4222,6 +4222,34 @@ void pci_dev_reset_iommu_done(struct pci_dev *pdev)
 }
 EXPORT_SYMBOL_GPL(pci_dev_reset_iommu_done);
 
+/*
+ * iommu_enable_cxl_ats() - Enable CXL source bit extension of ATS for the
+ * given device.
+ * @dev: CXL.cache-capable PCIe device to enable capability for
+ *
+ * Returns:
+ * + -ENODEV if no IOMMU present
+ * + -EOPNOTSUPP if IOMMU isn't CXL-aware or doesn't need extra set up
+ * + Result of enablement callback otherwise
+ *
+ * Required by devices looking to use CXL.cache with non-identity IOMMU domais:
+ * see CXL 4.0 specification, section 3.1.6 "Memory Type Indication on ATS"
+ */
+int iommu_enable_cxl_ats(struct device *dev)
+{
+	const struct iommu_ops *ops;
+
+	if (!dev_has_iommu(dev))
+		return -ENODEV;
+
+	ops = dev_iommu_ops(dev);
+	if (!ops->enable_cxl_ats)
+		return -EOPNOTSUPP;
+
+	return ops->enable_cxl_ats(dev);
+}
+EXPORT_SYMBOL_GPL(iommu_enable_cxl_ats);
+
 #if IS_ENABLED(CONFIG_IRQ_MSI_IOMMU)
 /**
  * iommu_dma_prepare_msi() - Map the MSI page in the IOMMU domain
diff --git a/include/cxl/cxl.h b/include/cxl/cxl.h
index 5d550dd70870..0471d5c00144 100644
--- a/include/cxl/cxl.h
+++ b/include/cxl/cxl.h
@@ -8,6 +8,7 @@
 #include <linux/node.h>
 #include <linux/ioport.h>
 #include <cxl/mailbox.h>
+#include <linux/pci.h>
 
 /**
  * enum cxl_devtype - delineate type-2 from a generic type-3 device
@@ -274,6 +275,23 @@ int cxl_set_capacity(struct cxl_dev_state *cxlds, u64 capacity);
 struct cxl_cachedev *devm_cxl_add_cachedev(struct cxl_dev_state *cxlds);
 int devm_cxl_cachedev_alloc_snoop_capacity(struct cxl_cachedev *cxlcd,
 					   u64 size);
+int cxl_cache_configure_iommu(struct cxl_dev_state *cxlds);
+
+static inline bool cxl_cache_supported(struct pci_dev *pdev)
+{
+	int offset;
+	u16 cap;
+
+	offset = pci_find_dvsec_capability(pdev, PCI_VENDOR_ID_CXL,
+					   PCI_DVSEC_CXL_DEVICE);
+	if (!offset)
+		return false;
+
+	if (pci_read_config_word(pdev, offset + PCI_DVSEC_CXL_CAP, &cap))
+		return false;
+
+	return cap & PCI_DVSEC_CXL_CACHE_CAPABLE;
+}
 #else
 static inline struct cxl_cachedev *
 devm_cxl_add_cachedev(struct cxl_dev_state *cxlds)
@@ -281,5 +299,9 @@ devm_cxl_add_cachedev(struct cxl_dev_state *cxlds)
 static inline int
 devm_cxl_cachedev_alloc_snoop_capacity(struct cxl_cachedev *cxlcd, u64 size)
 { return -ENXIO; }
+static inline int cxl_cache_configure_iommu(struct cxl_dev_state *cxlds)
+{ return -ENXIO; }
+static inline bool cxl_cache_supported(struct pci_dev *pdev)
+{ return false; }
 #endif /* CONFIG_CXL_CACHE */
 #endif /* __CXL_CXL_H__ */
diff --git a/include/linux/iommu.h b/include/linux/iommu.h
index ac43b8b93f14..aa1e42ec5e38 100644
--- a/include/linux/iommu.h
+++ b/include/linux/iommu.h
@@ -734,6 +734,7 @@ struct iommu_ops {
 	int (*viommu_init)(struct iommufd_viommu *viommu,
 			   struct iommu_domain *parent_domain,
 			   const struct iommu_user_data *user_data);
+	int (*enable_cxl_ats)(struct device *dev);
 
 	const struct iommu_domain_ops *default_domain_ops;
 	struct module *owner;
@@ -1222,6 +1223,8 @@ void iommu_detach_device_pasid(struct iommu_domain *domain,
 ioasid_t iommu_alloc_global_pasid(struct device *dev);
 void iommu_free_global_pasid(ioasid_t pasid);
 
+int iommu_enable_cxl_ats(struct device *dev);
+
 /* PCI device reset functions */
 int pci_dev_reset_iommu_prepare(struct pci_dev *pdev);
 void pci_dev_reset_iommu_done(struct pci_dev *pdev);
@@ -1549,6 +1552,11 @@ static inline ioasid_t iommu_alloc_global_pasid(struct device *dev)
 
 static inline void iommu_free_global_pasid(ioasid_t pasid) {}
 
+static inline int iommu_enable_cxl_ats(struct device *dev)
+{
+	return -ENODEV;
+}
+
 static inline int pci_dev_reset_iommu_prepare(struct pci_dev *pdev)
 {
 	return 0;
-- 
2.53.0