add a try_alloc api for memregions, fix examples and docs

pnnl · Nov 8, 2024 · 3872700 · 3872700
1 parent 59dde95
commit 3872700
Show file tree

Hide file tree

Showing 61 changed files with 542 additions and 398 deletions.
diff --git a/examples/active_message_examples/am_local_memregions.rs b/examples/active_message_examples/am_local_memregions.rs
@@ -50,7 +50,7 @@ fn main() {
     let world = lamellar::LamellarWorldBuilder::new().build();
     let my_pe = world.my_pe();
     let num_pes = world.num_pes();
-    let array = world.alloc_one_sided_mem_region::<u8>(10).expect("Enough memory should exist");
+    let array = world.alloc_one_sided_mem_region::<u8>(10);
 
     let mut rng = rand::thread_rng();
     let pes = Uniform::from(0..num_pes);

diff --git a/examples/array_examples/array_am.rs b/examples/array_examples/array_am.rs
@@ -32,7 +32,7 @@ impl LamellarAM for RdmaAM {
         });
 
         //get the original nodes data
-        let local = lamellar::world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+        let local = lamellar::world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
         let local_slice = unsafe { local.as_mut_slice().unwrap() };
         local_slice[ARRAY_LEN - 1] = num_pes as u8;
         unsafe {
@@ -67,7 +67,7 @@ fn main() {
     println!("creating array");
     let array = UnsafeArray::<u8>::new(world.team(), ARRAY_LEN, Distribution::Block).block();
     println!("creating memregion");
-    let local_mem_region = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+    let local_mem_region = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     println!("about to initialize array");
     array.print();
     if my_pe == 0 {

diff --git a/examples/array_examples/array_put_get.rs b/examples/array_examples/array_put_get.rs
@@ -34,11 +34,11 @@ fn main() {
         let shared_mem_region = world
             .alloc_shared_mem_region(total_len)
             .await
-            .expect("Enough memory should exist")
+
             .into(); //Convert into abstract LamellarMemoryRegion
         let local_mem_region = world
             .alloc_one_sided_mem_region(total_len)
-            .expect("Enough memory should exist")
+
             .into();
         initialize_array(&block_array).await;
         initialize_array(&cyclic_array).await;

diff --git a/examples/array_examples/dist_array_reduce.rs b/examples/array_examples/dist_array_reduce.rs
@@ -40,7 +40,7 @@ fn main() {
         UnsafeArray::<usize>::new(world.team(), total_len, Distribution::Block).block();
     let cyclic_array =
         UnsafeArray::<usize>::new(world.team(), total_len, Distribution::Cyclic).block();
-    let local_mem_region = world.alloc_one_sided_mem_region(total_len).expect("Enough memory should exist");
+    let local_mem_region = world.alloc_one_sided_mem_region(total_len);
     world.barrier();
     if my_pe == 0 {
         unsafe {

diff --git a/examples/bandwidths/am_bw_get.rs b/examples/bandwidths/am_bw_get.rs
@@ -25,7 +25,7 @@ impl LamellarAM for DataAM {
     async fn exec(&self) {
         unsafe {
             // let local = lamellar::team.local_array::<u8>(self.length, 255u8);
-            let local = lamellar::team.alloc_one_sided_mem_region::<u8>(self.length).expect("Enough memory should exist");
+            let local = lamellar::team.alloc_one_sided_mem_region::<u8>(self.length);
             let local_slice = local.as_mut_slice().unwrap();
             local_slice[self.length - 1] = 255u8;
             self.array.get_unchecked(self.index, local.clone());
@@ -42,8 +42,8 @@ fn main() {
     let world = lamellar::LamellarWorldBuilder::new().build();
     let my_pe = world.my_pe();
     let num_pes = world.num_pes();
-    let array = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
-    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+    let array = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
+    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     unsafe {
         for i in data.as_mut_slice().unwrap() {
             *i = my_pe as u8;

diff --git a/examples/bandwidths/am_group_bw_get.rs b/examples/bandwidths/am_group_bw_get.rs
@@ -25,7 +25,7 @@ impl LamellarAM for DataAM {
     async fn exec(&self) {
         unsafe {
             // let local = lamellar::team.local_array::<u8>(self.length, 255u8);
-            let local = lamellar::team.alloc_one_sided_mem_region::<u8>(self.length).expect("Enough memory should exist");
+            let local = lamellar::team.alloc_one_sided_mem_region::<u8>(self.length);
             let local_slice = local.as_mut_slice().unwrap();
             local_slice[self.length - 1] = 255u8;
             self.array.get_unchecked(self.index, local.clone());
@@ -42,8 +42,8 @@ fn main() {
     let world = lamellar::LamellarWorldBuilder::new().build();
     let my_pe = world.my_pe();
     let num_pes = world.num_pes();
-    let array = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
-    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+    let array = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
+    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     unsafe {
         for i in data.as_mut_slice().unwrap() {
             *i = my_pe as u8;

diff --git a/examples/bandwidths/atomic_array_get_bw.rs b/examples/bandwidths/atomic_array_get_bw.rs
@@ -14,7 +14,7 @@ fn main() {
     let num_pes = world.num_pes();
     let array: LocalLockArray<u8> =
         LocalLockArray::new(&world, ARRAY_LEN * num_pes, Distribution::Block).block();
-    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     unsafe {
         for i in data.as_mut_slice().unwrap() {
             *i = my_pe as u8;

diff --git a/examples/bandwidths/atomic_array_put_bw.rs b/examples/bandwidths/atomic_array_put_bw.rs
@@ -14,7 +14,7 @@ fn main() {
     let num_pes = world.num_pes();
     let array: LocalLockArray<u8> =
         LocalLockArray::new(&world, ARRAY_LEN * num_pes, Distribution::Block).block();
-    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     unsafe {
         for i in data.as_mut_slice().unwrap() {
             *i = my_pe as u8;

diff --git a/examples/bandwidths/get_bw.rs b/examples/bandwidths/get_bw.rs
@@ -14,9 +14,8 @@ fn main() {
     let num_pes = world.num_pes();
     let mem_reg = world
         .alloc_shared_mem_region::<u8>(MEMREG_LEN)
-        .block()
-        .unwrap();
-    let data = world.alloc_one_sided_mem_region::<u8>(MEMREG_LEN).expect("Enough memory should exist");
+        .block();
+    let data = world.alloc_one_sided_mem_region::<u8>(MEMREG_LEN);
     for j in 0..MEMREG_LEN as usize {
         unsafe {
             data.as_mut_slice().unwrap()[j] = my_pe as u8;

diff --git a/examples/bandwidths/global_lock_atomic_array_get_bw.rs b/examples/bandwidths/global_lock_atomic_array_get_bw.rs
@@ -15,7 +15,7 @@ fn main() {
 
     let array: GlobalLockArray<u8> =
         GlobalLockArray::new(&world, ARRAY_LEN * num_pes, Distribution::Block).block();
-    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     unsafe {
         for i in data.as_mut_slice().unwrap() {
             *i = my_pe as u8;

diff --git a/examples/bandwidths/global_lock_atomic_array_put_bw.rs b/examples/bandwidths/global_lock_atomic_array_put_bw.rs
@@ -15,7 +15,7 @@ fn main() {
     let num_pes = world.num_pes();
     let array: GlobalLockArray<u8> =
         GlobalLockArray::new(&world, ARRAY_LEN * num_pes, Distribution::Block).block();
-    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     unsafe {
         for i in data.as_mut_slice().unwrap() {
             *i = my_pe as u8;

diff --git a/examples/bandwidths/local_lock_atomic_array_get_bw.rs b/examples/bandwidths/local_lock_atomic_array_get_bw.rs
@@ -15,7 +15,7 @@ fn main() {
 
     let array: LocalLockArray<u8> =
         LocalLockArray::new(&world, ARRAY_LEN * num_pes, Distribution::Block).block();
-    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     unsafe {
         for i in data.as_mut_slice().unwrap() {
             *i = my_pe as u8;

diff --git a/examples/bandwidths/local_lock_atomic_array_put_bw.rs b/examples/bandwidths/local_lock_atomic_array_put_bw.rs
@@ -15,7 +15,7 @@ fn main() {
     let num_pes = world.num_pes();
     let array: LocalLockArray<u8> =
         LocalLockArray::new(&world, ARRAY_LEN * num_pes, Distribution::Block).block();
-    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     unsafe {
         for i in data.as_mut_slice().unwrap() {
             *i = my_pe as u8;

diff --git a/examples/bandwidths/put_bw.rs b/examples/bandwidths/put_bw.rs
@@ -13,11 +13,10 @@ fn main() {
     let num_pes = world.num_pes();
     let array = world
         .alloc_shared_mem_region::<u8>(ARRAY_LEN)
-        .block()
-        .unwrap();
+        .block();
     let data = world
         .alloc_one_sided_mem_region::<u8>(ARRAY_LEN)
-        .expect("Enough memory should exist");
+        ;
     unsafe {
         for i in data.as_mut_slice().unwrap() {
             *i = my_pe as u8;

diff --git a/examples/bandwidths/readonly_array_get_bw.rs b/examples/bandwidths/readonly_array_get_bw.rs
@@ -14,7 +14,7 @@ fn main() {
     let my_pe = world.my_pe();
     let num_pes = world.num_pes();
     let array: UnsafeArray<u8> = UnsafeArray::new(&world, ARRAY_LEN * num_pes, Distribution::Block).block();
-    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     unsafe {
         for i in data.as_mut_slice().unwrap() {
             *i = my_pe as u8;

diff --git a/examples/bandwidths/readonly_array_get_unchecked_bw.rs b/examples/bandwidths/readonly_array_get_unchecked_bw.rs
@@ -13,7 +13,7 @@ fn main() {
     let my_pe = world.my_pe();
     let num_pes = world.num_pes();
     let array: UnsafeArray<u8> = UnsafeArray::new(&world, ARRAY_LEN * num_pes, Distribution::Block).block();
-    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     unsafe {
         for i in data.as_mut_slice().unwrap() {
             *i = my_pe as u8;

diff --git a/examples/bandwidths/unsafe_array_get_bw.rs b/examples/bandwidths/unsafe_array_get_bw.rs
@@ -14,7 +14,7 @@ fn main() {
     let my_pe = world.my_pe();
     let num_pes = world.num_pes();
     let array: UnsafeArray<u8> = UnsafeArray::new(&world, ARRAY_LEN * num_pes, Distribution::Block).block();
-    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     unsafe {
         for i in data.as_mut_slice().unwrap() {
             *i = my_pe as u8;

diff --git a/examples/bandwidths/unsafe_array_get_unchecked_bw.rs b/examples/bandwidths/unsafe_array_get_unchecked_bw.rs
@@ -13,7 +13,7 @@ fn main() {
     let my_pe = world.my_pe();
     let num_pes = world.num_pes();
     let array: UnsafeArray<u8> = UnsafeArray::new(&world, ARRAY_LEN * num_pes, Distribution::Block).block();
-    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     unsafe {
         for i in data.as_mut_slice().unwrap() {
             *i = my_pe as u8;

diff --git a/examples/bandwidths/unsafe_array_put_bw.rs b/examples/bandwidths/unsafe_array_put_bw.rs
@@ -13,7 +13,7 @@ fn main() {
     let my_pe = world.my_pe();
     let num_pes = world.num_pes();
     let array: UnsafeArray<u8> = UnsafeArray::new(&world, ARRAY_LEN * num_pes, Distribution::Block).block();
-    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     unsafe {
         for i in data.as_mut_slice().unwrap() {
             *i = my_pe as u8;

diff --git a/examples/bandwidths/unsafe_array_put_unchecked_bw.rs b/examples/bandwidths/unsafe_array_put_unchecked_bw.rs
@@ -13,7 +13,7 @@ fn main() {
     let my_pe = world.my_pe();
     let num_pes = world.num_pes();
     let array: UnsafeArray<u8> = UnsafeArray::new(&world, ARRAY_LEN * num_pes, Distribution::Block).block();
-    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     unsafe {
         for i in data.as_mut_slice().unwrap() {
             *i = my_pe as u8;

diff --git a/examples/bandwidths/unsafe_array_store_bw.rs b/examples/bandwidths/unsafe_array_store_bw.rs
@@ -14,7 +14,7 @@ fn main() {
     let my_pe = world.my_pe();
     let num_pes = world.num_pes();
     let array: UnsafeArray<u8> = UnsafeArray::new(&world, ARRAY_LEN * num_pes, Distribution::Block).block();
-    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+    let data = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     unsafe {
         for i in data.as_mut_slice().unwrap() {
             *i = my_pe as u8;

diff --git a/examples/kernels/am_gemm.rs b/examples/kernels/am_gemm.rs
@@ -97,8 +97,8 @@ struct NaiveMM {
 #[lamellar::am]
 impl LamellarAM for NaiveMM {
     async fn exec() {
-        let a = lamellar::world.alloc_one_sided_mem_region(self.a.block_size * self.a.block_size).expect("Enough memory should exist"); //the tile for the A matrix
-        let b = lamellar::world.alloc_one_sided_mem_region(self.b.block_size * self.b.block_size).expect("Enough memory should exist"); //the tile for the B matrix
+        let a = lamellar::world.alloc_one_sided_mem_region(self.a.block_size * self.a.block_size); //the tile for the A matrix
+        let b = lamellar::world.alloc_one_sided_mem_region(self.b.block_size * self.b.block_size); //the tile for the B matrix
         let b_fut = get_sub_mat(&self.b, &b); //b is remote so we will launch "gets" for this data first
         let a_fut = get_sub_mat(&self.a, &a);
         let a_b_fut = future::join(a_fut, b_fut);
@@ -164,16 +164,13 @@ fn main() {
 
     let a = world
         .alloc_shared_mem_region::<f32>((m * n) / num_pes)
-        .block()
-        .unwrap();
+        .block();
     let b = world
         .alloc_shared_mem_region::<f32>((n * p) / num_pes)
-        .block()
-        .unwrap();
+        .block();
     let c = world
         .alloc_shared_mem_region::<f32>((m * p) / num_pes)
-        .block()
-        .unwrap();
+        .block();
     unsafe {
         let mut cnt = (((m * n) / num_pes) * my_pe) as f32;
         for elem in a.as_mut_slice().unwrap() {

diff --git a/examples/kernels/cached_am_gemm.rs b/examples/kernels/cached_am_gemm.rs
@@ -105,7 +105,7 @@ impl LamellarAM for MatMulAM {
     async fn exec() {
         let b = lamellar::world
             .alloc_one_sided_mem_region::<f32>(self.b.block_size * self.b.block_size)
-            .expect("enough memory exists");
+            ;
         get_sub_mat(&self.b, &b).await;
         // we dont actually want to alloc a shared memory region as there is an implicit barrier here
         // introduces sync point and potential for deadlock
@@ -122,7 +122,7 @@ impl LamellarAM for MatMulAM {
             c.row_block = row;
             let sub_a = lamellar::world
                 .alloc_one_sided_mem_region::<f32>(a.block_size * a.block_size)
-                .expect("enough memory exists");
+                ;
             get_sub_mat(&a, &sub_a).await; //this should be local copy so returns immediately
             do_gemm(&sub_a, &b, c, self.block_size);
         }
@@ -179,15 +179,15 @@ fn main() {
     let a = world
         .alloc_shared_mem_region::<f32>((m * n) / num_pes)
         .block()
-        .expect("enough memory exists");
+        ;
     let b = world
         .alloc_shared_mem_region::<f32>((n * p) / num_pes)
         .block()
-        .expect("enough memory exists");
+        ;
     let c = world
         .alloc_shared_mem_region::<f32>((m * p) / num_pes)
         .block()
-        .expect("enough memory exists");
+        ;
     // let c2 = world.alloc_shared_mem_region::<f32>((m * p) / num_pes);
     unsafe {
         let mut cnt = my_pe as f32 * ((m * n) / num_pes) as f32;

diff --git a/examples/kernels/dft_proxy.rs b/examples/kernels/dft_proxy.rs
@@ -145,8 +145,7 @@ fn dft_lamellar(
     let spectrum_slice = unsafe { spectrum.as_slice().unwrap() };
     let add_spec = world
         .alloc_shared_mem_region::<f64>(spectrum_slice.len())
-        .block()
-        .unwrap();
+        .block();
 
     let timer = Instant::now();
     for pe in 0..num_pes {
@@ -640,18 +639,18 @@ fn main() {
         let partial_sum = world
             .alloc_shared_mem_region::<f64>(num_pes)
             .block()
-            .expect("Enough memory should exist");
+            ;
         let partial_spectrum = world
             .alloc_shared_mem_region::<f64>(array_len)
             .block()
-            .expect("Enough memory should exist");
+            ;
         let partial_signal = world
             .alloc_shared_mem_region::<f64>(array_len)
             .block()
-            .expect("Enough memory should exist");
-        let full_signal = world.alloc_one_sided_mem_region::<f64>(global_len).expect("Enough memory should exist");
-        let full_spectrum = world.alloc_one_sided_mem_region::<f64>(global_len).expect("Enough memory should exist");
-        let magic = world.alloc_one_sided_mem_region::<f64>(num_pes).expect("Enough memory should exist");
+            ;
+        let full_signal = world.alloc_one_sided_mem_region::<f64>(global_len);
+        let full_spectrum = world.alloc_one_sided_mem_region::<f64>(global_len);
+        let magic = world.alloc_one_sided_mem_region::<f64>(num_pes);
 
         let full_spectrum_array =
             UnsafeArray::<f64>::new(world.team(), global_len, Distribution::Block).block();

diff --git a/examples/misc/ping_pong.rs b/examples/misc/ping_pong.rs
@@ -328,28 +328,23 @@ fn main() {
 
     let indices = world
         .alloc_shared_mem_region::<usize>(UPDATES_PER_CORE * world.num_threads_per_pe())
-        .block()
-        .unwrap();
+        .block();
 
     let index_send_buffers = world
         .alloc_shared_mem_region::<usize>(buffer_size * num_pes)
-        .block()
-        .unwrap();
+        .block();
     world.barrier();
     let index_recv_buffers = world
         .alloc_shared_mem_region::<usize>(buffer_size * num_pes)
-        .block()
-        .unwrap();
+        .block();
     world.barrier();
     let result_send_buffers = world
         .alloc_shared_mem_region::<usize>(buffer_size * num_pes)
-        .block()
-        .unwrap();
+        .block();
     world.barrier();
     let result_recv_buffers = world
         .alloc_shared_mem_region::<usize>(buffer_size * num_pes)
-        .block()
-        .unwrap();
+        .block();
     world.barrier();
     let mut rng: StdRng = SeedableRng::seed_from_u64(my_pe as u64);
     let table_size_per_pe = 100000 * world.num_threads_per_pe();

diff --git a/examples/rdma_examples/rdma_am.rs b/examples/rdma_examples/rdma_am.rs
@@ -34,7 +34,7 @@ impl LamellarAM for RdmaAM {
         //get the original nodes data
         let local = lamellar::world
             .alloc_one_sided_mem_region::<u8>(ARRAY_LEN)
-            .expect("Enough memory should exist");
+            ;
         let local_slice = unsafe { local.as_mut_slice().unwrap() };
         local_slice[ARRAY_LEN - 1] = lamellar::num_pes as u8;
         unsafe {
@@ -68,7 +68,7 @@ impl LamellarAM for RdmaLocalMRAM {
         );
 
         //get the original nodes data
-        let local = lamellar::world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+        let local = lamellar::world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
         let local_slice = unsafe { local.as_mut_slice().unwrap() };
         local_slice[ARRAY_LEN - 1] = lamellar::num_pes as u8;
         unsafe {
@@ -110,8 +110,8 @@ fn main() {
     let array = world
         .alloc_shared_mem_region::<u8>(ARRAY_LEN)
         .block()
-        .expect("Enough memory should exist");
-    let local_array = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN).expect("Enough memory should exist");
+        ;
+    let local_array = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
     unsafe {
         for i in array.as_mut_slice().unwrap() {
             *i = 255_u8;