From 9d5cc67ad4e94b63cfccbfcecc08706769469dd5 Mon Sep 17 00:00:00 2001
From: Jessie Yang <jiaxiyan@amazon.com>
Date: Mon, 9 Nov 2026 12:15:55 -0800
Subject: [PATCH 1/4] prov/shm: Size the SAR pool independently of
 SMR_MAX_PEERS

The SAR pool size was tied to SMR_MAX_PEERS, so raising the peer limit
would grow every region by 32 KiB per peer. Give the pool its own size.

The per-peer SAR quota becomes 0 once peers outnumber SAR buffers,
which blocks SAR for every peer. Clamp it to at least 1, and don't
start a SAR transfer when the pool is empty.

Signed-off-by: Jessie Yang <jiaxiyan@amazon.com>
---
 prov/shm/src/smr_av.c       | 12 +++---------
 prov/shm/src/smr_ep.c       |  3 ++-
 prov/shm/src/smr_progress.c |  5 ++---
 prov/shm/src/smr_util.c     |  4 ++--
 prov/shm/src/smr_util.h     |  8 ++++++++
 5 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/prov/shm/src/smr_av.c b/prov/shm/src/smr_av.c
index 50781a7c1d8..2d84d5302f4 100644
--- a/prov/shm/src/smr_av.c
+++ b/prov/shm/src/smr_av.c
@@ -441,8 +441,7 @@ static int smr_av_insert(struct fid_av *av_fid, const void *addr, size_t count,
 					       av_entry);
         		smr_ep = container_of(util_ep, struct smr_ep, util_ep);
 			smr_ep->region->max_sar_buf_per_peer =
-				MIN(SMR_BUF_BATCH_MAX,
-				    SMR_MAX_PEERS / smr_av->smr_map.num_peers);
+				smr_sar_bufs_per_peer(smr_av->smr_map.num_peers);
 			smr_ep->srx->owner_ops->foreach_unspec_addr(
 						smr_ep->srx, &smr_get_addr);
 		}
@@ -482,13 +481,8 @@ static int smr_av_remove(struct fid_av *av_fid, fi_addr_t *fi_addr,
 			util_ep = container_of(av_entry, struct util_ep,
 					       av_entry);
 			smr_ep = container_of(util_ep, struct smr_ep, util_ep);
-			if (smr_av->smr_map.num_peers > 0)
-				smr_ep->region->max_sar_buf_per_peer =
-						SMR_MAX_PEERS /
-						smr_av->smr_map.num_peers;
-			else
-				smr_ep->region->max_sar_buf_per_peer =
-							SMR_BUF_BATCH_MAX;
+			smr_ep->region->max_sar_buf_per_peer =
+				smr_sar_bufs_per_peer(smr_av->smr_map.num_peers);
 		}
 		smr_av->used--;
 	}
