[PATCH v2] x86/events/amd/iommu: Use a per-IOMMU cpumask

[PATCH v2] x86/events/amd/iommu: Use a per-IOMMU cpumask

[Vasant Hegde]

IOMMU performance counters are accessed via MMIO. The AMD IOMMU PMUs
share one global cpumask, fixed to CPU0, so every IOMMU's events are
read through the same CPU. That prevents collecting per-NUMA-node
numbers (for example perf stat --per-socket). The mask is also never
updated, so counting stops if that CPU goes offline.

Give each IOMMU PMU its own cpumask. The CPU is chosen from the online
CPUs of the NUMA node of the IOMMU's PCI device, via
amd_iommu_idx_to_dev(). If the node is unknown or has no online CPU,
any online CPU is used. If no online CPU exists at all, that PMU is
not registered.

A CPU hotplug callback moves the reader when that CPU goes offline,
preferring another online CPU on the same node. If the reader had to
leave the node, it moves back when a CPU on that node comes online.
Events are pinned to the current reader so they share one context and
migrate together.

Assisted-by: Cursor
Signed-off-by: Vasant Hegde <vasant.hegde@amd.com>
---
V1: https://lore.kernel.org/linux-iommu/20260517114026.6817-1-vasant.hegde@amd.com/

Changes since v1:
  - Rebased on top of v7.4-rc2
  - Added hotplug support

-Vasant

 arch/x86/events/amd/iommu.c | 130 ++++++++++++++++++++++++++++++++++--
 drivers/iommu/amd/init.c    |   7 ++
 include/linux/amd-iommu.h   |   1 +
 include/linux/cpuhotplug.h  |   1 +
 4 files changed, 134 insertions(+), 5 deletions(-)

diff --git a/arch/x86/events/amd/iommu.c b/arch/x86/events/amd/iommu.c
index f332c7089bd5..3b35702aa702 100644
--- a/arch/x86/events/amd/iommu.c
+++ b/arch/x86/events/amd/iommu.c
@@ -10,6 +10,7 @@
 
 #define pr_fmt(fmt)	"perf/amd_iommu: " fmt
 
+#include <linux/device.h>
 #include <linux/perf_event.h>
 #include <linux/init.h>
 #include <linux/cpumask.h>
@@ -43,6 +44,8 @@ struct perf_amd_iommu {
 	u8 max_counters;
 	u64 cntr_assign_mask;
 	raw_spinlock_t lock;
+	int node;
+	cpumask_t cpumask;
 };
 
 static LIST_HEAD(perf_amd_iommu_list);
