[PATCH v3 0/2] support large align and nid in Rust allocators

[PATCH v3 0/2] support large align and nid in Rust allocators

[Vitaly Wool]

The coming patches provide the ability for Rust allocators to set
NUMA node and large alignment.

Changelog:
v2 -> v3: fixed the build breakage for non-MMU configs

Signed-off-by: Vitaly Wool <vitaly.wool@konsulko.se>

[PATCH v3 1/2] mm/vmalloc: allow to set node and align in vrealloc

[Vitaly Wool]

Reimplement vrealloc() to be able to set node and alignment should
a user need to do so. Rename the function to vrealloc_node() to
better match what it actually does now and introduce a macro for
vrealloc() for backward compatibility.

With that change we also provide the ability for the Rust part of
the kernel to set node and aligmnent in its allocations.

Signed-off-by: Vitaly Wool <vitaly.wool@konsulko.se>
---
 include/linux/vmalloc.h |  8 +++++---
 mm/nommu.c              |  3 ++-
 mm/vmalloc.c            | 16 +++++++++++++---
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/include/linux/vmalloc.h b/include/linux/vmalloc.h
index fdc9aeb74a44..7d5251287687 100644
--- a/include/linux/vmalloc.h
+++ b/include/linux/vmalloc.h
@@ -197,9 +197,11 @@ extern void *__vcalloc_noprof(size_t n, size_t size, gfp_t flags) __alloc_size(1
 extern void *vcalloc_noprof(size_t n, size_t size) __alloc_size(1, 2);
 #define vcalloc(...)		alloc_hooks(vcalloc_noprof(__VA_ARGS__))
 
-void * __must_check vrealloc_noprof(const void *p, size_t size, gfp_t flags)
-		__realloc_size(2);
-#define vrealloc(...)		alloc_hooks(vrealloc_noprof(__VA_ARGS__))
+void *__must_check vrealloc_node_noprof(const void *p, size_t size,
+		unsigned long align, gfp_t flags, int nid) __realloc_size(2);
+#define vrealloc_noprof(p, s, f)	vrealloc_node_noprof(p, s, 1, f, NUMA_NO_NODE)
+#define vrealloc_node(...)		alloc_hooks(vrealloc_node_noprof(__VA_ARGS__))
+#define vrealloc(...)			alloc_hooks(vrealloc_noprof(__VA_ARGS__))
 
 extern void vfree(const void *addr);
 extern void vfree_atomic(const void *addr);
diff --git a/mm/nommu.c b/mm/nommu.c
index b624acec6d2e..ae240115607f 100644
--- a/mm/nommu.c
+++ b/mm/nommu.c
@@ -119,7 +119,8 @@ void *__vmalloc_noprof(unsigned long size, gfp_t gfp_mask)
 }
 EXPORT_SYMBOL(__vmalloc_noprof);
 
-void *vrealloc_noprof(const void *p, size_t size, gfp_t flags)
+void *vrealloc_node_noprof(const void *p, size_t size, unsigned long align,
+			   gfp_t flags, int nid)
 {
 	return krealloc_noprof(p, size, (flags | __GFP_COMP) & ~__GFP_HIGHMEM);
 }
diff --git a/mm/vmalloc.c b/mm/vmalloc.c
index ab986dd09b6a..117894301db1 100644
--- a/mm/vmalloc.c
+++ b/mm/vmalloc.c
@@ -4081,10 +4081,12 @@ void *vzalloc_node_noprof(unsigned long size, int node)
 EXPORT_SYMBOL(vzalloc_node_noprof);
 
 /**
- * vrealloc - reallocate virtually contiguous memory; contents remain unchanged
+ * vrealloc_node - reallocate virtually contiguous memory; contents remain unchanged
  * @p: object to reallocate memory for
  * @size: the size to reallocate
+ * @align: requested alignment
  * @flags: the flags for the page level allocator
+ * @nid: node id
  *
  * If @p is %NULL, vrealloc() behaves exactly like vmalloc(). If @size is 0 and
  * @p is not a %NULL pointer, the object pointed to is freed.
@@ -4103,7 +4105,7 @@ EXPORT_SYMBOL(vzalloc_node_noprof);
  * Return: pointer to the allocated memory; %NULL if @size is zero or in case of
  *         failure
  */
-void *vrealloc_noprof(const void *p, size_t size, gfp_t flags)
+void *vrealloc_node_noprof(const void *p, size_t size, unsigned long align, gfp_t flags, int nid)
 {
 	struct vm_struct *vm = NULL;
 	size_t alloced_size = 0;
@@ -4127,6 +4129,13 @@ void *vrealloc_noprof(const void *p, size_t size, gfp_t flags)
 		if (WARN(alloced_size < old_size,
 			 "vrealloc() has mismatched area vs requested sizes (%p)\n", p))
 			return NULL;
+		if (WARN(nid != NUMA_NO_NODE && nid != page_to_nid(vmalloc_to_page(p)),
+			 "vrealloc() has mismatched nids\n"))
+			return NULL;
+		if (WARN((uintptr_t)p & (align - 1),
+			 "will not reallocate with a bigger alignment (0x%lx)\n",
+			 align))
+			return NULL;
 	}
 
 	/*
@@ -4158,7 +4167,8 @@ void *vrealloc_noprof(const void *p, size_t size, gfp_t flags)
 	}
 
 	/* TODO: Grow the vm_area, i.e. allocate and map additional pages. */
-	n = __vmalloc_noprof(size, flags);
+	n = __vmalloc_node_noprof(size, align, flags, nid, __builtin_return_address(0));
+
 	if (!n)
 		return NULL;
 
-- 
2.39.2

[PATCH v3 2/2] rust: support align and NUMA id in allocations

[Vitaly Wool]

Add support for large (> PAGE_SIZE) alignments in Rust allocators
(Kmalloc support for large alignments is limited to the requested
size, which is a reasonable limitation anyway).
Besides, add support for NUMA id to Vmalloc.

Signed-off-by: Vitaly Wool <vitaly.wool@konsulko.se>
---
 rust/helpers/slab.c            |  8 +++++--
 rust/helpers/vmalloc.c         |  4 ++--
 rust/kernel/alloc.rs           | 28 ++++++++++++++++++++++--
 rust/kernel/alloc/allocator.rs | 40 +++++++++++++++++++---------------
 rust/kernel/alloc/kvec.rs      |  3 ++-
 5 files changed, 59 insertions(+), 24 deletions(-)

