diff --git a/.gitignore b/.gitignore index fcc3bbfec..61fa20e12 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/engines.py b/capytaine/bem/engines.py index c8eee507e..5c970a025 100644 --- a/capytaine/bem/engines.py +++ b/capytaine/bem/engines.py @@ -60,6 +60,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/bem/solver.py b/capytaine/bem/solver.py index 616520483..0e641ae34 100644 --- a/capytaine/bem/solver.py +++ b/capytaine/bem/solver.py @@ -23,10 +23,12 @@ 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 capytaine.tools.symbolic_multiplication import supporting_symbolic_multiplication LOG = logging.getLogger(__name__) + class BEMSolver: """ Solver for linear potential flow problems. @@ -160,7 +162,9 @@ def solve_all(self, problems, *, method='indirect', n_jobs=1, progress_bar=True, 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) groups_of_results = parallel(joblib.delayed(self.solve_all)(grp, method=method, n_jobs=1, progress_bar=False, _check_wavelength=False, **kwargs) for grp in groups_of_problems) if progress_bar: @@ -169,6 +173,8 @@ def solve_all(self, problems, *, method='indirect', n_jobs=1, progress_bar=True, description=f"Solving BEM problems with {n_jobs} threads:") results = [res for grp in groups_of_results for res in grp] # flatten the nested list return results + + @staticmethod def _check_wavelength(problems): diff --git a/capytaine/matrices/linear_solvers.py b/capytaine/matrices/linear_solvers.py index 72e0c25d2..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 @@ -149,3 +151,71 @@ def gmres_no_fft(A, b): LOG.warning(f"No convergence of the GMRES. Error code: {info}") return x + +##### GPU Solver ##### +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") + + a = torch.from_numpy(a).cfloat().to(device) + b = torch.from_numpy(b).cfloat().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}") + + 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 +