@@ -131,13 +134,15 @@ static struct amd_iommu_event_desc amd_iommu_v2_event_descs[] = {
 /*---------------------------------------------
  * sysfs cpumask attributes
  *---------------------------------------------*/
-static cpumask_t iommu_cpumask;
-
 static ssize_t _iommu_cpumask_show(struct device *dev,
 				   struct device_attribute *attr,
 				   char *buf)
 {
-	return sysfs_emit(buf, "%*pbl\n", cpumask_pr_args(&iommu_cpumask));
+	struct pmu *pmu = dev_get_drvdata(dev);
+	struct perf_amd_iommu *perf_iommu =
+			container_of(pmu, struct perf_amd_iommu, pmu);
+
+	return sysfs_emit(buf, "%*pbl\n", cpumask_pr_args(&perf_iommu->cpumask));
 }
 static DEVICE_ATTR(cpumask, S_IRUGO, _iommu_cpumask_show, NULL);
 
@@ -207,7 +212,11 @@ static int clear_avail_iommu_bnk_cntr(struct perf_amd_iommu *perf_iommu,
 
 static int perf_iommu_event_init(struct perf_event *event)
 {
+	struct perf_amd_iommu *perf_iommu =
+			container_of(event->pmu, struct perf_amd_iommu, pmu);
 	struct hw_perf_event *hwc = &event->hw;
+	unsigned long flags;
+	unsigned int cpu;
 
 	/* test the event attr type check for PMU enumeration */
 	if (event->attr.type != event->pmu->type)
@@ -224,6 +233,18 @@ static int perf_iommu_event_init(struct perf_event *event)
 	if (event->cpu < 0)
 		return -EINVAL;
 
+	/*
+	 * Counters are read via MMIO from one CPU. Pin every event there so
+	 * they share a context and move together when that CPU goes offline.
+	 */
+	raw_spin_lock_irqsave(&perf_iommu->lock, flags);
+	cpu = cpumask_first(&perf_iommu->cpumask);
+	raw_spin_unlock_irqrestore(&perf_iommu->lock, flags);
+	if (cpu >= nr_cpu_ids)
+		return -ENODEV;
+
+	event->cpu = cpu;
+
 	/* update the hw_perf_event struct with the iommu config data */
 	hwc->conf  = event->attr.config;
 	hwc->conf1 = event->attr.config1;
@@ -417,9 +438,93 @@ static const struct pmu iommu_pmu __initconst = {
 	.capabilities	= PERF_PMU_CAP_NO_EXCLUDE,
 };
 
+/*
+ * Pick a CPU to read this IOMMU's counters. Prefer an online CPU on
+ * @node. @except is skipped so a CPU going offline is not chosen again;
+ * pass -1 when any online CPU is acceptable.
+ */
+static unsigned int iommu_pmu_pick_cpu(int node, int except)
+{
+	unsigned int cpu = nr_cpu_ids;
+
+	if (node != NUMA_NO_NODE)
+		cpu = cpumask_any_and_but(cpumask_of_node(node),
+					  cpu_online_mask, except);
+	if (cpu >= nr_cpu_ids)
+		cpu = cpumask_any_but(cpu_online_mask, except);
+
+	return cpu;
+}
+
+static int iommu_pmu_cpu_offline(unsigned int cpu)
+{
+	struct perf_amd_iommu *perf_iommu;
+	unsigned long flags;
+	unsigned int target;
+
+	list_for_each_entry(perf_iommu, &perf_amd_iommu_list, list) {
+		raw_spin_lock_irqsave(&perf_iommu->lock, flags);
+		if (!cpumask_test_and_clear_cpu(cpu, &perf_iommu->cpumask)) {
+			raw_spin_unlock_irqrestore(&perf_iommu->lock, flags);
+			continue;
+		}
+
+		target = iommu_pmu_pick_cpu(perf_iommu->node, cpu);
+		if (target < nr_cpu_ids)
+			cpumask_set_cpu(target, &perf_iommu->cpumask);
+		raw_spin_unlock_irqrestore(&perf_iommu->lock, flags);
+
+		if (target < nr_cpu_ids)
+			perf_pmu_migrate_context(&perf_iommu->pmu, cpu, target);
+	}
+
+	return 0;
+}
+
+static int iommu_pmu_cpu_online(unsigned int cpu)
+{
+	struct perf_amd_iommu *perf_iommu;
+	unsigned long flags;
+	unsigned int target;
+
+	list_for_each_entry(perf_iommu, &perf_amd_iommu_list, list) {
+		if (perf_iommu->node != NUMA_NO_NODE &&
+		    cpu_to_node(cpu) != perf_iommu->node)
+			continue;
+
+		raw_spin_lock_irqsave(&perf_iommu->lock, flags);
+		if (cpumask_empty(&perf_iommu->cpumask)) {
+			cpumask_set_cpu(cpu, &perf_iommu->cpumask);
+			raw_spin_unlock_irqrestore(&perf_iommu->lock, flags);
+			continue;
+		}
+
+		/*
+		 * A reader that fell off this node, because every local CPU
+		 * was offline, moves back when a local CPU returns.
+		 */
+		target = cpumask_first(&perf_iommu->cpumask);
+		if (perf_iommu->node == NUMA_NO_NODE ||
+		    cpu_to_node(target) == perf_iommu->node) {
+			raw_spin_unlock_irqrestore(&perf_iommu->lock, flags);
+			continue;
+		}
+
+		cpumask_clear(&perf_iommu->cpumask);
+		cpumask_set_cpu(cpu, &perf_iommu->cpumask);
+		raw_spin_unlock_irqrestore(&perf_iommu->lock, flags);
+
+		perf_pmu_migrate_context(&perf_iommu->pmu, target, cpu);
+	}
+
+	return 0;
+}
+
 static __init int init_one_iommu(unsigned int idx)
 {
 	struct perf_amd_iommu *perf_iommu;
+	struct device *dev;
+	unsigned int cpu;
 	int ret;
 
 	perf_iommu = kzalloc_obj(struct perf_amd_iommu);
@@ -440,6 +545,17 @@ static __init int init_one_iommu(unsigned int idx)
 		return -EINVAL;
 	}
 
+	dev = amd_iommu_idx_to_dev(idx);
+	perf_iommu->node = dev ? dev_to_node(dev) : NUMA_NO_NODE;
+	cpu = iommu_pmu_pick_cpu(perf_iommu->node, -1);
+	if (cpu >= nr_cpu_ids) {
+		pr_warn("Failed to find online CPU for IOMMU %d.\n", idx);
+		kfree(perf_iommu);
+		return -ENODEV;
+	}
+
+	cpumask_set_cpu(cpu, &perf_iommu->cpumask);
+
 	snprintf(perf_iommu->name, IOMMU_NAME_SIZE, "amd_iommu_%u", idx);
 
 	ret = perf_pmu_register(&perf_iommu->pmu, perf_iommu->name, -1);
@@ -483,8 +599,12 @@ static __init int amd_iommu_pc_init(void)
 		return -ENODEV;
 	}
 
-	/* Init cpumask attributes to only core 0 */
-	cpumask_set_cpu(0, &iommu_cpumask);
+	ret = cpuhp_setup_state(CPUHP_AP_PERF_X86_AMD_IOMMU_ONLINE,
+				"perf/x86/amd/iommu:online",
+				iommu_pmu_cpu_online, iommu_pmu_cpu_offline);
+	if (ret)
+		pr_warn("Failed to install CPU hotplug callback (%d)\n", ret);
+
 	return 0;
 }
 
diff --git a/drivers/iommu/amd/init.c b/drivers/iommu/amd/init.c
index 3dd4ff145e55..f13aa858fd50 100644
--- a/drivers/iommu/amd/init.c
+++ b/drivers/iommu/amd/init.c
@@ -4077,6 +4077,13 @@ struct amd_iommu *get_amd_iommu(unsigned int idx)
 	return NULL;
 }
 
