From f765e62c443acb656108086ef312421abd82d457 Mon Sep 17 00:00:00 2001 From: SoundsSerious Date: Fri, 22 Dec 2023 13:58:45 -0500 Subject: [PATCH 1/5] GPU Direct Solver >use linear_solver='gpu_direct' >torch and cuda must be installed for this --- capytaine/bem/engines.py | 2 ++ capytaine/matrices/linear_solvers.py | 54 ++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/capytaine/bem/engines.py b/capytaine/bem/engines.py index 9c8605859..a7e367d87 100644 --- a/capytaine/bem/engines.py +++ b/capytaine/bem/engines.py @@ -62,6 +62,8 @@ class BasicMatrixEngine(MatrixEngine): 'lu_decomposition': linear_solvers.LUSolverWithCache().solve, 'gmres': linear_solvers.solve_gmres, } + if linear_solvers.gpu_direct: + available_linear_solvers['gpu_direct'] = linear_solvers.gpu_direct def __init__(self, *, linear_solver='lu_decomposition', matrix_cache_size=1): diff --git a/capytaine/matrices/linear_solvers.py b/capytaine/matrices/linear_solvers.py index 03fe402b7..f37d59d68 100644 --- a/capytaine/matrices/linear_solvers.py +++ b/capytaine/matrices/linear_solvers.py @@ -153,3 +153,57 @@ def gmres_no_fft(A, b): return x +##### GPU Solver ##### +try: + + import torch + def solve_gpu(a:np.ndarray,b:np.ndarray)->np.ndarray: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + a = torch.from_numpy(a).float().to(device) + b = torch.from_numpy(b).float().to(device) + + res = torch.linalg.solve(a, b) + + return res.cpu().numpy() + + def gpu_direct(A, b): + assert isinstance(b, np.ndarray) and A.ndim == b.ndim+1 and A.shape[-2] == b.shape[-1] + if isinstance(A, BlockCirculantMatrix): + LOG.debug("\tSolve linear circulant system %s", A) + blocks_of_diagonalization = A.block_diagonalize() + fft_of_rhs = np.fft.fft(np.reshape(b, (A.nb_blocks[0], A.block_shape[0])), axis=0) + try: # Try to run it as vectorized numpy arrays. + fft_of_result = solve_gpu(blocks_of_diagonalization, fft_of_rhs) + except np.linalg.LinAlgError: # Or do the same thing with list comprehension. + fft_of_result = np.array([gpu_direct(block, vec) for block, vec in zip(blocks_of_diagonalization, fft_of_rhs)]) + result = np.fft.ifft(fft_of_result, axis=0).reshape((A.shape[1],)) + return result + + elif isinstance(A, BlockSymmetricToeplitzMatrix): + if A.nb_blocks == (2, 2): + LOG.debug("\tSolve linear system %s", A) + A1, A2 = A._stored_blocks[0, :] + b1, b2 = b[:len(b)//2], b[len(b)//2:] + x_plus = gpu_direct(A1 + A2, b1 + b2) + x_minus = gpu_direct(A1 - A2, b1 - b2) + return np.concatenate([x_plus + x_minus, x_plus - x_minus])/2 + else: + # Not implemented + LOG.debug("\tSolve linear system %s", A) + return gpu_direct(A.full_matrix(), b) + + elif isinstance(A, BlockMatrix): + LOG.debug("\tSolve linear system %s", A) + return gpu_direct(A.full_matrix(), b) + + elif isinstance(A, np.ndarray): + LOG.debug(f"\tSolve linear system (size: {A.shape}) with numpy direct solver.") + return solve_gpu(A, b) + + else: + raise ValueError(f"Unrecognized type of matrix to solve: {A}") + +except ImportError as e: + LOG.debug(f'GPU solver not available. Install pytorch and cuda to enable. Error: {e}') + gpu_direct = None From 4b3788d6e4e36d0ef7b9ce17261a7e0b773fdece Mon Sep 17 00:00:00 2001 From: SoundsSerious Date: Fri, 22 Dec 2023 14:55:37 -0500 Subject: [PATCH 2/5] Add parallel solver --- .gitignore | 3 +++ capytaine/bem/solver.py | 29 ++++++++++++++++++++++++---- capytaine/matrices/linear_solvers.py | 1 + 3 files changed, 29 insertions(+), 4 deletions(-) diff --git a/.gitignore b/.gitignore index 95ad877ed..3d22a163d 100644 --- a/.gitignore +++ b/.gitignore @@ -139,6 +139,9 @@ ENV/ # Rope project settings .ropeproject +# meson build +.meson* + # mkdocs documentation /site diff --git a/capytaine/bem/solver.py b/capytaine/bem/solver.py index aa1b6f33a..8e7bc5b89 100644 --- a/capytaine/bem/solver.py +++ b/capytaine/bem/solver.py @@ -21,6 +21,7 @@ from capytaine.io.xarray import problems_from_dataset, assemble_dataset, kochin_data_array from capytaine.tools.optional_imports import silently_import_optional_dependency from capytaine.tools.lists_of_points import _normalize_points, _normalize_free_surface_points +from functools import partial LOG = logging.getLogger(__name__) @@ -126,13 +127,33 @@ def solve_all(self, problems, *, n_jobs=1, **kwargs): if n_jobs == 1: # force sequential resolution return [self.solve(pb, **kwargs) for pb in sorted(problems)] else: - joblib = silently_import_optional_dependency("joblib") - if joblib is None: - raise ImportError(f"Setting the `n_jobs` argument to {n_jobs} requires the missing optional dependency 'joblib'.") groups_of_problems = LinearPotentialFlowProblem._group_for_parallel_resolution(problems) - groups_of_results = joblib.Parallel(n_jobs=n_jobs)(joblib.delayed(self.solve_all)(grp, n_jobs=1, **kwargs) for grp in groups_of_problems) + groups_of_results = self.parallel(groups_of_problems,n_jobs,**kwargs) results = [res for grp in groups_of_results for res in grp] # flatten the nested list return results + + def parallel(self,problems,n_jobs=1,**kwargs)->list: + """handles the parallelization of solvers depending on linear solver""" + from capytaine.matrices import linear_solvers + print(self.engine.linear_solver,self.engine.linear_solver is linear_solvers.gpu_direct) + if self.engine.linear_solver is linear_solvers.gpu_direct: + LOG.info(f'using parallel gpu solver with {n_jobs} jobs') + from torch.multiprocessing import Pool, Process, set_start_method + try: + set_start_method('spawn') + except RuntimeError: + pass + pool = Pool(n_jobs) + return pool.map(partial(self.solve,**kwargs), problems) + + else: + LOG.info(f'using joblib parallel solver with {n_jobs} jobs') + joblib = silently_import_optional_dependency("joblib") + if joblib is None: + raise ImportError(f"Setting the `n_jobs` argument to {n_jobs} requires the missing optional dependency 'joblib'.") + return joblib.Parallel(n_jobs=n_jobs)(joblib.delayed(self.solve_all)(grp, n_jobs=1, **kwargs) for grp in problems) + + def fill_dataset(self, dataset, bodies, *, n_jobs=1, **kwargs): """Solve a set of problems defined by the coordinates of an xarray dataset. diff --git a/capytaine/matrices/linear_solvers.py b/capytaine/matrices/linear_solvers.py index f37d59d68..b8f0b73d3 100644 --- a/capytaine/matrices/linear_solvers.py +++ b/capytaine/matrices/linear_solvers.py @@ -204,6 +204,7 @@ def gpu_direct(A, b): else: raise ValueError(f"Unrecognized type of matrix to solve: {A}") + print(f'GPU available: {torch.cuda.get_device_name()}') except ImportError as e: LOG.debug(f'GPU solver not available. Install pytorch and cuda to enable. Error: {e}') gpu_direct = None From cc88c2e96cac1554dc0f5178cdb137859a40b649 Mon Sep 17 00:00:00 2001 From: SoundsSerious Date: Fri, 22 Dec 2023 15:26:06 -0500 Subject: [PATCH 3/5] Convert types from complex appropriately and depriciate parallel solver changes --- capytaine/bem/solver.py | 47 ++++++++++++++++++---------- capytaine/matrices/linear_solvers.py | 4 +-- 2 files changed, 32 insertions(+), 19 deletions(-) diff --git a/capytaine/bem/solver.py b/capytaine/bem/solver.py index 8e7bc5b89..aab7ba9a1 100644 --- a/capytaine/bem/solver.py +++ b/capytaine/bem/solver.py @@ -25,6 +25,7 @@ LOG = logging.getLogger(__name__) + class BEMSolver: """ Solver for linear potential flow problems. @@ -134,24 +135,31 @@ def solve_all(self, problems, *, n_jobs=1, **kwargs): def parallel(self,problems,n_jobs=1,**kwargs)->list: """handles the parallelization of solvers depending on linear solver""" - from capytaine.matrices import linear_solvers - print(self.engine.linear_solver,self.engine.linear_solver is linear_solvers.gpu_direct) - if self.engine.linear_solver is linear_solvers.gpu_direct: - LOG.info(f'using parallel gpu solver with {n_jobs} jobs') - from torch.multiprocessing import Pool, Process, set_start_method - try: - set_start_method('spawn') - except RuntimeError: - pass - pool = Pool(n_jobs) - return pool.map(partial(self.solve,**kwargs), problems) - else: - LOG.info(f'using joblib parallel solver with {n_jobs} jobs') - joblib = silently_import_optional_dependency("joblib") - if joblib is None: - raise ImportError(f"Setting the `n_jobs` argument to {n_jobs} requires the missing optional dependency 'joblib'.") - return joblib.Parallel(n_jobs=n_jobs)(joblib.delayed(self.solve_all)(grp, n_jobs=1, **kwargs) for grp in problems) + #FIXME: fix kwargs addition, getting pickle error +# from capytaine.matrices import linear_solvers +# print(self.engine.linear_solver,self.engine.linear_solver is linear_solvers.gpu_direct) +# if self.engine.linear_solver is linear_solvers.gpu_direct: +# LOG.info(f'using parallel gpu solver with {n_jobs} jobs') +# from torch.multiprocessing import Pool, Process, set_start_method +# try: +# set_start_method('spawn') +# except RuntimeError: +# pass +# +# pool = Pool(n_jobs) +# +# #probs = [(self,pb,kwargs) for pb in problems] +# +# return pool.map(self.solve,problems) +# +# else: + + LOG.info(f'using joblib parallel solver with {n_jobs} jobs') + joblib = silently_import_optional_dependency("joblib") + if joblib is None: + raise ImportError(f"Setting the `n_jobs` argument to {n_jobs} requires the missing optional dependency 'joblib'.") + return joblib.Parallel(n_jobs=n_jobs)(joblib.delayed(self.solve_all)(grp, n_jobs=1, **kwargs) for grp in problems) @@ -384,3 +392,8 @@ def get_free_surface_elevation(self, result, free_surface, keep_details=False): result.fs_elevation[free_surface] = fs_elevation return fs_elevation +# #GPU SOLVER +# def parallel_solve_function(inst,problem,kwargs=None): +# if kwargs: +# return inst.solve(problem, **kwargs) +# return inst.solve(problem) \ No newline at end of file diff --git a/capytaine/matrices/linear_solvers.py b/capytaine/matrices/linear_solvers.py index b8f0b73d3..213518c12 100644 --- a/capytaine/matrices/linear_solvers.py +++ b/capytaine/matrices/linear_solvers.py @@ -160,8 +160,8 @@ def gmres_no_fft(A, b): def solve_gpu(a:np.ndarray,b:np.ndarray)->np.ndarray: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - a = torch.from_numpy(a).float().to(device) - b = torch.from_numpy(b).float().to(device) + a = torch.from_numpy(a).cfloat().to(device) + b = torch.from_numpy(b).cfloat().to(device) res = torch.linalg.solve(a, b) From e973cc4761905b69bdfba98e5c40a775fc1a7531 Mon Sep 17 00:00:00 2001 From: SoundsSerious Date: Fri, 22 Dec 2023 15:57:09 -0500 Subject: [PATCH 4/5] updatejoblib --- capytaine/bem/solver.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/capytaine/bem/solver.py b/capytaine/bem/solver.py index 53eee0e8c..0e641ae34 100644 --- a/capytaine/bem/solver.py +++ b/capytaine/bem/solver.py @@ -159,6 +159,10 @@ def solve_all(self, problems, *, method='indirect', n_jobs=1, progress_bar=True, problems = track(problems, total=len(problems), description="Solving BEM problems") return [self.solve(pb, method=method, _check_wavelength=False, **kwargs) for pb in problems] else: + joblib = silently_import_optional_dependency("joblib") + if joblib is None: + raise ImportError(f"Setting the `n_jobs` argument to {n_jobs} requires the missing optional dependency 'joblib'.") + groups_of_problems = LinearPotentialFlowProblem._group_for_parallel_resolution(problems) parallel = joblib.Parallel(return_as="generator", n_jobs=n_jobs) From 4e747c74f7e5a37a05854f1dadd558355f302ca3 Mon Sep 17 00:00:00 2001 From: SoundsSerious Date: Fri, 26 Jan 2024 00:33:26 -0500 Subject: [PATCH 5/5] Add GPU Ignore Option --- capytaine/matrices/linear_solvers.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/capytaine/matrices/linear_solvers.py b/capytaine/matrices/linear_solvers.py index 95e94f418..b8ef1edbd 100644 --- a/capytaine/matrices/linear_solvers.py +++ b/capytaine/matrices/linear_solvers.py @@ -7,6 +7,8 @@ import logging +import os + import numpy as np from scipy import linalg as sl from scipy.sparse import linalg as ssl @@ -151,8 +153,13 @@ def gmres_no_fft(A, b): return x ##### GPU Solver ##### -try: +class IgnoreGPUException(Exception): + pass +try: + if not os.environ.get('CAPYTAINE_GPU',None) == 'True': + raise IgnoreGPUException() + import torch def solve_gpu(a:np.ndarray,b:np.ndarray)->np.ndarray: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -202,7 +209,13 @@ def gpu_direct(A, b): raise ValueError(f"Unrecognized type of matrix to solve: {A}") print(f'GPU available: {torch.cuda.get_device_name()}') +except IgnoreGPUException: + pass except ImportError as e: LOG.debug(f'GPU solver not available. Install pytorch and cuda to enable. Error: {e}') gpu_direct = None +except Exception as e: + LOG.debug(f'Unknown GPU Import Error: {e}') + gpu_direct = None +