[PATCH 0/8] nitro_enclaves: Support multi-NUMA CPU pools and per-node allocation

[PATCH 0/8] nitro_enclaves: Support multi-NUMA CPU pools and per-node allocation

[Alexander Graf]

An enclave draws its vCPU cores from the driver's single CPU pool, and
by design every CPU in that pool sits on one NUMA node: ne_cpus=
refuses a wider CPU list with -EINVAL and "CPUs with different NUMA
nodes" in dmesg. An operator who wants two enclaves on two nodes has
nowhere to put the second one's cores, and an enclave wider than one
node has to spare cannot be built at all.

Lift it, and let a caller name the node. A pool spanning nodes records
itself as node-agnostic. NE_SET_ALLOC_NUMA_NODE sets a sticky
allocation target on the enclave fd and NE_ALLOC_NUMA_NODE_ANY drops
it; NE_ADD_VCPU with vcpu_id 0 draws from the target and returns the id
it picked, so placing an enclave across two nodes is setting the target
once per node and draining the count. A named node does not spill,
yielding a core on that node or NE_ERR_NO_CPUS_AVAIL_IN_POOL. A
cpu_pool attribute group on the misc device reports mode, total, used
and avail as CPU-list strings; pool threads are offline, so the
per-node view comes from intersecting avail with the node<N>/cpu<M>
symlinks under /sys/devices/system/node.

I picked a per-fd target over a second NE_ADD_VCPU carrying the node on
every call. NE_ADD_VCPU already reports the id it chose, so a VMM
spreading an enclave over several nodes needs no new call, only the
target and the count it already drains; the variant would make that
same VMM learn a new ioctl for behaviour it already has.

Patches 1 to 3 are not the feature: a missing .compat_ioctl, a pool
mutex taken before mutex_init(), and a thread-mask index that collides
cores from different packages, the last preparation with no functional
change. Both fixes carry Fixes: and neither Cc: stable, the mutex splat
needing CONFIG_DEBUG_MUTEXES to appear.

Nothing consumes this yet; the kernel API goes first, and its shape is
what nitro-cli, or a VMM like QEMU, gets built against. A dynamic pool
mode and CMA-sourced enclave memory, both reusing this target, follow
in separate series.

I wrote this series with an AI coding assistant, which drafted the code
and the changelogs, this cover letter included; I reviewed and reworked
all of it before sending, and each commit carries an Assisted-by:
trailer. All eight commits were built detached in three configurations
(x86_64 defconfig with CONFIG_NITRO_ENCLAVES=m, the same with
CONFIG_NUMA=n, and allmodconfig) for 24 warning-free driver builds, and
scripts/checkpatch.pl --strict is clean on all eight.

Alexander Graf (8):
  nitro_enclaves: Let 32-bit processes use the ioctl interface
  nitro_enclaves: Initialise the CPU pool mutex statically
  nitro_enclaves: Slot pool cores by sibling mask
  nitro_enclaves: Allow the CPU pool to span NUMA nodes
  nitro_enclaves: Add NE_SET_ALLOC_NUMA_NODE
  nitro_enclaves: Expose CPU pool state under sysfs
  Documentation: ABI: Describe nitro_enclaves cpu_pool sysfs
  Documentation: virt: Describe multi-NUMA Nitro Enclaves CPU pools

 .../sysfs-devices-virtual-misc-nitro_enclaves |  44 ++
 Documentation/virt/ne_overview.rst            |  25 +-
 MAINTAINERS                                   |   1 +
 drivers/virt/nitro_enclaves/ne_misc_dev.c     | 489 +++++++++++++-----
 drivers/virt/nitro_enclaves/ne_misc_dev.h     |  16 +-
 include/uapi/linux/nitro_enclaves.h           |  62 ++-
 6 files changed, 508 insertions(+), 129 deletions(-)
 create mode 100644 Documentation/ABI/testing/sysfs-devices-virtual-misc-nitro_enclaves


base-commit: f5098b6bae761e346ebcd9da7f95622c04733cff
-- 
2.47.1

[PATCH 1/8] nitro_enclaves: Let 32-bit processes use the ioctl interface

[Alexander Graf]

Neither the misc device nor an enclave fd sets .compat_ioctl, so every
NE ioctl answers -ENOTTY to a 32-bit process on a 64-bit kernel. Nothing
about the interface asks for that. Each ioctl argument is a fixed-width
type or a struct built out of fixed-width types, laid out the same
either way, and the one thing that does differ between a 32-bit and a
64-bit caller is the argument pointer, which is what compat_ptr_ioctl()
converts.

So point both file_operations at it. A 32-bit program that drives
enclaves works with no translation layer on either side, and the ioctl
added later in this series is reachable from such a program the day it
lands, instead of waiting for somebody to notice that it is not.