diff --git a/rust/helpers/slab.c b/rust/helpers/slab.c
index a842bfbddcba..221c517f57a1 100644
--- a/rust/helpers/slab.c
+++ b/rust/helpers/slab.c
@@ -3,13 +3,17 @@
 #include <linux/slab.h>
 
 void * __must_check __realloc_size(2)
-rust_helper_krealloc(const void *objp, size_t new_size, gfp_t flags)
+rust_helper_krealloc(const void *objp, size_t new_size, unsigned long align, gfp_t flags, int nid)
 {
+	if (WARN_ON(new_size & (align - 1)))
+		return NULL;
 	return krealloc(objp, new_size, flags);
 }
 
 void * __must_check __realloc_size(2)
-rust_helper_kvrealloc(const void *p, size_t size, gfp_t flags)
+rust_helper_kvrealloc(const void *p, size_t size, unsigned long align, gfp_t flags, int nid)
 {
+	if (WARN_ON(size & (align - 1)))
+		return NULL;
 	return kvrealloc(p, size, flags);
 }
diff --git a/rust/helpers/vmalloc.c b/rust/helpers/vmalloc.c
index 80d34501bbc0..9131279222fa 100644
--- a/rust/helpers/vmalloc.c
+++ b/rust/helpers/vmalloc.c
@@ -3,7 +3,7 @@
 #include <linux/vmalloc.h>
 
 void * __must_check __realloc_size(2)
-rust_helper_vrealloc(const void *p, size_t size, gfp_t flags)
+rust_helper_vrealloc_node(const void *p, size_t size, unsigned long align, gfp_t flags, int node)
 {
-	return vrealloc(p, size, flags);
+	return vrealloc_node(p, size, align, flags, node);
 }
diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
index 2e377c52fa07..12a723bf6092 100644
--- a/rust/kernel/alloc.rs
+++ b/rust/kernel/alloc.rs
@@ -161,7 +161,30 @@ pub unsafe trait Allocator {
     fn alloc(layout: Layout, flags: Flags) -> Result<NonNull<[u8]>, AllocError> {
         // SAFETY: Passing `None` to `realloc` is valid by its safety requirements and asks for a
         // new memory allocation.
-        unsafe { Self::realloc(None, layout, Layout::new::<()>(), flags) }
+        unsafe { Self::realloc(None, layout, Layout::new::<()>(), flags, None) }
+    }
+
+    /// Allocate memory based on `layout`, `flags` and `nid`.
+    ///
+    /// On success, returns a buffer represented as `NonNull<[u8]>` that satisfies the layout
+    /// constraints (i.e. minimum size and alignment as specified by `layout`).
+    ///
+    /// This function is equivalent to `realloc` when called with `None`.
+    ///
+    /// # Guarantees
+    ///
+    /// When the return value is `Ok(ptr)`, then `ptr` is
+    /// - valid for reads and writes for `layout.size()` bytes, until it is passed to
+    ///   [`Allocator::free`] or [`Allocator::realloc`],
+    /// - aligned to `layout.align()`,
+    ///
+    /// Additionally, `Flags` are honored as documented in
+    /// <https://docs.kernel.org/core-api/mm-api.html#mm-api-gfp-flags>.
+    fn alloc_node(layout: Layout, flags: Flags, nid: Option<i32>)
+                -> Result<NonNull<[u8]>, AllocError> {
+        // SAFETY: Passing `None` to `realloc` is valid by its safety requirements and asks for a
+        // new memory allocation.
+        unsafe { Self::realloc(None, layout, Layout::new::<()>(), flags, nid) }
     }
 
     /// Re-allocate an existing memory allocation to satisfy the requested `layout`.
@@ -201,6 +224,7 @@ unsafe fn realloc(
         layout: Layout,
         old_layout: Layout,
         flags: Flags,
+        nid: Option<i32>,
     ) -> Result<NonNull<[u8]>, AllocError>;
 
     /// Free an existing memory allocation.
@@ -216,7 +240,7 @@ unsafe fn free(ptr: NonNull<u8>, layout: Layout) {
         // SAFETY: The caller guarantees that `ptr` points at a valid allocation created by this
         // allocator. We are passing a `Layout` with the smallest possible alignment, so it is
         // smaller than or equal to the alignment previously used with this allocation.
-        let _ = unsafe { Self::realloc(Some(ptr), Layout::new::<()>(), layout, Flags(0)) };
+        let _ = unsafe { Self::realloc(Some(ptr), Layout::new::<()>(), layout, Flags(0), None) };
     }
 }
 
