[kvm-unit-tests PATCH v1 4/4] lib/vmalloc: allow vm_memalign with alignment > PAGE_SIZE

[Claudio Imbrenda <imbrenda@xxxxxxxxxxxxx>]

Allow allocating aligned virtual memory with alignment larger than only
one page.

Add a check that the backing pages were actually allocated.

Export the alloc_vpages_aligned function to allow users to allocate
non-backed aligned virtual addresses.

Signed-off-by: Claudio Imbrenda <imbrenda@xxxxxxxxxxxxx>
---
 lib/vmalloc.h |  3 +++
 lib/vmalloc.c | 40 ++++++++++++++++++++++++++++++++--------
 2 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/lib/vmalloc.h b/lib/vmalloc.h
index 2b563f4..fa3fa22 100644
--- a/lib/vmalloc.h
+++ b/lib/vmalloc.h
@@ -5,6 +5,9 @@
 
 /* Allocate consecutive virtual pages (without backing) */
 extern void *alloc_vpages(ulong nr);
+/* Allocate consecutive and aligned virtual pages (without backing) */
+extern void *alloc_vpages_aligned(ulong nr, unsigned int alignment_pages);
+
 /* Allocate one virtual page (without backing) */
 extern void *alloc_vpage(void);
 /* Set the top of the virtual address space */
diff --git a/lib/vmalloc.c b/lib/vmalloc.c
index 9237a0f..2c482aa 100644
--- a/lib/vmalloc.c
+++ b/lib/vmalloc.c
@@ -12,19 +12,28 @@
 #include "alloc.h"
 #include "alloc_phys.h"
 #include "alloc_page.h"
+#include <bitops.h>
 #include "vmalloc.h"
 
 static struct spinlock lock;
 static void *vfree_top = 0;
 static void *page_root;
 
-void *alloc_vpages(ulong nr)
+/*
+ * Allocate a certain number of pages from the virtual address space (without
+ * physical backing).
+ *
+ * nr is the number of pages to allocate
+ * alignment_pages is the alignment of the allocation *in pages*
+ */
+static void *alloc_vpages_intern(ulong nr, unsigned int alignment_pages)
 {
 	uintptr_t ptr;
 
 	spin_lock(&lock);
 	ptr = (uintptr_t)vfree_top;
 	ptr -= PAGE_SIZE * nr;
+	ptr &= GENMASK_ULL(63, PAGE_SHIFT + get_order(alignment_pages));
 	vfree_top = (void *)ptr;
 	spin_unlock(&lock);
 
@@ -32,6 +41,16 @@ void *alloc_vpages(ulong nr)
 	return (void *)ptr;
 }
 
+void *alloc_vpages(ulong nr)
+{
+	return alloc_vpages_intern(nr, 1);
+}
+
+void *alloc_vpages_aligned(ulong nr, unsigned int alignment_pages)
+{
+	return alloc_vpages_intern(nr, alignment_pages);
+}
+
 void *alloc_vpage(void)
 {
 	return alloc_vpages(1);
@@ -55,17 +74,22 @@ void *vmap(phys_addr_t phys, size_t size)
 	return mem;
 }
 
+/*
+ * Allocate virtual memory, with the specified minimum alignment.
+ */
 static void *vm_memalign(size_t alignment, size_t size)
 {
+	phys_addr_t pa;
 	void *mem, *p;
-	size_t pages;
 
-	assert(alignment <= PAGE_SIZE);
-	size = PAGE_ALIGN(size);
-	pages = size / PAGE_SIZE;
-	mem = p = alloc_vpages(pages);
-	while (pages--) {
-		phys_addr_t pa = virt_to_phys(alloc_page());
+	assert(is_power_of_2(alignment));
+
+	size = PAGE_ALIGN(size) / PAGE_SIZE;
+	alignment = PAGE_ALIGN(alignment) / PAGE_SIZE;
+	mem = p = alloc_vpages_intern(size, alignment);
+	while (size--) {
+		pa = virt_to_phys(alloc_page());
+		assert(pa);
 		install_page(page_root, pa, p);
 		p += PAGE_SIZE;
 	}
-- 
2.26.2