fix(rocr/thunk): revoke only the deregistered SVM-API extent

projects/rocr-runtime/libhsakmt/src/fmm.c

@@ -258,12 +258,22 @@ typedef struct {
  * base and size here, refcounted to mirror the registration_count handling of
  * regular userptr vm_objects, so that deregistration can issue the inverse SVM
  * SET_ATTR (NO_ACCESS + clear coherency flags) instead of being a silent no-op.
+ *
+ * Keyed by the exact user pointer, not by the page-aligned base. Distinct
+ * sub-page buffers routinely share a page-aligned base and round to different
+ * spans, and register/deregister both carry the user pointer, so it is the only
+ * key that identifies one registration. Keying by base instead conflates those
+ * registrations and forces one shared extent, which then has to cover the
+ * largest of them - so deregistering a small buffer revokes GPU access across a
+ * span it never owned.
  */
 struct svm_api_range {
-	void *start;			/* page-aligned base address */
-	uint64_t size;			/* page-aligned size */
-	uint32_t refcount;		/* number of outstanding registrations */
-	rbtree_node_t node;		/* svm_api_range_tree, key addr only (size field 0) */
+	void *start;			/* page-aligned base, rounded out */
+	uint64_t size;			/* page-aligned size, rounded out */
+	void *revoke_start;		/* pages this registration owns outright */
+	uint64_t revoke_size;		/* 0 if it fills no whole page */
+	uint32_t refcount;		/* registrations of this same user pointer */
+	rbtree_node_t node;		/* svm_api_range_tree, key user VA (size field 0) */
 };
 typedef struct svm_api_range svm_api_range_t;
 
@@ -286,10 +296,13 @@ struct hsa_kfd_fmm_context
 	svm_t svm;
 
 	/* RB tree of host ranges registered via the SVM API (see svm_api_range).
-	 * Protected by svm_api_mutex. Keys use aligned VA only (LKP_ADDR).
+	 * Protected by svm_api_mutex. Keys use the user VA only (LKP_ADDR).
+	 * svm_api_max_extent is the largest tracked extent, used to bound the
+	 * leftward scan for registrations that overlap a range being revoked.
 	 */
 	rbtree_t svm_api_range_tree;
 	pthread_mutex_t svm_api_mutex;
+	uint64_t svm_api_max_extent;
 
 	/* On APU, for memory allocated on the system memory that GPU doesn't
 	 * access via GPU driver, they are not managed by GPUVM. cpuvm_aperture
@@ -350,6 +363,7 @@ int hsakmt_kfdcontext_init_fmm_context(HsaKFDContext *ctx)
 	/* Initialize SVM-API range tracking */
 	rbtree_init(&ctx->fmm_context->svm_api_range_tree);
 	pthread_mutex_init(&ctx->fmm_context->svm_api_mutex, NULL);
+	ctx->fmm_context->svm_api_max_extent = 0;
 
 	/* Initialize cpuvm_aperture */
 	ctx->fmm_context->cpuvm_aperture = init_aperture;
@@ -1136,9 +1150,9 @@ static HsaMemFlags fmm_translate_ioc_to_hsa_flags(uint32_t ioc_flags)
  * internally; callers must hold the mutex only when calling svm_api_range_find.
  */
 static svm_api_range_t *svm_api_range_find(struct hsa_kfd_fmm_context *fmm_ctx,
-					   void *aligned_addr)
+					   void *user_addr)
 {
-	rbtree_key_t key = rbtree_key((unsigned long)aligned_addr, 0);
+	rbtree_key_t key = rbtree_key((unsigned long)user_addr, 0);
 	rbtree_node_t *n = rbtree_lookup(&fmm_ctx->svm_api_range_tree, &key, LKP_ADDR);
 
 	if (!n)
@@ -1147,39 +1161,43 @@ static svm_api_range_t *svm_api_range_find(struct hsa_kfd_fmm_context *fmm_ctx,
 }
 
 /*
- * Add a new SVM-API range, or refcount an existing one sharing the page-aligned
- * base.
+ * Track one SVM-API registration of @user_addr covering the page-rounded range
+ * [@aligned_addr, @aligned_addr + @aligned_size).
  *
- * Several distinct registrations can share a page-aligned base - small buffers
- * packed into one page, one of which may straddle into the next page and so
- * round to a larger span. They are refcounted together on that base and the
- * tracked extent grows to the largest span seen, so deregister revokes the full
- * extent once the last owner is dropped. The exact sub-range to revoke is then
- * computed by interval subtraction in svm_api_range_put (which subtracts the
- * coverage of any surviving registration, at this base or a neighbouring one).
+ * Each registration keeps its own extent, so buffers that merely share a page
+ * no longer share a revoke extent. Repeat registrations of the same pointer are
+ * refcounted, and only there can the extent grow to the largest span seen - an
+ * allocator cannot hand out the same address twice while the first registration
+ * is still live, so in practice that is the same buffer being re-registered.
+ * Sub-ranges still owned by another registration are preserved by the interval
+ * subtraction in svm_api_range_put.
  */
-static HSAKMT_STATUS svm_api_range_get(struct hsa_kfd_fmm_context *fmm_ctx,
-				       void *aligned_addr, uint64_t aligned_size)
+static void svm_api_range_get(struct hsa_kfd_fmm_context *fmm_ctx, void *user_addr,
+			      void *aligned_addr, uint64_t aligned_size,
+			      uint64_t user_size)
 {
 	svm_api_range_t *r;
+	/* Register rounds out to map whole pages, revoke rounds in. The bytes
+	 * past either end belong to whatever else shares those pages.
+	 */
+	HSAuint64 in_start = PAGE_ALIGN_UP((HSAuint64)user_addr);
+	HSAuint64 in_end = ((HSAuint64)user_addr + user_size) & ~(HSAuint64)(PAGE_SIZE - 1);
+	HSAuint64 in_size = in_end > in_start ? in_end - in_start : 0;
 
 	pthread_mutex_lock(&fmm_ctx->svm_api_mutex);
-	r = svm_api_range_find(fmm_ctx, aligned_addr);
+	r = svm_api_range_find(fmm_ctx, user_addr);
 	if (r) {
-		/* Distinct sub-page buffers can share a page-aligned base with
-		 * different page-aligned spans (e.g. small calloc'd buffers
-		 * packed into one page, one of which straddles into the next
-		 * page). They are independent registrations, so refcount them
-		 * together and grow the tracked extent to the largest span seen
-		 * at this base; deregister then revokes the full extent once the
-		 * last owner is gone, and overlap with other bases is resolved by
-		 * the interval subtraction in svm_api_range_put.
-		 */
 		++r->refcount;
 		if (aligned_size > r->size)
 			r->size = aligned_size;
+		if (in_size > r->revoke_size) {
+			r->revoke_start = (void *)in_start;
+			r->revoke_size = in_size;
+		}
+		if (r->size > fmm_ctx->svm_api_max_extent)
+			fmm_ctx->svm_api_max_extent = r->size;
 		pthread_mutex_unlock(&fmm_ctx->svm_api_mutex);
-		return HSAKMT_STATUS_SUCCESS;
+		return;
 	}
 
 	r = calloc(1, sizeof(*r));
@@ -1188,17 +1206,20 @@ static HSAKMT_STATUS svm_api_range_get(struct hsa_kfd_fmm_context *fmm_ctx,
 		 * falls back to the previous no-op behavior for this range.
 		 */
 		pr_warn("Failed to track SVM-API range %p, deregister will be a no-op\n",
-			aligned_addr);
+			user_addr);
 		pthread_mutex_unlock(&fmm_ctx->svm_api_mutex);
-		return HSAKMT_STATUS_SUCCESS;
+		return;
 	}
 	r->start = aligned_addr;
 	r->size = aligned_size;
+	r->revoke_start = (void *)in_start;
+	r->revoke_size = in_size;
 	r->refcount = 1;
-	r->node.key = rbtree_key((unsigned long)aligned_addr, 0);
+	r->node.key = rbtree_key((unsigned long)user_addr, 0);
 	hsakmt_rbtree_insert(&fmm_ctx->svm_api_range_tree, &r->node);
+	if (aligned_size > fmm_ctx->svm_api_max_extent)
+		fmm_ctx->svm_api_max_extent = aligned_size;
 	pthread_mutex_unlock(&fmm_ctx->svm_api_mutex);
-	return HSAKMT_STATUS_SUCCESS;
 }
 
 /* A page range whose GPU access must be revoked on deregister. */
