Add optimized subset_ops

examples/heat_1d_ap_fft.rs

@@ -13,7 +13,7 @@ fn main() {
     // Create AP Solver
     let solver_params = args.solver_parameters();
     let mut solver =
-        generate_ap_solver(&stencil, direct_solver, &solver_params);
+        generate_ap_solver_1d(&stencil, direct_solver, &solver_params);
 
     args.run_solver(&mut solver);
 }

examples/heat_2d_ap_fft.rs

@@ -13,7 +13,7 @@ fn main() {
     // Create AP Solver
     let solver_params = args.solver_parameters();
     let mut solver =
-        generate_ap_solver(&stencil, direct_solver, &solver_params);
+        generate_ap_solver_2d(&stencil, direct_solver, &solver_params);
 
     args.run_solver(&mut solver);
 }

examples/heat_3d_ap_fft.rs

@@ -21,9 +21,9 @@ fn main() {
     };
 
     // Create AP Solver
-    let solver_parameters = args.solver_parameters();
+    let solver_params = args.solver_parameters();
     let mut solver =
-        generate_ap_solver(&stencil, direct_solver, &solver_parameters);
+        generate_ap_solver_3d(&stencil, direct_solver, &solver_params);
 
     args.run_solver(&mut solver);
 }

examples/tv_heat_1d_ap_fft.rs

@@ -11,7 +11,7 @@ fn main() {
     // Create AP Solver
     let solver_params = args.solver_parameters();
     let mut solver =
-        generate_tv_ap_solver(&stencil, direct_solver, &solver_params);
+        generate_tv_ap_solver_1d(&stencil, direct_solver, &solver_params);
 
     args.run_solver(&mut solver);
 }

examples/tv_heat_2d_ap_fft.rs

@@ -13,7 +13,7 @@ fn main() {
     // Create AP Solver
     let solver_params = args.solver_parameters();
     let mut solver =
-        generate_tv_ap_solver(&stencil, direct_solver, &solver_params);
+        generate_tv_ap_solver_2d(&stencil, direct_solver, &solver_params);
 
     args.run_solver(&mut solver);
 }

examples/tv_rotating_2d.rs

@@ -14,7 +14,7 @@ fn main() {
     // Create AP Solver
     let solver_params = args.solver_parameters();
     let mut solver =
-        generate_tv_ap_solver(&stencil, direct_solver, &solver_params);
+        generate_tv_ap_solver_2d(&stencil, direct_solver, &solver_params);
 
     args.run_solver(&mut solver);
 }

src/ap_solver/generate_solver.rs

@@ -5,37 +5,122 @@ use crate::ap_solver::solver::*;
 use crate::ap_solver::solver_parameters::*;
 use crate::ap_solver::tv_periodic_ops_collector::*;
 use crate::direct_solver::*;
+use crate::domain::*;
 use crate::stencil::*;
 use crate::SolverInterface;
 
-pub fn generate_ap_solver<
-    const GRID_DIMENSION: usize,
+pub fn generate_ap_solver_1d<
     const NEIGHBORHOOD_SIZE: usize,
-    DirectSolverType: DirectSolverInterface<GRID_DIMENSION>,
+    DirectSolverType: DirectSolverInterface<1>,
 >(
-    stencil: &Stencil<GRID_DIMENSION, NEIGHBORHOOD_SIZE>,
+    stencil: &Stencil<1, NEIGHBORHOOD_SIZE>,
     direct_solver: DirectSolverType,
-    params: &SolverParameters<GRID_DIMENSION>,
-) -> impl SolverInterface<GRID_DIMENSION> {
+    params: &SolverParameters<1>,
+) -> impl SolverInterface<1> {
     let create_ops_builder = || ApPeriodicOpsBuilder::new(stencil, params);
     let planner_result = generate_plan(stencil, create_ops_builder, params);
     let complex_buffer_type = ComplexBufferType::DomainOnly;
-    Solver::new(direct_solver, params, planner_result, complex_buffer_type)
+    let subset_ops = SubsetOps1d {};
+    Solver::new(
+        direct_solver,
+        subset_ops,
+        params,
+        planner_result,
+        complex_buffer_type,
+    )
 }
 