diff --git a/rust/kernel/alloc/allocator.rs b/rust/kernel/alloc/allocator.rs
index aa2dfa9dca4c..91b36e128b92 100644
--- a/rust/kernel/alloc/allocator.rs
+++ b/rust/kernel/alloc/allocator.rs
@@ -58,7 +58,8 @@ fn aligned_size(new_layout: Layout) -> usize {
 ///
 /// One of the following: `krealloc`, `vrealloc`, `kvrealloc`.
 struct ReallocFunc(
-    unsafe extern "C" fn(*const crate::ffi::c_void, usize, u32) -> *mut crate::ffi::c_void,
+    unsafe extern "C" fn(*const crate::ffi::c_void, usize, usize, u32, i32)
+                        -> *mut crate::ffi::c_void,
 );
 
 impl ReallocFunc {
@@ -66,7 +67,7 @@ impl ReallocFunc {
     const KREALLOC: Self = Self(bindings::krealloc);
 
     // INVARIANT: `vrealloc` satisfies the type invariants.
-    const VREALLOC: Self = Self(bindings::vrealloc);
+    const VREALLOC: Self = Self(bindings::vrealloc_node);
 
     // INVARIANT: `kvrealloc` satisfies the type invariants.
     const KVREALLOC: Self = Self(bindings::kvrealloc);
@@ -87,6 +88,7 @@ unsafe fn call(
         layout: Layout,
         old_layout: Layout,
         flags: Flags,
+        nid: Option<i32>,
     ) -> Result<NonNull<[u8]>, AllocError> {
         let size = aligned_size(layout);
         let ptr = match ptr {
@@ -100,6 +102,11 @@ unsafe fn call(
             None => ptr::null(),
         };
 
+        let c_nid = match nid {
+            None => bindings::NUMA_NO_NODE,
+            Some(n) => n,
+        };
+
         // SAFETY:
         // - `self.0` is one of `krealloc`, `vrealloc`, `kvrealloc` and thus only requires that
         //   `ptr` is NULL or valid.
@@ -110,7 +117,7 @@ unsafe fn call(
         // - Those functions provide the guarantees of this function.
         let raw_ptr = unsafe {
             // If `size == 0` and `ptr != NULL` the memory behind the pointer is freed.
-            self.0(ptr.cast(), size, flags.0).cast()
+            self.0(ptr.cast(), size, layout.align(), flags.0, c_nid).cast()
         };
 
         let ptr = if size == 0 {
@@ -134,9 +141,10 @@ unsafe fn realloc(
         layout: Layout,
         old_layout: Layout,
         flags: Flags,
+        _nid: Option<i32>,
     ) -> Result<NonNull<[u8]>, AllocError> {
         // SAFETY: `ReallocFunc::call` has the same safety requirements as `Allocator::realloc`.
-        unsafe { ReallocFunc::KREALLOC.call(ptr, layout, old_layout, flags) }
+        unsafe { ReallocFunc::KREALLOC.call(ptr, layout, old_layout, flags, None) }
     }
 }
 
@@ -151,16 +159,11 @@ unsafe fn realloc(
         layout: Layout,
         old_layout: Layout,
         flags: Flags,
+        nid: Option<i32>,
     ) -> Result<NonNull<[u8]>, AllocError> {
-        // TODO: Support alignments larger than PAGE_SIZE.
-        if layout.align() > bindings::PAGE_SIZE {
-            pr_warn!("Vmalloc does not support alignments larger than PAGE_SIZE yet.\n");
-            return Err(AllocError);
-        }
-
         // SAFETY: If not `None`, `ptr` is guaranteed to point to valid memory, which was previously
         // allocated with this `Allocator`.
-        unsafe { ReallocFunc::VREALLOC.call(ptr, layout, old_layout, flags) }
+        unsafe { ReallocFunc::VREALLOC.call(ptr, layout, old_layout, flags, nid) }
     }
 }
 
@@ -175,15 +178,18 @@ unsafe fn realloc(
         layout: Layout,
         old_layout: Layout,
         flags: Flags,
+        _nid: Option<i32>,
     ) -> Result<NonNull<[u8]>, AllocError> {
-        // TODO: Support alignments larger than PAGE_SIZE.
-        if layout.align() > bindings::PAGE_SIZE {
-            pr_warn!("KVmalloc does not support alignments larger than PAGE_SIZE yet.\n");
-            return Err(AllocError);
-        }
+        // if the caller wants to have alignment bigger than PAGE_SIZE
+        // it's only vmalloc we can offer to satisfy that request
+        let alloc_func = if layout.align() > bindings::PAGE_SIZE {
+            ReallocFunc::VREALLOC
+        } else {
+            ReallocFunc::KVREALLOC
+        };
 
         // SAFETY: If not `None`, `ptr` is guaranteed to point to valid memory, which was previously
         // allocated with this `Allocator`.
-        unsafe { ReallocFunc::KVREALLOC.call(ptr, layout, old_layout, flags) }
+        unsafe { alloc_func.call(ptr, layout, old_layout, flags, None) }
     }
 }
diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index 1a0dd852a468..ef4f977ba012 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -633,6 +633,7 @@ pub fn reserve(&mut self, additional: usize, flags: Flags) -> Result<(), AllocEr
                 layout.into(),
                 self.layout.into(),
                 flags,
+                None,
             )?
         };
 
@@ -1058,7 +1059,7 @@ pub fn collect(self, flags: Flags) -> Vec<T, A> {
             // the type invariant to be smaller than `cap`. Depending on `realloc` this operation
             // may shrink the buffer or leave it as it is.
             ptr = match unsafe {
-                A::realloc(Some(buf.cast()), layout.into(), old_layout.into(), flags)
+                A::realloc(Some(buf.cast()), layout.into(), old_layout.into(), flags, None)
             } {
                 // If we fail to shrink, which likely can't even happen, continue with the existing
                 // buffer.
-- 
2.39.2

Re: [PATCH v3 1/2] mm/vmalloc: allow to set node and align in vrealloc

[Uladzislau Rezki]

On Wed, Jun 25, 2025 at 08:30:13AM +0200, Vitaly Wool wrote:
> Reimplement vrealloc() to be able to set node and alignment should
> a user need to do so. Rename the function to vrealloc_node() to
> better match what it actually does now and introduce a macro for
> vrealloc() for backward compatibility.
> 
> With that change we also provide the ability for the Rust part of
> the kernel to set node and aligmnent in its allocations.
> 
> Signed-off-by: Vitaly Wool <vitaly.wool@konsulko.se>
> ---
>  include/linux/vmalloc.h |  8 +++++---
>  mm/nommu.c              |  3 ++-
>  mm/vmalloc.c            | 16 +++++++++++++---
>  3 files changed, 20 insertions(+), 7 deletions(-)
> 
> diff --git a/include/linux/vmalloc.h b/include/linux/vmalloc.h
> index fdc9aeb74a44..7d5251287687 100644
> --- a/include/linux/vmalloc.h
> +++ b/include/linux/vmalloc.h
> @@ -197,9 +197,11 @@ extern void *__vcalloc_noprof(size_t n, size_t size, gfp_t flags) __alloc_size(1
>  extern void *vcalloc_noprof(size_t n, size_t size) __alloc_size(1, 2);
>  #define vcalloc(...)		alloc_hooks(vcalloc_noprof(__VA_ARGS__))
>  
> -void * __must_check vrealloc_noprof(const void *p, size_t size, gfp_t flags)
> -		__realloc_size(2);
> -#define vrealloc(...)		alloc_hooks(vrealloc_noprof(__VA_ARGS__))
> +void *__must_check vrealloc_node_noprof(const void *p, size_t size,
> +		unsigned long align, gfp_t flags, int nid) __realloc_size(2);
> +#define vrealloc_noprof(p, s, f)	vrealloc_node_noprof(p, s, 1, f, NUMA_NO_NODE)
> +#define vrealloc_node(...)		alloc_hooks(vrealloc_node_noprof(__VA_ARGS__))
> +#define vrealloc(...)			alloc_hooks(vrealloc_noprof(__VA_ARGS__))
>  
>  extern void vfree(const void *addr);
>  extern void vfree_atomic(const void *addr);
> diff --git a/mm/nommu.c b/mm/nommu.c
> index b624acec6d2e..ae240115607f 100644
> --- a/mm/nommu.c
> +++ b/mm/nommu.c
> @@ -119,7 +119,8 @@ void *__vmalloc_noprof(unsigned long size, gfp_t gfp_mask)
>  }
>  EXPORT_SYMBOL(__vmalloc_noprof);
>  
> -void *vrealloc_noprof(const void *p, size_t size, gfp_t flags)
> +void *vrealloc_node_noprof(const void *p, size_t size, unsigned long align,
> +			   gfp_t flags, int nid)
>  {
>  	return krealloc_noprof(p, size, (flags | __GFP_COMP) & ~__GFP_HIGHMEM);
>  }
> diff --git a/mm/vmalloc.c b/mm/vmalloc.c
> index ab986dd09b6a..117894301db1 100644
> --- a/mm/vmalloc.c
> +++ b/mm/vmalloc.c
> @@ -4081,10 +4081,12 @@ void *vzalloc_node_noprof(unsigned long size, int node)
>  EXPORT_SYMBOL(vzalloc_node_noprof);
>  
>  /**
> - * vrealloc - reallocate virtually contiguous memory; contents remain unchanged
> + * vrealloc_node - reallocate virtually contiguous memory; contents remain unchanged
>   * @p: object to reallocate memory for
>   * @size: the size to reallocate
> + * @align: requested alignment
>   * @flags: the flags for the page level allocator
> + * @nid: node id
>   *
>   * If @p is %NULL, vrealloc() behaves exactly like vmalloc(). If @size is 0 and
>   * @p is not a %NULL pointer, the object pointed to is freed.
> @@ -4103,7 +4105,7 @@ EXPORT_SYMBOL(vzalloc_node_noprof);
>   * Return: pointer to the allocated memory; %NULL if @size is zero or in case of
>   *         failure
>   */
> -void *vrealloc_noprof(const void *p, size_t size, gfp_t flags)
> +void *vrealloc_node_noprof(const void *p, size_t size, unsigned long align, gfp_t flags, int nid)
>  {
>  	struct vm_struct *vm = NULL;
>  	size_t alloced_size = 0;
> @@ -4127,6 +4129,13 @@ void *vrealloc_noprof(const void *p, size_t size, gfp_t flags)
>  		if (WARN(alloced_size < old_size,
>  			 "vrealloc() has mismatched area vs requested sizes (%p)\n", p))
>  			return NULL;
> +		if (WARN(nid != NUMA_NO_NODE && nid != page_to_nid(vmalloc_to_page(p)),
> +			 "vrealloc() has mismatched nids\n"))
> +			return NULL;
> +		if (WARN((uintptr_t)p & (align - 1),
> +			 "will not reallocate with a bigger alignment (0x%lx)\n",
> +			 align))
> +			return NULL;
>  	}
>  
>  	/*
> @@ -4158,7 +4167,8 @@ void *vrealloc_noprof(const void *p, size_t size, gfp_t flags)
>  	}
>  
>  	/* TODO: Grow the vm_area, i.e. allocate and map additional pages. */
> -	n = __vmalloc_noprof(size, flags);
> +	n = __vmalloc_node_noprof(size, align, flags, nid, __builtin_return_address(0));
> +
>  	if (!n)
>  		return NULL;
>  
> -- 
> 2.39.2
> 
Reviewed-by: "Uladzislau Rezki (Sony)" <urezki@gmail.com>

--
Uladzislau Rezki

Re: [PATCH v3 2/2] rust: support align and NUMA id in allocations

[Boqun Feng]

On Wed, Jun 25, 2025 at 08:30:26AM +0200, Vitaly Wool wrote:
[...]
> @@ -151,16 +159,11 @@ unsafe fn realloc(
>          layout: Layout,
>          old_layout: Layout,
>          flags: Flags,
> +        nid: Option<i32>,
>      ) -> Result<NonNull<[u8]>, AllocError> {
> -        // TODO: Support alignments larger than PAGE_SIZE.

Thanks a lot for doing this! While you're at it, maybe we can add a few
tests for various alignments of allocation? I'm thinking:

#[repr(align(65536)]
pub struct Test64k {
    a: i32
}

#[kunit_tests(rust_vbox)]
mod tests {
    #[test]
    fn large_allocation() -> Result {
        // Better use `new_uninit()` to avoid allocation on the stack.
        let x = VBox::<Test64k>::new_uninit(...)?;

	assert!(x.as_ptr().addr() & (kernel::sizes::SZ_64K - 1) == 0);
    }
}

Thoughts?

Regards,
Boqun

> -        if layout.align() > bindings::PAGE_SIZE {
> -            pr_warn!("Vmalloc does not support alignments larger than PAGE_SIZE yet.\n");
> -            return Err(AllocError);
> -        }
> -
>          // SAFETY: If not `None`, `ptr` is guaranteed to point to valid memory, which was previously
>          // allocated with this `Allocator`.
> -        unsafe { ReallocFunc::VREALLOC.call(ptr, layout, old_layout, flags) }
> +        unsafe { ReallocFunc::VREALLOC.call(ptr, layout, old_layout, flags, nid) }
>      }
>  }
>  
[...]

Re: [PATCH v3 2/2] rust: support align and NUMA id in allocations

[Uladzislau Rezki]

On Wed, Jun 25, 2025 at 06:12:52AM -0700, Boqun Feng wrote:
> On Wed, Jun 25, 2025 at 08:30:26AM +0200, Vitaly Wool wrote:
> [...]
> > @@ -151,16 +159,11 @@ unsafe fn realloc(
> >          layout: Layout,
> >          old_layout: Layout,
> >          flags: Flags,
> > +        nid: Option<i32>,
> >      ) -> Result<NonNull<[u8]>, AllocError> {
> > -        // TODO: Support alignments larger than PAGE_SIZE.
> 
> Thanks a lot for doing this! While you're at it, maybe we can add a few
> tests for various alignments of allocation? I'm thinking:
> 
> #[repr(align(65536)]
> pub struct Test64k {
>     a: i32
> }
> 
> #[kunit_tests(rust_vbox)]
> mod tests {
>     #[test]
>     fn large_allocation() -> Result {
>         // Better use `new_uninit()` to avoid allocation on the stack.
>         let x = VBox::<Test64k>::new_uninit(...)?;
> 
> 	assert!(x.as_ptr().addr() & (kernel::sizes::SZ_64K - 1) == 0);
>     }
> }
> 
> Thoughts?
> 
> Regards,
> Boqun
> 
> > -        if layout.align() > bindings::PAGE_SIZE {
> > -            pr_warn!("Vmalloc does not support alignments larger than PAGE_SIZE yet.\n");
> > -            return Err(AllocError);
> > -        }
> > -
> >          // SAFETY: If not `None`, `ptr` is guaranteed to point to valid memory, which was previously
> >          // allocated with this `Allocator`.
> > -        unsafe { ReallocFunc::VREALLOC.call(ptr, layout, old_layout, flags) }
> > +        unsafe { ReallocFunc::VREALLOC.call(ptr, layout, old_layout, flags, nid) }
> >      }
> >  }
> >  
> [...]
>
At least we are lacking of vrealloc() exercising in the vmalloc-test suite.
I am not sure it makes a lot of sense to add a kunit test on top of rust-wrapper
around vrealloc().

From my side, i will add the test case to the test_vmalloc.c test-suite.

--
Uladzislau Rezki

Re: [PATCH v3 2/2] rust: support align and NUMA id in allocations

[Boqun Feng]

On Wed, Jun 25, 2025 at 06:07:06PM +0200, Uladzislau Rezki wrote:
> On Wed, Jun 25, 2025 at 06:12:52AM -0700, Boqun Feng wrote:
> > On Wed, Jun 25, 2025 at 08:30:26AM +0200, Vitaly Wool wrote:
> > [...]
> > > @@ -151,16 +159,11 @@ unsafe fn realloc(
> > >          layout: Layout,
> > >          old_layout: Layout,
> > >          flags: Flags,
> > > +        nid: Option<i32>,
> > >      ) -> Result<NonNull<[u8]>, AllocError> {
> > > -        // TODO: Support alignments larger than PAGE_SIZE.
> > 
> > Thanks a lot for doing this! While you're at it, maybe we can add a few
> > tests for various alignments of allocation? I'm thinking:
> > 
> > #[repr(align(65536)]
> > pub struct Test64k {
> >     a: i32
> > }
> > 
> > #[kunit_tests(rust_vbox)]
> > mod tests {
> >     #[test]
> >     fn large_allocation() -> Result {
> >         // Better use `new_uninit()` to avoid allocation on the stack.
> >         let x = VBox::<Test64k>::new_uninit(...)?;
> > 
> > 	assert!(x.as_ptr().addr() & (kernel::sizes::SZ_64K - 1) == 0);
> >     }
> > }
> > 
> > Thoughts?
> > 
> > Regards,
> > Boqun
> > 
> > > -        if layout.align() > bindings::PAGE_SIZE {
> > > -            pr_warn!("Vmalloc does not support alignments larger than PAGE_SIZE yet.\n");
> > > -            return Err(AllocError);
> > > -        }
> > > -
> > >          // SAFETY: If not `None`, `ptr` is guaranteed to point to valid memory, which was previously
> > >          // allocated with this `Allocator`.
> > > -        unsafe { ReallocFunc::VREALLOC.call(ptr, layout, old_layout, flags) }
> > > +        unsafe { ReallocFunc::VREALLOC.call(ptr, layout, old_layout, flags, nid) }
> > >      }
> > >  }
> > >  
> > [...]
> >
> At least we are lacking of vrealloc() exercising in the vmalloc-test suite.
> I am not sure it makes a lot of sense to add a kunit test on top of rust-wrapper
> around vrealloc().
> 
> From my side, i will add the test case to the test_vmalloc.c test-suite.
> 

Thanks! But we will need these tests from Rust side anyway, to test
1) whether the Rust wrapper does the right thing, and 2) whether any C
change cause the behavior changes on the API that Rust wrapper rely on.

