Add chunksize to 3pt1Dopt solver (#80)

Author: SallySoul

examples/heat_1d_ap_direct.rs

@@ -1,5 +1,4 @@
 use nhls::direct_solver::*;
-use nhls::domain::*;
 use nhls::image_1d_example::*;
 
 fn main() {
@@ -8,8 +7,12 @@ fn main() {
     let stencil = nhls::standard_stencils::heat_1d(1.0, 1.0, 0.5);
 
     // Create BC
-    let mut solver =
-        Direct3Pt1DSolver::new(&stencil, args.steps_per_line, args.threads);
+    let mut solver = Direct3Pt1DSolver::new(
+        &stencil,
+        args.steps_per_line,
+        args.threads,
+        args.chunk_size,
+    );
 
     args.run_solver(&mut solver);
 }

examples/heat_1d_ap_fft.rs

@@ -8,7 +8,7 @@ fn main() {
     let stencil = nhls::standard_stencils::heat_1d(1.0, 1.0, 0.5);
 
     // This optimized direct solver implement a uniform boundary condition of 0.0
-    let direct_solver = DirectSolver3Pt1DOpt::new(&stencil);
+    let direct_solver = DirectSolver3Pt1DOpt::new(&stencil, args.chunk_size);
 
     // Create AP Solver
     let solver_params = args.solver_parameters();

examples/tv_heat_1d_ap_fft.rs

@@ -7,7 +7,7 @@ fn main() {
 
     let stencil = nhls::standard_stencils::TVHeat1D::new();
 
-    let direct_solver = DirectSolver3Pt1DOpt::new(&stencil);
+    let direct_solver = DirectSolver3Pt1DOpt::new(&stencil, args.chunk_size);
     // Create AP Solver
     let solver_params = args.solver_parameters();
     let mut solver =

src/direct_solver/direct_3pt1d_opt.rs

@@ -8,17 +8,21 @@ use crate::SolverInterface;
 /// Implements a constant zero boundary condition.
 pub struct DirectSolver3Pt1DOpt<'a, StencilType: TVStencil<1, 3>> {
     stencil: &'a StencilType,
+    chunk_size: usize,
 }
 
 impl<'a, StencilType: TVStencil<1, 3>> DirectSolver3Pt1DOpt<'a, StencilType> {
-    pub fn new(stencil: &'a StencilType) -> Self {
+    pub fn new(stencil: &'a StencilType, chunk_size: usize) -> Self {
         let expected_offsets = [
             vector![1],  // 0
             vector![-1], // 1
             vector![0],  // 4
         ];
         assert_eq!(&expected_offsets, stencil.offsets());
-        DirectSolver3Pt1DOpt { stencil }
+        DirectSolver3Pt1DOpt {
+            stencil,
+            chunk_size,
+        }
     }
 
     fn apply_step<DomainType: DomainView<1> + Send>(
@@ -51,12 +55,12 @@ impl<'a, StencilType: TVStencil<1, 3>> DirectSolver3Pt1DOpt<'a, StencilType> {
 
             let const_output: &DomainType = output;
             rayon::scope(|s| {
-                profiling::scope!("direct_solver: Thread Callback");
-                let chunk_size = (n_r - 2) / (threads * 2);
+                let chunk_size = ((n_r - 2) / threads).max(self.chunk_size);
                 let mut start: usize = 1;
                 while start < n_r - 1 {
                     let end = (start + chunk_size).min(n_r - 1);
                     s.spawn(move |_| {
+                        profiling::scope!("direct_solver: Thread Callback");
                         let mut o = const_output.unsafe_mut_access();
                         for i in start..end {
                             *o.buffer_mut().get_unchecked_mut(i) = w
@@ -104,9 +108,17 @@ pub struct Direct3Pt1DSolver<'a, StencilType: TVStencil<1, 3>> {
 }
 
 impl<'a, StencilType: TVStencil<1, 3>> Direct3Pt1DSolver<'a, StencilType> {
-    pub fn new(stencil: &'a StencilType, steps: usize, threads: usize) -> Self {
+    pub fn new(
+        stencil: &'a StencilType,
+        steps: usize,
+        threads: usize,
+        chunk_size: usize,
+    ) -> Self {
         Direct3Pt1DSolver {
-            solver: DirectSolver3Pt1DOpt { stencil },
+            solver: DirectSolver3Pt1DOpt {
+                stencil,
+                chunk_size,
+            },
             steps,
             threads,
         }

src/domain/view/mod.rs

@@ -81,6 +81,7 @@ pub trait DomainView<const GRID_DIMENSION: usize>: Sync {
         other: &DomainType,
         chunk_size: usize,
     ) {
+        profiling::scope!("domain::par_set_subdomain");
         let const_self_ref: &Self = self;
         other.buffer()[0..other.aabb().buffer_size()]
             .par_chunks(chunk_size)
@@ -109,6 +110,7 @@ pub trait DomainView<const GRID_DIMENSION: usize>: Sync {
         other: &DomainType,
         chunk_size: usize,
     ) {
+        profiling::scope!("domain::par_from_subdomain");
         self.par_set_values(|world_coord| other.view(&world_coord), chunk_size);
     }