@@ -1230,10 +1251,10 @@ static bool svm_revoke_add(struct svm_revoke_range **arr, int *n, int *cap,
 }
 
 /*
- * Drop a reference on the SVM-API registration based at @addr. If this was the
- * last reference, remove it and compute which sub-ranges of [base, base+size)
- * no *other* surviving registration still covers - those are the only ranges
- * whose GPU access must be revoked. Overlapping or page-rounded-neighbor
+ * Drop a reference on the SVM-API registration of @addr. If this was the last
+ * reference, remove it and compute which sub-ranges of [base, base+size) no
+ * *other* surviving registration still covers - those are the only ranges whose
+ * GPU access must be revoked. Overlapping or page-rounded-neighbor
  * registrations keep coverage of the parts they still own, so a shared region
  * survives until its last owner is dropped.
  *
@@ -1244,38 +1265,46 @@ static bool svm_revoke_add(struct svm_revoke_range **arr, int *n, int *cap,
 static int svm_api_range_put(struct hsa_kfd_fmm_context *fmm_ctx, void *addr,
 			     struct svm_revoke_range **out)
 {
-	HSAuint64 page_offset = (HSAuint64)addr & (PAGE_SIZE - 1);
-	void *aligned_addr = (void *)((HSAuint64)addr - page_offset);
 	struct svm_revoke_range *ranges = NULL;
 	int n = 0, cap = 0;
-	HSAuint64 base, end, cur;
+	HSAuint64 base, end, cur, scan;
 	rbtree_key_t key;
 	rbtree_node_t *node;
 	svm_api_range_t *r;
 
 	*out = NULL;
 
 	pthread_mutex_lock(&fmm_ctx->svm_api_mutex);
-	r = svm_api_range_find(fmm_ctx, aligned_addr);
+	r = svm_api_range_find(fmm_ctx, addr);
 	if (!r || --r->refcount > 0) {
 		pthread_mutex_unlock(&fmm_ctx->svm_api_mutex);
 		return 0;
 	}
 
-	base = (HSAuint64)r->start;
-	end = base + r->size;
+	base = (HSAuint64)r->revoke_start;
+	end = base + r->revoke_size;
 	hsakmt_rbtree_delete(&fmm_ctx->svm_api_range_tree, &r->node);
 	free(r);
 
+	/* Fills no whole page, so nothing here is ours alone. */
+	if (end <= base) {
+		pthread_mutex_unlock(&fmm_ctx->svm_api_mutex);
+		return 0;
+	}
+
 	/* Emit the gaps in [base, end) not covered by any surviving registration.
-	 * Walk the tree in address order starting from the nearest range at or
-	 * before base (it may extend into us), else from the smallest.
+	 *
+	 * Any registration that can reach into [base, end) must start after
+	 * base - svm_api_max_extent, so the sweep starts at the first range from
+	 * there on. Starting at the immediate predecessor of base is not enough:
+	 * a longer range further left still covers base, and skipping it revokes
+	 * GPU access to memory that registration still owns.
 	 */
 	cur = base;