Regards,
Boqun

> --
> Uladzislau Rezki

Re: [PATCH v3 2/2] rust: support align and NUMA id in allocations

[Uladzislau Rezki]

On Wed, Jun 25, 2025 at 09:12:21AM -0700, Boqun Feng wrote:
> On Wed, Jun 25, 2025 at 06:07:06PM +0200, Uladzislau Rezki wrote:
> > On Wed, Jun 25, 2025 at 06:12:52AM -0700, Boqun Feng wrote:
> > > On Wed, Jun 25, 2025 at 08:30:26AM +0200, Vitaly Wool wrote:
> > > [...]
> > > > @@ -151,16 +159,11 @@ unsafe fn realloc(
> > > >          layout: Layout,
> > > >          old_layout: Layout,
> > > >          flags: Flags,
> > > > +        nid: Option<i32>,
> > > >      ) -> Result<NonNull<[u8]>, AllocError> {
> > > > -        // TODO: Support alignments larger than PAGE_SIZE.
> > > 
> > > Thanks a lot for doing this! While you're at it, maybe we can add a few
> > > tests for various alignments of allocation? I'm thinking:
> > > 
> > > #[repr(align(65536)]
> > > pub struct Test64k {
> > >     a: i32
> > > }
> > > 
> > > #[kunit_tests(rust_vbox)]
> > > mod tests {
> > >     #[test]
> > >     fn large_allocation() -> Result {
> > >         // Better use `new_uninit()` to avoid allocation on the stack.
> > >         let x = VBox::<Test64k>::new_uninit(...)?;
> > > 
> > > 	assert!(x.as_ptr().addr() & (kernel::sizes::SZ_64K - 1) == 0);
> > >     }
> > > }
> > > 
> > > Thoughts?
> > > 
> > > Regards,
> > > Boqun
> > > 
> > > > -        if layout.align() > bindings::PAGE_SIZE {
> > > > -            pr_warn!("Vmalloc does not support alignments larger than PAGE_SIZE yet.\n");
> > > > -            return Err(AllocError);
> > > > -        }
> > > > -
> > > >          // SAFETY: If not `None`, `ptr` is guaranteed to point to valid memory, which was previously
> > > >          // allocated with this `Allocator`.
> > > > -        unsafe { ReallocFunc::VREALLOC.call(ptr, layout, old_layout, flags) }
> > > > +        unsafe { ReallocFunc::VREALLOC.call(ptr, layout, old_layout, flags, nid) }
> > > >      }
> > > >  }
> > > >  
> > > [...]
> > >
> > At least we are lacking of vrealloc() exercising in the vmalloc-test suite.
> > I am not sure it makes a lot of sense to add a kunit test on top of rust-wrapper
> > around vrealloc().
> > 
> > From my side, i will add the test case to the test_vmalloc.c test-suite.
> > 
> 
> Thanks! But we will need these tests from Rust side anyway, to test
> 1) whether the Rust wrapper does the right thing, and 2) whether any C
> change cause the behavior changes on the API that Rust wrapper rely on.
> 
Ah. Got it :)