-pub fn generate_tv_ap_solver<
+pub fn generate_ap_solver_2d<
+    const NEIGHBORHOOD_SIZE: usize,
+    DirectSolverType: DirectSolverInterface<2>,
+>(
+    stencil: &Stencil<2, NEIGHBORHOOD_SIZE>,
+    direct_solver: DirectSolverType,
+    params: &SolverParameters<2>,
+) -> impl SolverInterface<2> {
+    let create_ops_builder = || ApPeriodicOpsBuilder::new(stencil, params);
+    let planner_result = generate_plan(stencil, create_ops_builder, params);
+    let complex_buffer_type = ComplexBufferType::DomainOnly;
+    let subset_ops = SubsetOps2d {
+        chunk_size: params.chunk_size,
+    };
+    Solver::new(
+        direct_solver,
+        subset_ops,
+        params,
+        planner_result,
+        complex_buffer_type,
+    )
+}
+
+pub fn generate_ap_solver_3d<
+    const NEIGHBORHOOD_SIZE: usize,
+    DirectSolverType: DirectSolverInterface<3>,
+>(
+    stencil: &Stencil<3, NEIGHBORHOOD_SIZE>,
+    direct_solver: DirectSolverType,
+    params: &SolverParameters<3>,
+) -> impl SolverInterface<3> {
+    let create_ops_builder = || ApPeriodicOpsBuilder::new(stencil, params);
+    let planner_result = generate_plan(stencil, create_ops_builder, params);
+    let complex_buffer_type = ComplexBufferType::DomainOnly;
+    let subset_ops = SubsetOps3d {
+        chunk_size: params.chunk_size,
+        task_min: params.task_min,
+    };
+    Solver::new(
+        direct_solver,
+        subset_ops,
+        params,
+        planner_result,
+        complex_buffer_type,
+    )
+}
+
+pub fn generate_tv_ap_solver_1d<
+    'a,
+    const NEIGHBORHOOD_SIZE: usize,
+    StencilType: TVStencil<1, NEIGHBORHOOD_SIZE>,
+    DirectSolverType: DirectSolverInterface<1> + 'a,
+>(
+    stencil: &'a StencilType,
+    direct_solver: DirectSolverType,
+    params: &'a SolverParameters<1>,
+) -> impl SolverInterface<1> + 'a {
+    let create_ops_builder = || TvPeriodicOpsCollector::new(stencil, params);
+    let planner_result = generate_plan(stencil, create_ops_builder, params);
+    let complex_buffer_type = ComplexBufferType::DomainAndOp;
+    let subset_ops = SubsetOps1d {};
+    Solver::new(
+        direct_solver,
+        subset_ops,
+        params,
+        planner_result,
+        complex_buffer_type,
+    )
+}
+
+pub fn generate_tv_ap_solver_2d<
     'a,
-    const GRID_DIMENSION: usize,
     const NEIGHBORHOOD_SIZE: usize,
-    StencilType: TVStencil<GRID_DIMENSION, NEIGHBORHOOD_SIZE>,
-    DirectSolverType: DirectSolverInterface<GRID_DIMENSION> + 'a,
+    StencilType: TVStencil<2, NEIGHBORHOOD_SIZE>,
+    DirectSolverType: DirectSolverInterface<2> + 'a,
 >(
     stencil: &'a StencilType,
     direct_solver: DirectSolverType,
-    params: &'a SolverParameters<GRID_DIMENSION>,
-) -> impl SolverInterface<GRID_DIMENSION> + 'a {
+    params: &'a SolverParameters<2>,
+) -> impl SolverInterface<2> + 'a {
     let create_ops_builder = || TvPeriodicOpsCollector::new(stencil, params);
     let planner_result = generate_plan(stencil, create_ops_builder, params);
     let complex_buffer_type = ComplexBufferType::DomainAndOp;
-    Solver::new(direct_solver, params, planner_result, complex_buffer_type)
+    let subset_ops = SubsetOps2d {
+        chunk_size: params.chunk_size,
+    };
+    Solver::new(
+        direct_solver,
+        subset_ops,
+        params,
+        planner_result,
+        complex_buffer_type,
+    )
 }