修复流Malloc和Free为Context操作

Cargo.toml

@@ -3,6 +3,6 @@ members = ["mudrv", "search-musa-tools"]
 resolver = "2"
 
 [workspace.dependencies]
-bindgen = "0.69"
+bindgen = "0.70"
 build-script-cfg = "0.0"
 half = "2.4"

mudrv/src/device.rs

@@ -132,6 +132,13 @@ impl Device {
     }
 }
 
+pub fn mem_get_info() -> (usize, usize) {
+    let mut free = 0;
+    let mut total = 0;
+    muruntime!(musaMemGetInfo(&mut free, &mut total));
+    (free, total)
+}
+
 pub struct InfoFmt<'a>(&'a Device);
 
 impl fmt::Display for InfoFmt<'_> {

mudrv/src/lib.rs

@@ -49,7 +49,7 @@ pub fn init() -> Result<(), NoDevice> {
 
 pub use context::{Context, CurrentCtx};
 pub use context_spore::{impl_spore, AsRaw, ContextResource, ContextSpore, RawContainer};
-pub use device::{BlockLimit, Device, SMLimit};
+pub use device::{mem_get_info, BlockLimit, Device, SMLimit};
 pub use event::{Event, EventSpore};
 pub use memory::{
     memcpy_d2d, memcpy_d2h, memcpy_h2d, DevByte, DevMem, DevMemSpore, HostMem, HostMemSpore,

mudrv/src/memory.rs

@@ -3,7 +3,7 @@ use context_spore::{impl_spore, AsRaw};
 use std::{
     alloc::Layout,
     marker::PhantomData,
-    mem::{forget, size_of_val},
+    mem::size_of_val,
     ops::{Deref, DerefMut},
     os::raw::c_void,
     ptr::null_mut,
@@ -66,14 +66,14 @@ impl Stream<'_> {
 impl_spore!(DevMem and DevMemSpore by (CurrentCtx, Blob<MUdeviceptr>));
 
 impl CurrentCtx {
-    pub fn malloc<T: Copy>(&self, len: usize) -> DevMem<'_> {
+    pub fn malloc<'ctx, T: Copy>(&self, len: usize) -> DevMem<'ctx> {
         let len = Layout::array::<T>(len).unwrap().size();
         let mut ptr = 0;
         mudrv!(muMemAlloc_v2(&mut ptr, len));
         DevMem(unsafe { self.wrap_raw(Blob { ptr, len }) }, PhantomData)
     }
 
-    pub fn from_host<T: Copy>(&self, slice: &[T]) -> DevMem<'_> {
+    pub fn from_host<'ctx, T: Copy>(&self, slice: &[T]) -> DevMem<'ctx> {
         let len = size_of_val(slice);
         let src = slice.as_ptr().cast();
         let mut ptr = 0;
@@ -85,34 +85,11 @@ impl CurrentCtx {
 
 impl<'ctx> Stream<'ctx> {
     pub fn malloc<T: Copy>(&self, len: usize) -> DevMem<'ctx> {
-        let len = Layout::array::<T>(len).unwrap().size();
-        let mut ptr = 0;
-        mudrv!(muMemAlloc_v2(&mut ptr, len));
-        DevMem(
-            unsafe { self.ctx().wrap_raw(Blob { ptr, len }) },
-            PhantomData,
-        )
+        self.ctx().malloc::<T>(len)
     }
 
     pub fn from_host<T: Copy>(&self, slice: &[T]) -> DevMem<'ctx> {
-        let stream = unsafe { self.as_raw() };
-        let len = size_of_val(slice);
-        let src = slice.as_ptr().cast();
-        let mut ptr = 0;
-        mudrv!(muMemAlloc_v2(&mut ptr, len));
-        mudrv!(muMemcpyHtoDAsync_v2(ptr, src, len, stream));
-        DevMem(
-            unsafe { self.ctx().wrap_raw(Blob { ptr, len }) },
-            PhantomData,
-        )
-    }
-}
-
-impl DevMem<'_> {
-    #[inline]
-    pub fn drop_on(self) {
-        mudrv!(muMemFree_v2(self.0.rss.ptr));
-        forget(self);
+        self.ctx().from_host(slice)
     }
 }