--
Uladzislau Rezki

Re: [PATCH v3 2/2] rust: support align and NUMA id in allocations

[Danilo Krummrich]

On Wed, Jun 25, 2025 at 08:30:26AM +0200, Vitaly Wool wrote:
> Add support for large (> PAGE_SIZE) alignments in Rust allocators
> (Kmalloc support for large alignments is limited to the requested
> size, which is a reasonable limitation anyway).

Please split this..

> Besides, add support for NUMA id to Vmalloc.

and this into separate patches.

Please also add some information to the commit message what you need node
support for. Do you also have patches to add node support to Box and Vec?

> 
> Signed-off-by: Vitaly Wool <vitaly.wool@konsulko.se>
> ---
>  rust/helpers/slab.c            |  8 +++++--
>  rust/helpers/vmalloc.c         |  4 ++--
>  rust/kernel/alloc.rs           | 28 ++++++++++++++++++++++--
>  rust/kernel/alloc/allocator.rs | 40 +++++++++++++++++++---------------
>  rust/kernel/alloc/kvec.rs      |  3 ++-
>  5 files changed, 59 insertions(+), 24 deletions(-)
> 
> diff --git a/rust/helpers/slab.c b/rust/helpers/slab.c
> index a842bfbddcba..221c517f57a1 100644
> --- a/rust/helpers/slab.c
> +++ b/rust/helpers/slab.c
> @@ -3,13 +3,17 @@
>  #include <linux/slab.h>
>  
>  void * __must_check __realloc_size(2)
> -rust_helper_krealloc(const void *objp, size_t new_size, gfp_t flags)
> +rust_helper_krealloc(const void *objp, size_t new_size, unsigned long align, gfp_t flags, int nid)