+struct device *amd_iommu_idx_to_dev(unsigned int idx)
+{
+	struct amd_iommu *iommu = get_amd_iommu(idx);
+
+	return iommu ? &iommu->dev->dev : NULL;
+}
+
 /****************************************************************************
  *
  * IOMMU EFR Performance Counter support functionality. This code allows
diff --git a/include/linux/amd-iommu.h b/include/linux/amd-iommu.h
index 5c9dcf720f84..e3bddbe1c21b 100644
--- a/include/linux/amd-iommu.h
+++ b/include/linux/amd-iommu.h
@@ -67,6 +67,7 @@ int amd_iommu_pc_set_reg(struct amd_iommu *iommu, u8 bank, u8 cntr, u8 fxn,
 int amd_iommu_pc_get_reg(struct amd_iommu *iommu, u8 bank, u8 cntr, u8 fxn,
 		u64 *value);
 struct amd_iommu *get_amd_iommu(unsigned int idx);
+struct device *amd_iommu_idx_to_dev(unsigned int idx);
 
 #ifdef CONFIG_KVM_AMD_SEV
 int amd_iommu_snp_disable(void);
diff --git a/include/linux/cpuhotplug.h b/include/linux/cpuhotplug.h
index 1b6dba159bb2..468e61c98e80 100644
--- a/include/linux/cpuhotplug.h
+++ b/include/linux/cpuhotplug.h
@@ -205,6 +205,7 @@ enum cpuhp_state {
 	CPUHP_AP_PERF_X86_UNCORE_ONLINE,
 	CPUHP_AP_PERF_X86_AMD_UNCORE_ONLINE,
 	CPUHP_AP_PERF_X86_AMD_POWER_ONLINE,
+	CPUHP_AP_PERF_X86_AMD_IOMMU_ONLINE,
 	CPUHP_AP_PERF_S390_CF_ONLINE,
 	CPUHP_AP_PERF_S390_SF_ONLINE,
 	CPUHP_AP_PERF_ARM_CCI_ONLINE,
-- 
2.31.1