Fixes: 38907e124088 ("nitro_enclaves: Add logic for creating an enclave VM")
Assisted-by: Kiro:claude-opus-5
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index c91300a73f50..8222fbe75800 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -67,6 +67,7 @@ static const struct file_operations ne_fops = {
 	.owner		= THIS_MODULE,
 	.llseek		= noop_llseek,
 	.unlocked_ioctl	= ne_ioctl,
+	.compat_ioctl	= compat_ptr_ioctl,
 };
 
 static struct miscdevice ne_misc_dev = {
@@ -1569,6 +1570,7 @@ static const struct file_operations ne_enclave_fops = {
 	.llseek		= noop_llseek,
 	.poll		= ne_enclave_poll,
 	.unlocked_ioctl	= ne_enclave_ioctl,
+	.compat_ioctl	= compat_ptr_ioctl,
 	.release	= ne_enclave_release,
 };
 
-- 
2.47.1

[PATCH 2/8] nitro_enclaves: Initialise the CPU pool mutex statically

[Alexander Graf]

ne_cpus is registered with module_param_cb(), so its setter runs while
the parameter section is parsed: from the kernel command line as
nitro_enclaves.ne_cpus= for a built-in driver, or from modprobe
nitro_enclaves ne_cpus=<cpu-list>. parse_args() precedes do_initcalls()
in start_kernel() and do_init_module() in load_module(), so the setter
runs before the driver's own code. It reaches ne_teardown_cpu_pool() and
then ne_setup_cpu_pool(), both of which lock ne_cpu_pool.mutex, while
the mutex_init() for that lock sits in ne_init(). An admin who
configures the pool at boot or at module load therefore takes a mutex
that nothing has initialised. Commit ff8a4d3e3a99 ("nitro_enclaves: Add
logic for setting an enclave vCPU") added those two locks; the parameter
and its runtime mutex_init() were already in place.

Under CONFIG_DEBUG_MUTEXES the first acquire trips a
DEBUG_LOCKS_WARN_ON() in kernel/locking/mutex.c, because ->magic is
still NULL rather than the lock's own address, and the splat carries
parse_args() and ne_set_kernel_param() in its backtrace. It is not a
crash: with lock debugging off nothing is reported and the pool comes
up, which is how this survived. A write to
/sys/module/nitro_enclaves/parameters/ne_cpus is not affected: the
initcall has run before user space exists.

So initialise the mutex where the pool is defined and drop the
mutex_init() call, leaving it valid from the moment the parameter
section is parsed.

Fixes: ff8a4d3e3a99 ("nitro_enclaves: Add logic for setting an enclave vCPU")
Assisted-by: Kiro:claude-opus-5
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 8222fbe75800..59e1794587e6 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -125,7 +125,9 @@ struct ne_cpu_pool {
 	int		numa_node;
 };
 
-static struct ne_cpu_pool ne_cpu_pool;
+static struct ne_cpu_pool ne_cpu_pool = {
+	.mutex = __MUTEX_INITIALIZER(ne_cpu_pool.mutex),
+};
 
 /**
  * struct ne_phys_contig_mem_regions - Contiguous physical memory regions.
@@ -1761,8 +1763,6 @@ static long ne_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
 
 static int __init ne_init(void)
 {
-	mutex_init(&ne_cpu_pool.mutex);
-
 	return pci_register_driver(&ne_pci_driver);
 }
 
-- 
2.47.1

[PATCH 3/8] nitro_enclaves: Slot pool cores by sibling mask

[Alexander Graf]

Preparation for letting the CPU pool span NUMA nodes; no functional
change intended. The pool records the threads of each core in
avail_threads_per_core[], indexed by topology_core_id(). That id is
package-local on x86, so two cores in different packages report the same
id and share one slot: an enclave that takes the slot is handed the
threads of both as though they came from one core, and one slot in the
array stays empty.

Walk each pool CPU's sibling cpumask instead and give every core the
first slot no core has taken yet. The per-CPU sibling mask identifies a
core on both architectures this driver builds for. Inside one package
the two ways partition the pool into the same cores, so an enclave is
still handed whole cores, and no ioctl hands a slot index to user space.
The bound moves with the index: a pool whose core ids reach nr_cpu_ids /
nr_threads_per_core used to be refused with -EINVAL and is now set up,
because the bound is now a count of cores.

The topology reads move under the CPU hotplug read lock, which the base
tree takes nowhere in this driver although it already read
topology_sibling_cpumask() and cpu_to_node(), so it closes a
pre-existing gap rather than protecting the new walk. It is dropped
before the offlining loop, because remove_cpu() ends up in
cpus_write_lock() on the same percpu-rwsem and would wait for a reader
that is the same thread. That is what draws the boundary of the new
ne_build_core_slots().

Assisted-by: Kiro:claude-opus-5
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 182 ++++++++++++++--------
 drivers/virt/nitro_enclaves/ne_misc_dev.h |   8 +-
 2 files changed, 125 insertions(+), 65 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 59e1794587e6..79cc53d8b465 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -105,11 +105,11 @@ MODULE_PARM_DESC(ne_cpus, "<cpu-list> - CPU pool used for Nitro Enclaves");
 /**
  * struct ne_cpu_pool - CPU pool used for Nitro Enclaves.
  * @avail_threads_per_core:	Available full CPU cores to be dedicated to
- *				enclave(s). The cpumasks from the array, indexed
- *				by core id, contain all the threads from the
- *				available cores, that are not set for created
- *				enclave(s). The full CPU cores are part of the
- *				NE CPU pool.
+ *				enclave(s). Each cpumask in the array holds the
+ *				threads of one core in the pool that are not set
+ *				for created enclave(s). Slots are handed out as
+ *				cores are discovered and carry no meaning
+ *				outside the driver.
  * @mutex:			Mutex for the access to the NE CPU pool.
  * @nr_parent_vm_cores :	The size of the available threads per core array.
  *				The total number of CPU cores available on the
@@ -166,6 +166,102 @@ static bool ne_check_enclaves_created(void)
 	return ret;
 }
 
+/**
+ * ne_free_core_slots() - Free the core slot arrays of the NE CPU pool.
+ * @void:	No parameters provided.
+ *
+ * Context: Process context. This function is called with the ne_cpu_pool mutex
+ *	    held.
+ */
+static void ne_free_core_slots(void)
+{
+	unsigned int i = 0;
+
+	if (!ne_cpu_pool.avail_threads_per_core)
+		return;
+
+	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
+		free_cpumask_var(ne_cpu_pool.avail_threads_per_core[i]);
+
+	kfree(ne_cpu_pool.avail_threads_per_core);
+	ne_cpu_pool.avail_threads_per_core = NULL;
+}
+
+/**
+ * ne_build_core_slots() - Group the CPU pool into full cores and give each one
+ *			   a slot of the NE CPU pool arrays.
+ * @cpu_pool:	The CPU pool as parsed out of the ne_cpus parameter.
+ *
+ * The slots preserve the CPU topology of the pool once its CPUs are offlined.
+ * topology_core_id() is package-local on x86, so two cores in different
+ * packages report the same id and would share a slot. The per-CPU sibling
+ * cpumask identifies a core on both architectures this driver builds for, so
+ * walk that and hand each core the next free slot.
+ *
+ * Context: Process context. This function is called with the ne_cpu_pool mutex
+ *	    and the CPU hotplug read lock held.
+ * Return:
+ * * 0 on success.
+ * * Negative return value on failure.
+ */
+static int ne_build_core_slots(const struct cpumask *cpu_pool)
+{
+	unsigned int cpu = 0;
+	unsigned int cpu_sibling = 0;
+	unsigned int i = 0;
+	unsigned int next_core_idx = 0;
+	cpumask_var_t processed;
+	int rc = -ENOMEM;
+
+	ne_cpu_pool.avail_threads_per_core = kzalloc_objs(*ne_cpu_pool.avail_threads_per_core,
+							  ne_cpu_pool.nr_parent_vm_cores);
+	if (!ne_cpu_pool.avail_threads_per_core)
+		return -ENOMEM;
+
+	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
+		if (!zalloc_cpumask_var(&ne_cpu_pool.avail_threads_per_core[i], GFP_KERNEL))
+			goto free_slots;
+
+	if (!zalloc_cpumask_var(&processed, GFP_KERNEL))
+		goto free_slots;
+
+	for_each_cpu(cpu, cpu_pool) {
+		if (cpumask_test_cpu(cpu, processed))
+			continue;
+
+		if (next_core_idx >= ne_cpu_pool.nr_parent_vm_cores) {
+			pr_err("%s: CPU pool has more cores than nr_parent_vm_cores=%d\n",
+			       ne_misc_dev.name, ne_cpu_pool.nr_parent_vm_cores);
+
+			rc = -EINVAL;
+
+			goto free_processed;
+		}
+
+		for_each_cpu(cpu_sibling, topology_sibling_cpumask(cpu)) {
+			if (!cpumask_test_cpu(cpu_sibling, cpu_pool))
+				continue;
+
+			cpumask_set_cpu(cpu_sibling,
+					ne_cpu_pool.avail_threads_per_core[next_core_idx]);
+			cpumask_set_cpu(cpu_sibling, processed);
+		}
+
+		next_core_idx++;
+	}
+
+	free_cpumask_var(processed);
+
+	return 0;
+
+free_processed:
+	free_cpumask_var(processed);
+free_slots:
+	ne_free_core_slots();
+
+	return rc;
+}
+
 /**
  * ne_setup_cpu_pool() - Set the NE CPU pool after handling sanity checks such
  *			 as not sharing CPU cores with the primary / parent VM
@@ -181,11 +277,9 @@ static bool ne_check_enclaves_created(void)
  */
 static int ne_setup_cpu_pool(const char *ne_cpu_list)
 {
-	int core_id = -1;
 	unsigned int cpu = 0;
 	cpumask_var_t cpu_pool;
 	unsigned int cpu_sibling = 0;
-	unsigned int i = 0;
 	int numa_node = -1;
 	int rc = -EINVAL;
 
@@ -193,12 +287,13 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 		return -ENOMEM;
 
 	mutex_lock(&ne_cpu_pool.mutex);
+	cpus_read_lock();
 
 	rc = cpulist_parse(ne_cpu_list, cpu_pool);
 	if (rc < 0) {
 		pr_err("%s: Error in cpulist parse [rc=%d]\n", ne_misc_dev.name, rc);
 
-		goto free_pool_cpumask;
+		goto unlock_hotplug;
 	}
 
 	cpu = cpumask_any(cpu_pool);
@@ -207,7 +302,7 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 
 		rc = -EINVAL;
 
-		goto free_pool_cpumask;
+		goto unlock_hotplug;
 	}
 
 	/*
@@ -221,7 +316,7 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 
 			rc = -EINVAL;
 
-			goto free_pool_cpumask;
+			goto unlock_hotplug;
 		}
 
 	/*
@@ -236,7 +331,7 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 
 				rc = -EINVAL;
 
-				goto free_pool_cpumask;
+				goto unlock_hotplug;
 			}
 		} else {
 			if (numa_node != cpu_to_node(cpu)) {
@@ -245,7 +340,7 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 
 				rc = -EINVAL;
 
-				goto free_pool_cpumask;
+				goto unlock_hotplug;
 			}
 		}
 
@@ -258,7 +353,7 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 
 		rc = -EINVAL;
 
-		goto free_pool_cpumask;
+		goto unlock_hotplug;
 	}
 
 	for_each_cpu(cpu_sibling, topology_sibling_cpumask(0)) {
@@ -268,7 +363,7 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 
 			rc = -EINVAL;
 
-			goto free_pool_cpumask;
+			goto unlock_hotplug;
 		}
 	}
 
@@ -285,7 +380,7 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 
 				rc = -EINVAL;
 
-				goto free_pool_cpumask;
+				goto unlock_hotplug;
 			}
 		}
 	}
@@ -297,38 +392,11 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 
 	ne_cpu_pool.nr_parent_vm_cores = nr_cpu_ids / ne_cpu_pool.nr_threads_per_core;
 
-	ne_cpu_pool.avail_threads_per_core = kzalloc_objs(*ne_cpu_pool.avail_threads_per_core,
-							  ne_cpu_pool.nr_parent_vm_cores);
-	if (!ne_cpu_pool.avail_threads_per_core) {
-		rc = -ENOMEM;
-
-		goto free_pool_cpumask;
-	}
-
-	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
-		if (!zalloc_cpumask_var(&ne_cpu_pool.avail_threads_per_core[i], GFP_KERNEL)) {
-			rc = -ENOMEM;
-
-			goto free_cores_cpumask;
-		}
-
-	/*
-	 * Split the NE CPU pool in threads per core to keep the CPU topology
-	 * after offlining the CPUs.
-	 */
-	for_each_cpu(cpu, cpu_pool) {
-		core_id = topology_core_id(cpu);
-		if (core_id < 0 || core_id >= ne_cpu_pool.nr_parent_vm_cores) {
-			pr_err("%s: Invalid core id  %d for CPU %d\n",
-			       ne_misc_dev.name, core_id, cpu);
-
-			rc = -EINVAL;
-
-			goto clear_cpumask;
-		}
+	rc = ne_build_core_slots(cpu_pool);
+	if (rc < 0)
+		goto unlock_hotplug;
 
-		cpumask_set_cpu(cpu, ne_cpu_pool.avail_threads_per_core[core_id]);
-	}
+	cpus_read_unlock();
 
 	/*
 	 * CPUs that are given to enclave(s) should not be considered online
@@ -361,14 +429,11 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 online_cpus:
 	for_each_cpu(cpu, cpu_pool)
 		add_cpu(cpu);
-clear_cpumask:
-	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
-		cpumask_clear(ne_cpu_pool.avail_threads_per_core[i]);
-free_cores_cpumask:
-	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
-		free_cpumask_var(ne_cpu_pool.avail_threads_per_core[i]);
-	kfree(ne_cpu_pool.avail_threads_per_core);
-free_pool_cpumask:
+	ne_free_core_slots();
+	goto reset_pool;
+unlock_hotplug:
+	cpus_read_unlock();
+reset_pool:
 	free_cpumask_var(cpu_pool);
 	ne_cpu_pool.nr_parent_vm_cores = 0;
 	ne_cpu_pool.nr_threads_per_core = 0;
@@ -399,7 +464,7 @@ static void ne_teardown_cpu_pool(void)
 		return;
 	}
 
-	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++) {
+	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
 		for_each_cpu(cpu, ne_cpu_pool.avail_threads_per_core[i]) {
 			rc = add_cpu(cpu);
 			if (rc != 0)
@@ -407,12 +472,7 @@ static void ne_teardown_cpu_pool(void)
 				       ne_misc_dev.name, cpu, rc);
 		}
 
-		cpumask_clear(ne_cpu_pool.avail_threads_per_core[i]);
-
-		free_cpumask_var(ne_cpu_pool.avail_threads_per_core[i]);
-	}
-
-	kfree(ne_cpu_pool.avail_threads_per_core);
+	ne_free_core_slots();
 	ne_cpu_pool.nr_parent_vm_cores = 0;
 	ne_cpu_pool.nr_threads_per_core = 0;
 	ne_cpu_pool.numa_node = -1;
diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.h b/drivers/virt/nitro_enclaves/ne_misc_dev.h
index 2a4d2224baba..94c6404bde22 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.h
@@ -55,10 +55,10 @@ struct ne_mem_region {
  * @numa_node:			NUMA node of the enclave memory and CPUs.
  * @slot_uid:			Slot unique id mapped to the enclave.
  * @state:			Enclave state, updated during enclave lifetime.
- * @threads_per_core:		Enclave full CPU cores array, indexed by core id,
- *				consisting of cpumasks with all their threads.
- *				Full CPU cores are taken from the NE CPU pool
- *				and are available to the enclave.
+ * @threads_per_core:		Enclave full CPU cores array. Each cpumask in the
+ *				array holds the threads of one core taken from
+ *				the NE CPU pool and available to the enclave.
+ *				The array shares the pool's slot indices.
  * @vcpu_ids:			Cpumask of the vCPUs that are set for the enclave.
  */
 struct ne_enclave {
-- 
2.47.1

[PATCH 4/8] nitro_enclaves: Allow the CPU pool to span NUMA nodes

[Alexander Graf]

There is one CPU pool, and every CPU in it has to sit on the same NUMA
node, so every enclave on the machine draws its cores from it. An
operator who wants enclaves on two nodes, or one enclave wider than a
node has to spare, has nothing to write into ne_cpus=: such a CPU list
is refused with -EINVAL and "CPUs with different NUMA nodes" in dmesg.
The restriction is deliberate and as old as the pool; it was the right
rule while an enclave sat on node 0 and stayed there.

Lift it. A pool whose CPUs report more than one node records its node as
NUMA_NO_NODE and counts as node-agnostic, and every enclave created from
it inherits that. That narrows the ABI: such an enclave skips the
memory-locality check, so NE_ERR_MEM_DIFFERENT_NUMA_NODE is not a
failure NE_SET_USER_MEMORY_REGION can return for it. Once an enclave's
cores come from more than one node there is no single node to hold its
memory to, and a caller that wants the check keeps its pool on one node.

The node then has to be part of picking a core.
ne_get_unused_core_from_cpu_pool() returns a free core from the target
node or nothing at all, and the target lives in ne_enclave->alloc_nid,
which NE_CREATE_VM sets to the node owning the pool's first core. The
picker does not spill onto another node when the target has no core
left: a wrong-node core would be stranded, because the sibling walk
honours the target too, and the mismatch would surface only at
NE_START_ENCLAVE, as NE_ERR_FULL_CORES_NOT_USED, which says nothing
about nodes.

Assisted-by: Kiro:claude-opus-5
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 133 +++++++++++++++-------
 drivers/virt/nitro_enclaves/ne_misc_dev.h |   8 ++
 include/uapi/linux/nitro_enclaves.h       |  20 ++--
 3 files changed, 110 insertions(+), 51 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 79cc53d8b465..e019bb1f6594 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -22,6 +22,7 @@
 #include <linux/module.h>
 #include <linux/mutex.h>
 #include <linux/nitro_enclaves.h>
+#include <linux/numa.h>
 #include <linux/pci.h>
 #include <linux/poll.h>
 #include <linux/range.h>
@@ -34,8 +35,8 @@
 
 /**
  * NE_CPUS_SIZE - Size for max 128 CPUs, for now, in a cpu-list string, comma
- *		  separated. The NE CPU pool includes CPUs from a single NUMA
- *		  node.
+ *		  separated. The NE CPU pool may include CPUs from more than
+ *		  one NUMA node.
  */
 #define NE_CPUS_SIZE		(512)
 
@@ -115,7 +116,20 @@ MODULE_PARM_DESC(ne_cpus, "<cpu-list> - CPU pool used for Nitro Enclaves");
  *				The total number of CPU cores available on the
  *				primary / parent VM.
  * @nr_threads_per_core:	The number of threads that a full CPU core has.
- * @numa_node:			NUMA node of the CPUs in the pool.
+ * @numa_node:			NUMA node of the CPUs in the pool, or
+ *				NUMA_NO_NODE if the pool spans multiple nodes.
+ * @first_nid:			The NUMA node that owns the first core in the
+ *				pool, or NUMA_NO_NODE if the pool is empty.
+ *				This is the node a new enclave targets until
+ *				user space says otherwise.
+ * @core_nid:			The NUMA node of each core slot, indexed the
+ *				same way as @avail_threads_per_core, or
+ *				NUMA_NO_NODE both for a slot that holds no
+ *				core and for a core whose CPUs report no node.
+ *				Taken while the pool CPUs are still online,
+ *				because from then on they are offline and
+ *				cpu_to_node() is not documented to keep
+ *				answering for them.
  */
 struct ne_cpu_pool {
 	cpumask_var_t	*avail_threads_per_core;
@@ -123,10 +137,13 @@ struct ne_cpu_pool {
 	unsigned int	nr_parent_vm_cores;
 	unsigned int	nr_threads_per_core;
 	int		numa_node;
+	int		first_nid;
+	int		*core_nid;
 };
 
 static struct ne_cpu_pool ne_cpu_pool = {
 	.mutex = __MUTEX_INITIALIZER(ne_cpu_pool.mutex),
+	.first_nid = NUMA_NO_NODE,
 };
 
 /**
@@ -185,6 +202,9 @@ static void ne_free_core_slots(void)
 
 	kfree(ne_cpu_pool.avail_threads_per_core);
 	ne_cpu_pool.avail_threads_per_core = NULL;
+
+	kfree(ne_cpu_pool.core_nid);
+	ne_cpu_pool.core_nid = NULL;
 }
 
 /**
@@ -218,10 +238,18 @@ static int ne_build_core_slots(const struct cpumask *cpu_pool)
 	if (!ne_cpu_pool.avail_threads_per_core)
 		return -ENOMEM;
 
-	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
+	ne_cpu_pool.core_nid = kmalloc_objs(*ne_cpu_pool.core_nid,
+					    ne_cpu_pool.nr_parent_vm_cores);
+	if (!ne_cpu_pool.core_nid)
+		goto free_slots;
+
+	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++) {
 		if (!zalloc_cpumask_var(&ne_cpu_pool.avail_threads_per_core[i], GFP_KERNEL))
 			goto free_slots;
 
+		ne_cpu_pool.core_nid[i] = NUMA_NO_NODE;
+	}
+
 	if (!zalloc_cpumask_var(&processed, GFP_KERNEL))
 		goto free_slots;
 
@@ -247,11 +275,15 @@ static int ne_build_core_slots(const struct cpumask *cpu_pool)
 			cpumask_set_cpu(cpu_sibling, processed);
 		}
 
+		ne_cpu_pool.core_nid[next_core_idx] = cpu_to_node(cpu);
+
 		next_core_idx++;
 	}
 
 	free_cpumask_var(processed);
 
+	ne_cpu_pool.first_nid = ne_cpu_pool.core_nid[0];
+
 	return 0;
 
 free_processed:
@@ -280,7 +312,9 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 	unsigned int cpu = 0;
 	cpumask_var_t cpu_pool;
 	unsigned int cpu_sibling = 0;
-	int numa_node = -1;
+	bool have_node = false;
+	bool multi_node = false;
+	int numa_node = NUMA_NO_NODE;
 	int rc = -EINVAL;
 
 	if (!zalloc_cpumask_var(&cpu_pool, GFP_KERNEL))
@@ -320,29 +354,21 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 		}
 
 	/*
-	 * Check if the CPUs from the NE CPU pool are from the same NUMA node.
+	 * Determine the NUMA node of the CPUs in the pool. If the pool spans
+	 * multiple NUMA nodes, set numa_node to NUMA_NO_NODE; the pool is then
+	 * treated as node-agnostic for placement and memory-locality checks.
 	 */
-	for_each_cpu(cpu, cpu_pool)
-		if (numa_node < 0) {
+	for_each_cpu(cpu, cpu_pool) {
+		if (!have_node) {
 			numa_node = cpu_to_node(cpu);
-			if (numa_node < 0) {
-				pr_err("%s: Invalid NUMA node %d\n",
-				       ne_misc_dev.name, numa_node);
-
-				rc = -EINVAL;
-
-				goto unlock_hotplug;
-			}
-		} else {
-			if (numa_node != cpu_to_node(cpu)) {
-				pr_err("%s: CPUs with different NUMA nodes\n",
-				       ne_misc_dev.name);
-
-				rc = -EINVAL;
-
-				goto unlock_hotplug;
-			}
+			have_node = true;
+		} else if (numa_node != cpu_to_node(cpu)) {
+			multi_node = true;
 		}
+	}
+
+	if (multi_node)
+		numa_node = NUMA_NO_NODE;
 
 	/*
 	 * Check if CPU 0 and its siblings are included in the provided CPU pool
@@ -437,7 +463,8 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 	free_cpumask_var(cpu_pool);
 	ne_cpu_pool.nr_parent_vm_cores = 0;
 	ne_cpu_pool.nr_threads_per_core = 0;
-	ne_cpu_pool.numa_node = -1;
+	ne_cpu_pool.numa_node = NUMA_NO_NODE;
+	ne_cpu_pool.first_nid = NUMA_NO_NODE;
 	mutex_unlock(&ne_cpu_pool.mutex);
 
 	return rc;
@@ -475,7 +502,8 @@ static void ne_teardown_cpu_pool(void)
 	ne_free_core_slots();
 	ne_cpu_pool.nr_parent_vm_cores = 0;
 	ne_cpu_pool.nr_threads_per_core = 0;
-	ne_cpu_pool.numa_node = -1;
+	ne_cpu_pool.numa_node = NUMA_NO_NODE;
+	ne_cpu_pool.first_nid = NUMA_NO_NODE;
 
 	mutex_unlock(&ne_cpu_pool.mutex);
 }
@@ -550,7 +578,13 @@ static bool ne_donated_cpu(struct ne_enclave *ne_enclave, unsigned int cpu)
 /**
  * ne_get_unused_core_from_cpu_pool() - Get the id of a full core from the
  *					NE CPU pool.
- * @void:	No parameters provided.
+ * @nid:	Required NUMA node for the core, or NUMA_NO_NODE for no
+ *		node constraint.
+ *
+ * When @nid names a specific node, only a free core whose CPUs are
+ * on that node is returned; if none is available the function returns -1
+ * rather than spilling onto another node. When @nid is NUMA_NO_NODE, any
+ * free core is returned.
  *
  * Context: Process context. This function is called with the ne_enclave and
  *	    ne_cpu_pool mutexes held.
@@ -558,19 +592,19 @@ static bool ne_donated_cpu(struct ne_enclave *ne_enclave, unsigned int cpu)
  * * Core id.
  * * -1 if no CPU core available in the pool.
  */
-static int ne_get_unused_core_from_cpu_pool(void)
+static int ne_get_unused_core_from_cpu_pool(int nid)
 {
-	int core_id = -1;
 	unsigned int i = 0;
 
-	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
-		if (!cpumask_empty(ne_cpu_pool.avail_threads_per_core[i])) {
-			core_id = i;
-
-			break;
-		}
+	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++) {
+		if (cpumask_empty(ne_cpu_pool.avail_threads_per_core[i]))
+			continue;
+		if (nid != NUMA_NO_NODE && ne_cpu_pool.core_nid[i] != nid)
+			continue;
+		return i;
+	}
 
-	return core_id;
+	return -1;
 }
 
 /**
@@ -637,28 +671,37 @@ static int ne_get_cpu_from_cpu_pool(struct ne_enclave *ne_enclave, u32 *vcpu_id)
 	int core_id = -1;
 	unsigned int cpu = 0;
 	unsigned int i = 0;
+	int nid = ne_enclave->alloc_nid;
 	int rc = -EINVAL;
 
+	mutex_lock(&ne_cpu_pool.mutex);
+
 	/*
 	 * If previously allocated a thread of a core to this enclave, first
 	 * check remaining sibling(s) for new CPU allocations, so that full
-	 * CPU cores are used for the enclave.
+	 * CPU cores are used for the enclave. Cores on a node other than the
+	 * allocation target are skipped: the target is a hard constraint, and
+	 * honouring it here is what keeps a sequence of targeted allocations
+	 * from handing back a core on the wrong node.
 	 */
-	for (i = 0; i < ne_enclave->nr_parent_vm_cores; i++)
+	for (i = 0; i < ne_enclave->nr_parent_vm_cores; i++) {
+		if (nid != NUMA_NO_NODE && ne_cpu_pool.core_nid[i] != nid)
+			continue;
+
 		for_each_cpu(cpu, ne_enclave->threads_per_core[i])
 			if (!ne_donated_cpu(ne_enclave, cpu)) {
 				*vcpu_id = cpu;
+				rc = 0;
 
-				return 0;
+				goto unlock_mutex;
 			}
-
-	mutex_lock(&ne_cpu_pool.mutex);
+	}
 
 	/*
 	 * If no remaining siblings, get a core from the NE CPU pool and keep
 	 * track of all the threads in the enclave threads per core data structure.
 	 */
-	core_id = ne_get_unused_core_from_cpu_pool();
+	core_id = ne_get_unused_core_from_cpu_pool(nid);
 
 	rc = ne_set_enclave_threads_per_core(ne_enclave, core_id, *vcpu_id);
 	if (rc < 0)
@@ -886,7 +929,8 @@ static int ne_sanity_check_user_mem_region_page(struct ne_enclave *ne_enclave,
 		return -NE_ERR_INVALID_PAGE_SIZE;
 	}
 
-	if (ne_enclave->numa_node != page_to_nid(mem_region_page)) {
+	if (ne_enclave->numa_node != NUMA_NO_NODE &&
+	    ne_enclave->numa_node != page_to_nid(mem_region_page)) {
 		dev_err_ratelimited(ne_misc_dev.this_device,
 				    "Page is not from NUMA node %d\n",
 				    ne_enclave->numa_node);
@@ -1687,6 +1731,7 @@ static int ne_create_vm_ioctl(struct ne_pci_dev *ne_pci_dev, u64 __user *slot_ui
 	ne_enclave->nr_parent_vm_cores = ne_cpu_pool.nr_parent_vm_cores;
 	ne_enclave->nr_threads_per_core = ne_cpu_pool.nr_threads_per_core;
 	ne_enclave->numa_node = ne_cpu_pool.numa_node;
+	ne_enclave->alloc_nid = ne_cpu_pool.first_nid;
 
 	mutex_unlock(&ne_cpu_pool.mutex);
 
diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.h b/drivers/virt/nitro_enclaves/ne_misc_dev.h
index 94c6404bde22..aed02183f0b9 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.h
@@ -53,6 +53,13 @@ struct ne_mem_region {
  * @nr_threads_per_core:	The number of threads that a full CPU core has.
  * @nr_vcpus:			Number of vcpus associated with the enclave.
  * @numa_node:			NUMA node of the enclave memory and CPUs.
+ * @alloc_nid:			NUMA node of the primary VM used for
+ *				kernel-side allocations on behalf of this
+ *				enclave (NE_ADD_VCPU auto-pick).
+ *				Set at NE_CREATE_VM to the first node that
+ *				owns a core in the CPU pool. User space
+ *				overrides it via NE_SET_ALLOC_NUMA_NODE.
+ *				NUMA_NO_NODE means "any node".
  * @slot_uid:			Slot unique id mapped to the enclave.
  * @state:			Enclave state, updated during enclave lifetime.
  * @threads_per_core:		Enclave full CPU cores array. Each cpumask in the
@@ -75,6 +82,7 @@ struct ne_enclave {
 	unsigned int		nr_threads_per_core;
 	unsigned int		nr_vcpus;
 	int			numa_node;
+	int			alloc_nid;
 	u64			slot_uid;
 	u16			state;
 	cpumask_var_t		*threads_per_core;
diff --git a/include/uapi/linux/nitro_enclaves.h b/include/uapi/linux/nitro_enclaves.h
index e808f5ba124d..7c6ec8dfe451 100644
--- a/include/uapi/linux/nitro_enclaves.h
+++ b/include/uapi/linux/nitro_enclaves.h
@@ -26,8 +26,8 @@
  *		  https://www.kernel.org/doc/html/latest/admin-guide/kernel-parameters.html
  *		  CPU 0 and its siblings have to remain available for the
  *		  primary / parent VM, so they cannot be set for enclaves. Full
- *		  CPU core(s), from the same NUMA node, need(s) to be included
- *		  in the CPU pool.
+ *		  CPU core(s) need(s) to be included in the CPU pool. The
+ *		  pool may span NUMA nodes.
  *
  * Context: Process context.
  * Return:
@@ -50,8 +50,9 @@
  * NE_ADD_VCPU - The command is used to set a vCPU for an enclave. The vCPU can
  *		 be auto-chosen from the NE CPU pool or it can be set by the
  *		 caller, with the note that it needs to be available in the NE
- *		 CPU pool. Full CPU core(s), from the same NUMA node, need(s) to
- *		 be associated with an enclave.
+ *		 CPU pool. Full CPU core(s) need(s) to be associated with an
+ *		 enclave. When the CPU pool spans NUMA nodes, the cores of
+ *		 one enclave can be from different nodes.
  *		 The vCPU id is an input / output parameter. If its value is 0,
  *		 then a CPU is chosen from the enclave CPU pool and returned via
  *		 this parameter.
@@ -108,8 +109,11 @@
 /**
  * NE_SET_USER_MEMORY_REGION - The command is used to set a memory region for an
  *			       enclave, given the allocated memory from the
- *			       userspace. Enclave memory needs to be from the
- *			       same NUMA node as the enclave CPUs.
+ *			       userspace. When the CPU pool sits on one
+ *			       NUMA node, enclave memory needs to be from
+ *			       that node. An enclave built out of a pool
+ *			       that spans nodes has no node of its own, and
+ *			       its memory can come from anywhere.
  *			       The user memory region is an input parameter. It
  *			       includes info provided by the caller - flags,
  *			       memory size and userspace address.
@@ -138,7 +142,9 @@
  * * NE_ERR_MEM_NOT_HUGE_PAGE		- The memory region is not backed by
  *					  huge pages.
  * * NE_ERR_MEM_DIFFERENT_NUMA_NODE	- The memory region is not from the same
- *					  NUMA node as the CPUs.
+ *					  NUMA node as the CPUs. Not returned
+ *					  for an enclave whose CPU pool spans
+ *					  nodes.
  * * NE_ERR_MEM_MAX_REGIONS		- The number of memory regions set for
  *					  the enclave reached maximum.
  * * NE_ERR_INVALID_PAGE_SIZE		- The memory region is not backed by
-- 
2.47.1

[PATCH 5/8] nitro_enclaves: Add NE_SET_ALLOC_NUMA_NODE

[Alexander Graf]

A CPU pool that spans NUMA nodes, which pool setup now permits, still
hands every enclave the same allocation target: the node that owns the
pool's first core. Allocation from that target does not spill to another
node, so once that node's contribution is exhausted NE_ADD_VCPU fails
with NE_ERR_NO_CPUS_AVAIL_IN_POOL (272) while cores on the pool's other
nodes sit unused. An enclave wider than one node's share of the pool
cannot be built, and nothing a caller passes to NE_ADD_VCPU changes
which node it draws from.

Give the enclave fd a sticky allocation target instead. NE_ADD_VCPU with
vcpu_id 0 draws its core from that node and returns the id it picked, so
a caller placing an enclave across two nodes sets the target once per
node and drains the count. NE_ALLOC_NUMA_NODE_ANY drops the constraint
and takes any free pool core. That constant is spelled (-1) in the uapi
header rather than reusing NUMA_NO_NODE, which is defined in
linux/nodemask_types.h and reachable from no uapi header at all.

I can think of two ways to let a caller name a node: a target that lives
on the fd, or a second NE_ADD_VCPU carrying the node on every call. I
picked the target. NE_ADD_VCPU already reports the id it chose, so a
caller spreading an enclave over several nodes needs no new call at all,
only the target and the existing count. The variant would make that same
caller learn a new ioctl to get behaviour it already has, and would
price every future input to a placement decision the same way: one more
ioctl number each.

Assisted-by: Kiro:claude-opus-5
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 38 ++++++++++++++++++++
 include/uapi/linux/nitro_enclaves.h       | 42 +++++++++++++++++++++++
 2 files changed, 80 insertions(+)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index e019bb1f6594..eb0091de1182 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -22,6 +22,7 @@
 #include <linux/module.h>
 #include <linux/mutex.h>
 #include <linux/nitro_enclaves.h>
+#include <linux/nodemask.h>
 #include <linux/numa.h>
 #include <linux/pci.h>
 #include <linux/poll.h>
@@ -1477,6 +1478,43 @@ static long ne_enclave_ioctl(struct file *file, unsigned int cmd, unsigned long
 		return 0;
 	}
 
+	case NE_SET_ALLOC_NUMA_NODE: {
+		struct ne_alloc_numa_node n;
+		int nid;
+
+		if (copy_from_user(&n, (void __user *)arg, sizeof(n)))
+			return -EFAULT;
+
+		if (n.flags)
+			return -EINVAL;
+
+		nid = n.numa_node;
+		if (nid != NUMA_NO_NODE &&
+		    (nid < 0 || nid >= nr_node_ids || !node_state(nid, N_POSSIBLE))) {
+			dev_err_ratelimited(ne_misc_dev.this_device,
+					    "Invalid NUMA node %d\n", nid);
+
+			return -EINVAL;
+		}
+
+		mutex_lock(&ne_enclave->enclave_info_mutex);
+
+		if (ne_enclave->state != NE_STATE_INIT) {
+			dev_err_ratelimited(ne_misc_dev.this_device,
+					    "Enclave is not in init state\n");
+
+			mutex_unlock(&ne_enclave->enclave_info_mutex);
+
+			return -NE_ERR_NOT_IN_INIT_STATE;
+		}
+
+		ne_enclave->alloc_nid = nid;
+
+		mutex_unlock(&ne_enclave->enclave_info_mutex);
+
+		return 0;
+	}
+
 	default:
 		return -ENOTTY;
 	}
diff --git a/include/uapi/linux/nitro_enclaves.h b/include/uapi/linux/nitro_enclaves.h
index 7c6ec8dfe451..8ddc4b150b7e 100644
--- a/include/uapi/linux/nitro_enclaves.h
+++ b/include/uapi/linux/nitro_enclaves.h
@@ -184,6 +184,31 @@
  */
 #define NE_START_ENCLAVE		_IOWR(0xAE, 0x24, struct ne_enclave_start_info)
 
+/**
+ * NE_SET_ALLOC_NUMA_NODE - Set the NUMA node of the primary VM used for
+ *			    subsequent kernel-side allocations on this enclave
+ *			    fd: NE_ADD_VCPU auto-pick (vcpu_id == 0) draws a
+ *			    core from this node. The setting is sticky
+ *			    until changed or the fd is closed. Pass
+ *			    %NE_ALLOC_NUMA_NODE_ANY for node-agnostic
+ *			    allocation (any pool core).
+ *
+ *			    Without this ioctl the target is the first node
+ *			    that owns a core in the CPU pool, and allocation
+ *			    never spills to another node, so a caller that
+ *			    cares which node it lands on names it here.
+ *
+ * Context: Process context.
+ * Return:
+ * * 0				- On success.
+ * * -1				- On failure, errno is set to:
+ * *   EFAULT			- copy_from_user() failed.
+ * *   EINVAL			- flags is non-zero, or node id is not a
+ *				  possible node.
+ * *   NE_ERR_NOT_IN_INIT_STATE	- Enclave is not in init state.
+ */
+#define NE_SET_ALLOC_NUMA_NODE		_IOW(0xAE, 0x25, struct ne_alloc_numa_node)
+
 /**
  * DOC: NE specific error codes
  */
@@ -297,6 +322,23 @@
 
 #define NE_IMAGE_LOAD_MAX_FLAG_VAL	(0x02)
 
+/**
+ * NE_ALLOC_NUMA_NODE_ANY - Node id for a node-agnostic allocation target.
+ */
+#define NE_ALLOC_NUMA_NODE_ANY		(-1)
+
+/**
+ * struct ne_alloc_numa_node - Argument for %NE_SET_ALLOC_NUMA_NODE.
+ * @numa_node:	NUMA node of the primary VM for subsequent kernel-side
+ *		allocations (NE_ADD_VCPU auto-pick), or
+ *		%NE_ALLOC_NUMA_NODE_ANY for node-agnostic allocation.
+ * @flags:	Must be 0.
+ */
+struct ne_alloc_numa_node {
+	__s32	numa_node;
+	__u32	flags;
+};
+
 /**
  * struct ne_image_load_info - Info necessary for in-memory enclave image
  *			       loading (in / out).
-- 
2.47.1

[PATCH 6/8] nitro_enclaves: Expose CPU pool state under sysfs

[Alexander Graf]

An orchestrator deciding where to place an enclave needs to know how
many cores of the pool are still free, and on which NUMA node. The total
is not the missing piece: /sys/module/nitro_enclaves/parameters/ne_cpus
is mode 0644 and reads back the CPU list that was accepted. Occupancy
is. NE_CREATE_VM returns NE_ERR_NO_CPUS_AVAIL_IN_POOL (272) before it
allocates a slot when no core anywhere in the pool is free, but that
check does not look at nodes, so the only per-node answer the kernel
offers arrives from NE_ADD_VCPU, as errno 272, once the slot exists and
has to be freed again.

So report the pool on the misc device, as a cpu_pool attribute group
under /sys/devices/virtual/misc/nitro_enclaves/: mode is none or static,
total is the pool as parsed from ne_cpus, used is what enclaves hold,
and avail is the rest. The group hangs off miscdevice.groups, so the
subtree comes and goes with the device and there is no kobject to manage
by hand.

The threads enclaves hold get a mask of their own in the pool rather
than being added up from the enclave list on every read. Walking that
list needs the enclave list mutex and each enclave's mutex, and
NE_ADD_VCPU, release and the event work handler each hold one of those
across a request to the device, so a read of used would sit behind an
enclave that is failing to stop. The four files take the pool mutex and
nothing else, and that mutex is not held across any of this driver's
seven calls into the device.

Assisted-by: Kiro:claude-opus-5
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 134 ++++++++++++++++++++--
 1 file changed, 122 insertions(+), 12 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index eb0091de1182..9b28d1f828b6 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -28,6 +28,7 @@
 #include <linux/poll.h>
 #include <linux/range.h>
 #include <linux/slab.h>
+#include <linux/sysfs.h>
 #include <linux/types.h>
 #include <uapi/linux/vm_sockets.h>
 
@@ -72,17 +73,6 @@ static const struct file_operations ne_fops = {
 	.compat_ioctl	= compat_ptr_ioctl,
 };
 
-static struct miscdevice ne_misc_dev = {
-	.minor	= MISC_DYNAMIC_MINOR,
-	.name	= "nitro_enclaves",
-	.fops	= &ne_fops,
-	.mode	= 0660,
-};
-
-struct ne_devs ne_devs = {
-	.ne_misc_dev	= &ne_misc_dev,
-};
-
 /*
  * TODO: Update logic to create new sysfs entries instead of using
  * a kernel parameter e.g. if multiple sysfs files needed.
@@ -131,6 +121,17 @@ MODULE_PARM_DESC(ne_cpus, "<cpu-list> - CPU pool used for Nitro Enclaves");
  *				because from then on they are offline and
  *				cpu_to_node() is not documented to keep
  *				answering for them.
+ * @pool_cpus:			Every CPU thread in the pool, as parsed from
+ *				ne_cpus. Empty if no pool is set up.
+ * @claimed_cpus:		The subset of @pool_cpus that enclaves hold.
+ *				A whole core enters when an enclave takes it
+ *				and leaves when the enclave is released, in
+ *				the same two places that move the core out of
+ *				and back into @avail_threads_per_core.
+ * @avail_cpus:			Scratch mask for the avail sysfs file:
+ *				@pool_cpus without @claimed_cpus, computed
+ *				and emitted under @mutex. It lives here so
+ *				that the read has nothing to allocate.
  */
 struct ne_cpu_pool {
 	cpumask_var_t	*avail_threads_per_core;
@@ -140,6 +141,9 @@ struct ne_cpu_pool {
 	int		numa_node;
 	int		first_nid;
 	int		*core_nid;
+	struct cpumask	pool_cpus;
+	struct cpumask	claimed_cpus;
+	struct cpumask	avail_cpus;
 };
 
 static struct ne_cpu_pool ne_cpu_pool = {
@@ -147,6 +151,101 @@ static struct ne_cpu_pool ne_cpu_pool = {
 	.first_nid = NUMA_NO_NODE,
 };
 
+/*
+ * cpu_pool/ sysfs group on the Nitro Enclaves misc device: report the CPU
+ * pool mode and occupancy.
+ *
+ *   /sys/devices/virtual/misc/nitro_enclaves/cpu_pool/mode
+ *   /sys/devices/virtual/misc/nitro_enclaves/cpu_pool/{total,used,avail}
+ *
+ * total, used and avail are CPU-list strings taken from the pool masks under
+ * the pool mutex. used is claimed_cpus and avail is what is left of the pool
+ * once claimed_cpus is taken out of it, so the two partition total whichever
+ * order they are read in. The three CPU-list files emit from a mask the pool
+ * already owns and mode emits a string constant, so none of the four show
+ * handlers can fail.
+ */
+
+static ssize_t mode_show(struct device *dev, struct device_attribute *attr,
+			 char *buf)
+{
+	const char *mode;
+
+	mutex_lock(&ne_cpu_pool.mutex);
+	mode = cpumask_empty(&ne_cpu_pool.pool_cpus) ? "none" : "static";
+	mutex_unlock(&ne_cpu_pool.mutex);
+
+	return sysfs_emit(buf, "%s\n", mode);
+}
+static DEVICE_ATTR_RO(mode);
+
+static ssize_t total_show(struct device *dev, struct device_attribute *attr,
+			  char *buf)
+{
+	ssize_t ret;
+
+	mutex_lock(&ne_cpu_pool.mutex);
+	ret = sysfs_emit(buf, "%*pbl\n", cpumask_pr_args(&ne_cpu_pool.pool_cpus));
+	mutex_unlock(&ne_cpu_pool.mutex);
+
+	return ret;
+}
+static DEVICE_ATTR_RO(total);
+
+static ssize_t used_show(struct device *dev, struct device_attribute *attr,
+			 char *buf)
+{
+	ssize_t ret;
+
+	mutex_lock(&ne_cpu_pool.mutex);
+	ret = sysfs_emit(buf, "%*pbl\n", cpumask_pr_args(&ne_cpu_pool.claimed_cpus));
+	mutex_unlock(&ne_cpu_pool.mutex);
+
+	return ret;
+}
+static DEVICE_ATTR_RO(used);
+
+static ssize_t avail_show(struct device *dev, struct device_attribute *attr,
+			  char *buf)
+{
+	ssize_t ret;
+
+	mutex_lock(&ne_cpu_pool.mutex);
+	cpumask_andnot(&ne_cpu_pool.avail_cpus, &ne_cpu_pool.pool_cpus,
+		       &ne_cpu_pool.claimed_cpus);
+	ret = sysfs_emit(buf, "%*pbl\n", cpumask_pr_args(&ne_cpu_pool.avail_cpus));
+	mutex_unlock(&ne_cpu_pool.mutex);
+
+	return ret;
+}
+static DEVICE_ATTR_RO(avail);
+
+static struct attribute *ne_cpu_pool_attrs[] = {
+	&dev_attr_mode.attr,
+	&dev_attr_total.attr,
+	&dev_attr_used.attr,
+	&dev_attr_avail.attr,
+	NULL,
+};
+
+static const struct attribute_group ne_cpu_pool_group = {
+	.name	= "cpu_pool",
+	.attrs	= ne_cpu_pool_attrs,
+};
+__ATTRIBUTE_GROUPS(ne_cpu_pool);
+
+static struct miscdevice ne_misc_dev = {
+	.minor	= MISC_DYNAMIC_MINOR,
+	.name	= "nitro_enclaves",
+	.fops	= &ne_fops,
+	.mode	= 0660,
+	.groups	= ne_cpu_pool_groups,
+};
+
+struct ne_devs ne_devs = {
+	.ne_misc_dev	= &ne_misc_dev,
+};
+
 /**
  * struct ne_phys_contig_mem_regions - Contiguous physical memory regions.
  * @num:	The number of regions that currently has.
@@ -445,6 +544,8 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 		}
 	}
 
+	cpumask_copy(&ne_cpu_pool.pool_cpus, cpu_pool);
+	cpumask_clear(&ne_cpu_pool.claimed_cpus);
 	free_cpumask_var(cpu_pool);
 
 	ne_cpu_pool.numa_node = numa_node;
@@ -462,6 +563,8 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 	cpus_read_unlock();
 reset_pool:
 	free_cpumask_var(cpu_pool);
+	cpumask_clear(&ne_cpu_pool.pool_cpus);
+	cpumask_clear(&ne_cpu_pool.claimed_cpus);
 	ne_cpu_pool.nr_parent_vm_cores = 0;
 	ne_cpu_pool.nr_threads_per_core = 0;
 	ne_cpu_pool.numa_node = NUMA_NO_NODE;
@@ -501,6 +604,8 @@ static void ne_teardown_cpu_pool(void)
 		}
 
 	ne_free_core_slots();
+	cpumask_clear(&ne_cpu_pool.pool_cpus);
+	cpumask_clear(&ne_cpu_pool.claimed_cpus);
 	ne_cpu_pool.nr_parent_vm_cores = 0;
 	ne_cpu_pool.nr_threads_per_core = 0;
 	ne_cpu_pool.numa_node = NUMA_NO_NODE;
@@ -650,6 +755,9 @@ static int ne_set_enclave_threads_per_core(struct ne_enclave *ne_enclave,
 	for_each_cpu(cpu, ne_cpu_pool.avail_threads_per_core[core_id])
 		cpumask_set_cpu(cpu, ne_enclave->threads_per_core[core_id]);
 
+	cpumask_or(&ne_cpu_pool.claimed_cpus, &ne_cpu_pool.claimed_cpus,
+		   ne_cpu_pool.avail_threads_per_core[core_id]);
+
 	cpumask_clear(ne_cpu_pool.avail_threads_per_core[core_id]);
 
 	return 0;
@@ -1564,9 +1672,11 @@ static void ne_enclave_remove_all_vcpu_id_entries(struct ne_enclave *ne_enclave)
 	mutex_lock(&ne_cpu_pool.mutex);
 
 	for (i = 0; i < ne_enclave->nr_parent_vm_cores; i++) {
-		for_each_cpu(cpu, ne_enclave->threads_per_core[i])
+		for_each_cpu(cpu, ne_enclave->threads_per_core[i]) {
 			/* Update the available NE CPU pool. */
 			cpumask_set_cpu(cpu, ne_cpu_pool.avail_threads_per_core[i]);
+			cpumask_clear_cpu(cpu, &ne_cpu_pool.claimed_cpus);
+		}
 
 		free_cpumask_var(ne_enclave->threads_per_core[i]);
 	}
-- 
2.47.1

[PATCH 7/8] Documentation: ABI: Describe nitro_enclaves cpu_pool sysfs

[Alexander Graf]

An operator or an orchestrator placing an enclave needs more than the
contents of these files: how far the partition between used and avail
goes, and which node the enclave allocates from when it does not name
one. Write both down.

Assisted-by: Kiro:claude-opus-5
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 .../sysfs-devices-virtual-misc-nitro_enclaves | 44 +++++++++++++++++++
 MAINTAINERS                                   |  1 +
 2 files changed, 45 insertions(+)
 create mode 100644 Documentation/ABI/testing/sysfs-devices-virtual-misc-nitro_enclaves

diff --git a/Documentation/ABI/testing/sysfs-devices-virtual-misc-nitro_enclaves b/Documentation/ABI/testing/sysfs-devices-virtual-misc-nitro_enclaves
new file mode 100644
index 000000000000..20b978c62244
--- /dev/null
+++ b/Documentation/ABI/testing/sysfs-devices-virtual-misc-nitro_enclaves
@@ -0,0 +1,44 @@
+What:		/sys/devices/virtual/misc/nitro_enclaves/cpu_pool/mode
+Date:		July 2026
+KernelVersion:	7.3
+Contact:	aws-nitro-enclaves-devel@amazon.com
+Description:
+		Read-only. Mode of the Nitro Enclaves CPU pool, one of:
+
+		  none     the pool is empty (ne_cpus is unset)
+		  static   ne_cpus is a CPU list; the listed CPUs are
+		           offlined and dedicated to the pool
+
+What:		/sys/devices/virtual/misc/nitro_enclaves/cpu_pool/total
+What:		/sys/devices/virtual/misc/nitro_enclaves/cpu_pool/used
+What:		/sys/devices/virtual/misc/nitro_enclaves/cpu_pool/avail
+Date:		July 2026
+KernelVersion:	7.3
+Contact:	aws-nitro-enclaves-devel@amazon.com
+Description:
+		Read-only. CPU-list strings describing the Nitro Enclaves CPU
+		pool: total is every CPU thread dedicated to the pool, used is
+		the threads held on behalf of enclaves, and avail is the
+		threads free for allocation. A thread leaves used when the
+		enclave holding it is released; if that release fails part way
+		through, the thread stays in used until the driver is unloaded
+		or the machine reboots. used and avail partition total at any
+		instant, and each read is a consistent snapshot, so a thread
+		that changes hands between two reads can be absent from both.
+		total changes only when the pool is reconfigured, which
+		requires that no enclave exists.
+
+		These files report no error of their own. A read does take the
+		pool lock, so it can wait for a write to ne_cpus to finish.
+
+		By default, allocation for a new enclave targets the NUMA
+		node that owns the first core in total and does not spill to
+		another node. NE_SET_ALLOC_NUMA_NODE changes the target, and
+		its node-agnostic value lets one enclave take cores from
+		more than one node.
+
+		Pool threads are offline, so the threads free on a given
+		node come from intersecting avail with the nodeN/cpuM
+		symlinks under /sys/devices/system/node, which outlive a CPU
+		going offline. nodeN/cpulist is filtered against
+		cpu_online_mask and lists none of them.
diff --git a/MAINTAINERS b/MAINTAINERS
index 716acfc3d7c1..e573764d9cc2 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -19184,6 +19184,7 @@ R:	The AWS Nitro Enclaves Team <aws-nitro-enclaves-devel@amazon.com>
 L:	linux-kernel@vger.kernel.org
 S:	Supported
 W:	https://aws.amazon.com/ec2/nitro/nitro-enclaves/
+F:	Documentation/ABI/testing/sysfs-devices-virtual-misc-nitro_enclaves
 F:	Documentation/virt/ne_overview.rst
 F:	drivers/virt/nitro_enclaves/
 F:	include/linux/nitro_enclaves.h
-- 
2.47.1

[PATCH 8/8] Documentation: virt: Describe multi-NUMA Nitro Enclaves CPU pools

[Alexander Graf]

The Nitro Enclaves overview tells an admin flatly that enclave memory
and CPUs come from the same NUMA node, and that is the page's only
mention of NUMA. The rule stops applying once the CPU pool spans
nodes. What replaces it lives in the NE_ADD_VCPU,
NE_SET_USER_MEMORY_REGION and NE_SET_ALLOC_NUMA_NODE kerneldoc and in
the cpu_pool sysfs ABI file, and the overview references none of them.
Qualify the sentence, and give the page a section covering which node
an allocation targets by default, how to move that target, and how to
read per-node occupancy out of the sysfs group.

A cross-reference to the header and the ABI file instead would have
left one copy of the behaviour to keep current, at the cost of sending
an admin who wanted only to know whether the rule still applied into a
uapi header. The section goes after the rest of the overview rather
than into the middle of it, so the paragraphs on vsock, the image
format and attestation stay part of the overview instead of ending up
underneath a CPU pool heading.

Assisted-by: Kiro:claude-opus-5
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 Documentation/virt/ne_overview.rst | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/Documentation/virt/ne_overview.rst b/Documentation/virt/ne_overview.rst
index 74c2f5919c88..157b632d7437 100644
--- a/Documentation/virt/ne_overview.rst
+++ b/Documentation/virt/ne_overview.rst
@@ -47,7 +47,8 @@ The memory regions carved out of the primary VM and given to an enclave need to
 be aligned 2 MiB / 1 GiB physically contiguous memory regions (or multiple of
 this size e.g. 8 MiB). The memory can be allocated e.g. by using hugetlbfs from
 user space [2][3][7]. The memory size for an enclave needs to be at least
-64 MiB. The enclave memory and CPUs need to be from the same NUMA node.
+64 MiB. The enclave memory and CPUs need to be from the same NUMA node,
+unless the CPU pool spans several of them; see `Multi-NUMA CPU pools`_.
 
 An enclave runs on dedicated cores. CPU 0 and its CPU siblings need to remain
 available for the primary VM. A CPU pool has to be set for NE purposes by an
@@ -90,6 +91,28 @@ the NE driver. This event is sent further to the user space enclave process
 running in the primary VM via a poll notification mechanism. Then the user space
 enclave process can exit.
 
+Multi-NUMA CPU pools
+--------------------
+
+The CPU pool may span multiple NUMA nodes. When it does, kernel-side
+allocations for a new enclave target the node that owns the first core in
+the pool: NE_ADD_VCPU with vcpu_id == 0 draws a core from that node, and
+does not spill to another node once that node has none left. User space
+moves the target with NE_SET_ALLOC_NUMA_NODE, and can move it between
+allocations to build an enclave that spans nodes. Passing
+NE_ALLOC_NUMA_NODE_ANY as the node lets an allocation come from any node
+in the pool.
+
+The target says where cores come from and nothing else. An enclave built
+out of a pool that spans nodes has no node of its own, so the rule that
+its memory be on that node does not apply to it and its memory can come
+from anywhere.
+
+The pool mode and occupancy are readable under
+/sys/devices/virtual/misc/nitro_enclaves/cpu_pool/. Intersect ``avail``
+with /sys/devices/system/node/nodeN/cpulist for the threads still free on
+one node.
+
 [1] https://aws.amazon.com/ec2/nitro/nitro-enclaves/
 [2] https://www.kernel.org/doc/html/latest/admin-guide/mm/hugetlbpage.html
 [3] https://lwn.net/Articles/807108/
-- 
2.47.1

Re: [PATCH 6/8] nitro_enclaves: Expose CPU pool state under sysfs

[Arnd Bergmann]

On Thu, Jul 30, 2026, at 14:53, Alexander Graf wrote:

> +static ssize_t total_show(struct device *dev, struct device_attribute *attr,
> +			  char *buf)
> +{
> +	ssize_t ret;
> +
> +	mutex_lock(&ne_cpu_pool.mutex);
> +	ret = sysfs_emit(buf, "%*pbl\n", cpumask_pr_args(&ne_cpu_pool.pool_cpus));
> +	mutex_unlock(&ne_cpu_pool.mutex);
> +
> +	return ret;
> +}
> +static DEVICE_ATTR_RO(total);
> +
> +static ssize_t used_show(struct device *dev, struct device_attribute *attr,
> +			 char *buf)
> +{
> +	ssize_t ret;
> +
> +	mutex_lock(&ne_cpu_pool.mutex);
> +	ret = sysfs_emit(buf, "%*pbl\n", cpumask_pr_args(&ne_cpu_pool.claimed_cpus));
> +	mutex_unlock(&ne_cpu_pool.mutex);
> +
> +	return ret;
> +}
> +static DEVICE_ATTR_RO(used);

The names are a bit inconsistent: you have pool/claimed for the internal
structure but total/used for the ABI. It might help to pick one set of
names here.

      Arnd

Re: [PATCH 5/8] nitro_enclaves: Add NE_SET_ALLOC_NUMA_NODE

[Arnd Bergmann]

On Thu, Jul 30, 2026, at 14:53, Alexander Graf wrote:
> +		if (ne_enclave->state != NE_STATE_INIT) {
> +			dev_err_ratelimited(ne_misc_dev.this_device,
> +					    "Enclave is not in init state\n");
> +
> +			mutex_unlock(&ne_enclave->enclave_info_mutex);
> +
> +			return -NE_ERR_NOT_IN_INIT_STATE;
> +		}

This does not look like a normal error code, what is going on here?

> +struct ne_alloc_numa_node {
> +	__s32	numa_node;
> +	__u32	flags;
> +};

I don't see the point of 'flags' here if that has to be zero.
Can you just drop that?

      Arnd

Re: [PATCH 0/8] nitro_enclaves: Support multi-NUMA CPU pools and per-node allocation

[Arnd Bergmann]

On Thu, Jul 30, 2026, at 14:53, Alexander Graf wrote:

Hi Alex

> I picked a per-fd target over a second NE_ADD_VCPU carrying the node on
> every call. NE_ADD_VCPU already reports the id it chose, so a VMM
> spreading an enclave over several nodes needs no new call, only the
> target and the count it already drains; the variant would make that
> same VMM learn a new ioctl for behaviour it already has.

I had to read this three times to understand what you are trying
to say, but still don't know why you picked one over the other.

> I wrote this series with an AI coding assistant, which drafted the code
> and the changelogs, this cover letter included; I reviewed and reworked
> all of it before sending, and each commit carries an Assisted-by:
> trailer. All eight commits were built detached in three configurations
> (x86_64 defconfig with CONFIG_NITRO_ENCLAVES=m, the same with
> CONFIG_NUMA=n, and allmodconfig) for 24 warning-free driver builds, and
> scripts/checkpatch.pl --strict is clean on all eight.

The patches look fine to me overall, I only have some very minor
comments on them. The autogenerated changelogs seem actively
counterproductive though, despite your attempts to to rewrite
them into something readable. This seems to have just made the
series harder to review. Instead of telling the reader a
consistent story about what each change is for, the generated
text has too much useless detail about the implementation.

     Arnd

Re: [PATCH 0/8] nitro_enclaves: Support multi-NUMA CPU pools and per-node allocation

[Graf (AWS), Alexander]

Hey Arnd!

On 30.07.26 17:43, Arnd Bergmann wrote:
> On Thu, Jul 30, 2026, at 14:53, Alexander Graf wrote:
>
> Hi Alex
>
>> I picked a per-fd target over a second NE_ADD_VCPU carrying the node on
>> every call. NE_ADD_VCPU already reports the id it chose, so a VMM
>> spreading an enclave over several nodes needs no new call, only the
>> target and the count it already drains; the variant would make that
>> same VMM learn a new ioctl for behaviour it already has.
> I had to read this three times to understand what you are trying
> to say, but still don't know why you picked one over the other.


Thanks a bunch for taking the time to do so. The message is: Both work. 
We can either have special ioctls per allocation (CPU, memory) that gets 
a special nid property or we can have a global "allocate from this nid" 
cookie behind the fd.

I don't have a super strong preference which way to pick. The main plus 
point for the cookie is that ADD_VCPUS is already an ioctl which we 
would otherwise have to add a new nid-aware variant for.

Do you have a preference?


>
>> I wrote this series with an AI coding assistant, which drafted the code
>> and the changelogs, this cover letter included; I reviewed and reworked
>> all of it before sending, and each commit carries an Assisted-by:
>> trailer. All eight commits were built detached in three configurations
>> (x86_64 defconfig with CONFIG_NITRO_ENCLAVES=m, the same with
>> CONFIG_NUMA=n, and allmodconfig) for 24 warning-free driver builds, and
>> scripts/checkpatch.pl --strict is clean on all eight.
> The patches look fine to me overall, I only have some very minor
> comments on them. The autogenerated changelogs seem actively
> counterproductive though, despite your attempts to to rewrite
> them into something readable. This seems to have just made the
> series harder to review. Instead of telling the reader a
> consistent story about what each change is for, the generated
> text has too much useless detail about the implementation.


I agree. And that's after I massaged it many times already :(. Maybe in 
the next set, I'll just write it myself again :). Or I'll find a way to 
generate actually readable ones.


Alex