This should have a comment making it obvious why the function has two arguments
that are discarded. I think we should even separate it with an additional inline
function.

I do agree with discarding the align argument, given that it's not exposed to
users though the Allocator API.

I do disagree with discarding the nid argument though, since you change the
generic Allocator::realloc() API to take a node argument, which for KREALLOC and
KVREALLOC is silently discarded. If we introduce it, we should do so for all
three allocators.

>  {
> +	if (WARN_ON(new_size & (align - 1)))
> +		return NULL;

I don't think we should have this WARN_ON(). If we want to warn about this, we
should already do so on the Rust side. The helper functions in this file should
not contain any logic.

>  	return krealloc(objp, new_size, flags);
>  }
>  
>  void * __must_check __realloc_size(2)
> -rust_helper_kvrealloc(const void *p, size_t size, gfp_t flags)
> +rust_helper_kvrealloc(const void *p, size_t size, unsigned long align, gfp_t flags, int nid)
>  {
> +	if (WARN_ON(size & (align - 1)))
> +		return NULL;
>  	return kvrealloc(p, size, flags);
>  }

Same as above.

> diff --git a/rust/helpers/vmalloc.c b/rust/helpers/vmalloc.c
> index 80d34501bbc0..9131279222fa 100644
> --- a/rust/helpers/vmalloc.c
> +++ b/rust/helpers/vmalloc.c
> @@ -3,7 +3,7 @@
>  #include <linux/vmalloc.h>
>  
>  void * __must_check __realloc_size(2)
> -rust_helper_vrealloc(const void *p, size_t size, gfp_t flags)
> +rust_helper_vrealloc_node(const void *p, size_t size, unsigned long align, gfp_t flags, int node)
>  {
> -	return vrealloc(p, size, flags);
> +	return vrealloc_node(p, size, align, flags, node);
>  }
> diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
> index 2e377c52fa07..12a723bf6092 100644
> --- a/rust/kernel/alloc.rs
> +++ b/rust/kernel/alloc.rs
> @@ -161,7 +161,30 @@ pub unsafe trait Allocator {
>      fn alloc(layout: Layout, flags: Flags) -> Result<NonNull<[u8]>, AllocError> {
>          // SAFETY: Passing `None` to `realloc` is valid by its safety requirements and asks for a
>          // new memory allocation.
> -        unsafe { Self::realloc(None, layout, Layout::new::<()>(), flags) }
> +        unsafe { Self::realloc(None, layout, Layout::new::<()>(), flags, None) }
> +    }
> +
> +    /// Allocate memory based on `layout`, `flags` and `nid`.
> +    ///
> +    /// On success, returns a buffer represented as `NonNull<[u8]>` that satisfies the layout
> +    /// constraints (i.e. minimum size and alignment as specified by `layout`).
> +    ///
> +    /// This function is equivalent to `realloc` when called with `None`.
> +    ///
> +    /// # Guarantees
> +    ///
> +    /// When the return value is `Ok(ptr)`, then `ptr` is
> +    /// - valid for reads and writes for `layout.size()` bytes, until it is passed to
> +    ///   [`Allocator::free`] or [`Allocator::realloc`],
> +    /// - aligned to `layout.align()`,
> +    ///
> +    /// Additionally, `Flags` are honored as documented in
> +    /// <https://docs.kernel.org/core-api/mm-api.html#mm-api-gfp-flags>.
> +    fn alloc_node(layout: Layout, flags: Flags, nid: Option<i32>)
> +                -> Result<NonNull<[u8]>, AllocError> {
> +        // SAFETY: Passing `None` to `realloc` is valid by its safety requirements and asks for a
> +        // new memory allocation.
> +        unsafe { Self::realloc(None, layout, Layout::new::<()>(), flags, nid) }
>      }
>  
>      /// Re-allocate an existing memory allocation to satisfy the requested `layout`.
> @@ -201,6 +224,7 @@ unsafe fn realloc(
>          layout: Layout,
>          old_layout: Layout,
>          flags: Flags,
> +        nid: Option<i32>,
>      ) -> Result<NonNull<[u8]>, AllocError>;

I think we should rename this to realloc_node() and introduce realloc(), which
just calls realloc_node() with None.

Re: [PATCH v3 2/2] rust: support align and NUMA id in allocations

[Danilo Krummrich]

On Wed, Jun 25, 2025 at 08:56:05PM +0200, Danilo Krummrich wrote:
> On Wed, Jun 25, 2025 at 08:30:26AM +0200, Vitaly Wool wrote:
> > Add support for large (> PAGE_SIZE) alignments in Rust allocators
> > (Kmalloc support for large alignments is limited to the requested
> > size, which is a reasonable limitation anyway).
> 
> Please split this..
> 
> > Besides, add support for NUMA id to Vmalloc.
> 
> and this into separate patches.
> 
> Please also add some information to the commit message what you need node
> support for. Do you also have patches to add node support to Box and Vec?
> 
> > 
> > Signed-off-by: Vitaly Wool <vitaly.wool@konsulko.se>
> > ---
> >  rust/helpers/slab.c            |  8 +++++--
> >  rust/helpers/vmalloc.c         |  4 ++--
> >  rust/kernel/alloc.rs           | 28 ++++++++++++++++++++++--
> >  rust/kernel/alloc/allocator.rs | 40 +++++++++++++++++++---------------
> >  rust/kernel/alloc/kvec.rs      |  3 ++-
> >  5 files changed, 59 insertions(+), 24 deletions(-)
> > 
> > diff --git a/rust/helpers/slab.c b/rust/helpers/slab.c
> > index a842bfbddcba..221c517f57a1 100644
> > --- a/rust/helpers/slab.c
> > +++ b/rust/helpers/slab.c
> > @@ -3,13 +3,17 @@
> >  #include <linux/slab.h>
> >  
> >  void * __must_check __realloc_size(2)
> > -rust_helper_krealloc(const void *objp, size_t new_size, gfp_t flags)
> > +rust_helper_krealloc(const void *objp, size_t new_size, unsigned long align, gfp_t flags, int nid)
> 
> This should have a comment making it obvious why the function has two arguments
> that are discarded. I think we should even separate it with an additional inline
> function.
> 
> I do agree with discarding the align argument, given that it's not exposed to
> users though the Allocator API.

What I meant is that proper alignment is implied when krealloc() succeeds.

> I do disagree with discarding the nid argument though, since you change the
> generic Allocator::realloc() API to take a node argument, which for KREALLOC and
> KVREALLOC is silently discarded. If we introduce it, we should do so for all
> three allocators.
> 
> >  {
> > +	if (WARN_ON(new_size & (align - 1)))
> > +		return NULL;
> 
> I don't think we should have this WARN_ON(). If we want to warn about this, we
> should already do so on the Rust side. The helper functions in this file should
> not contain any logic.
> 
> >  	return krealloc(objp, new_size, flags);
> >  }
> >  
> >  void * __must_check __realloc_size(2)
> > -rust_helper_kvrealloc(const void *p, size_t size, gfp_t flags)
> > +rust_helper_kvrealloc(const void *p, size_t size, unsigned long align, gfp_t flags, int nid)
> >  {
> > +	if (WARN_ON(size & (align - 1)))
> > +		return NULL;
> >  	return kvrealloc(p, size, flags);
> >  }
> 
> Same as above.

This is actually different though, here kvrealloc() may succeed even if the
requested alignment is not fulfilled, so this is incorrect.

Re: [PATCH v3 2/2] rust: support align and NUMA id in allocations

[Vitaly Wool]

> On Jun 25, 2025, at 9:07 PM, Danilo Krummrich <dakr@kernel.org> wrote:
> 
> On Wed, Jun 25, 2025 at 08:56:05PM +0200, Danilo Krummrich wrote:
>> On Wed, Jun 25, 2025 at 08:30:26AM +0200, Vitaly Wool wrote:
>>> Add support for large (> PAGE_SIZE) alignments in Rust allocators
>>> (Kmalloc support for large alignments is limited to the requested
>>> size, which is a reasonable limitation anyway).
>> 
>> Please split this..
>> 
>>> Besides, add support for NUMA id to Vmalloc.
>> 
>> and this into separate patches.
>> 
>> Please also add some information to the commit message what you need node
>> support for. Do you also have patches to add node support to Box and Vec?

No, but there is a zswap backend implementation written in Rust and it should be  NUMA id aware.
I’m planning on submitting that basically as soon as this piece gets accepted.

>> 
>>> 
>>> Signed-off-by: Vitaly Wool <vitaly.wool@konsulko.se>
>>> ---
>>> rust/helpers/slab.c            |  8 +++++--
>>> rust/helpers/vmalloc.c         |  4 ++--
>>> rust/kernel/alloc.rs           | 28 ++++++++++++++++++++++--
>>> rust/kernel/alloc/allocator.rs | 40 +++++++++++++++++++---------------
>>> rust/kernel/alloc/kvec.rs      |  3 ++-
>>> 5 files changed, 59 insertions(+), 24 deletions(-)
>>> 
>>> diff --git a/rust/helpers/slab.c b/rust/helpers/slab.c
>>> index a842bfbddcba..221c517f57a1 100644
>>> --- a/rust/helpers/slab.c
>>> +++ b/rust/helpers/slab.c
>>> @@ -3,13 +3,17 @@
>>> #include <linux/slab.h>
>>> 
>>> void * __must_check __realloc_size(2)
>>> -rust_helper_krealloc(const void *objp, size_t new_size, gfp_t flags)
>>> +rust_helper_krealloc(const void *objp, size_t new_size, unsigned long align, gfp_t flags, int nid)
>> 
>> This should have a comment making it obvious why the function has two arguments
>> that are discarded. I think we should even separate it with an additional inline
>> function.
>> 
>> I do agree with discarding the align argument, given that it's not exposed to
>> users though the Allocator API.
> 
> What I meant is that proper alignment is implied when krealloc() succeeds.

I agree, I need to add some comments explaining this.

> 
>> I do disagree with discarding the nid argument though, since you change the
>> generic Allocator::realloc() API to take a node argument, which for KREALLOC and
>> KVREALLOC is silently discarded. If we introduce it, we should do so for all
>> three allocators.
>> 
>>> {
>>> + if (WARN_ON(new_size & (align - 1)))
>>> + return NULL;
>> 
>> I don't think we should have this WARN_ON(). If we want to warn about this, we
>> should already do so on the Rust side. The helper functions in this file should
>> not contain any logic.

Agreed.

>> 
>>> return krealloc(objp, new_size, flags);
>>> }
>>> 
>>> void * __must_check __realloc_size(2)
>>> -rust_helper_kvrealloc(const void *p, size_t size, gfp_t flags)
>>> +rust_helper_kvrealloc(const void *p, size_t size, unsigned long align, gfp_t flags, int nid)
>>> {
>>> + if (WARN_ON(size & (align - 1)))
>>> + return NULL;
>>> return kvrealloc(p, size, flags);
>>> }
>> 
>> Same as above.
> 
> This is actually different though, here kvrealloc() may succeed even if the
> requested alignment is not fulfilled, so this is incorrect.

I can move this logic to the Rust part, too. My point here is, for Kvrealloc with a large alignment we’ll just make the decision to use vmalloc, period. We can indeed do that on the Rust side.

~Vitaly

Re: [PATCH v3 0/2] support large align and nid in Rust allocators

[Andrew Morton]

On Wed, 25 Jun 2025 08:29:17 +0200 Vitaly Wool <vitaly.wool@konsulko.se> wrote:

> The coming patches provide the ability for Rust allocators to set
> NUMA node and large alignment.

Including a diffstat would be nice.

 include/linux/vmalloc.h        |    8 +++---
 mm/nommu.c                     |    3 +-
 mm/vmalloc.c                   |   16 ++++++++++--
 rust/helpers/slab.c            |    8 ++++--
 rust/helpers/vmalloc.c         |    4 +--
 rust/kernel/alloc.rs           |   28 ++++++++++++++++++++-
 rust/kernel/alloc/allocator.rs |   40 +++++++++++++++++--------------
 rust/kernel/alloc/kvec.rs      |    3 +-
 8 files changed, 79 insertions(+), 31 deletions(-)

What is the preferred merge path for this series?

Re: [PATCH v3 0/2] support large align and nid in Rust allocators

[Danilo Krummrich]

On Wed, Jun 25, 2025 at 02:10:22PM -0700, Andrew Morton wrote:
> On Wed, 25 Jun 2025 08:29:17 +0200 Vitaly Wool <vitaly.wool@konsulko.se> wrote:
> 
> > The coming patches provide the ability for Rust allocators to set
> > NUMA node and large alignment.
> 
> Including a diffstat would be nice.
> 
>  include/linux/vmalloc.h        |    8 +++---
>  mm/nommu.c                     |    3 +-
>  mm/vmalloc.c                   |   16 ++++++++++--
>  rust/helpers/slab.c            |    8 ++++--
>  rust/helpers/vmalloc.c         |    4 +--
>  rust/kernel/alloc.rs           |   28 ++++++++++++++++++++-
>  rust/kernel/alloc/allocator.rs |   40 +++++++++++++++++--------------
>  rust/kernel/alloc/kvec.rs      |    3 +-
>  8 files changed, 79 insertions(+), 31 deletions(-)
> 
> What is the preferred merge path for this series?

I can take it through my Rust alloc tree, but I think it would probably be
better you take it through the mm tree, since I don't expect this to conflict
with anything else in my tree anyways.

The Rust changes gonna need some more discussion though.

- Danilo

Re: [PATCH v3 2/2] rust: support align and NUMA id in allocations

[Danilo Krummrich]

On Wed, Jun 25, 2025 at 10:22:36PM +0200, Vitaly Wool wrote:
> I can move this logic to the Rust part, too. My point here is, for Kvrealloc with a large alignment we’ll just make the decision to use vmalloc, period. We can indeed do that on the Rust side.

That's fine with me.

But please also make sure to properly support NUMA nodes for all allocator
backends.

Re: [PATCH v3 0/2] support large align and nid in Rust allocators

[Andrew Morton]

On Wed, 25 Jun 2025 23:15:46 +0200 Danilo Krummrich <dakr@kernel.org> wrote:

> > What is the preferred merge path for this series?
> 
> I can take it through my Rust alloc tree, but I think it would probably be
> better you take it through the mm tree, since I don't expect this to conflict
> with anything else in my tree anyways.

No probs.

> The Rust changes gonna need some more discussion though.

Great, I'll add a note-to-self that discussion is ongoing.

Re: [PATCH v3 0/2] support large align and nid in Rust allocators

[Vitaly Wool]

> On Jun 25, 2025, at 11:34 PM, Andrew Morton <akpm@linux-foundation.org> wrote:
> 
> On Wed, 25 Jun 2025 23:15:46 +0200 Danilo Krummrich <dakr@kernel.org> wrote:
> 
>>> What is the preferred merge path for this series?
>> 
>> I can take it through my Rust alloc tree, but I think it would probably be
>> better you take it through the mm tree, since I don't expect this to conflict
>> with anything else in my tree anyways.
> 
> No probs.
> 
>> The Rust changes gonna need some more discussion though.
> 
> Great, I'll add a note-to-self that discussion is ongoing.
> 

Awesome, thank you very much! The first patch is unlikely to need a change, the second one will be reworked as discussed with Danilo and the plan is that I post the new version tomorrow or on Friday.

~Vitaly