diff --git a/prov/shm/src/smr_ep.c b/prov/shm/src/smr_ep.c
index a279158dbe6..9a30aad62ea 100644
--- a/prov/shm/src/smr_ep.c
+++ b/prov/shm/src/smr_ep.c
@@ -335,7 +335,8 @@ static int smr_format_sar(struct smr_ep *ep, struct smr_cmd *cmd,
 	int i, ret;
 
 	if (ep->region->max_sar_buf_per_peer == 0 ||
-	    smr_peer_data(ep->region)[cmd->hdr.tx_id].sar_status)
+	    smr_peer_data(ep->region)[cmd->hdr.tx_id].sar_status ||
+	    !smr_freestack_avail(smr_sar_pool(ep->region)))
 		return -FI_EAGAIN;
 
 	cmd->data.buf_batch_size = MIN3(
diff --git a/prov/shm/src/smr_progress.c b/prov/shm/src/smr_progress.c
index 4508da255f3..4e7f21adf0b 100644
--- a/prov/shm/src/smr_progress.c
+++ b/prov/shm/src/smr_progress.c
@@ -884,9 +884,8 @@ static void smr_progress_connreq(struct smr_ep *ep, struct smr_cmd *cmd)
 	smr_peer_data(ep->region)[idx].local_region = (uintptr_t) peer_smr;
 
 	assert(ep->map->num_peers > 0);
-	ep->region->max_sar_buf_per_peer = MIN(
-		SMR_BUF_BATCH_MAX,
-		SMR_MAX_PEERS / ep->map->num_peers);
+	ep->region->max_sar_buf_per_peer =
+		smr_sar_bufs_per_peer(ep->map->num_peers);
 
 	//set last to indicate to peer that setup is complete
 	smr_peer_data(peer_smr)[cmd->hdr.tx_id].id = idx;
diff --git a/prov/shm/src/smr_util.c b/prov/shm/src/smr_util.c
index e6071fa8487..c3b0d841ccb 100644
--- a/prov/shm/src/smr_util.c
+++ b/prov/shm/src/smr_util.c
@@ -113,7 +113,7 @@ size_t smr_calculate_size_offsets(size_t tx_count, size_t rx_count,
 	sar_pool_offset = ret_queue_offset + sizeof(struct smr_return_queue) +
 		sizeof(struct smr_return_queue_entry) * tx_size;
 	peer_data_offset = sar_pool_offset +
-		freestack_size(sizeof(struct smr_sar_buf), SMR_MAX_PEERS);
+		freestack_size(sizeof(struct smr_sar_buf), SMR_SAR_POOL_SIZE);
 	ep_name_offset = peer_data_offset + sizeof(struct smr_peer_data) *
 		SMR_MAX_PEERS;
 
@@ -288,7 +288,7 @@ int smr_create(const struct fi_provider *prov, const struct smr_attr *attr,
 
 	smr_freestack_init(smr_cmd_stack(*smr), tx_size,
 			   sizeof(struct smr_cmd));
-	smr_freestack_init(smr_sar_pool(*smr), SMR_MAX_PEERS,
+	smr_freestack_init(smr_sar_pool(*smr), SMR_SAR_POOL_SIZE,
 			   sizeof(struct smr_sar_buf));
 	for (i = 0; i < SMR_MAX_PEERS; i++) {
 		smr_peer_data(*smr)[i].id = -1;
diff --git a/prov/shm/src/smr_util.h b/prov/shm/src/smr_util.h
index ffe9a518098..81204732cf4 100644
--- a/prov/shm/src/smr_util.h
+++ b/prov/shm/src/smr_util.h
@@ -234,6 +234,7 @@ struct smr_peer {
 };
 
 #define SMR_MAX_PEERS	256
+#define SMR_SAR_POOL_SIZE	256
 #define SMR_PREFETCH_SZ	128
 
 struct smr_region {
@@ -332,6 +333,13 @@ static inline struct smr_return_queue *smr_return_queue(struct smr_region *smr)
 	return (struct smr_return_queue *)
 			((char *) smr + smr->ret_queue_offset);
 }
+static inline uint16_t smr_sar_bufs_per_peer(int num_peers)
+{
+	if (num_peers <= 0)
+		return SMR_BUF_BATCH_MAX;
+	return MAX(1, MIN(SMR_BUF_BATCH_MAX, SMR_SAR_POOL_SIZE / num_peers));
+}
+
 static inline struct smr_peer_data *smr_peer_data(struct smr_region *smr)
 {
 	return (struct smr_peer_data *) ((char *) smr + smr->peer_data_offset);

From 5b28689c9495c0903a172107358ea3554d2c85a8 Mon Sep 17 00:00:00 2001
From: Jessie Yang <jiaxiyan@amazon.com>
Date: Tue, 10 Nov 2026 11:55:01 -0800
Subject: [PATCH 2/4] prov/shm: Size the peer limit from the number of online
 cores

shm is limited to 256 peers per AV, which is not enough for instances
with high core counts. Compute one node-wide limit at init,
max(256, SMR_PROCS_PER_CORE * online cores), so every process on the
node lays out its region the same way and the AV count stays a hint.

Store the limit in the region header and bump SMR_VERSION so peers with
a different layout are rejected. Also reset shm_id on every iteration of
smr_av_insert so a failed insert does not remove the previous peer.

Signed-off-by: Jessie Yang <jiaxiyan@amazon.com>
---
 prov/shm/src/smr.h      |  2 +-
 prov/shm/src/smr_av.c   | 41 +++++++++++++++++++++++++++++------------
 prov/shm/src/smr_ep.c   |  4 ++--
 prov/shm/src/smr_init.c |  1 +
 prov/shm/src/smr_util.c | 17 +++++++++++++++--
 prov/shm/src/smr_util.h | 11 +++++++++--
 6 files changed, 57 insertions(+), 19 deletions(-)

diff --git a/prov/shm/src/smr.h b/prov/shm/src/smr.h
index 32d16048d2c..6e158a6f93c 100644
--- a/prov/shm/src/smr.h
+++ b/prov/shm/src/smr.h
@@ -70,7 +70,7 @@ struct smr_map {
 	int 			num_peers;
 	uint16_t		flags;
 	struct ofi_rbmap	rbmap;
-	struct smr_peer		peers[SMR_MAX_PEERS];
+	struct smr_peer		*peers;
 };
 
 struct smr_av {
diff --git a/prov/shm/src/smr_av.c b/prov/shm/src/smr_av.c
index 2d84d5302f4..32c1c88be8e 100644
--- a/prov/shm/src/smr_av.c
+++ b/prov/shm/src/smr_av.c
@@ -143,6 +143,16 @@ int smr_map_to_region(struct smr_map *map, int64_t id)
 		goto out;
 	}
 
+	if (peer->version != SMR_VERSION || peer->max_peers != smr_max_peers) {
+		FI_WARN(&smr_prov, FI_LOG_AV,
+			"peer %s region version %d max peers %d, expected %d %d\n",
+			name, peer->version, peer->max_peers, SMR_VERSION,
+			smr_max_peers);
+		munmap(peer, sizeof(*peer));
+		ret = -FI_EINVAL;
+		goto out;
+	}
+
 	size = peer->total_size;
 	munmap(peer, sizeof(*peer));
 
@@ -261,18 +271,18 @@ void smr_map_add(struct smr_map *map, const char *name, int64_t *id)
 		return;
 	}
 
-	while (map->peers[map->cur_id].id_assigned && tries < SMR_MAX_PEERS) {
-		if (++map->cur_id == SMR_MAX_PEERS)
+	while (map->peers[map->cur_id].id_assigned && tries < smr_max_peers) {
+		if (++map->cur_id == smr_max_peers)
 			map->cur_id = 0;
 		tries++;
 	}
 
-	assert(map->cur_id < SMR_MAX_PEERS && tries < SMR_MAX_PEERS);
+	assert(map->cur_id < smr_max_peers && tries < smr_max_peers);
 	*id = map->cur_id;
-	if (++map->cur_id == SMR_MAX_PEERS)
+	if (++map->cur_id == smr_max_peers)
 		map->cur_id = 0;
 	node->data = (void *) (intptr_t) *id;
-	strncpy(map->peers[*id].name, shm_name, SMR_NAME_MAX);
+	strncpy(map->peers[*id].name, shm_name, SMR_NAME_MAX - 1);
 	map->peers[*id].name[SMR_NAME_MAX - 1] = '\0';
 	map->peers[*id].region = NULL;
 	map->num_peers++;
@@ -287,7 +297,7 @@ static void smr_map_del(struct smr_map *map, int64_t id)
 	assert(ofi_genlock_held(&container_of(map, struct smr_av,
 					      smr_map)->util_av.lock));
 
-	assert(id >= 0 && id < SMR_MAX_PEERS);
+	assert(id >= 0 && id < smr_max_peers);
 	pthread_mutex_lock(&ep_list_lock);
 	dlist_foreach_container(&ep_name_list, struct smr_ep_name, name,
 				entry) {
@@ -308,7 +318,7 @@ static void smr_map_del(struct smr_map *map, int64_t id)
 
 struct smr_region *smr_map_get(struct smr_map *map, int64_t id)
 {
-	if (id < 0 || id >= SMR_MAX_PEERS)
+	if (id < 0 || id >= smr_max_peers)
 		return NULL;
 
 	return map->peers[id].region;
@@ -328,6 +338,10 @@ static int smr_map_init(struct smr_map *map, int peer_count, uint16_t flags)
 {
 	int i;
 
+	map->peers = calloc(peer_count, sizeof(*map->peers));
+	if (!map->peers)
+		return -FI_ENOMEM;
+
 	for (i = 0; i < peer_count; i++) {
 		memset(&map->peers[i].name, 0, SMR_NAME_MAX);
 		map->peers[i].id_assigned = 0;
@@ -345,12 +359,14 @@ static void smr_map_cleanup(struct smr_av *av)
 	int64_t i;
 
 	ofi_genlock_lock(&av->util_av.lock);
-	for (i = 0; i < SMR_MAX_PEERS; i++) {
+	for (i = 0; i < smr_max_peers; i++) {
 		if (av->smr_map.peers[i].id_assigned)
 			smr_map_del(&av->smr_map, i);
 	}
 	ofi_rbmap_cleanup(&av->smr_map.rbmap);
 	ofi_genlock_unlock(&av->util_av.lock);
+	free(av->smr_map.peers);
+	av->smr_map.peers = NULL;
 }
 
 static int smr_av_close(struct fid *fid)
@@ -401,7 +417,8 @@ static int smr_av_insert(struct fid_av *av_fid, const void *addr, size_t count,
 		FI_INFO(&smr_prov, FI_LOG_AV, "%s\n", (const char *) addr);
 
 		util_addr = FI_ADDR_NOTAVAIL;
-		if (smr_av->used < SMR_MAX_PEERS) {
+		shm_id = -1;
+		if (smr_av->used < smr_max_peers) {
 			smr_map_add(&smr_av->smr_map, addr, &shm_id);
 			ret = ofi_av_insert_addr(util_av, &shm_id, &util_addr);
 		} else {
@@ -422,7 +439,7 @@ static int smr_av_insert(struct fid_av *av_fid, const void *addr, size_t count,
 			continue;
 		}
 
-		assert(shm_id >= 0 && shm_id < SMR_MAX_PEERS);
+		assert(shm_id >= 0 && shm_id < smr_max_peers);
 		if (flags & FI_AV_USER_ID) {
 			assert(fi_addr);
 			smr_av->smr_map.peers[shm_id].fiaddr = fi_addr[i];
@@ -572,7 +589,7 @@ int smr_av_open(struct fid_domain *domain, struct fi_av_attr *attr,
 	util_attr.addrlen = sizeof(int64_t);
 	util_attr.context_len = 0;
 	util_attr.flags = 0;
-	if (attr->count > SMR_MAX_PEERS) {
+	if (attr->count > smr_max_peers) {
 		FI_INFO(&smr_prov, FI_LOG_AV,
 			"count %d exceeds max peers\n", (int) attr->count);
 		ret = -FI_ENOSYS;
@@ -589,7 +606,7 @@ int smr_av_open(struct fid_domain *domain, struct fi_av_attr *attr,
 	(*av)->fid.ops = &smr_av_fi_ops;
 	(*av)->ops = &smr_av_ops;
 
-	ret = smr_map_init(&smr_av->smr_map, SMR_MAX_PEERS,
+	ret = smr_map_init(&smr_av->smr_map, smr_max_peers,
 			   util_domain->info_domain_caps & FI_HMEM ?
 			   SMR_FLAG_HMEM_ENABLED : 0);
 	if (ret)
diff --git a/prov/shm/src/smr_ep.c b/prov/shm/src/smr_ep.c
index 9a30aad62ea..f8693b4fc9f 100644
--- a/prov/shm/src/smr_ep.c
+++ b/prov/shm/src/smr_ep.c
@@ -203,7 +203,7 @@ int64_t smr_verify_peer(struct smr_ep *ep, fi_addr_t fi_addr)
 	int ret;
 
 	id = smr_addr_lookup(ep->util_ep.av, fi_addr);
-	assert(id < SMR_MAX_PEERS);
+	assert(id < smr_max_peers);
 	if (id < 0)
 		return -1;
 
@@ -797,7 +797,7 @@ static void smr_ep_map_all_peers(struct smr_ep *ep)
 	int64_t i;
 
 	ofi_genlock_lock(&ep->util_ep.av->lock);
-	for (i = 0; i < SMR_MAX_PEERS; i++)
+	for (i = 0; i < smr_max_peers; i++)
 		smr_map_to_endpoint(ep, i);
 
 	ofi_genlock_unlock(&ep->util_ep.av->lock);
diff --git a/prov/shm/src/smr_init.c b/prov/shm/src/smr_init.c
index f3df3f4a36d..55a44483962 100644
--- a/prov/shm/src/smr_init.c
+++ b/prov/shm/src/smr_init.c
@@ -233,6 +233,7 @@ SHM_INI
 			"buffering. (default: 1)");
 
 	smr_init_env();
+	smr_init_max_peers();
 
 	if (smr_env.use_dsa_sar)
 		smr_dsa_init();
diff --git a/prov/shm/src/smr_util.c b/prov/shm/src/smr_util.c
index c3b0d841ccb..cbfcfd9f704 100644
--- a/prov/shm/src/smr_util.c
+++ b/prov/shm/src/smr_util.c
@@ -88,6 +88,18 @@ void smr_cma_check(struct smr_region *smr, struct smr_region *peer_smr)
 	}
 }
 
+int smr_max_peers = SMR_MAX_PEERS;
+
+void smr_init_max_peers(void)
+{
+	long cores = ofi_sysconf(_SC_NPROCESSORS_ONLN);
+
+	if (cores > 0)
+		smr_max_peers = MIN(MAX((long) SMR_MAX_PEERS,
+					SMR_PROCS_PER_CORE * cores),
+				    INT16_MAX);
+}
+
 size_t smr_calculate_size_offsets(size_t tx_count, size_t rx_count,
 				  size_t *cmd_offset, size_t *cs_offset,
 				  size_t *inject_offset, size_t *rq_offset,
@@ -115,7 +127,7 @@ size_t smr_calculate_size_offsets(size_t tx_count, size_t rx_count,
 	peer_data_offset = sar_pool_offset +
 		freestack_size(sizeof(struct smr_sar_buf), SMR_SAR_POOL_SIZE);
 	ep_name_offset = peer_data_offset + sizeof(struct smr_peer_data) *
-		SMR_MAX_PEERS;
+		smr_max_peers;
 
 	total_size = ep_name_offset + SMR_NAME_MAX;
 
@@ -260,6 +272,7 @@ int smr_create(const struct fi_provider *prov, const struct smr_attr *attr,
 	*smr = mapped_addr;
 
 	(*smr)->version = SMR_VERSION;
+	(*smr)->max_peers = smr_max_peers;
 
 	(*smr)->flags = attr->flags;
 
@@ -290,7 +303,7 @@ int smr_create(const struct fi_provider *prov, const struct smr_attr *attr,
 			   sizeof(struct smr_cmd));
 	smr_freestack_init(smr_sar_pool(*smr), SMR_SAR_POOL_SIZE,
 			   sizeof(struct smr_sar_buf));
-	for (i = 0; i < SMR_MAX_PEERS; i++) {
+	for (i = 0; i < smr_max_peers; i++) {
 		smr_peer_data(*smr)[i].id = -1;
 		smr_peer_data(*smr)[i].sar_status = SMR_SAR_FREE;
 		smr_peer_data(*smr)[i].name_sent = 0;
diff --git a/prov/shm/src/smr_util.h b/prov/shm/src/smr_util.h
index 81204732cf4..3ea319ce758 100644
--- a/prov/shm/src/smr_util.h
+++ b/prov/shm/src/smr_util.h
@@ -42,7 +42,7 @@
 extern "C" {
 #endif
 
-#define SMR_VERSION	10
+#define SMR_VERSION	11
 
 struct smr_env {
 	int	disable_cma;
@@ -234,6 +234,12 @@ struct smr_peer {
 };
 
 #define SMR_MAX_PEERS	256
+#ifndef SMR_PROCS_PER_CORE
+#define SMR_PROCS_PER_CORE	1
+#endif
+
+extern int smr_max_peers;
+void smr_init_max_peers(void);
 #define SMR_SAR_POOL_SIZE	256
 #define SMR_PREFETCH_SZ	128
 
@@ -251,7 +257,8 @@ struct smr_region {
 			struct ofi_xpmem_pinfo	xpmem_peer;
 
 			int			pid;
-			int			resv2;
+			uint16_t		max_peers;
+			uint16_t		resv2;
 
 			uintptr_t		base_addr;
 

From b12134c178423b799641dc4cdc42cb1549b82ec0 Mon Sep 17 00:00:00 2001
From: Jessie Yang <jiaxiyan@amazon.com>
Date: Tue, 10 Nov 2026 14:16:45 -0800
Subject: [PATCH 3/4] prov/efa: Default FI_EFA_SHM_AV_SIZE to the number of
 online CPUs

FI_EFA_SHM_AV_SIZE defaulted to 256, so jobs on instances with high
core counts failed at the 257th local peer. Default it to the number of
online CPUs, with a minimum of 256, to match shm's peer limit. Setting
the variable still overrides the default.

shm now rejects AV counts above its limit, so drop EFA's own maximum
and warn when the shm AV cannot be opened. Add gtests for the default,
the limit, and an AV open above shm's limit.

The cmocka harness reset efa_env to a copy taken before provider init,
which no longer holds the shm AV size. Take the copy after init.

Signed-off-by: Jessie Yang <jiaxiyan@amazon.com>
---
 man/fi_efa.7.md                              |   1 +
 prov/efa/src/efa_av.h                        |   1 -
 prov/efa/src/efa_env.c                       |  13 +-
 prov/efa/src/rdm/efa_rdm_av.c                |  14 +--
 prov/efa/test/efa_unit_tests.c               |   4 +-
 prov/efa/test/gtest/efa_gtest_rdm_av.cc      | 123 +++++++++++++++++++
 prov/efa/test/gtest/efa_gtest_rdm_av_utils.c |  39 ++++++
 prov/efa/test/gtest/efa_gtest_rdm_av_utils.h |   9 ++
 8 files changed, 190 insertions(+), 14 deletions(-)

diff --git a/man/fi_efa.7.md b/man/fi_efa.7.md
index f2097e35276..379a0b425eb 100644
--- a/man/fi_efa.7.md
+++ b/man/fi_efa.7.md
@@ -1257,6 +1257,7 @@ These OFI runtime parameters apply only to the RDM endpoint.
 
 *FI_EFA_SHM_AV_SIZE*
 : Defines the maximum number of entries in SHM provider's address vector.
+  (Default: number of online CPUs, minimum 256)
 
 *FI_EFA_SHM_MAX_MEDIUM_SIZE*
 : Defines the switch point between small/medium message and large message. The message
diff --git a/prov/efa/src/efa_av.h b/prov/efa/src/efa_av.h
index 9744d9e958e..3c2abb3a9f0 100644
--- a/prov/efa/src/efa_av.h
+++ b/prov/efa/src/efa_av.h
@@ -22,7 +22,6 @@
 #endif
 
 #define EFA_MIN_AV_SIZE (16384)
-#define EFA_SHM_MAX_AV_COUNT       (256)
 
 struct efa_rdm_av;
 struct efa_rdm_pke;
diff --git a/prov/efa/src/efa_env.c b/prov/efa/src/efa_env.c
index 5ab09b43a37..cc97a6991d7 100644
--- a/prov/efa/src/efa_env.c
+++ b/prov/efa/src/efa_env.c
@@ -14,7 +14,6 @@ struct efa_env efa_env = {
 	.tx_queue_size = 0,
 	.enable_shm_transfer = 1,
 	.use_zcpy_rx = 1,
-	.shm_av_size = 256,
 	.recvwin_size = EFA_RDM_PEER_DEFAULT_REORDER_BUFFER_SIZE,
 	.ooo_pool_chunk_size = 64,
 	.unexp_pool_chunk_size = 1024,
@@ -75,6 +74,15 @@ static void efa_env_unregistered_param_get(void)
 		efa_env.enable_comp_action = atoi(tmp);
 }
 
+static int efa_env_default_shm_av_size(void)
+{
+	long cores = ofi_sysconf(_SC_NPROCESSORS_ONLN);
+
+	if (cores <= 256)
+		return 256;
+	return (int) cores;
+}
+
 /* @brief Read and store the FI_EFA_* environment variables.
  */
 void efa_env_param_get(void)
@@ -128,6 +136,7 @@ void efa_env_param_get(void)
 	fi_param_get_int(&efa_prov, "tx_queue_size", &efa_env.tx_queue_size);
 	fi_param_get_int(&efa_prov, "enable_shm_transfer", &efa_env.enable_shm_transfer);
 	fi_param_get_int(&efa_prov, "use_zcpy_rx", &efa_env.use_zcpy_rx);
+	efa_env.shm_av_size = efa_env_default_shm_av_size();
 	fi_param_get_int(&efa_prov, "shm_av_size", &efa_env.shm_av_size);
 	fi_param_get_int(&efa_prov, "recvwin_size", &efa_env.recvwin_size);
 	fi_param_get_int(&efa_prov, "readcopy_pool_size", &efa_env.readcopy_pool_size);
@@ -230,7 +239,7 @@ void efa_env_define()
 	fi_param_define(&efa_prov, "use_zcpy_rx", FI_PARAM_INT,
 			"Enables the use of application's receive buffers in place of bounce-buffers when feasible. (Default: 1)");
 	fi_param_define(&efa_prov, "shm_av_size", FI_PARAM_INT,
-			"Defines the maximum number of entries in SHM provider's address vector (Default 128).");
+			"Defines the maximum number of entries in SHM provider's address vector. (Default: number of online CPUs, minimum 256)");
 	fi_param_define(&efa_prov, "recvwin_size", FI_PARAM_INT,
 			"Defines the size of sliding receive window. (Default: 16384)");
 	fi_param_define(&efa_prov, "readcopy_pool_size", FI_PARAM_INT,
diff --git a/prov/efa/src/rdm/efa_rdm_av.c b/prov/efa/src/rdm/efa_rdm_av.c
index 16d2e1b69db..83badc16825 100644
--- a/prov/efa/src/rdm/efa_rdm_av.c
+++ b/prov/efa/src/rdm/efa_rdm_av.c
@@ -1880,20 +1880,16 @@ int efa_rdm_av_open(struct fid_domain *domain_fid, struct fi_av_attr *attr,
 		 * the need of the instances with more CPUs.
 		 */
 		av_attr = *attr;
-		if (efa_env.shm_av_size > EFA_SHM_MAX_AV_COUNT) {
-			ret = -FI_ENOSYS;
-			EFA_WARN(FI_LOG_AV,
-				 "The requested av size is beyond"
-				 " shm supported maximum av size: %s\n",
-				 fi_strerror(-ret));
-			goto err_close_util_av_implicit;
-		}
 		av_attr.count = efa_env.shm_av_size;
 		assert(av_attr.type == FI_AV_TABLE);
 		ret = fi_av_open(rdm_domain->shm_domain, &av_attr,
 				 &rdm_av->shm_rdm_av, context);
-		if (ret)
+		if (ret) {
+			EFA_WARN(FI_LOG_AV,
+				 "Failed to open shm AV with FI_EFA_SHM_AV_SIZE=%d: %s\n",
+				 efa_env.shm_av_size, fi_strerror(-ret));
 			goto err_close_util_av_implicit;
+		}
 	}
 
 	/* the shm AV above stays a table: its fi_addrs never reach the application */
diff --git a/prov/efa/test/efa_unit_tests.c b/prov/efa/test/efa_unit_tests.c
index 9d40ea67b09..2b50480b6d2 100644
--- a/prov/efa/test/efa_unit_tests.c
+++ b/prov/efa/test/efa_unit_tests.c
@@ -20,13 +20,13 @@ static int efa_unit_test_mocks_group_setup(void **state)
 	resource = calloc(1, sizeof(struct efa_resource));
 	*state = resource;
 
-	orig_efa_env = efa_env;
-
 	/* run fi_getinfo to populate g_efa_hmem_info and copy it */
 	fi_getinfo(FI_VERSION(2, 0), NULL, NULL, 0, NULL, &info);
 	memcpy(g_efa_hmem_info_backup, g_efa_hmem_info, sizeof(g_efa_hmem_info));
 	fi_freeinfo(info);
 
+	orig_efa_env = efa_env;
+
 	/* Initialize the random number generator */
 	srand(time(NULL));
 
diff --git a/prov/efa/test/gtest/efa_gtest_rdm_av.cc b/prov/efa/test/gtest/efa_gtest_rdm_av.cc
index 87b4887ce52..667ce7daeec 100644
--- a/prov/efa/test/gtest/efa_gtest_rdm_av.cc
+++ b/prov/efa/test/gtest/efa_gtest_rdm_av.cc
@@ -6,6 +6,10 @@
 #include "efa_gtest_common_resource.h"
 #include "efa_gtest_rdm_av_utils.h"
 #include <gtest/gtest.h>
+#include <algorithm>
+#include <cstdlib>
+#include <string>
+#include <unistd.h>
 
 using testing::Test;
 using testing::_;
@@ -226,3 +230,122 @@ TEST_F(EfaConnTest, promotion_publish_ordering_holds_on_recycled_entry)
 	EXPECT_TRUE(obs->peer_map_resolves);
 	EXPECT_TRUE(obs->peer_is_migrated_peer);
 }
+
+struct EfaShmAvSizeParam {
+	const char *name;
+	int env_value;
+};
+
+class EfaShmAvSizeTest : public testing::TestWithParam<EfaShmAvSizeParam>
+{
+	protected:
+	struct efa_resource resource = {};
+	int prev_shm_av_size;
+	int limit;
+
+	void SetUp() override
+	{
+		memset(&resource, 0, sizeof(resource));
+		prev_shm_av_size = efa_test_get_shm_av_size();
+		if (GetParam().env_value)
+			efa_test_set_shm_av_size(GetParam().env_value);
+		limit = efa_test_get_shm_av_size();
+		efa_test_resource_construct(
+			&resource,
+			efa_test_alloc_default_hints(FI_EP_RDM, EFA_FABRIC_NAME));
+		ASSERT_NE(resource.ep, nullptr);
+		ASSERT_NE(efa_test_get_shm_domain(resource.domain), nullptr);
+	}
+
+	void TearDown() override
+	{
+		efa_test_resource_destruct(&resource);
+		efa_test_set_shm_av_size(prev_shm_av_size);
+	}
+};
+
+TEST_P(EfaShmAvSizeTest, local_peers_up_to_shm_av_size_use_shm)
+{
+	fi_addr_t addr;
+	int i;
+
+	for (i = 0; i < limit; i++) {
+		ASSERT_EQ(efa_test_av_insert_local_peer(resource.ep, resource.av,
+							i + 1, &addr),
+			  1)
+			<< "local peer " << i;
+		ASSERT_TRUE(efa_test_av_peer_uses_shm(resource.ep, addr))
+			<< "local peer " << i;
+	}
+
+	EXPECT_EQ(efa_test_av_insert_local_peer(resource.ep, resource.av,
+						limit + 1, &addr),
+		  0);
+	EXPECT_EQ(addr, (fi_addr_t) FI_ADDR_NOTAVAIL);
+}
+
+INSTANTIATE_TEST_SUITE_P(
+	, EfaShmAvSizeTest,
+	testing::Values(EfaShmAvSizeParam{"default", 0},
+			EfaShmAvSizeParam{"lowered", 64}),
+	[](const testing::TestParamInfo<EfaShmAvSizeParam> &info) {
+		return std::string(info.param.name);
+	});
+
+class EfaShmAvSizeLimitTest : public Test
+{
+	protected:
+	struct efa_resource resource = {};
+
+	void SetUp() override
+	{
+		memset(&resource, 0, sizeof(resource));
+		efa_test_resource_construct(
+			&resource,
+			efa_test_alloc_default_hints(FI_EP_RDM, EFA_FABRIC_NAME));
+		ASSERT_NE(resource.ep, nullptr);
+		ASSERT_NE(efa_test_get_shm_domain(resource.domain), nullptr);
+	}
+
+	void TearDown() override
+	{
+		efa_test_resource_destruct(&resource);
+	}
+
+	int av_open_with_shm_av_size(int shm_av_size)
+	{
+		struct fi_av_attr av_attr = {};
+		struct fid_av *av = nullptr;
+		int prev, ret;
+
+		av_attr.type = FI_AV_TABLE;
+		prev = efa_test_set_shm_av_size(shm_av_size);
+		ret = fi_av_open(resource.domain, &av_attr, &av, NULL);
+		efa_test_set_shm_av_size(prev);
+		if (!ret)
+			fi_close(&av->fid);
+		return ret;
+	}
+};
+
+TEST_F(EfaShmAvSizeLimitTest, default_follows_online_cpus)
+{
+	const char *env = getenv("FI_EFA_SHM_AV_SIZE");
+	long cpus = sysconf(_SC_NPROCESSORS_ONLN);
+	long expected;
+
+	if (env)
+		expected = atol(env);
+	else
+		expected = std::max(256L, cpus);
+
+	EXPECT_EQ(efa_test_get_shm_av_size(), expected);
+}
+
+TEST_F(EfaShmAvSizeLimitTest, av_open_fails_above_shm_capacity)
+{
+	long cpus = sysconf(_SC_NPROCESSORS_ONLN);
+	int above_capacity = (int) std::max(256L, cpus) + 1;
+
+	EXPECT_EQ(av_open_with_shm_av_size(above_capacity), -FI_ENOSYS);
+}
diff --git a/prov/efa/test/gtest/efa_gtest_rdm_av_utils.c b/prov/efa/test/gtest/efa_gtest_rdm_av_utils.c
index cd8c7c6b18c..c7f9e542681 100644
--- a/prov/efa/test/gtest/efa_gtest_rdm_av_utils.c
+++ b/prov/efa/test/gtest/efa_gtest_rdm_av_utils.c
@@ -3,6 +3,7 @@
 
 #include "efa.h"
 #include "efa_av.h"
+#include "efa_env.h"
 #include "rdm/efa_rdm_av.h"
 #include "rdm/efa_rdm_ep.h"
 #include "rdm/efa_rdm_peer.h"
@@ -140,3 +141,41 @@ int efa_test_av_publish_ordering_recycle_slot(struct fid_ep *ep,
 	*recycled_fi_addr = fi_addr;
 	return 0;
 }
+
+int efa_test_get_shm_av_size(void)
+{
+	return efa_env.shm_av_size;
+}
+
+int efa_test_set_shm_av_size(int value)
+{
+	int prev = efa_env.shm_av_size;
+
+	efa_env.shm_av_size = value;
+	return prev;
+}
+
+int efa_test_av_insert_local_peer(struct fid_ep *ep, struct fid_av *av,
+				  uint32_t qpn, fi_addr_t *addr)
+{
+	struct efa_ep_addr raw_addr = {0};
+	size_t raw_addr_len = sizeof(raw_addr);
+	int ret;
+
+	ret = fi_getname(&ep->fid, &raw_addr, &raw_addr_len);
+	if (ret)
+		return ret;
+	raw_addr.qpn = qpn;
+	raw_addr.qkey = 0x1234;
+
+	return fi_av_insert(av, &raw_addr, 1, addr, 0, NULL);
+}
+
+bool efa_test_av_peer_uses_shm(struct fid_ep *ep, fi_addr_t addr)
+{
+	struct efa_rdm_ep *efa_rdm_ep =
+		container_of(ep, struct efa_rdm_ep, base_ep.util_ep.ep_fid);
+
+	return efa_rdm_ep_get_explicit_shm_fi_addr(efa_rdm_ep, addr) !=
+	       FI_ADDR_NOTAVAIL;
+}
diff --git a/prov/efa/test/gtest/efa_gtest_rdm_av_utils.h b/prov/efa/test/gtest/efa_gtest_rdm_av_utils.h
index 369cca34663..832f50a2a48 100644
--- a/prov/efa/test/gtest/efa_gtest_rdm_av_utils.h
+++ b/prov/efa/test/gtest/efa_gtest_rdm_av_utils.h
@@ -102,6 +102,15 @@ int efa_test_av_publish_ordering_recycle_slot(struct fid_ep *ep,
  */
 bool efa_test_av_publish_ordering_single_peer(fi_addr_t explicit_fi_addr);
 
+int efa_test_get_shm_av_size(void);
+
+int efa_test_set_shm_av_size(int value);
+
+int efa_test_av_insert_local_peer(struct fid_ep *ep, struct fid_av *av,
+				  uint32_t qpn, fi_addr_t *addr);
+
+bool efa_test_av_peer_uses_shm(struct fid_ep *ep, fi_addr_t addr);
+
 #ifdef __cplusplus
 }
 #endif

From 2d2f6b0b21689d2086eefdb63b31bcf3fa325a8f Mon Sep 17 00:00:00 2001
From: Jessie Yang <jiaxiyan@amazon.com>
Date: Tue, 10 Nov 2026 11:55:01 -0800
Subject: [PATCH 4/4] fabtests: Add fi_av_capacity_test

Insert addresses into an AV until it is full, check that it accepts at
least a given number, rejects the next insert, and accepts an insert
again after a remove. Run it for shm with the minimum set to
max(256, online CPUs).

Signed-off-by: Jessie Yang <jiaxiyan@amazon.com>
---
 fabtests/Makefile.am                    |   8 +
 fabtests/man/fabtests.7.md              |   5 +
 fabtests/man/man1/fi_av_capacity_test.1 |   1 +
 fabtests/pytest/shm/test_av_capacity.py |  12 ++
 fabtests/scripts/runfabtests.sh         |   1 +
 fabtests/unit/av_capacity_test.c        | 206 ++++++++++++++++++++++++
 6 files changed, 233 insertions(+)
 create mode 100644 fabtests/man/man1/fi_av_capacity_test.1
 create mode 100644 fabtests/pytest/shm/test_av_capacity.py
 create mode 100644 fabtests/unit/av_capacity_test.c

diff --git a/fabtests/Makefile.am b/fabtests/Makefile.am
index a4313b331db..e4d59e7789c 100644
--- a/fabtests/Makefile.am
+++ b/fabtests/Makefile.am
@@ -62,6 +62,7 @@ bin_PROGRAMS = \
 	unit/fi_mr_cache_evict \
 	unit/fi_cntr_test \
 	unit/fi_av_test \
+	unit/fi_av_capacity_test \
 	unit/fi_dom_test \
 	unit/fi_getinfo_test \
 	unit/fi_setopt_test \
@@ -208,6 +209,7 @@ nobase_dist_config_DATA = \
 	pytest/shm/conftest.py \
 	pytest/shm/shm_common.py \
 	pytest/shm/test_av.py \
+	pytest/shm/test_av_capacity.py \
 	pytest/shm/test_cntr.py \
 	pytest/shm/test_cq.py \
 	pytest/shm/test_dom.py \
@@ -483,6 +485,11 @@ unit_fi_av_test_SOURCES = \
 	$(unit_srcs)
 unit_fi_av_test_LDADD = libfabtests.la
 
+unit_fi_av_capacity_test_SOURCES = \
+	unit/av_capacity_test.c \
+	$(unit_srcs)
+unit_fi_av_capacity_test_LDADD = libfabtests.la
+
 unit_fi_dom_test_SOURCES = \
 	unit/dom_test.c \
 	$(unit_srcs)
@@ -691,6 +698,7 @@ dummy_man_pages = \
 	man/man1/fi_rdm_tagged_pingpong.1 \
 	man/man1/fi_rma_bw.1 \
 	man/man1/fi_av_test.1 \
+	man/man1/fi_av_capacity_test.1 \
 	man/man1/fi_cntr_test.1 \
 	man/man1/fi_cq_test.1 \
 	man/man1/fi_dom_test.1 \
diff --git a/fabtests/man/fabtests.7.md b/fabtests/man/fabtests.7.md
index d3c09914d47..8e97431b7d9 100644
--- a/fabtests/man/fabtests.7.md
+++ b/fabtests/man/fabtests.7.md
@@ -211,6 +211,11 @@ testing scope is limited.
 *fi_av_test*
 : Verify address vector interfaces.
 
+*fi_av_capacity_test*
+: Verify that an address vector accepts at least a given number of addresses,
+  rejects inserts beyond its limit, and accepts an insert again after a
+  remove. Supports providers using the FI_ADDR_STR address format.
+
 *fi_cntr_test*
 : Tests counter creation and destruction.
 
diff --git a/fabtests/man/man1/fi_av_capacity_test.1 b/fabtests/man/man1/fi_av_capacity_test.1
new file mode 100644
index 00000000000..3f6ccf96f11
--- /dev/null
+++ b/fabtests/man/man1/fi_av_capacity_test.1
@@ -0,0 +1 @@
+.so man7/fabtests.7
diff --git a/fabtests/pytest/shm/test_av_capacity.py b/fabtests/pytest/shm/test_av_capacity.py
new file mode 100644
index 00000000000..a4e2418c277
--- /dev/null
+++ b/fabtests/pytest/shm/test_av_capacity.py
@@ -0,0 +1,12 @@
+import os
+
+import pytest
+
+
+@pytest.mark.pr_ci
+@pytest.mark.unit
+def test_av_capacity(cmdline_args):
+    from common import UnitTest
+    min_addrs = max(256, os.cpu_count() or 0)
+    test = UnitTest(cmdline_args, "fi_av_capacity_test -n " + str(min_addrs))
+    test.run()
diff --git a/fabtests/scripts/runfabtests.sh b/fabtests/scripts/runfabtests.sh
index 9cae9eb5db2..648b428b0f9 100755
--- a/fabtests/scripts/runfabtests.sh
+++ b/fabtests/scripts/runfabtests.sh
@@ -240,6 +240,7 @@ unit_tests=(
 	"fi_getinfo_test -s SERVER_ADDR GOOD_ADDR"
 	"fi_av_test -g GOOD_ADDR -n 1 -s SERVER_ADDR -e rdm"
 	"fi_av_test -g GOOD_ADDR -n 1 -s SERVER_ADDR -e dgram"
+	"fi_av_capacity_test -n 256"
 	"fi_dom_test -n 2"
 	"fi_eq_test"
 	"fi_cq_test"
diff --git a/fabtests/unit/av_capacity_test.c b/fabtests/unit/av_capacity_test.c
new file mode 100644
index 00000000000..7251c69fbb9
--- /dev/null
+++ b/fabtests/unit/av_capacity_test.c
@@ -0,0 +1,206 @@
+/* SPDX-License-Identifier: BSD-2-Clause OR GPL-2.0-only */
+/* SPDX-FileCopyrightText: Copyright Amazon.com, Inc. or its affiliates. All rights reserved. */
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <getopt.h>
+#include <inttypes.h>
+
+#include <rdma/fi_cm.h>
+#include <rdma/fi_errno.h>
+
+#include "shared.h"
+#include "unit_common.h"
+
+#define AV_CAPACITY_ADDR_LEN 256
+
+/* 
+ * Distinct from FI_ADDR_NOTAVAIL so that a failed insert which leaves fi_addr
+ * untouched can be told apart from one where the provider correctly set it to
+ * FI_ADDR_NOTAVAIL.
+ */
+#define AV_CAPACITY_ADDR_POISON ((fi_addr_t) 0xdeadbeefdeadbeefULL)
+
+static char base_addr[AV_CAPACITY_ADDR_LEN];
+
+static int insert_one(struct fid_av *test_av, int index, fi_addr_t *fi_addr)
+{
+	char addr[AV_CAPACITY_ADDR_LEN];
+	int len;
+
+	len = snprintf(addr, sizeof(addr), "%s_%d", base_addr, index);
+	if (len < 0 || len >= (int) sizeof(addr))
+		return -FI_ETOOSMALL;
+
+	*fi_addr = AV_CAPACITY_ADDR_POISON;
+	return fi_av_insert(test_av, addr, 1, fi_addr, 0, NULL);
+}
+
+static int get_base_addr(void)
+{
+	size_t len = sizeof(base_addr) - 16;
+	int ret;
+
+	memset(base_addr, 0, sizeof(base_addr));
+	ret = fi_getname(&ep->fid, base_addr, &len);
+	if (ret) {
+		FT_PRINTERR("fi_getname", ret);
+		return ret;
+	}
+	base_addr[sizeof(base_addr) - 1] = '\0';
+	return 0;
+}
+
+static int av_capacity(int min_addrs, int max_addrs)
+{
+	struct fi_av_attr attr = { .type = FI_AV_TABLE };
+	struct fid_av *test_av;
+	fi_addr_t *fi_addrs;
+	fi_addr_t extra;
+	int inserted;
+	int ret, close_ret;
+
+	fi_addrs = calloc(max_addrs, sizeof(*fi_addrs));
+	if (!fi_addrs)
+		return -FI_ENOMEM;
+
+	ret = fi_av_open(domain, &attr, &test_av, NULL);
+	if (ret) {
+		FT_PRINTERR("fi_av_open", ret);
+		goto out;
+	}
+
+	for (inserted = 0; inserted < max_addrs; inserted++) {
+		ret = insert_one(test_av, inserted, &fi_addrs[inserted]);
+		if (ret != 1)
+			break;
+	}
+
+	if (inserted == 0) {
+		printf("cannot insert generated addresses: %d\n", ret);
+		ret = -FI_ENOSYS;
+		goto close;
+	}
+
+	printf("inserted %d addresses\n", inserted);
+	if (inserted < min_addrs) {
+		printf("AV accepted %d addresses, expected at least %d\n",
+		       inserted, min_addrs);
+		ret = -FI_EOTHER;
+		goto close;
+	}
+
+	if (inserted < max_addrs && fi_addrs[inserted] != FI_ADDR_NOTAVAIL) {
+		printf("insert beyond the limit returned %d, fi_addr %" PRIu64
+		       "\n", ret, fi_addrs[inserted]);
+		ret = -FI_EOTHER;
+		goto close;
+	}
+
+	ret = fi_av_remove(test_av, &fi_addrs[0], 1, 0);
+	if (ret) {
+		FT_PRINTERR("fi_av_remove", ret);
+		goto close;
+	}
+
+	ret = insert_one(test_av, inserted, &extra);
+	if (ret != 1 || extra == FI_ADDR_NOTAVAIL ||
+	    extra == AV_CAPACITY_ADDR_POISON) {
+		printf("insert after remove returned %d, fi_addr %" PRIu64 "\n",
+		       ret, extra);
+		ret = -FI_EOTHER;
+		goto close;
+	}
+	ret = 0;
+
+close:
+	close_ret = fi_close(&test_av->fid);
+	if (close_ret) {
+		FT_PRINTERR("fi_close", close_ret);
+		if (!ret)
+			ret = close_ret;
+	}
+out:
+	free(fi_addrs);
+	return ret;
+}
+
+static void usage(char *name)
+{
+	ft_unit_usage(name, "Unit test for AV capacity");
+	FT_PRINT_OPTS_USAGE("-n <num>",
+			    "minimum number of addresses the AV must accept (default 1)");
+	FT_PRINT_OPTS_USAGE("-m <num>",
+			    "maximum number of addresses to insert (default 65536)");
+}
+
+int main(int argc, char **argv)
+{
+	int min_addrs = 1, max_addrs = 65536;
+	int op, ret, cleanup_ret;
+
+	hints = fi_allocinfo();
+	if (!hints)
+		return EXIT_FAILURE;
+
+	while ((op = getopt(argc, argv, FAB_OPTS "n:m:h")) != -1) {
+		switch (op) {
+		case 'n':
+			min_addrs = atoi(optarg);
+			break;
+		case 'm':
+			max_addrs = atoi(optarg);
+			break;
+		default:
+			ft_parseinfo(op, optarg, hints, &opts);
+			break;
+		case '?':
+		case 'h':
+			usage(argv[0]);
+			return EXIT_FAILURE;
+		}
+	}
+
+	if (min_addrs <= 0 || max_addrs <= 0 || min_addrs > max_addrs) {
+		usage(argv[0]);
+		return EXIT_FAILURE;
+	}
+
+	hints->ep_attr->type = FI_EP_RDM;
+	hints->mode = ~0;
+	hints->domain_attr->mode = ~0;
+	hints->domain_attr->mr_mode = ~OFI_MR_DEPRECATED;
+
+	ret = ft_getinfo(hints, &fi);
+	if (ret)
+		goto out;
+
+	if (fi->addr_format != FI_ADDR_STR) {
+		printf("address format %s is not supported\n",
+		       fi_tostr(&fi->addr_format, FI_TYPE_ADDR_FORMAT));
+		ret = -FI_ENOSYS;
+		goto out;
+	}
+
+	ret = ft_open_fabric_res();
+	if (ret)
+		goto out;
+
+	ret = ft_alloc_active_res(fi);
+	if (ret)
+		goto out;
+
+	ret = ft_enable_ep(ep, eq, av, txcq, rxcq, txcntr, rxcntr, rma_cntr);
+	if (ret)
+		goto out;
+
+	ret = get_base_addr();
+	if (ret)
+		goto out;
+
+	ret = av_capacity(min_addrs, max_addrs);
+out:
+	cleanup_ret = ft_free_res();
+	return ft_exit_code(ret ? ret : cleanup_ret);
+}