-	key = rbtree_key(base, 0);
-	node = rbtree_lookup_nearest(&fmm_ctx->svm_api_range_tree, &key, LKP_ADDR, LEFT);
-	if (!node)
-		node = rbtree_min_max(&fmm_ctx->svm_api_range_tree, LEFT);
+	scan = base > fmm_ctx->svm_api_max_extent ?
+		base - fmm_ctx->svm_api_max_extent : 0;
+	key = rbtree_key(scan, 0);
+	node = rbtree_lookup_nearest(&fmm_ctx->svm_api_range_tree, &key, LKP_ADDR, RIGHT);
 
 	while (node && cur < end) {
 		svm_api_range_t *o = rb_entry(node, svm_api_range_t, node);
@@ -1301,6 +1330,13 @@ static int svm_api_range_put(struct hsa_kfd_fmm_context *fmm_ctx, void *addr,
 	if (cur < end)				/* uncovered tail */
 		svm_revoke_add(&ranges, &n, &cap, cur, end);
 
+	/* svm_api_max_extent only ever grows, which just widens the scan above.
+	 * Reset it once nothing is tracked so a single large registration does
+	 * not keep widening every later scan for the life of the process.
+	 */
+	if (fmm_ctx->svm_api_range_tree.root == &fmm_ctx->svm_api_range_tree.sentinel)
+		fmm_ctx->svm_api_max_extent = 0;
+
 	pthread_mutex_unlock(&fmm_ctx->svm_api_mutex);
 	*out = ranges;
 	return n;
@@ -1325,10 +1361,12 @@ static HSAKMT_STATUS fmm_unregister_mem_svm_api(HsaKFDContext *ctx,
 	if (!fmm_ctx->first_gpu_mem)
 		return HSAKMT_STATUS_ERROR;
 
-	/* One NO_ACCESS attribute per GPU plus two CLR_FLAGS attributes.
-	 * NO_ACCESS on a GPU that never had access is harmless.
+	/* NO_ACCESS only. The coherency flags are page granular and COHERENT is
+	 * the kernel default, so clearing them drops whatever else shares those
+	 * pages below default; on gfx942/950 that also turns a live mapping
+	 * cached. NO_ACCESS is what the unmap and eviction skip key off.
 	 */
-	nattr = num_gpus + 2;
+	nattr = num_gpus;
 	s_attr = nattr * sizeof(struct kfd_ioctl_svm_attribute);
 	args = alloca(sizeof(*args) + s_attr);
 	args->start_addr = (HSAuint64)aligned_addr;
@@ -1339,10 +1377,6 @@ static HSAKMT_STATUS fmm_unregister_mem_svm_api(HsaKFDContext *ctx,
 		args->attrs[i].type = HSA_SVM_ATTR_NO_ACCESS;
 		args->attrs[i].value = fmm_ctx->all_gpu_id_array[i];
 	}
-	args->attrs[num_gpus].type = HSA_SVM_ATTR_CLR_FLAGS;
-	args->attrs[num_gpus].value = HSA_SVM_FLAG_COHERENT;
-	args->attrs[num_gpus + 1].type = HSA_SVM_ATTR_CLR_FLAGS;
-	args->attrs[num_gpus + 1].value = HSA_SVM_FLAG_EXT_COHERENT;
 
 	pr_debug("Deregistering from SVM %p size: %" PRIu64 "\n", aligned_addr,
 		 aligned_size);
@@ -1413,10 +1447,11 @@ static HSAKMT_STATUS fmm_register_mem_svm_api(HsaKFDContext *ctx,
 	}
 
 	/* Track only once the kernel has the attributes, so a failed register
-	 * leaves nothing behind and cannot grow a same-base entry's extent to a
-	 * span that never got GPU access.
+	 * leaves nothing behind and there is no reserved extent to roll back.
 	 */
-	ret = svm_api_range_get(fmm_ctx, (void *)aligned_addr, aligned_size);
+	svm_api_range_get(fmm_ctx, address, (void *)aligned_addr, aligned_size, size);
+
+	ret = HSAKMT_STATUS_SUCCESS;
 
 out:
 	free(args);
@@ -3510,6 +3545,7 @@ void hsakmt_fmm_destroy_process_apertures(HsaKFDContext *ctx)
 		hsakmt_rbtree_delete(&fmm_ctx->svm_api_range_tree, n);
 		free(r);
 	}
+	fmm_ctx->svm_api_max_extent = 0;
 	pthread_mutex_unlock(&fmm_ctx->svm_api_mutex);
 
 	if (fmm_ctx->all_gpu_id_array) {

projects/rocr-runtime/libhsakmt/tests/kfdtest/CMakeLists.txt

@@ -274,6 +274,7 @@ set (SRC_FILES gtest-1.6.0/gtest-all.cpp
   src/KFDPMTest.cpp
   src/KFDSVMRangeTest.cpp
   src/KFDSVMRangeReGrant_test.cpp
+  src/KFDSVMRangePartialPage_test.cpp
   src/KFDSVMEvictTest.cpp
   src/KFDRASTest.cpp
   src/KFDPCSamplingTest.cpp

projects/rocr-runtime/libhsakmt/tests/kfdtest/src/KFDSVMRangePartialPage_test.cpp

@@ -0,0 +1,68 @@
+// Copyright © Advanced Micro Devices, Inc., or its affiliates.
+//
+// SPDX-License-Identifier: MIT
+
+#include "KFDSVMRangeTest.hpp"
+#include <sys/mman.h>
+#include <string.h>
+
+/*
+ * Registration rounds out to page boundaries, so a small buffer drags in bytes
+ * an allocator handed to something else. Deregister must not revoke those: on a
+ * kernel that unmaps on no-access the GPU faults on the co-tenant data.
+ */
+TEST_P(KFDSVMRangeTest, SVMApiPartialPageNotRevoked) {
+    TEST_REQUIRE_ENV_CAPABILITIES(ENVCAPS_64BITLINUX);
+    TEST_START(TESTPROFILE_RUNALL);
+
+    ASSERT_SUCCESS(KFDTestLaunch([this](int gpuNode) {
+        if (!SVMAPISupported_GPU(gpuNode)) {
+            LOG() << "Skipping test: SVM API not supported on gpuNode " << gpuNode << std::endl;
+            return;
+        }
+        if (GetFamilyIdFromNodeId(gpuNode) < FAMILY_AI) {
+            LOG() << "Skipping test: no svm range support on gpuNode " << gpuNode << std::endl;
+            return;
+        }
+        const HsaNodeProperties *pNodeProperties =
+            Get_NodeInfo()->GetNodeProperties(gpuNode);
+        if (pNodeProperties->Integrated || Get_NodeInfo()->IsAppAPU(gpuNode)) {
+            LOG() << "Skipping test on (App)APU gpuNode " << gpuNode << std::endl;
+            return;
+        }
+
+        const HSAuint64 PG = PAGE_SIZE;
+
+        auto access = [&](void *addr) -> HSAuint32 {
+            HSA_SVM_ATTRIBUTE attr = { HSA_SVM_ATTR_ACCESS, (HSAuint32)gpuNode };
+            EXPECT_SUCCESS_GPU(HSAKMT_CALL(hsaKmtSVMGetAttr, m_hsakmt_current_ctx,
+                                           addr, PG, 1, &attr), gpuNode);
+            return attr.type;
+        };
+
+        void *page = mmap(0, 2 * PG, PROT_READ | PROT_WRITE,
+                          MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
+        ASSERT_NE_GPU(MAP_FAILED, page, gpuNode);
+        memset(page, 0, 2 * PG);
+
+        /* Only 16 bytes are ever registered. The rest of that page stands in
+         * for the neighbours an allocator packs alongside it.
+         */
+        EXPECT_SUCCESS_GPU(HSAKMT_CALL(hsaKmtRegisterMemory, m_hsakmt_current_ctx,
+                                       page, 16), gpuNode);
+        EXPECT_SUCCESS_GPU(HSAKMT_CALL(hsaKmtDeregisterMemory, m_hsakmt_current_ctx,
+                                       page), gpuNode);
+        EXPECT_NE_GPU(HSA_SVM_ATTR_NO_ACCESS, access(page), gpuNode);
+
+        /* A whole page is ours, so that one does get revoked. */
+        EXPECT_SUCCESS_GPU(HSAKMT_CALL(hsaKmtRegisterMemory, m_hsakmt_current_ctx,
+                                       page, PG), gpuNode);
+        EXPECT_SUCCESS_GPU(HSAKMT_CALL(hsaKmtDeregisterMemory, m_hsakmt_current_ctx,
+                                       page), gpuNode);
+        EXPECT_EQ_GPU(HSA_SVM_ATTR_NO_ACCESS, access(page), gpuNode);
+
+        munmap(page, 2 * PG);
+    }));
+
+    TEST_END
+}