From 8fbe0f38f1831380fe901a39aa85a0d2fd9ef888 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Fri, 6 Feb 2026 11:49:02 +0100 Subject: [PATCH 01/98] Add init vector for cg --- pykeops/pykeops/common/operations.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/pykeops/pykeops/common/operations.py b/pykeops/pykeops/common/operations.py index 950c1c7c9..37ad5bd24 100644 --- a/pykeops/pykeops/common/operations.py +++ b/pykeops/pykeops/common/operations.py @@ -87,14 +87,14 @@ def postprocess(out, binding, reduction_op, nout, opt_arg, dtype): return out -def ConjugateGradientSolver(binding, linop, b, eps=1e-6): +def ConjugateGradientSolver(binding, linop, b, x0, eps=1e-6): # Conjugate gradient algorithm to solve linear system of the form # Ma=b where linop is a linear operation corresponding # to a symmetric and positive definite matrix tools = get_tools(binding) - delta = tools.size(b) * eps**2 - a = 0 - r = tools.copy(b) + delta = tools.size(b) * eps**2 # FIXME: See scipy cg implementation + a = 0 if x0 is None else tools.copy(x0) + r = tools.copy(b) if x0 is None else b - linop(x0) nr2 = (r**2).sum() if nr2 < delta: return 0 * r @@ -105,7 +105,7 @@ def ConjugateGradientSolver(binding, linop, b, eps=1e-6): alp = nr2 / (p * Mp).sum() a += alp * p r -= alp * Mp - nr2new = (r**2).sum() + nr2new = (r ** 2).sum() if nr2new < delta: break p = r + (nr2new / nr2) * p @@ -115,18 +115,18 @@ def ConjugateGradientSolver(binding, linop, b, eps=1e-6): def KernelLinearSolver( - binding, K, x, b, alpha=0, eps=1e-6, precond=False, precondKernel=None + binding, K, x, b, x0=None, alpha=0, eps=1e-6, precond=False, precondKernel=None ): tools = get_tools(binding) dtype = tools.dtype(x) - def PreconditionedConjugateGradientSolver(linop, b, invprecondop, eps=1e-6): + def PreconditionedConjugateGradientSolver(linop, b, invprecondop, x0, eps=1e-6): # Preconditioned conjugate gradient algorithm to solve linear system of the form # Ma=b where linop is a linear operation corresponding # to a symmetric and positive definite matrix # invprecondop is linear operation corresponding to the inverse of the preconditioner matrix - a = 0 - r = tools.copy(b) + a = 0 if x0 is None else tools.copy(x0) + r = tools.copy(b) if x0 is None else b - linop(x0) z = invprecondop(r) p = tools.copy(z) rz = (r * z).sum() @@ -228,6 +228,6 @@ def f(x, y): invprecondop = NystromInversePreconditioner(K, precondKernel, x, alpha) a = PreconditionedConjugateGradientSolver(KernelLinOp, b, invprecondop, eps) else: - a = ConjugateGradientSolver(binding, KernelLinOp, b, eps=eps) + a = ConjugateGradientSolver(binding, KernelLinOp, b, x0, eps=eps) return a From 67b289ef453a7f23f0e5f2256b779986c788ea4c Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Wed, 8 Apr 2026 14:41:26 +0200 Subject: [PATCH 02/98] Fix backend helper function binding --- pykeops/pykeops/numpy/utils.py | 10 +++++----- pykeops/pykeops/test/test_numpy.py | 14 ++++++++++++++ pykeops/pykeops/test/test_torch.py | 13 +++++++++++++ pykeops/pykeops/torch/utils.py | 10 +++++----- 4 files changed, 37 insertions(+), 10 deletions(-) diff --git a/pykeops/pykeops/numpy/utils.py b/pykeops/pykeops/numpy/utils.py index cc0d015d4..f5683910a 100644 --- a/pykeops/pykeops/numpy/utils.py +++ b/pykeops/pykeops/numpy/utils.py @@ -6,13 +6,13 @@ class numpytools: - norm = np.linalg.norm - arraysum = np.sum - exp = np.exp - log = np.log + norm = staticmethod(np.linalg.norm) + arraysum = staticmethod(np.sum) + exp = staticmethod(np.exp) + log = staticmethod(np.log) Genred = Genred KernelSolve = KernelSolve - swap_axes = np_swap_axes + swap_axes = staticmethod(np_swap_axes) arraytype = np.ndarray float_types = [float, np.float16, np.float32, np.float64] diff --git a/pykeops/pykeops/test/test_numpy.py b/pykeops/pykeops/test/test_numpy.py index 3b66ed731..028926d3b 100644 --- a/pykeops/pykeops/test/test_numpy.py +++ b/pykeops/pykeops/test/test_numpy.py @@ -54,6 +54,20 @@ class NumpyUnitTestCase(unittest.TestCase): type_to_test = ["float32", "float64"] + ############################################################ + def test_numpytools_function_binding(self): + ############################################################ + from pykeops.numpy.utils import numpytools + + tools = numpytools() + x = self.x.astype(self.type_to_test[0]) + + self.assertTrue(np.array_equal(tools.copy(x), x)) + self.assertTrue(np.allclose(tools.exp(x), np.exp(x))) + self.assertTrue(np.allclose(tools.log(x + 1), np.log(x + 1))) + self.assertTrue(np.allclose(tools.norm(x), np.linalg.norm(x))) + self.assertTrue(np.allclose(tools.arraysum(x, axis=0), np.sum(x, axis=0))) + ############################################################ def test_generic_syntax_sum(self): ############################################################ diff --git a/pykeops/pykeops/test/test_torch.py b/pykeops/pykeops/test/test_torch.py index d6402cb00..d7d8aad72 100644 --- a/pykeops/pykeops/test/test_torch.py +++ b/pykeops/pykeops/test/test_torch.py @@ -100,6 +100,19 @@ class PytorchUnitTestCase(unittest.TestCase): print("Pytorch could not be loaded. Skip tests.") pass + ############################################################ + def test_torchtools_function_binding(self): + ############################################################ + from pykeops.torch.utils import torchtools + + tools = torchtools() + x = self.xc.detach() + + self.assertTrue(torch.equal(tools.copy(x), x)) + self.assertTrue(torch.allclose(tools.exp(x), torch.exp(x))) + self.assertTrue(torch.allclose(tools.log(x + 1), torch.log(x + 1))) + self.assertTrue(torch.allclose(tools.norm(x), torch.norm(x))) + ############################################################ def test_generic_syntax_float(self): ############################################################ diff --git a/pykeops/pykeops/torch/utils.py b/pykeops/pykeops/torch/utils.py index 45fb1d509..36749cfae 100644 --- a/pykeops/pykeops/torch/utils.py +++ b/pykeops/pykeops/torch/utils.py @@ -18,12 +18,12 @@ def is_on_device(x): class torchtools: - copy = torch.clone - exp = torch.exp - log = torch.log - norm = torch.norm + copy = staticmethod(torch.clone) + exp = staticmethod(torch.exp) + log = staticmethod(torch.log) + norm = staticmethod(torch.norm) - swap_axes = torch_swap_axes + swap_axes = staticmethod(torch_swap_axes) Genred = Genred KernelSolve = KernelSolve From 641eb1a7973fa2b9f091a55baa2b6dcab1883255 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Wed, 8 Apr 2026 14:49:36 +0200 Subject: [PATCH 03/98] Make Keops_Print compatible with multiple messages --- keopscore/keopscore/utils/misc_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/keopscore/keopscore/utils/misc_utils.py b/keopscore/keopscore/utils/misc_utils.py index 887127dba..bced3371d 100644 --- a/keopscore/keopscore/utils/misc_utils.py +++ b/keopscore/keopscore/utils/misc_utils.py @@ -8,9 +8,9 @@ import platform -def KeOps_Print(message, force_print=False, **kwargs): +def KeOps_Print(*messages, force_print=False, **kwargs): if keopscore.verbose or force_print: - print(message, **kwargs) + print(*messages, **kwargs) def KeOps_Message(message, use_tag=True, **kwargs): From b6b1a087526a2cd6660155b1140d81f1dd6e3c98 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Wed, 8 Apr 2026 14:51:56 +0200 Subject: [PATCH 04/98] Add default formula --- keopscore/keopscore/sandbox/do_test_formula.py | 8 +++++++- keopscore/keopscore/utils/TestFormula.py | 4 ++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/keopscore/keopscore/sandbox/do_test_formula.py b/keopscore/keopscore/sandbox/do_test_formula.py index 85833e601..678b5e570 100644 --- a/keopscore/keopscore/sandbox/do_test_formula.py +++ b/keopscore/keopscore/sandbox/do_test_formula.py @@ -2,13 +2,19 @@ from keopscore.utils.TestFormula import TestFormula if __name__ == "__main__": - if len(sys.argv) == 1: + + if len(sys.argv) == 0: + # Add a default formula + res = TestFormula("Exp(-Sum((Var(0,3,0)-Var(1,3,1))**2))*Var(2,1,1)") + + elif len(sys.argv) == 1: print( """ you should provide a formula to test, e.g. 'python do_test_formula.py "Exp(-Sum((Var(0,3,0)-Var(1,3,1))**2))*Var(2,1,1)"' """ ) + else: if len(sys.argv) == 2: res = TestFormula(sys.argv[1]) diff --git a/keopscore/keopscore/utils/TestFormula.py b/keopscore/keopscore/utils/TestFormula.py index a3dd253fc..d249ce666 100644 --- a/keopscore/keopscore/utils/TestFormula.py +++ b/keopscore/keopscore/utils/TestFormula.py @@ -80,3 +80,7 @@ def TestFormula(formula, tol=1e-4, dtype="float32", test_grad=False, randseed=No ) return c, g return c + + +if __name__ == "__main__": + res = TestFormula("Exp(-Sum((Var(0,3,0)-Var(1,3,1))**2))*Var(2,1,1)") \ No newline at end of file From 86062ac1d6f78f80bbc7f57d4bb8fa5836464c15 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Wed, 8 Apr 2026 17:20:27 +0200 Subject: [PATCH 05/98] add precommit --- .pre-commit-config.yaml | 14 +++++++++ keopscore/keopscore/utils/TestFormula.py | 2 +- pykeops/pykeops/common/operations.py | 36 ++++++++++++++++++------ pykeops/pykeops/test/test_torch.py | 1 + pykeops/pykeops/torch/operations.py | 2 +- 5 files changed, 44 insertions(+), 11 deletions(-) create mode 100644 .pre-commit-config.yaml diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 000000000..79294ee1c --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,14 @@ +repos: + - repo: https://github.com/psf/black + rev: 24.10.0 + hooks: + - id: black + + - repo: local + hooks: + - id: pytest + name: pytest + entry: ./pytest.sh + language: system + pass_filenames: false + always_run: true diff --git a/keopscore/keopscore/utils/TestFormula.py b/keopscore/keopscore/utils/TestFormula.py index d249ce666..21857d45d 100644 --- a/keopscore/keopscore/utils/TestFormula.py +++ b/keopscore/keopscore/utils/TestFormula.py @@ -83,4 +83,4 @@ def TestFormula(formula, tol=1e-4, dtype="float32", test_grad=False, randseed=No if __name__ == "__main__": - res = TestFormula("Exp(-Sum((Var(0,3,0)-Var(1,3,1))**2))*Var(2,1,1)") \ No newline at end of file + res = TestFormula("Exp(-Sum((Var(0,3,0)-Var(1,3,1))**2))*Var(2,1,1)") diff --git a/pykeops/pykeops/common/operations.py b/pykeops/pykeops/common/operations.py index 37ad5bd24..41c78377c 100644 --- a/pykeops/pykeops/common/operations.py +++ b/pykeops/pykeops/common/operations.py @@ -87,16 +87,18 @@ def postprocess(out, binding, reduction_op, nout, opt_arg, dtype): return out -def ConjugateGradientSolver(binding, linop, b, x0, eps=1e-6): +def ConjugateGradientSolver(binding, linop, b, x0=None, eps=1e-6): # Conjugate gradient algorithm to solve linear system of the form # Ma=b where linop is a linear operation corresponding # to a symmetric and positive definite matrix + # Follow the scipy conjugate gradient implementation tools = get_tools(binding) - delta = tools.size(b) * eps**2 # FIXME: See scipy cg implementation + atol, _ = _get_atol_rtol(tools.norm(b), eps) + a = 0 if x0 is None else tools.copy(x0) r = tools.copy(b) if x0 is None else b - linop(x0) nr2 = (r**2).sum() - if nr2 < delta: + if nr2 < atol * atol: return 0 * r p = tools.copy(r) k = 0 @@ -105,8 +107,8 @@ def ConjugateGradientSolver(binding, linop, b, x0, eps=1e-6): alp = nr2 / (p * Mp).sum() a += alp * p r -= alp * Mp - nr2new = (r ** 2).sum() - if nr2new < delta: + nr2new = (r**2).sum() + if nr2new < atol * atol: break p = r + (nr2new / nr2) * p nr2 = nr2new @@ -114,17 +116,31 @@ def ConjugateGradientSolver(binding, linop, b, x0, eps=1e-6): return a +def _get_atol_rtol(b_norm, atol=0.0, rtol=1e-5): + """ + A helper function to handle tolerance normalization. See scipy.linalg.cg. + """ + atol = max(float(atol), float(rtol) * float(b_norm)) + + return atol, rtol + + def KernelLinearSolver( binding, K, x, b, x0=None, alpha=0, eps=1e-6, precond=False, precondKernel=None ): tools = get_tools(binding) dtype = tools.dtype(x) - def PreconditionedConjugateGradientSolver(linop, b, invprecondop, x0, eps=1e-6): + def PreconditionedConjugateGradientSolver( + linop, b, invprecondop, x0=None, eps=1e-6 + ): # Preconditioned conjugate gradient algorithm to solve linear system of the form # Ma=b where linop is a linear operation corresponding # to a symmetric and positive definite matrix # invprecondop is linear operation corresponding to the inverse of the preconditioner matrix + + atol, _ = _get_atol_rtol(tools.norm(b), eps) + a = 0 if x0 is None else tools.copy(x0) r = tools.copy(b) if x0 is None else b - linop(x0) z = invprecondop(r) @@ -135,7 +151,7 @@ def PreconditionedConjugateGradientSolver(linop, b, invprecondop, x0, eps=1e-6): alp = rz / (p * linop(p)).sum() a += alp * p r -= alp * linop(p) - if (r**2).sum() < eps**2: + if (r**2).sum() < atol * atol: break z = invprecondop(r) rznew = (r * z).sum() @@ -226,8 +242,10 @@ def f(x, y): if precond: invprecondop = NystromInversePreconditioner(K, precondKernel, x, alpha) - a = PreconditionedConjugateGradientSolver(KernelLinOp, b, invprecondop, eps) + a = PreconditionedConjugateGradientSolver( + KernelLinOp, b, invprecondop, x0=x0, eps=eps + ) else: - a = ConjugateGradientSolver(binding, KernelLinOp, b, x0, eps=eps) + a = ConjugateGradientSolver(binding, KernelLinOp, b, x0=x0, eps=eps) return a diff --git a/pykeops/pykeops/test/test_torch.py b/pykeops/pykeops/test/test_torch.py index d7d8aad72..85df1b21b 100644 --- a/pykeops/pykeops/test/test_torch.py +++ b/pykeops/pykeops/test/test_torch.py @@ -104,6 +104,7 @@ class PytorchUnitTestCase(unittest.TestCase): def test_torchtools_function_binding(self): ############################################################ from pykeops.torch.utils import torchtools + import torch tools = torchtools() x = self.xc.detach() diff --git a/pykeops/pykeops/torch/operations.py b/pykeops/pykeops/torch/operations.py index e6232e88b..5a784fe41 100644 --- a/pykeops/pykeops/torch/operations.py +++ b/pykeops/pykeops/torch/operations.py @@ -78,7 +78,7 @@ def linop(var): res += params.alpha * var return res - result = ConjugateGradientSolver("torch", linop, varinv.data, params.eps) + result = ConjugateGradientSolver("torch", linop, varinv.data, eps=params.eps) # relying on the 'ctx.saved_variables' attribute is necessary if you want to be able to differentiate the output # of the backward once again. It helps pytorch to keep track of 'who is who'. From 9e49563dea2fb1b6f9bfbd03332be814dac8b385 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Thu, 9 Apr 2026 16:19:04 +0200 Subject: [PATCH 06/98] black files --- benchmarks/PyTorch_GPU.py | 1 - benchmarks/TVM.py | 1 - doc/tools/nb_to_doc.py | 1 + .../keopscore/formulas/LinearOperators/AdjointOperator.py | 1 - keopscore/keopscore/formulas/LinearOperators/LinearOperator.py | 1 - keopscore/keopscore/formulas/LinearOperators/SumLinOperator.py | 1 - keopscore/keopscore/formulas/LinearOperators/TraceOperator.py | 1 - keopscore/keopscore/formulas/complex/ComplexAbs.py | 1 - keopscore/keopscore/formulas/complex/ComplexExp.py | 1 - keopscore/keopscore/formulas/maths/Add.py | 1 - keopscore/keopscore/formulas/maths/Atan2.py | 1 - keopscore/keopscore/formulas/maths/Concat.py | 1 - keopscore/keopscore/formulas/maths/Exp.py | 1 - keopscore/keopscore/formulas/maths/GradMatrix.py | 1 - keopscore/keopscore/formulas/maths/Inv.py | 1 - keopscore/keopscore/formulas/maths/Minus.py | 1 - keopscore/keopscore/formulas/maths/Rsqrt.py | 1 - keopscore/keopscore/formulas/maths/Sign.py | 1 - keopscore/keopscore/formulas/maths/Sqrt.py | 1 - keopscore/keopscore/formulas/maths/Square.py | 1 - keopscore/keopscore/formulas/maths/Sum.py | 1 - keopscore/keopscore/formulas/maths/TensorProd.py | 1 - keopscore/keopscore/formulas/variables/Var.py | 1 - keopscore/keopscore/sandbox/lazytensor_tensordot.py | 1 - pykeops/pykeops/benchmarks/plot_benchmark_convolutions.py | 1 - pykeops/pykeops/benchmarks/plot_benchmark_grad1convolutions.py | 1 - pykeops/pykeops/examples/pytorch/plot_anisotropic_kernels.py | 1 - pykeops/pykeops/sandbox/issue_327.py | 1 - pykeops/pykeops/sandbox/issue_353.py | 1 - pykeops/pykeops/sandbox/test_ifelse.py | 1 - pykeops/pykeops/test/test_block_sparse_reduction.py | 1 - pykeops/pykeops/torch/lazytensor/LazyTensor.py | 1 - pykeops/pykeops/torch/utils.py | 1 - .../tutorials/gaussian_mixture/plot_gaussian_mixture.py | 3 +-- 34 files changed, 2 insertions(+), 34 deletions(-) diff --git a/benchmarks/PyTorch_GPU.py b/benchmarks/PyTorch_GPU.py index bfecaa956..1836f7f9a 100644 --- a/benchmarks/PyTorch_GPU.py +++ b/benchmarks/PyTorch_GPU.py @@ -12,7 +12,6 @@ import numpy as np from time import time - nits = 100 Ns, D = [10000, 100000, 1000000], 3 diff --git a/benchmarks/TVM.py b/benchmarks/TVM.py index 837021116..ec4be25cd 100644 --- a/benchmarks/TVM.py +++ b/benchmarks/TVM.py @@ -41,7 +41,6 @@ import numpy as np from time import time - # Global declarations of environment. tgt_host = "llvm" tgt = "cuda" diff --git a/doc/tools/nb_to_doc.py b/doc/tools/nb_to_doc.py index 7c594e837..0b1fe76e3 100755 --- a/doc/tools/nb_to_doc.py +++ b/doc/tools/nb_to_doc.py @@ -3,6 +3,7 @@ Convert empty IPython notebook to a sphinx doc page. """ + import sys from subprocess import check_call as sh diff --git a/keopscore/keopscore/formulas/LinearOperators/AdjointOperator.py b/keopscore/keopscore/formulas/LinearOperators/AdjointOperator.py index 9411ca848..f6b980fac 100644 --- a/keopscore/keopscore/formulas/LinearOperators/AdjointOperator.py +++ b/keopscore/keopscore/formulas/LinearOperators/AdjointOperator.py @@ -21,7 +21,6 @@ from keopscore.formulas.maths.Scalprod import Scalprod_Impl from keopscore.formulas.Operation import Broadcast - # ///////////////////////////////////////////////////////////// # /// ADJOINT OPERATOR //// # ///////////////////////////////////////////////////////////// diff --git a/keopscore/keopscore/formulas/LinearOperators/LinearOperator.py b/keopscore/keopscore/formulas/LinearOperators/LinearOperator.py index bd667da26..503993799 100644 --- a/keopscore/keopscore/formulas/LinearOperators/LinearOperator.py +++ b/keopscore/keopscore/formulas/LinearOperators/LinearOperator.py @@ -15,7 +15,6 @@ from keopscore.formulas.maths.Scalprod import Scalprod_Impl from keopscore.formulas.Operation import Broadcast - # ///////////////////////////////////////////////////////////// # /// LINEAR OPERATORS //// # ///////////////////////////////////////////////////////////// diff --git a/keopscore/keopscore/formulas/LinearOperators/SumLinOperator.py b/keopscore/keopscore/formulas/LinearOperators/SumLinOperator.py index 548077a41..337db5301 100644 --- a/keopscore/keopscore/formulas/LinearOperators/SumLinOperator.py +++ b/keopscore/keopscore/formulas/LinearOperators/SumLinOperator.py @@ -17,7 +17,6 @@ from keopscore.formulas.maths.Scalprod import Scalprod_Impl from keopscore.formulas.Operation import Broadcast - # ///////////////////////////////////////////////////////////// # /// SUM_LIN OPERATOR //// # ///////////////////////////////////////////////////////////// diff --git a/keopscore/keopscore/formulas/LinearOperators/TraceOperator.py b/keopscore/keopscore/formulas/LinearOperators/TraceOperator.py index d5ab12ac4..ee659857b 100644 --- a/keopscore/keopscore/formulas/LinearOperators/TraceOperator.py +++ b/keopscore/keopscore/formulas/LinearOperators/TraceOperator.py @@ -18,7 +18,6 @@ from keopscore.formulas.maths.Scalprod import Scalprod_Impl from keopscore.formulas.Operation import Broadcast - # ///////////////////////////////////////////////////////////// # /// TRACE OPERATOR //// # ///////////////////////////////////////////////////////////// diff --git a/keopscore/keopscore/formulas/complex/ComplexAbs.py b/keopscore/keopscore/formulas/complex/ComplexAbs.py index 8f0c1e41c..758050cbb 100644 --- a/keopscore/keopscore/formulas/complex/ComplexAbs.py +++ b/keopscore/keopscore/formulas/complex/ComplexAbs.py @@ -1,7 +1,6 @@ from keopscore.formulas.complex.ComplexSquareAbs import ComplexSquareAbs from keopscore.formulas.maths.Sqrt import Sqrt - # ///////////////////////////////////////////////////////////////////////// # //// ComplexAbs //// # ///////////////////////////////////////////////////////////////////////// diff --git a/keopscore/keopscore/formulas/complex/ComplexExp.py b/keopscore/keopscore/formulas/complex/ComplexExp.py index c14d6c1b5..8f3a19340 100644 --- a/keopscore/keopscore/formulas/complex/ComplexExp.py +++ b/keopscore/keopscore/formulas/complex/ComplexExp.py @@ -14,7 +14,6 @@ from keopscore.formulas.maths.Cos import Cos from keopscore.formulas.maths.Sin import Sin - # ///////////////////////////////////////////////////////////////////////// # //// ComplexExp //// # ///////////////////////////////////////////////////////////////////////// diff --git a/keopscore/keopscore/formulas/maths/Add.py b/keopscore/keopscore/formulas/maths/Add.py index ecb42f534..7be4aafd8 100644 --- a/keopscore/keopscore/formulas/maths/Add.py +++ b/keopscore/keopscore/formulas/maths/Add.py @@ -7,7 +7,6 @@ from keopscore.formulas.variables.RatCst import RatCst, RatCst_Impl from keopscore.formulas.variables.Zero import Zero - ########################## ###### Add ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/Atan2.py b/keopscore/keopscore/formulas/maths/Atan2.py index c3cb4fb8e..b9f3d7086 100644 --- a/keopscore/keopscore/formulas/maths/Atan2.py +++ b/keopscore/keopscore/formulas/maths/Atan2.py @@ -1,7 +1,6 @@ from keopscore.formulas.VectorizedScalarOp import VectorizedScalarOp from keopscore.utils.math_functions import keops_atan2 - # ////////////////////////////////////////////////////////////// # //// ATAN2 : Atan2< F, G > //// # ////////////////////////////////////////////////////////////// diff --git a/keopscore/keopscore/formulas/maths/Concat.py b/keopscore/keopscore/formulas/maths/Concat.py index b10f55886..a104733bf 100644 --- a/keopscore/keopscore/formulas/maths/Concat.py +++ b/keopscore/keopscore/formulas/maths/Concat.py @@ -3,7 +3,6 @@ from keopscore.utils.code_gen_utils import VectCopy from keopscore.utils.misc_utils import KeOps_Error - ############################ ###### Concat ##### ############################ diff --git a/keopscore/keopscore/formulas/maths/Exp.py b/keopscore/keopscore/formulas/maths/Exp.py index b52514cbd..a6b1a7be3 100644 --- a/keopscore/keopscore/formulas/maths/Exp.py +++ b/keopscore/keopscore/formulas/maths/Exp.py @@ -1,7 +1,6 @@ from keopscore.formulas.VectorizedScalarOp import VectorizedScalarOp from keopscore.utils.math_functions import keops_exp - ########################## ###### Exp ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/GradMatrix.py b/keopscore/keopscore/formulas/maths/GradMatrix.py index a6c8495ef..215b5474e 100644 --- a/keopscore/keopscore/formulas/maths/GradMatrix.py +++ b/keopscore/keopscore/formulas/maths/GradMatrix.py @@ -4,7 +4,6 @@ from keopscore.formulas.variables.Var import Var from keopscore.utils.code_gen_utils import GetInds - # ////////////////////////////////////////////////////////////////////////////////////////////// # //// Standard basis of R^DIM : < (1,0,0,...) , (0,1,0,...) , ... , (0,...,0,1) > //// # ////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/keopscore/keopscore/formulas/maths/Inv.py b/keopscore/keopscore/formulas/maths/Inv.py index a0355ffec..0e3905017 100644 --- a/keopscore/keopscore/formulas/maths/Inv.py +++ b/keopscore/keopscore/formulas/maths/Inv.py @@ -1,7 +1,6 @@ from keopscore.formulas.VectorizedScalarOp import VectorizedScalarOp from keopscore.utils.math_functions import keops_rcp - ########################## ###### INVERSE : Inv ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/Minus.py b/keopscore/keopscore/formulas/maths/Minus.py index f374d1867..5a59c193b 100644 --- a/keopscore/keopscore/formulas/maths/Minus.py +++ b/keopscore/keopscore/formulas/maths/Minus.py @@ -4,7 +4,6 @@ from keopscore.formulas.variables.IntCst import IntCst_Impl, IntCst from keopscore.formulas.variables.RatCst import RatCst_Impl, RatCst - ########################## ###### Minus ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/Rsqrt.py b/keopscore/keopscore/formulas/maths/Rsqrt.py index 62454ccc6..f8f930831 100644 --- a/keopscore/keopscore/formulas/maths/Rsqrt.py +++ b/keopscore/keopscore/formulas/maths/Rsqrt.py @@ -1,7 +1,6 @@ from keopscore.formulas.VectorizedScalarOp import VectorizedScalarOp from keopscore.formulas.maths.IntInv import IntInv - ########################## ###### Rsqrt ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/Sign.py b/keopscore/keopscore/formulas/maths/Sign.py index 760e38b10..122faea48 100644 --- a/keopscore/keopscore/formulas/maths/Sign.py +++ b/keopscore/keopscore/formulas/maths/Sign.py @@ -2,7 +2,6 @@ from keopscore.formulas.variables.Zero import Zero from keopscore.utils.math_functions import keops_sign - ########################## ###### Sign ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/Sqrt.py b/keopscore/keopscore/formulas/maths/Sqrt.py index ee00b1b51..12696cfde 100644 --- a/keopscore/keopscore/formulas/maths/Sqrt.py +++ b/keopscore/keopscore/formulas/maths/Sqrt.py @@ -3,7 +3,6 @@ from keopscore.formulas.maths.Rsqrt import Rsqrt from keopscore.utils.math_functions import keops_sqrt - ########################## ###### Sqrt ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/Square.py b/keopscore/keopscore/formulas/maths/Square.py index 808bc8698..1ca9ced5b 100644 --- a/keopscore/keopscore/formulas/maths/Square.py +++ b/keopscore/keopscore/formulas/maths/Square.py @@ -2,7 +2,6 @@ from keopscore.formulas.variables.Zero import Zero - ########################## ###### Square ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/Sum.py b/keopscore/keopscore/formulas/maths/Sum.py index f0483ae6a..45632508a 100644 --- a/keopscore/keopscore/formulas/maths/Sum.py +++ b/keopscore/keopscore/formulas/maths/Sum.py @@ -3,7 +3,6 @@ from keopscore.utils.code_gen_utils import c_zero_float, VectApply from keopscore.formulas.maths.Square import Square_Impl - ########################## ###### Sum ##### ########################## diff --git a/keopscore/keopscore/formulas/maths/TensorProd.py b/keopscore/keopscore/formulas/maths/TensorProd.py index f03001e86..44f5e72eb 100644 --- a/keopscore/keopscore/formulas/maths/TensorProd.py +++ b/keopscore/keopscore/formulas/maths/TensorProd.py @@ -5,7 +5,6 @@ ) from keopscore.utils.misc_utils import KeOps_Error - #################################### ###### Tensor product ##### #################################### diff --git a/keopscore/keopscore/formulas/variables/Var.py b/keopscore/keopscore/formulas/variables/Var.py index ae299df4e..2d009ba67 100644 --- a/keopscore/keopscore/formulas/variables/Var.py +++ b/keopscore/keopscore/formulas/variables/Var.py @@ -1,7 +1,6 @@ from keopscore.utils.code_gen_utils import VectCopy from keopscore.formulas.Operation import Operation - ####################### ## Var operation ####################### diff --git a/keopscore/keopscore/sandbox/lazytensor_tensordot.py b/keopscore/keopscore/sandbox/lazytensor_tensordot.py index 9a9d44a6f..d43213667 100644 --- a/keopscore/keopscore/sandbox/lazytensor_tensordot.py +++ b/keopscore/keopscore/sandbox/lazytensor_tensordot.py @@ -11,7 +11,6 @@ from pykeops.torch import LazyTensor - M, N = 2, 10 ####################################################################################################################### diff --git a/pykeops/pykeops/benchmarks/plot_benchmark_convolutions.py b/pykeops/pykeops/benchmarks/plot_benchmark_convolutions.py index 2ecdf516a..04a88f253 100644 --- a/pykeops/pykeops/benchmarks/plot_benchmark_convolutions.py +++ b/pykeops/pykeops/benchmarks/plot_benchmark_convolutions.py @@ -27,7 +27,6 @@ from pykeops.torch.utils import torch_kernel from pykeops.torch import Vi, Vj, Pm - ###################################################################### # Benchmark specifications: # diff --git a/pykeops/pykeops/benchmarks/plot_benchmark_grad1convolutions.py b/pykeops/pykeops/benchmarks/plot_benchmark_grad1convolutions.py index cbe438cbc..a6922dc96 100644 --- a/pykeops/pykeops/benchmarks/plot_benchmark_grad1convolutions.py +++ b/pykeops/pykeops/benchmarks/plot_benchmark_grad1convolutions.py @@ -28,7 +28,6 @@ from pykeops.torch.utils import torch_kernel from pykeops.torch import Vi, Vj, Pm - ###################################################################### # Benchmark specifications: # diff --git a/pykeops/pykeops/examples/pytorch/plot_anisotropic_kernels.py b/pykeops/pykeops/examples/pytorch/plot_anisotropic_kernels.py index 837cf964e..26d0bc743 100644 --- a/pykeops/pykeops/examples/pytorch/plot_anisotropic_kernels.py +++ b/pykeops/pykeops/examples/pytorch/plot_anisotropic_kernels.py @@ -20,7 +20,6 @@ import torch from pykeops.torch import Vi, Vj, Pm, LazyTensor - ############################################## # Dataset: diff --git a/pykeops/pykeops/sandbox/issue_327.py b/pykeops/pykeops/sandbox/issue_327.py index c37742a7d..9c6c01544 100644 --- a/pykeops/pykeops/sandbox/issue_327.py +++ b/pykeops/pykeops/sandbox/issue_327.py @@ -1,7 +1,6 @@ from pykeops.torch import Vi, Vj, LazyTensor import torch - xc = torch.randn(256, 5) xc.requires_grad_(True) diff --git a/pykeops/pykeops/sandbox/issue_353.py b/pykeops/pykeops/sandbox/issue_353.py index 7245ea04a..1017b1c1e 100644 --- a/pykeops/pykeops/sandbox/issue_353.py +++ b/pykeops/pykeops/sandbox/issue_353.py @@ -2,7 +2,6 @@ from pykeops.torch import Vi, Vj, LazyTensor, Genred import pykeops - d = 1 n, m = 1, 1 q_points = torch.rand((n, d)).requires_grad_(True).to("cuda") diff --git a/pykeops/pykeops/sandbox/test_ifelse.py b/pykeops/pykeops/sandbox/test_ifelse.py index 3593a0ef6..58f8f0557 100644 --- a/pykeops/pykeops/sandbox/test_ifelse.py +++ b/pykeops/pykeops/sandbox/test_ifelse.py @@ -2,7 +2,6 @@ import pykeops from pykeops.torch import LazyTensor - ttypes = ( (torch.cuda.FloatTensor,) if torch.cuda.is_available() else (torch.FloatTensor,) ) diff --git a/pykeops/pykeops/test/test_block_sparse_reduction.py b/pykeops/pykeops/test/test_block_sparse_reduction.py index 3ca6fe1e9..90e1590d5 100644 --- a/pykeops/pykeops/test/test_block_sparse_reduction.py +++ b/pykeops/pykeops/test/test_block_sparse_reduction.py @@ -3,7 +3,6 @@ import torch from pykeops.torch import LazyTensor - # Import clustering functions from KeOps from pykeops.torch.cluster import ( grid_cluster, diff --git a/pykeops/pykeops/torch/lazytensor/LazyTensor.py b/pykeops/pykeops/torch/lazytensor/LazyTensor.py index 656b7cbc2..69785283a 100644 --- a/pykeops/pykeops/torch/lazytensor/LazyTensor.py +++ b/pykeops/pykeops/torch/lazytensor/LazyTensor.py @@ -3,7 +3,6 @@ from pykeops.common.lazy_tensor import GenericLazyTensor, ComplexGenericLazyTensor from pykeops.torch.utils import torchtools - # Convenient aliases: diff --git a/pykeops/pykeops/torch/utils.py b/pykeops/pykeops/torch/utils.py index 36749cfae..1d0182a4f 100644 --- a/pykeops/pykeops/torch/utils.py +++ b/pykeops/pykeops/torch/utils.py @@ -3,7 +3,6 @@ from pykeops.torch import Genred, KernelSolve from pykeops.torch.cluster import swap_axes as torch_swap_axes - # from pykeops.torch.generic.generic_red import GenredLowlevel diff --git a/pykeops/pykeops/tutorials/gaussian_mixture/plot_gaussian_mixture.py b/pykeops/pykeops/tutorials/gaussian_mixture/plot_gaussian_mixture.py index faa502585..f90ad5e05 100644 --- a/pykeops/pykeops/tutorials/gaussian_mixture/plot_gaussian_mixture.py +++ b/pykeops/pykeops/tutorials/gaussian_mixture/plot_gaussian_mixture.py @@ -23,7 +23,6 @@ from pykeops.torch import Vi, Vj, LazyTensor - #################################################################### # Define our dataset: a collection of points :math:`(x_i)_{i\in[1,N]}` which describe a # spiral in the unit square. @@ -114,7 +113,7 @@ def __init__(self, M, sparsity=0, D=2): def update_covariances(self): """Computes the full covariance matrices from the model's parameters.""" - (M, D, _) = self.A.shape + M, D, _ = self.A.shape self.params["gamma"] = (torch.matmul(self.A, self.A.transpose(1, 2))).view( M, D * D ) / 2 From ccd5b1dbd8f0b1991beb2e300cc97dffd65a97db Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Thu, 9 Apr 2026 16:20:12 +0200 Subject: [PATCH 07/98] black files --- keopscore/keopscore/sandbox/do_test_formula.py | 6 ++---- keopscore/keopscore/utils/gpu_utils.py | 12 ++++-------- 2 files changed, 6 insertions(+), 12 deletions(-) diff --git a/keopscore/keopscore/sandbox/do_test_formula.py b/keopscore/keopscore/sandbox/do_test_formula.py index 678b5e570..21f359de0 100644 --- a/keopscore/keopscore/sandbox/do_test_formula.py +++ b/keopscore/keopscore/sandbox/do_test_formula.py @@ -8,12 +8,10 @@ res = TestFormula("Exp(-Sum((Var(0,3,0)-Var(1,3,1))**2))*Var(2,1,1)") elif len(sys.argv) == 1: - print( - """ + print(""" you should provide a formula to test, e.g. 'python do_test_formula.py "Exp(-Sum((Var(0,3,0)-Var(1,3,1))**2))*Var(2,1,1)"' - """ - ) + """) else: if len(sys.argv) == 2: diff --git a/keopscore/keopscore/utils/gpu_utils.py b/keopscore/keopscore/utils/gpu_utils.py index df5b1debc..878e94f4a 100644 --- a/keopscore/keopscore/utils/gpu_utils.py +++ b/keopscore/keopscore/utils/gpu_utils.py @@ -88,16 +88,14 @@ def get_cuda_include_path(): return path # finally nothing found, so we display a warning asking the user to do something - KeOps_Warning( - """ + KeOps_Warning(""" The location of Cuda header files cuda.h and nvrtc.h could not be detected on your system. You must determine their location and then define the environment variable CUDA_PATH, either before launching Python or using os.environ before importing keops. For example if these files are in /vol/cuda/10.2.89-cudnn7.6.4.38/include you can do : import os os.environ['CUDA_PATH'] = '/vol/cuda/10.2.89-cudnn7.6.4.38' - """ - ) + """) def orig_cuda_include_fp16_path(): @@ -191,13 +189,11 @@ def get_gpu_props(): def safe_call(d, result): test = result == CUDA_SUCCESS if not test: - KeOps_Warning( - f""" + KeOps_Warning(f""" cuda was detected, driver API has been initialized, but there was an error for detecting properties of GPU device nr {d}. Switching to cpu only. - """ - ) + """) return test test = True From bd3acb240c12c973572e23d17d653a823b313fe9 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Tue, 14 Apr 2026 15:34:46 +0200 Subject: [PATCH 08/98] add ConjugateGradientSolver with x0, maxiter, and cv_info support --- pykeops/pykeops/common/lazy_tensor.py | 43 ++++--- pykeops/pykeops/common/operations.py | 116 +++++++++++++++--- pykeops/pykeops/numpy/operations.py | 18 ++- pykeops/pykeops/numpy/utils.py | 11 +- pykeops/pykeops/torch/operations.py | 52 ++++++-- pykeops/pykeops/torch/utils.py | 8 ++ .../plot_RBF_interpolation_numpy.py | 20 ++- .../plot_RBF_interpolation_torch.py | 18 ++- 8 files changed, 212 insertions(+), 74 deletions(-) diff --git a/pykeops/pykeops/common/lazy_tensor.py b/pykeops/pykeops/common/lazy_tensor.py index 791c5a2da..46765497c 100644 --- a/pykeops/pykeops/common/lazy_tensor.py +++ b/pykeops/pykeops/common/lazy_tensor.py @@ -328,28 +328,27 @@ def fixvariables(self): self.variables = newvars def separate_kwargs(self, kwargs): - # separating keyword arguments for Genred init vs Genred call... - # Currently the only additional optional keyword arguments that are passed to Genred init - # are accuracy options: dtype_acc, use_double_acc and sum_scheme, - # chunk mode option enable_chunks, - # use_fast_math option, - # and compiler option optional_flags. - kwargs_init = [] - kwargs_call = [] - for key in kwargs: - if key in ( - "dtype_acc", - "use_double_acc", - "sum_scheme", - "enable_chunks", - "use_fast_math", - "optional_flags", - ): - kwargs_init += [(key, kwargs[key])] - else: - kwargs_call += [(key, kwargs[key])] - kwargs_init = dict(kwargs_init) - kwargs_call = dict(kwargs_call) + """ + separating keyword arguments for Genred init vs Genred call... + Currently the only additional optional keyword arguments that are passed to Genred init are + accuracy options: dtype_acc, use_double_acc and sum_scheme, + chunk mode option enable_chunks, + use_fast_math option, + and compiler option optional_flags. + """ + init_keys = { + "dtype_acc", + "use_double_acc", + "sum_scheme", + "enable_chunks", + "use_fast_math", + "optional_flags", + } + + kwargs_init = {key: kwargs.get(key) for key in init_keys if key in kwargs} + kwargs_call = { + key: value for key, value in kwargs.items() if key not in init_keys + } return kwargs_init, kwargs_call def promote(self, other, props, is_complex=False): diff --git a/pykeops/pykeops/common/operations.py b/pykeops/pykeops/common/operations.py index 41c78377c..a26404ed5 100644 --- a/pykeops/pykeops/common/operations.py +++ b/pykeops/pykeops/common/operations.py @@ -87,33 +87,109 @@ def postprocess(out, binding, reduction_op, nout, opt_arg, dtype): return out -def ConjugateGradientSolver(binding, linop, b, x0=None, eps=1e-6): - # Conjugate gradient algorithm to solve linear system of the form - # Ma=b where linop is a linear operation corresponding - # to a symmetric and positive definite matrix - # Follow the scipy conjugate gradient implementation +def ConjugateGradientSolver( + binding, linop, b, x0=None, eps=1e-6, maxiter=None, cv_info=False +): + """ + Conjugate gradient algorithm to solve a linear system of the form + ``linop(x) = b``, where ``linop`` is a symmetric positive definite + linear operator. + + This implementation mostly follows SciPy's conjugate gradient solver. + + Parameters + ---------- + binding : str + Backend used by :func:`get_tools`, typically ``"numpy"`` or + ``"torch"``. + + linop : function + Function implementing the matrix-vector product associated with + the linear operator. Typically a PyKeOps routine. + + b : tensor + Right-hand side of the linear system. + + x0 : tensor, optional + Initial guess for the solution of the linear system. + + eps : float, optional + Relative tolerance used to define the absolute stopping threshold. + Internally, the solver stops when ``||r|| <= max(eps * ||b||, eps)``, + where ``r = b - linop(x)`` is the current residual. + + maxiter : int, optional + Maximum number of conjugate gradient iterations. Defaults to + ``10 * b.shape[0]``. + + cv_info : bool, optional + If ``False`` (default), return only the estimated solution ``x``. + If ``True``, return a tuple ``(x, info)`` where ``info`` is a + dictionary containing convergence diagnostics + Returns + ------- + x : tensor + Approximate solution returned by the conjugate gradient iterations. + + info : dict, optional + Returned only when ``cv_info=True``. Contains the convergence + information with the following + keys: + + - ``"status"``: ``"Converged"`` or ``"Maximum iterations reached"``. + - ``"niter"``: number of iterations performed. + - ``"residual_norm"``: final residual norm ``||r||``. + - ``"relative_residual_norm"``: final residual norm divided by + ``||b||``. + - ``"atol"``: absolute tolerance used internally for stopping. + - ``"maxiter"``: effective maximum number of iterations. + - ``"x0_provided"``: whether a non-``None`` initial guess was given. + - ``"dtype"``: data type of the solution ``x``. + """ + tools = get_tools(binding) + + # stopping criterion atol, _ = _get_atol_rtol(tools.norm(b), eps) + maxiter = maxiter if maxiter is not None else 10 * b.shape[0] - a = 0 if x0 is None else tools.copy(x0) + x = 0 if x0 is None else tools.copy(x0) r = tools.copy(b) if x0 is None else b - linop(x0) nr2 = (r**2).sum() if nr2 < atol * atol: - return 0 * r + return tools.zeros_like(x) + p = tools.copy(r) - k = 0 - while True: + + for it in range(maxiter): Mp = linop(p) alp = nr2 / (p * Mp).sum() - a += alp * p + x += alp * p r -= alp * Mp nr2new = (r**2).sum() if nr2new < atol * atol: break p = r + (nr2new / nr2) * p nr2 = nr2new - k += 1 - return a + + else: # for loop exhausted + # Return incomplete progress + it = -maxiter - 1 + + if cv_info: + nr = tools.sqrt(nr2new) + return x, { + "status": ("Converged" if nr <= atol else "Maximum iterations reached"), + "niter": it + 1, + "dtype": tools.dtypename(x.dtype), + "residual_norm": float(nr), + "relative_residual_norm": float(nr / tools.norm(b)), + "atol": atol, + "maxiter": maxiter, + "x0_provided": x0 is not None, + } + else: + return x def _get_atol_rtol(b_norm, atol=0.0, rtol=1e-5): @@ -126,7 +202,17 @@ def _get_atol_rtol(b_norm, atol=0.0, rtol=1e-5): def KernelLinearSolver( - binding, K, x, b, x0=None, alpha=0, eps=1e-6, precond=False, precondKernel=None + binding, + K, + x, + b, + x0=None, + alpha=0, + eps=1e-6, + maxiter=None, + precond=False, + precondKernel=None, + cv_info=False, ): tools = get_tools(binding) dtype = tools.dtype(x) @@ -246,6 +332,8 @@ def f(x, y): KernelLinOp, b, invprecondop, x0=x0, eps=eps ) else: - a = ConjugateGradientSolver(binding, KernelLinOp, b, x0=x0, eps=eps) + a = ConjugateGradientSolver( + binding, KernelLinOp, b, eps=eps, maxiter=maxiter, x0=x0, cv_info=cv_info + ) return a diff --git a/pykeops/pykeops/numpy/operations.py b/pykeops/pykeops/numpy/operations.py index b12876c0d..1c7246c82 100644 --- a/pykeops/pykeops/numpy/operations.py +++ b/pykeops/pykeops/numpy/operations.py @@ -1,11 +1,10 @@ import numpy as np - +from pykeops import default_device_id from pykeops.common.get_options import get_tag_backend from pykeops.common.keops_io import keops_binder from pykeops.common.operations import ConjugateGradientSolver from pykeops.common.parse_type import get_sizes, complete_aliases, get_optional_flags from pykeops.common.utils import axis2cat -from pykeops import default_device_id from pykeops.common.utils import pyKeOps_Warning @@ -177,7 +176,16 @@ def __init__( self.optional_flags = optional_flags def __call__( - self, *args, backend="auto", device_id=-1, alpha=1e-10, eps=1e-6, ranges=None + self, + *args, + backend="auto", + device_id=-1, + ranges=None, + alpha=1e-10, + eps=1e-6, + x0=None, + maxiter=None, + cv_info=False ): r""" To apply the routine on arbitrary NumPy arrays. @@ -272,4 +280,6 @@ def linop(var): res += alpha * var return res - return ConjugateGradientSolver("numpy", linop, varinv, eps=eps) + return ConjugateGradientSolver( + "numpy", linop, varinv, x0=x0, eps=eps, cv_info=cv_info + ) diff --git a/pykeops/pykeops/numpy/utils.py b/pykeops/pykeops/numpy/utils.py index f5683910a..0322596e1 100644 --- a/pykeops/pykeops/numpy/utils.py +++ b/pykeops/pykeops/numpy/utils.py @@ -1,8 +1,7 @@ import numpy as np - +import pykeops.config from pykeops.numpy import Genred, KernelSolve from pykeops.numpy.cluster import swap_axes as np_swap_axes -import pykeops.config class numpytools: @@ -107,6 +106,10 @@ def randn(m, n, dtype): def zeros(shape, dtype, device=None, requires_grad=None): return np.zeros(shape).astype(dtype) + @staticmethod + def zeros_like(x): + return np.zeros_like(x) + @staticmethod def empty(shape, dtype, device=None, requires_grad=None): return np.empty(shape, dtype=dtype) @@ -135,6 +138,10 @@ def device_type_index(x): def device_dict(x): return dict(cat="cpu") + @staticmethod + def sqrt(x): + return np.sqrt(x) + def squared_distances(x, y): x_norm = (x**2).sum(1).reshape(-1, 1) diff --git a/pykeops/pykeops/torch/operations.py b/pykeops/pykeops/torch/operations.py index 5a784fe41..8a26173ae 100644 --- a/pykeops/pykeops/torch/operations.py +++ b/pykeops/pykeops/torch/operations.py @@ -1,6 +1,6 @@ -import torch import copy +import torch from pykeops.common.get_options import get_tag_backend from pykeops.common.keops_io import keops_binder from pykeops.common.operations import ConjugateGradientSolver @@ -11,13 +11,12 @@ get_optional_flags, ) from pykeops.common.utils import axis2cat +from pykeops.common.utils import pyKeOps_Warning from pykeops.torch.generic.generic_red import ( GenredAutograd_fun, Genred_parameters, set_device, ) -from pykeops import default_device_id -from pykeops.common.utils import pyKeOps_Warning class KernelSolveAutograd(torch.autograd.Function): @@ -78,7 +77,15 @@ def linop(var): res += params.alpha * var return res - result = ConjugateGradientSolver("torch", linop, varinv.data, eps=params.eps) + result = ConjugateGradientSolver( + "torch", + linop, + varinv.data, + eps=params.eps, + x0=params.x0, + maxiter=params.maxiter, + cv_info=params.cv_info, + ) # relying on the 'ctx.saved_variables' attribute is necessary if you want to be able to differentiate the output # of the backward once again. It helps pytorch to keep track of 'who is who'. @@ -236,6 +243,7 @@ def __init__( that should be computed and reduced. The correct syntax is described in the :doc:`documentation <../../Genred>`, using appropriate :doc:`mathematical operations <../../../api/math-operations>`. + aliases (list of strings): A list of identifiers of the form ``"AL = TYPE(DIM)"`` that specify the categories and dimensions of the input variables. Here: @@ -250,16 +258,13 @@ def __init__( As described below, :meth:`__call__` will expect input Tensors whose shape are compatible with **aliases**. + varinvalias (string): The alphanumerical **alias** of the variable with respect to which we shall perform our conjugate gradient descent. **formula** is supposed to be linear with respect to **varinvalias**, but may be more sophisticated than a mere ``"K(x,y) * {varinvalias}"``. Keyword Args: - alpha (float, default = 1e-10): Non-negative - **ridge regularization** parameter, added to the diagonal - of the Kernel matrix :math:`K_{xx}`. - axis (int, default = 0): Specifies the dimension of the kernel matrix :math:`K_{x_ix_j}` that is reduced by our routine. The supported values are: @@ -332,7 +337,16 @@ def __init__( self.axis = axis def __call__( - self, *args, backend="auto", device_id=-1, alpha=1e-10, eps=1e-6, ranges=None + self, + *args, + backend="auto", + device_id=-1, + ranges=None, + alpha=1e-10, + eps=1e-6, + x0=None, + maxiter=None, + cv_info=False, ): r""" Apply the routine on arbitrary torch Tensors. @@ -373,6 +387,15 @@ def __call__( If **None** (default), we simply use a **dense Kernel matrix** as we loop over all indices :math:`i\in[0,M)` and :math:`j\in[0,N)`. + alpha (float, default = 1e-10): Non-negative floating-point + **ridge regularization** parameter, added to the diagonal + of the Kernel matrix :math:`K_{xx}`. + + eps (float, default = 1e-6): Stopping criterion for the conjugate gradient algorithm. + + x0 (2d Tensor, default = None): Initial guess for the solution of the linear system. + should be of the same shape as b. + Returns: (M,D) or (N,D) Tensor: @@ -389,18 +412,25 @@ def __call__( nx, ny = get_sizes(self.aliases, *args) params = Genred_parameters() + params.formula = self.formula params.aliases = self.aliases params.varinvpos = self.varinvpos - params.alpha = alpha params.backend = backend params.dtype = dtype + params.device_id_request = device_id - params.eps = eps params.ranges = ranges params.optional_flags = self.optional_flags params.rec_multVar_highdim = self.rec_multVar_highdim params.nx = nx params.ny = ny + # parameter of the Conjugate Gradient + params.alpha = alpha + params.eps = eps + params.x0 = x0 + params.cv_info = cv_info + params.maxiter = maxiter + return KernelSolveAutograd.apply(params, *args) diff --git a/pykeops/pykeops/torch/utils.py b/pykeops/pykeops/torch/utils.py index 1d0182a4f..c5df4a81d 100644 --- a/pykeops/pykeops/torch/utils.py +++ b/pykeops/pykeops/torch/utils.py @@ -153,6 +153,10 @@ def zeros( *shape, dtype=dtype, device=device, requires_grad=requires_grad ) + @staticmethod + def zeros_like(x): + return torch.zeros_like(x) + @staticmethod def empty( shape, @@ -209,6 +213,10 @@ def device_type_index(x): def pointer(x): return x.data.data_ptr() + @staticmethod + def sqrt(x): + return torch.sqrt(x) + def squared_distances(x, y): x_norm = (x**2).sum(1).reshape(-1, 1) diff --git a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py index 0109c3988..ea5ba15ec 100644 --- a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py +++ b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py @@ -29,11 +29,11 @@ import time -import numpy as np from matplotlib import pyplot as plt -from pykeops.numpy import LazyTensor +import numpy as np import pykeops.config +from pykeops.numpy import LazyTensor ####################################################################### # Generate some data: @@ -79,15 +79,14 @@ def gaussian_kernel(x, y, sigma=0.1): start = time.time() K_xx = gaussian_kernel(x, x) -a = K_xx.solve(b, alpha=alpha) +a, cv_info = K_xx.solve(b, alpha=alpha, cv_info=True) end = time.time() print( - "Time to perform an RBF interpolation with {:,} samples in 1D: {:.5f}s".format( - N, end - start - ) + f"Time to perform an RBF interpolation with {N:,} samples in 1D: {end - start:.5f}s" ) +print(f"Conjugate gradient infos: {cv_info}\n\n") ####################################################################### # Display the (fitted) model on the unit interval: @@ -146,20 +145,19 @@ def laplacian_kernel(x, y, sigma=0.1): # between a perfect fit (**alpha** = 0) and a # smooth interpolation (**alpha** = :math:`+\infty`): -alpha = 10 # Ridge regularization +alpha = 5 # Ridge regularization start = time.time() K_xx = laplacian_kernel(x, x) -a = K_xx.solve(b, alpha=alpha) +a, cv_info = K_xx.solve(b, alpha=alpha, cv_info=True) end = time.time() print( - "Time to perform an RBF interpolation with {:,} samples in 2D: {:.5f}s".format( - N, end - start - ) + f"Time to perform an RBF interpolation with {N:,} samples in 2D: {end - start:.5f}s" ) +print(f"Conjugate gradient infos: {cv_info}\n\n") ######################################################################## # Display the (fitted) model on the unit square: diff --git a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py index ea0a91b59..0b32c8bcb 100644 --- a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py +++ b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py @@ -29,9 +29,9 @@ import time -import torch from matplotlib import pyplot as plt +import torch from pykeops.torch import LazyTensor ############################################################################################### @@ -80,15 +80,14 @@ def gaussian_kernel(x, y, sigma=0.1): start = time.time() K_xx = gaussian_kernel(x, x) -a = K_xx.solve(b, alpha=alpha) +a, cv_info = K_xx.solve(b, alpha=alpha, cv_info=True) end = time.time() print( - "Time to perform an RBF interpolation with {:,} samples in 1D: {:.5f}s".format( - N, end - start - ) + f"Time to perform an RBF interpolation with {N:,} samples in 1D: {end - start:.5f}s" ) +print(f"Conjugate gradient infos: {cv_info}\n\n") ############################################################################################### # Display the (fitted) model on the unit interval: @@ -148,20 +147,19 @@ def laplacian_kernel(x, y, sigma=0.1): # between a perfect fit (**alpha** = 0) and a # smooth interpolation (**alpha** = :math:`+\infty`): -alpha = 10 # Ridge regularization +alpha = 5 # Ridge regularization start = time.time() K_xx = laplacian_kernel(x, x) -a = K_xx.solve(b, alpha=alpha) +a, cv_info = K_xx.solve(b, alpha=alpha, cv_info=True) end = time.time() print( - "Time to perform an RBF interpolation with {:,} samples in 2D: {:.5f}s".format( - N, end - start - ) + f"Time to perform an RBF interpolation with {N:,} samples in 2D: {end - start:.5f}s" ) +print(f"Conjugate gradient infos: {cv_info}\n\n") ############################################################################################### From 3296b0b02faddb42e041f4b542733ae7dff96bd0 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Tue, 14 Apr 2026 15:39:36 +0200 Subject: [PATCH 09/98] Add explicit indexing to torch.meshgrid call in doc --- .../tutorials/interpolation/plot_RBF_interpolation_torch.py | 2 +- pykeops/pykeops/tutorials/knn/plot_knn_torch.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py index 0b32c8bcb..657c9db1f 100644 --- a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py +++ b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py @@ -168,7 +168,7 @@ def laplacian_kernel(x, y, sigma=0.1): # Extrapolate on a uniform sample: X = Y = torch.linspace(0, 1, 101).type(dtype) -X, Y = torch.meshgrid(X, Y) +X, Y = torch.meshgrid(X, Y, indexing="ij") t = torch.stack((X.contiguous().view(-1), Y.contiguous().view(-1)), dim=1) K_tx = laplacian_kernel(t, x) diff --git a/pykeops/pykeops/tutorials/knn/plot_knn_torch.py b/pykeops/pykeops/tutorials/knn/plot_knn_torch.py index 610cb48a5..34428b3e8 100644 --- a/pykeops/pykeops/tutorials/knn/plot_knn_torch.py +++ b/pykeops/pykeops/tutorials/knn/plot_knn_torch.py @@ -48,7 +48,7 @@ def fth(x): M = 1000 if use_cuda else 100 tmp = torch.linspace(0, 1, M).type(dtype) -g2, g1 = torch.meshgrid(tmp, tmp) +g2, g1 = torch.meshgrid(tmp, tmp, indexing="ij") g = torch.cat((g1.contiguous().view(-1, 1), g2.contiguous().view(-1, 1)), dim=1) From 11a2c179a0b47e708d8780af1e353fb364b60bf9 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Tue, 14 Apr 2026 15:57:16 +0200 Subject: [PATCH 10/98] Remove Jenkinsfile --- Jenkinsfile | 25 ------------------------- 1 file changed, 25 deletions(-) delete mode 100644 Jenkinsfile diff --git a/Jenkinsfile b/Jenkinsfile deleted file mode 100644 index 01c1a4661..000000000 --- a/Jenkinsfile +++ /dev/null @@ -1,25 +0,0 @@ -pipeline { - agent none - stages { - stage('Test PyKeOps') { - parallel { - - stage('Test Cuda') { - agent { - label 'cuda' - } - steps { - echo 'Clean KeOps Cache...' - sh 'rm -rf $HOME/.cache/keops*' - echo 'Testing..' - sh '''#!/bin/bash - srun -n 1 -c 16 --mem=8G --gpus=1 --gres-flags=enforce-binding -J keops_ci pytest.sh -v - ''' - } - } - - } - } - - } -} From 4d13f20a6c0afb2663822480629be3346e138d89 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Wed, 15 Apr 2026 09:57:01 +0200 Subject: [PATCH 11/98] fix setup.py and doc building --- keopscore/setup.py | 13 +- pykeops/pykeops/benchmarks/plot_accuracy.py | 7 +- pykeops/pykeops/common/lazy_tensor.py | 12 +- pykeops/setup.py | 19 +- pytest.sh | 183 ++++++++------------ 5 files changed, 104 insertions(+), 130 deletions(-) diff --git a/keopscore/setup.py b/keopscore/setup.py index a0854b4a2..5a4f04b16 100644 --- a/keopscore/setup.py +++ b/keopscore/setup.py @@ -1,7 +1,7 @@ # Always prefer setuptools over distutils # To use a consistent encoding -from codecs import open import os +from codecs import open from os import path from setuptools import setup @@ -26,18 +26,20 @@ "Bug Reports": "https://github.com/getkeops/keops/issues", "Source": "https://github.com/getkeops/keops", }, - author="B. Charlier, J. Feydy, J. Glaunes", - author_email="benjamin.charlier@umontpellier.fr, jean.feydy@gmail.com, alexis.glaunes@parisdescartes.fr", + author="B. Charlier, J. Feydy, J. Glaunès", + author_email="benjamin.charlier@inrae.fr, jean.feydy@inria.com, alexis.glaunes@parisdescartes.fr", python_requires=">=3.8", classifiers=[ + "Topic :: Scientific/Engineering", "Development Status :: 5 - Production/Stable", "Intended Audience :: Developers", - "Topic :: Scientific/Engineering", - "License :: OSI Approved :: MIT License", + "Intended Audience :: Science/Research", "Operating System :: POSIX :: Linux", "Operating System :: MacOS :: MacOS X", "Programming Language :: C++", "Programming Language :: Python :: 3 :: Only", + "Environment :: GPU :: NVIDIA CUDA", + "License :: OSI Approved :: MIT License", ], keywords="kernels gpu autodiff", packages=[ @@ -65,7 +67,6 @@ "readme.md", "licence.txt", "keops_version", - "config/libiomp5.dylib", "binders/nvrtc/keops_nvrtc.cpp", "binders/nvrtc/nvrtc_jit.cpp", "include/CudaSizes.h", diff --git a/pykeops/pykeops/benchmarks/plot_accuracy.py b/pykeops/pykeops/benchmarks/plot_accuracy.py index c5c014a4e..30db3093e 100644 --- a/pykeops/pykeops/benchmarks/plot_accuracy.py +++ b/pykeops/pykeops/benchmarks/plot_accuracy.py @@ -13,7 +13,6 @@ output_filename = "accuracy" -import importlib import os import time @@ -145,7 +144,7 @@ def benchmark( ) ) else: - elapsed = np.NaN + elapsed = np.nan # accuracy ind = torch.randperm(y.shape[0]) @@ -159,11 +158,11 @@ def benchmark( mean_err = ( (out.double() - ref_out.double()).abs().mean() / ref_out.double().abs().mean() ).item() - mean_err = float("NaN") if mean_err == 0 else mean_err + mean_err = float("nan") if mean_err == 0 else mean_err max_err = ( (out.double() - ref_out.double()).abs().max() / ref_out.double().abs().mean() ).item() - max_err = float("NaN") if max_err == 0 else max_err + max_err = float("nan") if max_err == 0 else max_err print( "accuracy of an MxN convolution, with M = {}, N ={:7}: mean err={:.1e}, max err={:.1e}".format( M, N, mean_err, max_err diff --git a/pykeops/pykeops/common/lazy_tensor.py b/pykeops/pykeops/common/lazy_tensor.py index 46765497c..8c381d881 100644 --- a/pykeops/pykeops/common/lazy_tensor.py +++ b/pykeops/pykeops/common/lazy_tensor.py @@ -328,13 +328,11 @@ def fixvariables(self): self.variables = newvars def separate_kwargs(self, kwargs): - """ - separating keyword arguments for Genred init vs Genred call... - Currently the only additional optional keyword arguments that are passed to Genred init are - accuracy options: dtype_acc, use_double_acc and sum_scheme, - chunk mode option enable_chunks, - use_fast_math option, - and compiler option optional_flags. + """Split keyword arguments between ``Genred`` initialization and call. + + The initialization keywords are ``dtype_acc``, ``use_double_acc``, + ``sum_scheme``, ``enable_chunks``, ``use_fast_math`` and + ``optional_flags``. """ init_keys = { "dtype_acc", diff --git a/pykeops/setup.py b/pykeops/setup.py index d8154d91f..822c6d599 100644 --- a/pykeops/setup.py +++ b/pykeops/setup.py @@ -1,7 +1,7 @@ # Always prefer setuptools over distutils # To use a consistent encoding -from codecs import open import os +from codecs import open from os import path from setuptools import setup @@ -28,18 +28,21 @@ "Bug Reports": "https://github.com/getkeops/keops/issues", "Source": "https://github.com/getkeops/keops", }, - author="B. Charlier, J. Feydy, J. Glaunes", - author_email="benjamin.charlier@umontpellier.fr, jean.feydy@gmail.com, alexis.glaunes@parisdescartes.fr", + author="B. Charlier, J. Feydy, J. Glaunès", + author_email="benjamin.charlier@inrae.fr, jean.feydy@inria.com, alexis.glaunes@parisdescartes.fr", python_requires=">=3.8", classifiers=[ + "Topic :: Scientific/Engineering", "Development Status :: 5 - Production/Stable", "Intended Audience :: Developers", - "Topic :: Scientific/Engineering", - "License :: OSI Approved :: MIT License", + "Intended Audience :: Science/Research", "Operating System :: POSIX :: Linux", "Operating System :: MacOS :: MacOS X", + "Programming Language :: C", "Programming Language :: C++", "Programming Language :: Python :: 3 :: Only", + "Environment :: GPU :: NVIDIA CUDA", + "License :: OSI Approved :: MIT License", ], keywords="kernels gpu autodiff", packages=[ @@ -70,20 +73,22 @@ "sphinx", "sphinx-gallery", "recommonmark", + "myst-parser", "sphinxcontrib-httpdomain", "sphinx_rtd_theme", "sphinx-prompt", - "breathe", "matplotlib", "imageio", "torch", "gpytorch", "scikit-learn", "multiprocess", - "faiss", "h5py", "jaxlib", "jax", + "plotly", + "si_prefix", + "pandas", ], "test:": ["pytest", "numpy", "torch"], }, diff --git a/pytest.sh b/pytest.sh index 3b56ed1e3..a4a733fe3 100755 --- a/pytest.sh +++ b/pytest.sh @@ -1,128 +1,99 @@ -#!/bin/bash - -# exit in case of any errors -set -e - -################################################################################ -# help # -################################################################################ -function print_help() { - # Display Help - echo "Test script for keopscore/pykeops packages." - echo - echo "Usage: $0 [option...]" - echo - echo " -h Print the help" - echo " -v Verbose mode" - echo - exit 1 -} - -################################################################################ -# utils # -################################################################################ +#!/usr/bin/env bash -# log with verbosity management -function logging() { - if [[ ${PYTEST_VERBOSE} == 1 ]]; then - echo -e $1 - fi -} +set -euo pipefail -################################################################################ -# process script options # -################################################################################ +readonly PYTHON_BIN="python3" +readonly SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +readonly TEST_VENV="${SCRIPT_DIR}/.test_venv_pytest" +readonly TEST_REQUIREMENTS=(pip) -# default options PYTEST_VERBOSE=0 -# Get the options -while getopts 'hv' option; do - case $option in - h) # display Help - print_help - ;; - v) # enable verbosity - PYTEST_VERBOSE=1 - logging "## verbose mode" - ;; - \?) # Invalid option - echo "Error: Invalid option" - exit 1 - ;; - esac -done - -################################################################################ -# script setup # -################################################################################ - -# project root directory -PROJDIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) - -# python exec -PYTHON="python3" - -# python environment for test -TEST_VENV=${PROJDIR}/.test_venv_pytest - -# python test requirements (names of packages to be installed with pip) -TEST_REQ="pip" - -################################################################################ -# prepare python environment # -################################################################################ - -logging "-- Preparing python environment for test..." +print_help() { + cat < Date: Wed, 15 Apr 2026 13:59:11 +0200 Subject: [PATCH 12/98] Update pydoc with parallel jobs and fix --- doc/conf.py | 1 + pydoc.sh | 254 ++++++++++++++++++++++++++++------------------------ 2 files changed, 140 insertions(+), 115 deletions(-) diff --git a/doc/conf.py b/doc/conf.py index 1c12a72cd..365c591e3 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -108,6 +108,7 @@ def find_source(): # Add patterns # 'filename_pattern': r'../pykeops/pykeops/tutorials/*', "ignore_pattern": r"__init__\.py|benchmark_utils\.py|dataset_utils\.py", + "parallel": int(os.environ.get("SPHINX_GALLERY_JOBS", "1")), } # Generate the API documentation when building diff --git a/pydoc.sh b/pydoc.sh index 13563fc58..8f4e544d8 100755 --- a/pydoc.sh +++ b/pydoc.sh @@ -1,115 +1,139 @@ -#! /bin/sh -# -# This script build the doc and fix some links - -# do not exit in case of errors -set +e - - -################################################################################ -# process script options # -################################################################################ - -fix_link=false -noplot=false - -while getopts "l:n" opt; do - case ${opt} in - l ) fix_link=true - ;; - n ) noplot=true - ;; - \? ) echo "Usage: generate_doc [-l] [-n] - - -l : make correction on links - -n : no plot generation (html-noplot) - " - exit 255 - ;; - esac -done - - - -################################################################################ -# script setup # -################################################################################ - -# project root directory -PROJDIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) - -# python exec -PYTHON="python3" - -# python environment for test -DOC_VENV=${PROJDIR}/.doc_venv - -# python test requirements (names of packages to be installed with pip) -DOC_REQ="pip" - - -################################################################################ -# prepare python environment # -################################################################################ - -logging "-- Preparing python environment for doc build..." - -${PYTHON} -m venv --clear ${DOC_VENV} -source ${DOC_VENV}/bin/activate - -logging "---- Python version = $(python -V)" - -pip install -U ${DOC_REQ} - - -################################################################################ -# Installing keopscore # -################################################################################ - -logging "-- Installing keopscore..." - -pip install -e ${PROJDIR}/keopscore - -################################################################################ -# Installing pykeops # -################################################################################ - -logging "-- Installing pykeops..." - -pip install -e "${PROJDIR}/pykeops[full]" - - - -################################################################################ -# Building the doc # -################################################################################ - -printf "\n----------------------\n Building the doc \n----------------------\n\n" - -# go to the doc directory -CURRENT_DIR=$(pwd) -cd $PROJDIR/doc - -make clean -if [ $noplot = true ]; then - make html-noplot -else - make html -fi - -################################################################################ -# fixing doc link # -################################################################################ - -if [ $fix_link = true ]; then - printf "\n----------------------\n Fixing doc links \n----------------------\n\n" - # Fix some bad links due interaction between rtd-theme and sphinx-gallery - find . -path "*_auto_*" -name "plot_*.html" -exec sed -i "s/doc\/_auto_\(.*\)rst/pykeops\/pykeops\/\1py/" {} \; - find . -path "*_auto_*" -name "index.html" -exec sed -i "s/doc\/_auto_\(.*\)\/index\.rst/pykeops\/pykeops\/\1\//" {} \; -fi - -set -e - -# comes back to directory of -cd $CURRENT_DIR - +#!/usr/bin/env bash + +set -euo pipefail + +readonly PYTHON_BIN="python3" +readonly SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +readonly DOC_VENV="${SCRIPT_DIR}/.doc_venv" +readonly DOC_REQUIREMENTS=(pip) +readonly DOC_DIR="${SCRIPT_DIR}/doc" + +FIX_LINKS=0 +NO_PLOT=0 +DOC_JOBS="" + +print_help() { + cat </dev/null + make clean + if [[ "${NO_PLOT}" -eq 1 ]]; then + make_target="html-noplot" + else + make_target="html" + fi + make "${make_args[@]}" "${make_target}" + popd >/dev/null +} + +fix_doc_links() { + if [[ "${FIX_LINKS}" -ne 1 ]]; then + return + fi + + log_step "" + log_step "----------------------" + log_step " Fixing doc links" + log_step "----------------------" + log_step "" + + pushd "${DOC_DIR}" >/dev/null + find . -path "*_auto_*" -name "plot_*.html" -exec \ + sed -i "s/doc\/_auto_\(.*\)rst/pykeops\/pykeops\/\1py/" {} \; + find . -path "*_auto_*" -name "index.html" -exec \ + sed -i "s/doc\/_auto_\(.*\)\/index\.rst/pykeops\/pykeops\/\1\//" {} \; + popd >/dev/null +} + +main() { + parse_options "$@" + prepare_python_environment + install_editable_package "keopscore" "${SCRIPT_DIR}/keopscore" + install_editable_package "pykeops" "${SCRIPT_DIR}/pykeops[full]" + build_doc + fix_doc_links +} + +main "$@" From 47308ea5c84981e4011c93d052ac778355c1e7e0 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Wed, 22 Apr 2026 11:59:36 +0200 Subject: [PATCH 13/98] add verbose flag to CG --- pykeops/pykeops/common/operations.py | 106 ++++++++++-------- pykeops/pykeops/numpy/operations.py | 25 ++++- pykeops/pykeops/test/test_numpy.py | 51 ++++++++- pykeops/pykeops/test/test_torch.py | 34 +++++- pykeops/pykeops/torch/operations.py | 17 ++- .../plot_RBF_interpolation_numpy.py | 6 +- .../plot_RBF_interpolation_torch.py | 26 ++--- 7 files changed, 183 insertions(+), 82 deletions(-) diff --git a/pykeops/pykeops/common/operations.py b/pykeops/pykeops/common/operations.py index a26404ed5..80dce1978 100644 --- a/pykeops/pykeops/common/operations.py +++ b/pykeops/pykeops/common/operations.py @@ -88,7 +88,7 @@ def postprocess(out, binding, reduction_op, nout, opt_arg, dtype): def ConjugateGradientSolver( - binding, linop, b, x0=None, eps=1e-6, maxiter=None, cv_info=False + binding, linop, b, x0=None, eps=1e-6, maxiter=None, verbose=False ): """ Conjugate gradient algorithm to solve a linear system of the form @@ -122,63 +122,63 @@ def ConjugateGradientSolver( Maximum number of conjugate gradient iterations. Defaults to ``10 * b.shape[0]``. - cv_info : bool, optional - If ``False`` (default), return only the estimated solution ``x``. - If ``True``, return a tuple ``(x, info)`` where ``info`` is a - dictionary containing convergence diagnostics + verbose : bool, optional + If ``True``, print the convergence information dictionary before + returning. Defaults to ``False``. It contains the convergence + information with the following keys: + + - ``"status"``: ``"Converged"`` or ``"Maximum iterations reached"``. + - ``"niter"``: number of iterations performed. + - ``"residual_norm"``: final residual norm ``||r||``. + - ``"relative_residual_norm"``: final residual norm divided by + ``||b||``. + - ``"atol"``: absolute tolerance used internally for stopping. + - ``"maxiter"``: effective maximum number of iterations. + - ``"x0_provided"``: whether a non-``None`` initial guess was given. + - ``"dtype"``: data type of the solution ``x``. + Returns ------- x : tensor Approximate solution returned by the conjugate gradient iterations. - info : dict, optional - Returned only when ``cv_info=True``. Contains the convergence - information with the following - keys: - - - ``"status"``: ``"Converged"`` or ``"Maximum iterations reached"``. - - ``"niter"``: number of iterations performed. - - ``"residual_norm"``: final residual norm ``||r||``. - - ``"relative_residual_norm"``: final residual norm divided by - ``||b||``. - - ``"atol"``: absolute tolerance used internally for stopping. - - ``"maxiter"``: effective maximum number of iterations. - - ``"x0_provided"``: whether a non-``None`` initial guess was given. - - ``"dtype"``: data type of the solution ``x``. + """ tools = get_tools(binding) # stopping criterion atol, _ = _get_atol_rtol(tools.norm(b), eps) - maxiter = maxiter if maxiter is not None else 10 * b.shape[0] + maxiter = 10 * b.shape[0] if (maxiter is None) else maxiter - x = 0 if x0 is None else tools.copy(x0) - r = tools.copy(b) if x0 is None else b - linop(x0) + x = tools.zeros_like(b) if (x0 is None) else tools.copy(x0) + r = tools.copy(b) if (x0 is None) else b - linop(x0) nr2 = (r**2).sum() - if nr2 < atol * atol: - return tools.zeros_like(x) - - p = tools.copy(r) - - for it in range(maxiter): - Mp = linop(p) - alp = nr2 / (p * Mp).sum() - x += alp * p - r -= alp * Mp - nr2new = (r**2).sum() - if nr2new < atol * atol: - break - p = r + (nr2new / nr2) * p - nr2 = nr2new - - else: # for loop exhausted - # Return incomplete progress - it = -maxiter - 1 - - if cv_info: + nr2new = nr2 + + if nr2 <= atol * atol: + it = -1 + else: + p = tools.copy(r) + + for it in range(maxiter): + Mp = linop(p) + alp = nr2 / (p * Mp).sum() + x += alp * p + r -= alp * Mp + nr2new = (r**2).sum() + if nr2new < atol * atol: + break + p = r + (nr2new / nr2) * p + nr2 = nr2new + + else: # for loop exhausted + # Return incomplete progress + it = -maxiter - 1 + + if verbose: nr = tools.sqrt(nr2new) - return x, { + info = { "status": ("Converged" if nr <= atol else "Maximum iterations reached"), "niter": it + 1, "dtype": tools.dtypename(x.dtype), @@ -188,9 +188,11 @@ def ConjugateGradientSolver( "maxiter": maxiter, "x0_provided": x0 is not None, } - else: - return x + + print(f"[KeOps CG]: {info}") + return x + def _get_atol_rtol(b_norm, atol=0.0, rtol=1e-5): """ @@ -206,13 +208,13 @@ def KernelLinearSolver( K, x, b, - x0=None, alpha=0, eps=1e-6, + x0=None, maxiter=None, precond=False, precondKernel=None, - cv_info=False, + verbose=False, ): tools = get_tools(binding) dtype = tools.dtype(x) @@ -333,7 +335,13 @@ def f(x, y): ) else: a = ConjugateGradientSolver( - binding, KernelLinOp, b, eps=eps, maxiter=maxiter, x0=x0, cv_info=cv_info + binding, + KernelLinOp, + b, + eps=eps, + x0=x0, + maxiter=maxiter, + verbose=verbose, ) return a diff --git a/pykeops/pykeops/numpy/operations.py b/pykeops/pykeops/numpy/operations.py index 1c7246c82..383d25eea 100644 --- a/pykeops/pykeops/numpy/operations.py +++ b/pykeops/pykeops/numpy/operations.py @@ -185,7 +185,7 @@ def __call__( eps=1e-6, x0=None, maxiter=None, - cv_info=False + verbose=False, ): r""" To apply the routine on arbitrary NumPy arrays. @@ -212,6 +212,21 @@ def __call__( **ridge regularization** parameter, added to the diagonal of the Kernel matrix :math:`K_{xx}`. + eps (float, default = 1e-6): Stopping criterion for the + conjugate gradient algorithm. + + x0 (2d array, default = None): Initial guess for the solution of + the linear system. Should be of the same shape as ``b``. + + maxiter (int, default = None): Maximum number of conjugate + gradient iterations. If ``None``, uses the default from + :func:`pykeops.common.operations.ConjugateGradientSolver`. + + verbose (bool, default = False): If ``True``, prints the + conjugate gradient convergence information dictionary + produced by + :func:`pykeops.common.operations.ConjugateGradientSolver`. + backend (string): Specifies the map-reduce scheme, as detailed in the documentation of the :class:`numpy.Genred ` module. @@ -281,5 +296,11 @@ def linop(var): return res return ConjugateGradientSolver( - "numpy", linop, varinv, x0=x0, eps=eps, cv_info=cv_info + "numpy", + linop, + varinv, + eps=eps, + x0=x0, + maxiter=maxiter, + verbose=verbose, ) diff --git a/pykeops/pykeops/test/test_numpy.py b/pykeops/pykeops/test/test_numpy.py index 028926d3b..c529ef1d6 100644 --- a/pykeops/pykeops/test/test_numpy.py +++ b/pykeops/pykeops/test/test_numpy.py @@ -1,5 +1,7 @@ import os.path import sys +from contextlib import redirect_stdout +import io sys.path.append( os.path.join( @@ -21,12 +23,8 @@ import pykeops import pykeops.config from pykeops.numpy.utils import ( - np_kernel, - grad_np_kernel, - differences, squared_distances, log_sum_exp, - np_kernel_sphere, ) @@ -68,6 +66,51 @@ def test_numpytools_function_binding(self): self.assertTrue(np.allclose(tools.norm(x), np.linalg.norm(x))) self.assertTrue(np.allclose(tools.arraysum(x, axis=0), np.sum(x, axis=0))) + ############################################################ + def test_cg_solver_stops_immediately_when_x0_is_good(self): + ############################################################ + + from pykeops.numpy import LazyTensor + + alpha = 2.0 + + x_i = LazyTensor(self.x[:, None, :]) + x_j = LazyTensor(self.x[None, :, :]) + K_xx = (((x_i - x_j).abs()).sum(-1)).exp() + + b = K_xx @ self.f + alpha * self.f + + x = K_xx.solve(b, alpha=alpha, x0=self.f, eps=1e-12) + self.assertTrue(np.allclose(self.f, x)) + + _, info = K_xx.solve(b, alpha=alpha, x0=self.f, eps=1e-12) + self.assertEqual(info["status"], "Converged") + self.assertEqual(info["niter"], 0) + self.assertLessEqual(info["residual_norm"], info["atol"]) + + ############################################################ + def test_cg_solver_verbose_prints_info(self): + ############################################################ + + from pykeops.numpy import LazyTensor + + alpha = 2.0 + + x_i = LazyTensor(self.x[:, None, :]) + x_j = LazyTensor(self.x[None, :, :]) + K_xx = (((x_i - x_j).abs()).sum(-1)).exp() + + b = K_xx @ self.f + alpha * self.f + + stream = io.StringIO() + with redirect_stdout(stream): + x = K_xx.solve(b, alpha=alpha, x0=self.f, eps=1e-12, verbose=True) + + self.assertTrue(np.allclose(self.f, x)) + output = stream.getvalue() + self.assertIn("'status': 'Converged'", output) + self.assertIn("'x0_provided': True", output) + ############################################################ def test_generic_syntax_sum(self): ############################################################ diff --git a/pykeops/pykeops/test/test_torch.py b/pykeops/pykeops/test/test_torch.py index 85df1b21b..4d272264c 100644 --- a/pykeops/pykeops/test/test_torch.py +++ b/pykeops/pykeops/test/test_torch.py @@ -1,5 +1,7 @@ import os.path import sys +from contextlib import redirect_stdout +import io sys.path.append( os.path.join( @@ -15,17 +17,13 @@ ) import unittest -import itertools import numpy as np import pykeops import pykeops.config from pykeops.numpy.utils import ( squared_distances, - np_kernel, log_np_kernel, - grad_np_kernel, - differences, log_sum_exp, ) @@ -427,6 +425,34 @@ def test_invkernel(self): ) ) + ############################################################ + def test_cg_solver_stops_immediately_when_x0_is_good(self): + ############################################################ + + import torch + from pykeops.torch import LazyTensor + + alpha = 2.0 + + x_i = LazyTensor(self.xc[:, None, :]) + x_j = LazyTensor(self.xc[None, :, :]) + K_xx = (((x_i - x_j).abs()).sum(-1)).exp() + + b = K_xx @ self.fc + alpha * self.fc + + x = K_xx.solve(b, alpha=alpha, x0=self.fc, eps=1e-12) + self.assertTrue(torch.allclose(self.fc, x)) + + stream = io.StringIO() + with redirect_stdout(stream): + x = K_xx.solve(b, alpha=alpha, x0=self.fc, eps=1e-12, verbose=True) + + self.assertTrue(torch.allclose(self.fc, x)) + output = stream.getvalue() + self.assertIn("'status': 'Converged'", output) + self.assertIn("'niter': 0", output) + self.assertIn("'x0_provided': True", output) + ############################################################ def test_softmax(self): ############################################################ diff --git a/pykeops/pykeops/torch/operations.py b/pykeops/pykeops/torch/operations.py index 8a26173ae..9434abda6 100644 --- a/pykeops/pykeops/torch/operations.py +++ b/pykeops/pykeops/torch/operations.py @@ -84,11 +84,12 @@ def linop(var): eps=params.eps, x0=params.x0, maxiter=params.maxiter, - cv_info=params.cv_info, + verbose=params.verbose, ) # relying on the 'ctx.saved_variables' attribute is necessary if you want to be able to differentiate the output # of the backward once again. It helps pytorch to keep track of 'who is who'. + # Here we do not save the extra info dict if it is in the output ctx.save_for_backward(*args, result) return result @@ -96,7 +97,6 @@ def linop(var): @staticmethod def backward(ctx, G): params = ctx.params - device_id = ctx.params myconv = ctx.myconv args = ctx.saved_tensors[:-1] # Unwrap the saved variables @@ -346,7 +346,7 @@ def __call__( eps=1e-6, x0=None, maxiter=None, - cv_info=False, + verbose=False, ): r""" Apply the routine on arbitrary torch Tensors. @@ -396,6 +396,15 @@ def __call__( x0 (2d Tensor, default = None): Initial guess for the solution of the linear system. should be of the same shape as b. + maxiter (int, default = None): Maximum number of conjugate + gradient iterations. If ``None``, uses the default from + :func:`pykeops.common.operations.ConjugateGradientSolver`. + + verbose (bool, default = False): If ``True``, prints the + conjugate gradient convergence information dictionary + produced by + :func:`pykeops.common.operations.ConjugateGradientSolver`. + Returns: (M,D) or (N,D) Tensor: @@ -430,7 +439,7 @@ def __call__( params.alpha = alpha params.eps = eps params.x0 = x0 - params.cv_info = cv_info params.maxiter = maxiter + params.verbose = verbose return KernelSolveAutograd.apply(params, *args) diff --git a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py index ea5ba15ec..cac02ea26 100644 --- a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py +++ b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py @@ -79,14 +79,13 @@ def gaussian_kernel(x, y, sigma=0.1): start = time.time() K_xx = gaussian_kernel(x, x) -a, cv_info = K_xx.solve(b, alpha=alpha, cv_info=True) +a = K_xx.solve(b, alpha=alpha, verbose=True) end = time.time() print( f"Time to perform an RBF interpolation with {N:,} samples in 1D: {end - start:.5f}s" ) -print(f"Conjugate gradient infos: {cv_info}\n\n") ####################################################################### # Display the (fitted) model on the unit interval: @@ -150,14 +149,13 @@ def laplacian_kernel(x, y, sigma=0.1): start = time.time() K_xx = laplacian_kernel(x, x) -a, cv_info = K_xx.solve(b, alpha=alpha, cv_info=True) +a = K_xx.solve(b, alpha=alpha, verbose=True) end = time.time() print( f"Time to perform an RBF interpolation with {N:,} samples in 2D: {end - start:.5f}s" ) -print(f"Conjugate gradient infos: {cv_info}\n\n") ######################################################################## # Display the (fitted) model on the unit square: diff --git a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py index 657c9db1f..fb7f460d1 100644 --- a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py +++ b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py @@ -29,9 +29,8 @@ import time -from matplotlib import pyplot as plt - import torch +from matplotlib import pyplot as plt from pykeops.torch import LazyTensor ############################################################################################### @@ -47,10 +46,10 @@ # Some random-ish 1D signal: b = ( - x - + 0.5 * (6 * x).sin() - + 0.1 * (20 * x).sin() - + 0.05 * torch.randn(N, 1).type(dtype) + x + + 0.5 * (6 * x).sin() + + 0.1 * (20 * x).sin() + + 0.05 * torch.randn(N, 1).type(dtype) ) @@ -66,7 +65,7 @@ def gaussian_kernel(x, y, sigma=0.1): x_i = LazyTensor(x[:, None, :]) # (M, 1, 1) y_j = LazyTensor(y[None, :, :]) # (1, N, 1) D_ij = ((x_i - y_j) ** 2).sum(-1) # (M, N) symbolic matrix of squared distances - return (-D_ij / (2 * sigma**2)).exp() # (M, N) symbolic Gaussian kernel matrix + return (-D_ij / (2 * sigma ** 2)).exp() # (M, N) symbolic Gaussian kernel matrix ####################################################################### @@ -80,14 +79,13 @@ def gaussian_kernel(x, y, sigma=0.1): start = time.time() K_xx = gaussian_kernel(x, x) -a, cv_info = K_xx.solve(b, alpha=alpha, cv_info=True) +a = K_xx.solve(b, alpha=alpha, verbose=True) end = time.time() print( f"Time to perform an RBF interpolation with {N:,} samples in 1D: {end - start:.5f}s" ) -print(f"Conjugate gradient infos: {cv_info}\n\n") ############################################################################################### # Display the (fitted) model on the unit interval: @@ -119,9 +117,9 @@ def gaussian_kernel(x, y, sigma=0.1): # Some random-ish 2D signal: b = ((x - 0.5) ** 2).sum(1, keepdim=True) -b[b > 0.4**2] = 0 -b[b < 0.3**2] = 0 -b[b >= 0.3**2] = 1 +b[b > 0.4 ** 2] = 0 +b[b < 0.3 ** 2] = 0 +b[b >= 0.3 ** 2] = 1 b = b + 0.05 * torch.randn(N, 1).type(dtype) # Add 25% of outliers: @@ -152,15 +150,13 @@ def laplacian_kernel(x, y, sigma=0.1): start = time.time() K_xx = laplacian_kernel(x, x) -a, cv_info = K_xx.solve(b, alpha=alpha, cv_info=True) +a = K_xx.solve(b, alpha=alpha, verbose=True) end = time.time() print( f"Time to perform an RBF interpolation with {N:,} samples in 2D: {end - start:.5f}s" ) -print(f"Conjugate gradient infos: {cv_info}\n\n") - ############################################################################################### # Display the (fitted) model on the unit square: From ca25ca3d5769bb37c27217696de0122a81725d4e Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Wed, 22 Apr 2026 13:13:22 +0200 Subject: [PATCH 14/98] fix KernelSolve calls --- pykeops/pykeops/common/operations.py | 11 ++++++----- pykeops/pykeops/test/test_numpy.py | 14 ++++++++++---- .../plot_RBF_interpolation_torch.py | 16 ++++++++-------- 3 files changed, 24 insertions(+), 17 deletions(-) diff --git a/pykeops/pykeops/common/operations.py b/pykeops/pykeops/common/operations.py index 80dce1978..9b51bb9f1 100644 --- a/pykeops/pykeops/common/operations.py +++ b/pykeops/pykeops/common/operations.py @@ -1,5 +1,6 @@ import numpy as np +from keopscore.utils.misc_utils import KeOps_Print, KeOps_Warning from pykeops.common.utils import get_tools @@ -130,8 +131,7 @@ def ConjugateGradientSolver( - ``"status"``: ``"Converged"`` or ``"Maximum iterations reached"``. - ``"niter"``: number of iterations performed. - ``"residual_norm"``: final residual norm ``||r||``. - - ``"relative_residual_norm"``: final residual norm divided by - ``||b||``. + - ``"relative_residual_norm"``: final residual norm divided by ``||b||``. - ``"atol"``: absolute tolerance used internally for stopping. - ``"maxiter"``: effective maximum number of iterations. - ``"x0_provided"``: whether a non-``None`` initial guess was given. @@ -142,7 +142,7 @@ def ConjugateGradientSolver( x : tensor Approximate solution returned by the conjugate gradient iterations. - + """ tools = get_tools(binding) @@ -174,6 +174,7 @@ def ConjugateGradientSolver( else: # for loop exhausted # Return incomplete progress + KeOps_Warning("[KeOps CG]: Maximum iterations reached. Check convergence...") it = -maxiter - 1 if verbose: @@ -189,10 +190,10 @@ def ConjugateGradientSolver( "x0_provided": x0 is not None, } - print(f"[KeOps CG]: {info}") + KeOps_Print(f"[KeOps CG]: {info}") return x - + def _get_atol_rtol(b_norm, atol=0.0, rtol=1e-5): """ diff --git a/pykeops/pykeops/test/test_numpy.py b/pykeops/pykeops/test/test_numpy.py index c529ef1d6..dbc08ab73 100644 --- a/pykeops/pykeops/test/test_numpy.py +++ b/pykeops/pykeops/test/test_numpy.py @@ -82,11 +82,17 @@ def test_cg_solver_stops_immediately_when_x0_is_good(self): x = K_xx.solve(b, alpha=alpha, x0=self.f, eps=1e-12) self.assertTrue(np.allclose(self.f, x)) + + stream = io.StringIO() + with redirect_stdout(stream): + x = K_xx.solve(b, alpha=alpha, x0=self.f, eps=1e-12, verbose=True) + + self.assertTrue(np.allclose(self.f, x)) + output = stream.getvalue() + self.assertIn("'status': 'Converged'", output) + self.assertIn("'niter': 0", output) + self.assertIn("'x0_provided': True", output) - _, info = K_xx.solve(b, alpha=alpha, x0=self.f, eps=1e-12) - self.assertEqual(info["status"], "Converged") - self.assertEqual(info["niter"], 0) - self.assertLessEqual(info["residual_norm"], info["atol"]) ############################################################ def test_cg_solver_verbose_prints_info(self): diff --git a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py index fb7f460d1..636ecf55a 100644 --- a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py +++ b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py @@ -46,10 +46,10 @@ # Some random-ish 1D signal: b = ( - x - + 0.5 * (6 * x).sin() - + 0.1 * (20 * x).sin() - + 0.05 * torch.randn(N, 1).type(dtype) + x + + 0.5 * (6 * x).sin() + + 0.1 * (20 * x).sin() + + 0.05 * torch.randn(N, 1).type(dtype) ) @@ -65,7 +65,7 @@ def gaussian_kernel(x, y, sigma=0.1): x_i = LazyTensor(x[:, None, :]) # (M, 1, 1) y_j = LazyTensor(y[None, :, :]) # (1, N, 1) D_ij = ((x_i - y_j) ** 2).sum(-1) # (M, N) symbolic matrix of squared distances - return (-D_ij / (2 * sigma ** 2)).exp() # (M, N) symbolic Gaussian kernel matrix + return (-D_ij / (2 * sigma**2)).exp() # (M, N) symbolic Gaussian kernel matrix ####################################################################### @@ -117,9 +117,9 @@ def gaussian_kernel(x, y, sigma=0.1): # Some random-ish 2D signal: b = ((x - 0.5) ** 2).sum(1, keepdim=True) -b[b > 0.4 ** 2] = 0 -b[b < 0.3 ** 2] = 0 -b[b >= 0.3 ** 2] = 1 +b[b > 0.4**2] = 0 +b[b < 0.3**2] = 0 +b[b >= 0.3**2] = 1 b = b + 0.05 * torch.randn(N, 1).type(dtype) # Add 25% of outliers: From 7fa45dabf471d77f4add21eb96389d04dd68a1a6 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Tue, 28 Apr 2026 14:55:41 +0200 Subject: [PATCH 15/98] WIP. Refactor config stack. Platform; Cpp; OpenMP --- keopscore/keopscore/config/CppConfig.py | 145 +++++++++++++ keopscore/keopscore/config/Platform.py | 155 +++++++++----- keopscore/keopscore/config/_shared.py | 155 ++++++++++++++ keopscore/keopscore/config/openmp.py | 264 +++++++++++++++--------- 4 files changed, 577 insertions(+), 142 deletions(-) create mode 100644 keopscore/keopscore/config/CppConfig.py create mode 100644 keopscore/keopscore/config/_shared.py diff --git a/keopscore/keopscore/config/CppConfig.py b/keopscore/keopscore/config/CppConfig.py new file mode 100644 index 000000000..1033d184f --- /dev/null +++ b/keopscore/keopscore/config/CppConfig.py @@ -0,0 +1,145 @@ +import os, shutil + +from keopscore.utils.misc_utils import KeOps_Warning, KeOps_OS_Run +from keopscore.config.Platform import Platform +from keopscore.config._shared import print_envs, not_found_str + + +class CppConfig(Platform): + """ + Common platform and C++ compiler configuration. + """ + + _cxx_compiler = None + _cpp_env_flags = None + _compile_options = None + _disable_pragma_unrolls = None + cpp_envs = ["CXX", "CXXFLAGS"] + + + def __init__(self): + super().__init__() + self.set_cxx_compiler() + self.set_cpp_env_flags() + self.set_compile_options() + self.set_disable_pragma_unrolls() + + # C++ Compiler Detection + def set_cxx_compiler(self): + """Set the C++ compiler.""" + self._cxx_compiler = self.detect_cxx_compiler() + + def get_cxx_compiler(self): + """Get the C++ compiler.""" + return self._cxx_compiler + + def print_cxx_compiler(self): + """Print the C++ compiler information.""" + print(f"C++ Compiler Path: {self.get_cxx_compiler() or not_found_str}") + + def detect_cxx_compiler(self): + """Return the best available C++ compiler for the current platform.""" + preferred_compilers = [] + + env_cxx = os.getenv("CXX") + if env_cxx: + preferred_compilers.append(env_cxx) + + # Common alias for the default C++ compiler + preferred_compilers.append("c++") + + if self.get_platform() == "Darwin": + # On macOS, prioritize clang++ over g++ + preferred_compilers.extend(("clang++", "g++")) + else: + preferred_compilers.append(["g++", "clang++"]) + + # Return the first available compiler from the preferred list. + for compiler in preferred_compilers: + cxx_compiler_path = shutil.which(compiler) + if cxx_compiler_path: + return cxx_compiler_path + else: + KeOps_Warning("No C++ compiler found. Define CXX environment variable or install g++.") + + return None + + def get_cxx_compiler_version(self): + """Detect if using Apple Clang.""" + compiler_info = KeOps_OS_Run(f"{self.get_cxx_compiler()} --version").stdout.decode("utf-8").splitlines()[0].strip() + return compiler_info + + def print_cxx_compiler_version(self): + """Print the C++ compiler version information.""" + print(f"C++ Compiler Version: {self.get_cxx_compiler_version() or not_found_str}") + + @property + def use_Apple_clang(self): + """Detect if using Apple Clang.""" + return "Apple clang" in self.get_cxx_compiler_version() if self.get_cxx_compiler_version() else False + + # Disable Pragma Unrolls + def set_disable_pragma_unrolls(self): + """Set the flag for disabling pragma unrolls.""" + self._disable_pragma_unrolls = True + + def get_disable_pragma_unrolls(self): + """Get the flag for disabling pragma unrolls.""" + return self._disable_pragma_unrolls + + def print_disable_pragma_unrolls(self): + """Print the flag for disabling pragma unrolls.""" + status = "no (default)" if self.get_disable_pragma_unrolls() else "yes" + print(f"Pragma unroll: {status}") + + # Compile Options + def set_compile_options(self): + """Set the compile options.""" + self._compile_options = " -shared -fPIC -O3 -std=c++11 -flto=auto" + + # Specific Options for Apple Clang (maybe unnecessary in recent MacOs versions) + if self.get_platform() == "Darwin" and self.use_Apple_clang: + self.cpp_flags += " -undefined dynamic_lookup" + + # ... and Silicon chips + if self.get_platform() == "Darwin" and self.get_machine() in ["arm64", "arm64e"]: + self.cpp_flags += " -arch arm64" + + def get_compile_options(self): + """Get the compile options.""" + return self._compile_options + + def print_compile_options(self): + """Print the compile options.""" + print(f"Compile Options: {self.get_compile_options()}") + + # C++ Environment Flags. Unused yet... + def set_cpp_env_flags(self): + """Recover the C++ environment flags.""" + self._cpp_env_flags = os.getenv("CXXFLAGS") if "CXXFLAGS" in os.environ else "" + + def get_cpp_env_flags(self): + """Get the C++ environment flags.""" + return self._cpp_env_flags + + # Comprehensive Platform and Compiler Information + def print_cpp(self): + """Print all platform and C++ compiler information.""" + + print("=" * 60) + print("C/C++ Compiler Information") + print("=" * 60) + + self.print_cxx_compiler() + self.print_cxx_compiler_version() + self.print_compile_options() + self.print_disable_pragma_unrolls() + + # Print relevant environment variables. + print_envs(self.cpp_envs) + + +if __name__ == "__main__": + cpp_config = CppConfig() + cpp_config.print_platform() + cpp_config.print_cpp() \ No newline at end of file diff --git a/keopscore/keopscore/config/Platform.py b/keopscore/keopscore/config/Platform.py index 238b44042..35687a4ce 100644 --- a/keopscore/keopscore/config/Platform.py +++ b/keopscore/keopscore/config/Platform.py @@ -1,91 +1,154 @@ +import os import platform import sys -import os +from keopscore.config._shared import print_envs -class DetectPlatform: + +class Platform: """ Class for detecting the operating system, Python version, and environment type. """ - + _os = None + _platform = None + _machine = None + _uname = None + _python_version = None + _python_executable = None + _env_type = None + platform_envs = [ + "PYTHONPATH", + "PATH", + "VIRTUAL_ENV", + "CONDA_DEFAULT_ENV", + "CONDA_PREFIX", + ] + def __init__(self): - self.os = None - self.python_version = None - self.env_type = None self.set_os() + self.set_platform() + self.set_machine() + self.set_uname() self.set_python_version() + self.set_python_executable() self.set_env_type() + # OS Detection (Distribution Name and Version) def set_os(self): """Set the operating system.""" + self._os = self.detect_os() + + def get_os(self): + return self._os + + def print_os(self): + print(f"Operating System: {self.get_os()}") + + @staticmethod + def detect_os(): + """Return a human-readable operating system description.""" if platform.system() == "Linux": try: with open("/etc/os-release") as f: info = dict(line.strip().split("=", 1) for line in f if "=" in line) name = info.get("NAME", "Linux").strip('"') version = info.get("VERSION_ID", "").strip('"') - self.os = f"{platform.system()} {name} {version}" + return f"{platform.system()} {name} {version}" except FileNotFoundError: - self.os = "Linux (distribution info not found)" - else: - self.os = platform.system() + " " + platform.version() + return "Linux (distribution info not found)" + + return platform.system() + " " + platform.version() - def get_os(self): - return self.os + # Platform detection (Darwin, Windows, Linux, etc.) + def set_platform(self): + self._platform = platform.system() + + def get_platform(self): + return self._platform - def print_os(self): - print(f"Operating System: {self.os}") + # Machine architecture detection (x86_64, arm64, etc.) + def set_machine(self): + self._machine = platform.machine() + + def get_machine(self): + return self._machine + + def print_machine(self): + print(f"Machine Architecture: {self.get_machine()}") + # uname detection + def set_uname(self): + self._uname = platform.uname() + + def get_uname(self): + return self._uname + + # Python Version Detection def set_python_version(self): """Set the Python version.""" - self.python_version = platform.python_version() + self._python_version = platform.python_version() def get_python_version(self): - return self.python_version + return self._python_version def print_python_version(self): - print(f"Python Version: {self.python_version}") + print(f"Python Version: {self.get_python_version()}") + # Python Executable Detection + def set_python_executable(self): + """Set the Python executable path.""" + self._python_executable = sys.executable + + def get_python_executable(self): + return self._python_executable + + def print_python_executable(self): + print(f"Python Executable: {self.get_python_executable()}") + + # Environment Type Detection def set_env_type(self): """Set the environment type (conda, virtualenv, or system).""" - if "CONDA_DEFAULT_ENV" in os.environ: - self.env_type = f"conda ({os.environ['CONDA_DEFAULT_ENV']})" - elif hasattr(sys, "real_prefix") or ( - hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix - ): - self.env_type = "virtualenv" - else: - self.env_type = "system" + self._env_type = self.detect_env_type() def get_env_type(self): - return self.env_type + return self._env_type def print_env_type(self): - print(f"Environment Type: {self.env_type}") + print(f"Environment Type: {self.get_env_type()} {sys.prefix}", end="") + if self.get_env_type() != 'system': + print(f' (base at {sys.base_prefix})', end="") + print() - def print_all(self): + @staticmethod + def detect_env_type(): + """Return whether Python runs in conda, virtualenv, or the system env.""" + if "CONDA_DEFAULT_ENV" in os.environ: + return f"conda ({os.environ['CONDA_DEFAULT_ENV']})" + if hasattr(sys, "real_prefix") or ( + hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix + ): + return "virtualenv" + return "system" + + # Comprehensive Platform Information + def print_platform(self): """ - Print all platform-related information, including relevant environment variables. + Print all platform-related information. """ - print("\nPlatform Information") + print("=" * 60) + print("Platform Information") + print("=" * 60) + self.print_os() self.print_python_version() self.print_env_type() + self.print_python_executable() - # Print relevant environment variables - print("\nRelevant Environment Variables") - print("-" * 60) - env_vars = [ - "CONDA_DEFAULT_ENV", - "VIRTUAL_ENV", - "PATH", - "PYTHONPATH", - "CXX", - ] - for var in env_vars: - value = os.environ.get(var) - if value: - print(f"{var} = {value}") - else: - print(f"{var} is not set") + # Print relevant environment variables. + print_envs(self.platform_envs) + + +if __name__ == "__main__": + platform_info = Platform() + platform_info.print_platform() \ No newline at end of file diff --git a/keopscore/keopscore/config/_shared.py b/keopscore/keopscore/config/_shared.py new file mode 100644 index 000000000..8283040ca --- /dev/null +++ b/keopscore/keopscore/config/_shared.py @@ -0,0 +1,155 @@ +import os +import site +import sys +import sysconfig +from pathlib import Path + +from keopscore.utils.misc_utils import CHECK_MARK, CROSS_MARK, find_library_abspath + +not_found_str = f"Not Found. {CROSS_MARK}" +enabled_dict = {True: f"Enabled {CHECK_MARK}", False: f"Disabled {CROSS_MARK}"} + +def _unique_paths(paths): + """Return paths in first-seen order with duplicates and None values removed.""" + unique = [] + seen = set() + for path in paths: + if path is None: + continue + path = Path(path) + key = str(path) + if key not in seen: + seen.add(key) + unique.append(path) + return unique + + +def _env_roots(env_vars): + """Return existing environment variable values as normalized Path objects.""" + return _unique_paths(os.getenv(env_var) for env_var in env_vars) + + +def _path_candidates(roots, suffixes): + """Expand root directories with relative suffixes while preserving order.""" + candidates = [] + for root in _unique_paths(roots): + for suffix in suffixes: + candidates.append(root / suffix if suffix else root) + return _unique_paths(candidates) + + +def _first_existing_file(paths): + """Return the first path that points to an existing file.""" + for path in _unique_paths(paths): + if path.is_file(): + return str(path) + return None + + +def _first_matching_file(directories, patterns): + """Return the first file matching glob patterns inside ordered directories.""" + for directory in _unique_paths(directories): + if not directory.is_dir(): + continue + for pattern in patterns: + matches = sorted(directory.glob(pattern)) + for match in matches: + if match.is_file(): + return str(match) + return None + + +def _first_existing_dir_with_files(directories, required_filenames): + """Return the first directory containing all required filenames.""" + for directory in _unique_paths(directories): + if all((directory / filename).is_file() for filename in required_filenames): + return str(directory) + return None + + +def _find_library_by_names(library_names): + """Return the first library path resolved by ctypes for known library names.""" + from ctypes.util import find_library + + # ctypes delegates to the platform loader, so this is the final fallback. + for library_name in library_names: + library_path = find_library(library_name) + if library_path: + return find_library_abspath(library_name) + """ + library_path = find_library(library_name) + + if library_path: + full_path = KeOps_OS_Run("ldconfig -p").stdout.decode("utf-8") + + #print(f"Checking for OpenMP library '{full_path}' in ldconfig output...") + for line in full_path.splitlines(): + if library_path in line: + library_full_path = line.split("=>", 1)[1].strip() + return library_path, os.path.dirname(library_full_path) + """ + + return None + + +def _python_package_roots(): + """Return Python package roots where pip-installed wheels may live.""" + package_roots = [] + getters = [ + lambda: getattr(site, "getsitepackages", lambda: [])(), + lambda: [site.getusersitepackages()], + lambda: [sysconfig.get_path("purelib")], + lambda: [sysconfig.get_path("platlib")], + lambda: sys.path, + ] + for getter in getters: + try: + # Some site helpers are unavailable in embedded or non-standard Python builds. + package_roots.extend(getter() or []) + except Exception: + continue + return _unique_paths(package_roots) + + +def _ordered_search_roots(env_vars=(), pip_suffixes=(), conda_root=None, system_roots=()): + """Return roots ordered by explicit env vars, pip, conda, then system paths. + + :argument + - env_vars (list): environment variable names to search for paths + - pip_suffixes (list): suffixes to append to pip-installed wheel paths + - conda_root (str): environment variable name for conda root + - system_roots (list): explicit system paths to search + """ + roots = [] + + if env_vars: + roots.extend(_env_roots(env_vars)) + + if pip_suffixes: + # Pip wheels such as nvidia-cuda-runtime expose libraries under site-packages. + roots.extend(_path_candidates(_python_package_roots(), pip_suffixes)) + + if conda_root: + roots.extend(_env_roots((conda_root,))) + roots.extend(system_roots) + return _unique_paths(roots) + + +def ensure_directory(path, add_to_syspath=False): + """Create a directory and optionally register it in sys.path.""" + os.makedirs(path, exist_ok=True) + if add_to_syspath and path not in sys.path: + sys.path.append(path) + return path + + +def print_envs(env_vars): + """Print the values of specified environment variables.""" + print("\nRelevant Environment Variables") + print("-" * 60) + for var in env_vars: + value = os.environ.get(var) + if value: + print(f"{var} = {value}") + else: + print(f"{var} is not set") \ No newline at end of file diff --git a/keopscore/keopscore/config/openmp.py b/keopscore/keopscore/config/openmp.py index 66ab72317..69492decf 100644 --- a/keopscore/keopscore/config/openmp.py +++ b/keopscore/keopscore/config/openmp.py @@ -1,56 +1,174 @@ import os -import shutil -import tempfile import subprocess -import platform -from ctypes.util import find_library +import tempfile -from keopscore.utils.misc_utils import KeOps_Warning, KeOps_OS_Run, get_brew_prefix -from keopscore.utils.misc_utils import CHECK_MARK, CROSS_MARK +from keopscore.config.CppConfig import CppConfig +from keopscore.config._shared import ( + _find_library_by_names, + _first_existing_dir_with_files, + _first_existing_file, + _ordered_search_roots, + _path_candidates, + print_envs, + enabled_dict, + not_found_str, +) +from keopscore.utils.misc_utils import KeOps_Warning, get_brew_prefix -class OpenMPConfig: +class OpenMPConfig(CppConfig): """ Class for OpenMP detection and configuration. """ + _use_OpenMP = None + _openmp_lib_name = None + _openmp_lib_include_dir = None + _openmp_header_name = "omp.h" + openmp_env_vars = ( + "OMP_PATH", + "LIBOMP_PATH", + "OpenMP_ROOT", + "OpenMP_ROOT_DIR", + ) + + openmp_dir_prefixes = ( + get_brew_prefix(), + os.path.join(os.path.sep, "usr", "local", "opt", "libomp"), + os.path.join(os.path.sep, "opt", "local"), + ) + def __init__(self): - self._use_OpenMP = None - self.openmp_lib_path = None - self.os = platform.system() - self.set_cxx_compiler() + super().__init__() + self.set_openmplib_path() self.set_use_OpenMP() - def set_cxx_compiler(self): - """Set the C++ compiler.""" - env_cxx = os.getenv("CXX") - if env_cxx and shutil.which(env_cxx): - self.cxx_compiler = env_cxx - elif shutil.which("g++"): - self.cxx_compiler = "g++" + # OpenMP library path + def set_openmplib_path(self): + """try to locate OpenMP libraries""" + openmp_install = self.find_openmp_install() + openmp_lib = openmp_install["library"] + if openmp_lib: + self._openmp_lib_include_dir = openmp_install["include_dir"] + self._openmp_lib_lib_dir = os.path.dirname(openmp_lib) + self._openmp_lib_name = openmp_lib else: - self.cxx_compiler = None KeOps_Warning( - "No C++ compiler found. You need to either define the CXX environment variable pointing to a valid compiler, or ensure that 'g++' is installed and in your PATH." + "OpenMP runtime library not found. " + "Set OMP_PATH, LIBOMP_PATH, or OpenMP_ROOT if it is installed in a non-standard location." + ) + + def print_openmplib_path(self): + if self.get_openmp_lib_name() and self.get_openmp_lib_dir(): + full_path = os.path.join( + self.get_openmp_lib_dir(), self.get_openmp_lib_name() + ) + elif self.get_openmp_lib_name(): + full_path = self.get_openmp_lib_name() + else: + full_path = None + + print(f"OpenMP Library Path: {full_path or not_found_str}") + + def find_openmp_install(self): + """ + Locate OpenMP runtime files without assuming a specific package manager. + + Returns a dict with optional ``library`` and ``include_dir`` entries. + """ + result = {"library": None, "include_dir": None} + + # First try to find OpenMP library using standard names via ctypes. + result["library"] = _find_library_by_names(("gomp", "omp")) + p = subprocess.run( + [ + self.get_cxx_compiler(), + f"-print-file-name=include/{self._openmp_header_name}", + ], + capture_output=True, + text=True, + ).stdout.strip() + result["include_dir"] = p if os.path.isabs(p) and os.path.exists(p) else None + + # If that fails, search for OpenMP headers and libraries in common locations. + if not result["library"]: + candidate_roots = _ordered_search_roots( + env_vars=self.openmp_env_vars, + conda_root="CONDA_PREFIX", + system_roots=self.openmp_dir_prefixes, ) + result["library"] = _first_existing_file( + _path_candidates( + candidate_roots, + ( + "lib/libomp.dylib", + "lib/libgomp.dylib", + "lib/libomp.so", + "opt/libomp/lib/libomp.dylib", + ), + ) + ) + + # Finally, search for OpenMP headers in common locations. + result["include_dir"] = _first_existing_dir_with_files( + _path_candidates(candidate_roots, ("include", "opt/libomp/include")), + (self._openmp_header_name,), + ) + + return result + + # OpenMP library name and directory getters/setters + def set_openmp_lib_name(self): + """Set the OpenMP library name (e.g., libomp.so or libgomp.dylib).""" + # This is set in set_openmplib_path if the library is found, otherwise it remains None. + pass + + def get_openmp_lib_name(self): + """Get the OpenMP library name (e.g., libomp.so or libgomp.dylib).""" + return self._openmp_lib_name + + def set_openmp_lib_dir(self): + """Set the OpenMP library directory (containing .so or .dylib files).""" + # This is set in set_openmplib_path if the library is found, otherwise it remains None. + pass + + def get_openmp_lib_dir(self): + """Get the OpenMP library directory (containing .so or .dylib files).""" + return self._openmp_lib_lib_dir + + # OpenMP header path + def set_openmp_lib_include_dir(self): + """Set the OpenMP include directory (containing headers).""" + # This is set in set_openmplib_path if the library is found, otherwise it remains None. + pass + + def get_openmp_include_dir(self): + """Get the OpenMP include directory (containing headers).""" + return self._openmp_lib_include_dir + + def print_openmp_include_dir(self): + if self.get_openmp_include_dir(): + print(f"OpenMP Header Path: {self.get_openmp_include_dir()}") + + # OpenMP use detection def set_use_OpenMP(self): """Determine and set whether to use OpenMP.""" compiler_supports_openmp = self.check_compiler_for_openmp() - openmp_libs_available = self.check_openmp_libraries() - self._use_OpenMP = compiler_supports_openmp or openmp_libs_available - if not self._use_OpenMP: - KeOps_Warning("OpenMP support is not available. Disabling OpenMP.") + self._use_OpenMP = compiler_supports_openmp and ( + self.get_openmp_lib_name() is not None + ) def get_use_OpenMP(self): return self._use_OpenMP def print_use_OpenMP(self): - status = "Enabled ✅" if self._use_OpenMP else "Disabled ❌" - print(f"OpenMP Support: {status}") + print(f"OpenMP Support: {enabled_dict[self.get_use_OpenMP() or False]}") def check_compiler_for_openmp(self): - if not self.cxx_compiler: + """Attempt to compile a simple OpenMP program to check if the compiler supports OpenMP.""" + + if not self.get_cxx_compiler(): KeOps_Warning("No C++ compiler available to check for OpenMP support.") return False @@ -67,12 +185,14 @@ def check_compiler_for_openmp(self): test_file = f.name compile_command = [ - self.cxx_compiler, + self.get_cxx_compiler(), test_file, "-fopenmp", - "-o", - test_file + ".out", ] + if self.get_openmp_lib_dir(): + compile_command.append(f"-L{self.get_openmp_lib_dir()}") + compile_command.extend(["-o", f"{test_file}.out"]) + try: # Warning : subprocess is used below to compile the test program (using subprocess.check_output to capture stderr) subprocess.check_output(compile_command, stderr=subprocess.STDOUT) @@ -83,74 +203,26 @@ def check_compiler_for_openmp(self): os.remove(test_file) return False - def check_openmp_libraries(self): - if self.os.startswith("Linux"): - openmp_lib = find_library("gomp") - if not openmp_lib: - KeOps_Warning("OpenMP library 'libgomp' not found.") - return False - else: - self.openmp_lib_path = openmp_lib - return True - # Specific check for M1/M2/M3 apple Silicon chips - elif self.os.startswith("Darwin") and platform.machine() in ["arm64", "arm64e"]: - brew_prefix = get_brew_prefix() - if brew_prefix is not None: - openmp_path = f"{brew_prefix}/opt/libomp/lib/libomp.dylib" - openmp_lib = openmp_path if os.path.exists(openmp_path) else None - else: - openmp_lib = None - if openmp_lib is None: - KeOps_Warning( - "OpenMP library not found, it must be downloaded through Homebrew for apple Silicon chips" - ) - return False - else: - self.openmp_lib_path = openmp_lib - return True - elif self.os.startswith("Darwin"): - openmp_lib = find_library("omp") - if not openmp_lib: - KeOps_Warning("OpenMP library 'libomp' not found.") - return False - else: - self.openmp_lib_path = openmp_lib - return True - else: - self.openmp_lib_path = None - return False - - def print_all(self): + # OpenMP configuration printing + def print_openmp(self): """ Print all OpenMP-related configuration and system health status. """ - # OpenMP Support - openmp_status = CHECK_MARK if self.get_use_OpenMP() else CROSS_MARK - print(f"\nOpenMP Support") - print("-" * 60) + + print("=" * 60) + print(f"OpenMP Configuration") + print("=" * 60) + self.print_use_OpenMP() - if self.get_use_OpenMP(): - openmp_lib_path = self.openmp_lib_path or "Not Found" - print(f"OpenMP Library Path: {openmp_lib_path}") - # Compiler path - compiler_path = ( - shutil.which(self.cxx_compiler) if self.cxx_compiler else None - ) - print(f"C++ Compiler: {self.cxx_compiler}") - if not compiler_path: - print( - f"Compiler '{self.cxx_compiler}' not found on the system.{CROSS_MARK}" - ) - else: - print(f"OpenMP support is disabled or not available.{CROSS_MARK}") + self.print_openmplib_path() + self.print_openmp_include_dir() + # Print relevant environment variables. - print("\nRelevant Environment Variables:") - env_vars = [ - "OMP_PATH", - ] - for var in env_vars: - value = os.environ.get(var, None) - if value: - print(f"{var} = {value}") - else: - print(f"{var} is not set") + print_envs(self.openmp_env_vars) + + +if __name__ == "__main__": + openmp_config = OpenMPConfig() + openmp_config.print_platform() + openmp_config.print_cpp() + openmp_config.print_openmp() From 7a42bed73b314551f3b4b8defc7fe705c794a2a5 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Tue, 28 Apr 2026 20:39:59 +0200 Subject: [PATCH 16/98] WIP. Cuda Co-authored-by: Copilot --- keopscore/keopscore/config/cuda.py | 637 ++++++++++++++++------------- 1 file changed, 349 insertions(+), 288 deletions(-) diff --git a/keopscore/keopscore/config/cuda.py b/keopscore/keopscore/config/cuda.py index 320f700e9..4252ad78d 100644 --- a/keopscore/keopscore/config/cuda.py +++ b/keopscore/keopscore/config/cuda.py @@ -1,30 +1,31 @@ -import os import ctypes -from ctypes.util import find_library +import os from ctypes import ( + DEFAULT_MODE, c_int, - c_void_p, - c_char_p, CDLL, byref, - cast, - POINTER, - Structure, RTLD_GLOBAL, ) -from pathlib import Path -import shutil -from os.path import join -import platform -import subprocess -import sys -import keopscore + +from keopscore.config.CppConfig import CppConfig +from keopscore.config._shared import ( + _find_library_by_names, + _first_existing_dir_with_files, + _first_matching_file, + _ordered_search_roots, + _path_candidates, + print_envs, + not_found_str, +) +from keopscore.utils.misc_utils import ( + CHECK_MARK, + CROSS_MARK, +) from keopscore.utils.misc_utils import KeOps_Warning -from keopscore.utils.misc_utils import KeOps_OS_Run -from keopscore.utils.misc_utils import CHECK_MARK, CROSS_MARK, get_include_file_abspath -class CUDAConfig: +class CUDAConfig(CppConfig): """ Class for CUDA detection and configuration. """ @@ -34,115 +35,252 @@ class CUDAConfig: CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK = 1 CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK = 8 - # Cuda attributes - libcuda_folder = None - libnvrtc_folder = None - cuda_include_path = None - nvrtc_flags = None - cuda_version = None + # Cuda detection variables + _use_cuda = None + _specific_gpus = None + + _cuda_include_path = None + _nvrtc_flags = None + _cuda_version = None n_gpus = 0 - gpu_compile_flags = "" + _gpu_compile_flags = "" cuda_message = "" - specific_gpus = None cuda_block_size = None + cuda_install_info = None + + # ------------------------ # + # Search location # + # ------------------------ # + + cuda_env_vars = [ + "CUDA_VISIBLE_DEVICES", + "CUDA_PATH", + "CUDA_HOME", + "CUDA_ROOT", + "CUDA_TOOLKIT_ROOT_DIR", + ] + + pip_suffixes = ( + "nvidia/cuda_runtime", + "nvidia/cuda_nvrtc", + ) + + system_suffixes = ( + os.path.join(os.path.sep, "usr", "local", "cuda"), + os.path.join(os.path.sep, "usr", "local"), + os.path.join(os.path.sep, "opt", "cuda"), + os.path.join(os.path.sep, "usr"), + os.path.join(os.path.sep, "lib"), + ) + + library_suffixes = ( + "lib64", + "lib", + "lib/x86_64-linux-gnu/", + ) + + include_suffixes = ( + "include", + "targets/x86_64-linux/include", + "targets/sbsa-linux/include", + "targets/aarch64-linux/include", + ) + + # ------------------------- # + # Library info dicts # + # ------------------------- # + + _libcuda_info = { + "name": "cuda", + "lib_file_name_candidate": ["libcuda.so.*", "libcuda.dylib", "cuda.lib"], + "header_file_name": "cuda.h", + "library": None, # to be filled later + "include_dir": None, # to be filled later + } + _libnvrtc_info = { + "name": "nvrtc", + "lib_file_name_candidate": ["libnvrtc.so.*", "libnvrtc.dylib", "nvrtc.lib"], + "header_file_name": "nvrtc.h", + "library": None, # to be filled later + "include_dir": None, # to be filled later + } + _cudart_info = { + "name": "cudart", + "lib_file_name_candidate": ["libcudart.so.*", "libcudart.dylib", "cudart.lib"], + "header_file_name": None, + "library": None, # to be filled later + "include_dir": None, # not needed + } def __init__(self): - self.set_keops_cache_folder() - self.set_default_build_folder_name() + self.set_specific_gpus() - self.set_build_folder() - self.set_cxx_compiler() + + super().__init__() + self.set_use_cuda() + # If cuda is enabled, then we finalize the rest of the config - if self._use_cuda: - self.set_libcuda_folder() - self.set_libnvrtc_folder() + if self.get_use_cuda(): + self.get_cuda_version() + self.get_cuda_include_path() self.set_nvrtc_flags() self.set_cuda_block_size() - def _try_load_library(self, lib_name): + def find_cuda_install(self, lib_dict_info, warn=None): """ - Attempt to locate and load libraties. + Locate a cuda and headers using an explicit ordered search. + + Arguments: + lib_dict_info (dict): A dictionary containing at least the keys 'name', 'lib_file_name_candidate', and 'header_file_name' for the library to find. This allows the function to be used for finding libcuda, libcudart, or nvrtc by passing the appropriate info dict. + + Returns: + result (dict): a copy of lib_dict_info completed with the ``library`` and ``include_dir`` keys containing the absolute paths to the library file and include directory, or None if not found. + """ + result = ( + lib_dict_info.copy() + ) # Start with the provided info, which may contain names and file patterns + + candidate_roots = _ordered_search_roots( + env_vars=self.cuda_env_vars, + pip_suffixes=self.pip_suffixes, + conda_root="CONDA_PREFIX", + system_roots=self.system_suffixes, + ) + + # ------------------------ # + # Search for library file # + # ------------------------ # + + # First try to find the library file using the candidate roots and library suffixes + candidate_library_dirs = _path_candidates(candidate_roots, self.library_suffixes) + result["library"] = _first_matching_file( + candidate_library_dirs, result["lib_file_name_candidate"] + ) + if result["library"] is None: + result["library"] = _find_library_by_names((result["name"],)) + + if result["library"] is None and warn: + KeOps_Warning(f"lib{result['name']} not found.") + + # ------------------------ # + # Search for header files # + # ------------------------ # + + result["include_dir"] = _first_existing_dir_with_files( + _path_candidates(candidate_roots, self.include_suffixes), + (result["header_file_name"],), + ) + + if ( + result["include_dir"] is None + and warn + and result["header_file_name"] is not None + ): + KeOps_Warning(f"{result['name']} header files not found.") + + return result + + def _try_load_library(self, lib_full_path_path, mode=DEFAULT_MODE): + """ + Attempt to load a shared library. Returns: success (bool): True if the library was found and loaded. - abspath (str): Absolute path to the loaded library if success==True, else "". error_msg (str): Contains error details if success==False, else "". """ - # Find library - found_path = find_library(lib_name) - if not found_path: - return (False, "", f"Library '{lib_name}' not found on this system.") - - # Try to load it + if not lib_full_path_path: + return False, "Library path not found" try: - lib_handle = CDLL(found_path, mode=RTLD_GLOBAL) + CDLL(lib_full_path_path, mode=mode) except OSError as e: - return (False, "", f"Failed to load library '{lib_name}': {e}") + return ( + False, + f"Failed to load library '{os.path.basename(lib_full_path_path)}': {e}", + ) - class LINKMAP(Structure): - _fields_ = [("l_addr", c_void_p), ("l_name", c_char_p)] + return True, "" + def _try_find_cuda_version(self, libcudart_full_path_path): + """ + Attempt to find the CUDA version by loading the CUDA runtime library and querying its version. + Returns: + success (bool): True if the version was successfully determined, False otherwise. + error_msg (str): Contains error details if success==False, else "". + + """ try: - # Attempt to load libdl to use dlinfo - libdl_path = find_library("dl") - if not libdl_path: - # If we can't find libdl, we can't do dlinfo; fallback - return (True, found_path, "") - - libdl = CDLL(libdl_path) - dlinfo = libdl.dlinfo - dlinfo.argtypes = (c_void_p, c_int, c_void_p) - dlinfo.restype = c_int - - lmptr = c_void_p() - # RTLD_DI_LINKMAP = 2 - result = dlinfo(lib_handle._handle, 2, byref(lmptr)) - if result != 0: - # dlinfo call failed, fallback - return (True, found_path, "") - - abspath_bytes = cast(lmptr, POINTER(LINKMAP)).contents.l_name - abspath_str = abspath_bytes.decode("utf-8") - if abspath_str: - return (True, abspath_str, "") - else: - return (True, found_path, "") - except Exception as err: - # If anything goes wrong, fallback to found_path - return (True, found_path, "") + libcudart = ctypes.CDLL(libcudart_full_path_path) + cuda_version = ctypes.c_int() + libcudart.cudaRuntimeGetVersion(ctypes.byref(cuda_version)) + self._cuda_version = int(cuda_version.value) + return True, "" + + except OSError as e: + self._cuda_version = None + return ( + False, + f"Failed to load '{os.path.basename(libcudart_full_path_path)}': {e}", + ) def _cuda_libraries_available(self): """ - Check if both cuda and nvrtc libraries are available. + Check if libcuda (driver); libcudart (cudatoolkit) and nvrtc (cudatoolkit) libraries are available. + This is where ```_cuda_include_path``` is set. + Returns: - True if both cuda and nvrtc are loadable, False otherwise. + True if all three are loadable, False otherwise. This is also where we handle one single warning if needed. """ - # This step loads "libcuda.so (driver) and libnvrtc (cuda tool kit) **Globaly** to - # make cuda avalaible to keops shared objects - success_cuda, cuda_path, err_cuda = self._try_load_library("cuda") - success_nvrtc, nvrtc_path, err_nvrtc = self._try_load_library("nvrtc") + # Libcuda (driver) loaded globally so they are available to KeOps shared objects. + self._libcuda_info = self.find_cuda_install(self._libcuda_info, warn=True) + success_cuda = False + err_cuda = "libcuda not found" + if self._libcuda_info["library"] is not None: + success_cuda, err_cuda = self._try_load_library( + self._libcuda_info["library"], mode=RTLD_GLOBAL + ) + if not success_cuda: + KeOps_Warning(f"{err_cuda}. Switching to CPU only.") + return False - if not success_cuda or not success_nvrtc: - self.cuda_message = "CUDA libraries not found or could not be loaded; Switching to CPU only." - KeOps_Warning(self.cuda_message) + # libnvrtc as well, since it's required for the runtime compilation of CUDA code. + self._libnvrtc_info = self.find_cuda_install(self._libnvrtc_info, warn=True) + success_nvrtc = False + err_nvrtc = "libnvrtc not found" + if self._libnvrtc_info["library"] is not None: + success_nvrtc, err_nvrtc = self._try_load_library( + self._libnvrtc_info["library"], mode=RTLD_GLOBAL + ) + if not success_nvrtc: + KeOps_Warning(f"{err_nvrtc}. Switching to CPU only.") + return False + # Populate the cuda_install_info which is needed for include path and cuda version detection + self._cuda_include_path = list(set([ + self._libnvrtc_info["include_dir"], + self._libcuda_info["include_dir"], + ])) + + # Finally check cudart + self._cudart_info = self.find_cuda_install(self._cudart_info, warn=False) + success_cudart = False + err_cudart = "libcudart not found" + if self._cudart_info["library"] is not None: + success_cudart, err_cudart = self._try_find_cuda_version( + self._cudart_info["library"] + ) + if not success_cudart: + KeOps_Warning(f"{err_cudart}. Switching to CPU only.") return False - # If both succeeded, store their folder paths - self.libcuda_folder = os.path.dirname(cuda_path) - self.libnvrtc_folder = os.path.dirname(nvrtc_path) return True + # CUDA Support def set_use_cuda(self): """Determine and set whether to use CUDA.""" - self._use_cuda = True - if not self._cuda_libraries_available(): - self._use_cuda = False + self._use_cuda = self._cuda_libraries_available() - self.get_cuda_version() - self.get_cuda_include_path() self.get_gpu_props() if self.n_gpus == 0 and self._use_cuda: self._use_cuda = False @@ -153,9 +291,10 @@ def get_use_cuda(self): return self._use_cuda def print_use_cuda(self): - status = "Enabled ✅" if self._use_cuda else "Disabled ❌" + status = f"Enabled {CHECK_MARK}" if self._use_cuda else f"Disabled {CROSS_MARK}" print(f"CUDA Support: {status}") + # CUDA Block Size def set_cuda_block_size(self, cuda_block_size=192): """Sets default cuda block size.""" self.cuda_block_size = cuda_block_size @@ -166,205 +305,133 @@ def get_cuda_block_size(self): def print_cuda_block_size(self): print(f"CUDA Block Size: {self.cuda_block_size}") + # Specific GPUs def set_specific_gpus(self): """Set specific GPUs from CUDA_VISIBLE_DEVICES.""" - self.specific_gpus = os.getenv("CUDA_VISIBLE_DEVICES") - if self.specific_gpus: + if os.getenv("CUDA_VISIBLE_DEVICES"): + self._specific_gpus = os.getenv("CUDA_VISIBLE_DEVICES") # Modify the build folder name to include GPU specifics - gpu_suffix = self.specific_gpus.replace(",", "_") - self.default_build_folder_name += f"_CUDA_VISIBLE_DEVICES_{gpu_suffix}" + gpu_suffix = self._specific_gpus.replace(",", "_") + self.set_default_build_folder_name( + appended_name=f"CUDA_VISIBLE_DEVICES_{gpu_suffix}" + ) def get_specific_gpus(self): """Get the specific GPUs.""" - return self.specific_gpus + return self._specific_gpus def print_specific_gpus(self): """Print the specific GPUs.""" - if self.specific_gpus: - print(f"Specific GPUs (CUDA_VISIBLE_DEVICES): {self.specific_gpus}") - else: - print("Specific GPUs (CUDA_VISIBLE_DEVICES): Not Set") - - def set_cxx_compiler(self): - """Set the C++ compiler.""" - env_cxx = os.getenv("CXX") - if env_cxx and shutil.which(env_cxx): - self.cxx_compiler = env_cxx - elif shutil.which("g++"): - self.cxx_compiler = "g++" - else: - self.cxx_compiler = None - KeOps_Warning( - "No C++ compiler found. You need to either define the CXX environment variable pointing to a valid compiler, or ensure that 'g++' is installed and in your PATH." - ) - - def set_keops_cache_folder(self): - """Set the KeOps cache folder.""" - self.keops_cache_folder = os.getenv("KEOPS_CACHE_FOLDER") - if self.keops_cache_folder is None: - self.keops_cache_folder = join( - os.path.expanduser("~"), ".cache", f"keops{keopscore.__version__}" - ) - # Ensure the cache folder exists - os.makedirs(self.keops_cache_folder, exist_ok=True) - - def set_default_build_folder_name(self): - """Set the default build folder name.""" - uname = platform.uname() - self.default_build_folder_name = ( - "_".join(uname[:3]) + f"_p{sys.version.split(' ')[0]}" + print( + f"Specific GPUs (CUDA_VISIBLE_DEVICES): {self.get_specific_gpus() or "Not set"}" ) - def set_build_folder(self): - self.build_folder = join( - self.keops_cache_folder, self.default_build_folder_name - ) - - def get_build_folder(self): - return self.build_folder - + # Libcuda folder def set_libcuda_folder(self): """ - Return nothing if not using cuda - self.libcuda_folder is already set in _cuda_libraries_available. + Is set in _cuda_libraries_available. """ - if not self._use_cuda: - return + pass def get_libcuda_folder(self): - return self.libcuda_folder + return self._libcuda_info["library"] and os.path.dirname( + self._libcuda_info["library"] + ) + def print_libcuda_folder(self): + print(f"Libcuda Folder: {self.get_libcuda_folder() or not_found_str}") + + # Libnvrtc folder def set_libnvrtc_folder(self): """ Return nothing if not using cuda self.libnvrtc_folder is already set in _cuda_libraries_available. """ - if not self._use_cuda: - return + pass def get_libnvrtc_folder(self): - return self.libnvrtc_folder + return self._libnvrtc_info["library"] and os.path.dirname( + self._libnvrtc_info["library"] + ) + + def print_libnvrtc_folder(self): + print(f"Libnvrtc Folder: {self.get_libnvrtc_folder() or not_found_str}") + + # CUDA Version + def set_cuda_version(self, warn=True): + """Set the CUDA version by querying the CUDA runtime library.""" + self._cuda_version = self.find_cuda_version(warn=warn) def get_cuda_version(self, out_type="single_value"): - if not self._use_cuda: - self.cuda_version = None - return None - try: - libcudart_path = find_library("cudart") - if not libcudart_path: - self.cuda_version = None - return None - libcudart = ctypes.CDLL(libcudart_path) - cuda_version = ctypes.c_int() - libcudart.cudaRuntimeGetVersion(ctypes.byref(cuda_version)) - cuda_version_value = int(cuda_version.value) - - if out_type == "single_value": - self.cuda_version = cuda_version_value - return cuda_version_value - - major = cuda_version_value // 1000 - minor = (cuda_version_value % 1000) // 10 - - if out_type == "major,minor": - return major, minor - elif out_type == "string": - return f"{major}.{minor}" - except Exception as e: - KeOps_Warning(f"Could not determine CUDA version: {e}") - self.cuda_version = None - return None + major = self._cuda_version // 1000 + minor = (self._cuda_version % 1000) // 10 + + if out_type == "major,minor": + return major, minor + elif out_type == "string": + return f"{major}.{minor}" + else: + return self._cuda_version + + def print_cuda_version(self): + str = f"CUDA Version: {self.get_cuda_version(out_type="string") if self.get_cuda_version() else not_found_str}" + print(str) + + # CUDA Include Path + def set_cuda_include_path(self): + """Set the CUDA include path by searching for cuda.h and nvrtc.h.""" + # This is done in find_cuda_install since it relies on the cuda installation info which is only available after checking library availability. + pass def get_cuda_include_path(self): """ - Attempt to find CUDA headers (cuda.h, nvrtc.h) in standard - places or environment variables. + Attempt to find CUDA headers (cuda.h, nvrtc.h) using an explicit + ordered search over environment variables and standard locations. """ if not self._use_cuda: - self.cuda_include_path = None return None - # Check the CUDA_PATH and CUDA_HOME environment variables - for env_var in ["CUDA_PATH", "CUDA_HOME"]: - path = os.getenv(env_var) - if path: - include_path = Path(path) / "include" - if (include_path / "cuda.h").is_file() and ( - include_path / "nvrtc.h" - ).is_file(): - self.cuda_include_path = str(include_path) - return self.cuda_include_path - - # Check if CUDA is installed via conda - conda_prefix = os.getenv("CONDA_PREFIX") - if conda_prefix: - include_path = Path(conda_prefix) / "include" - if (include_path / "cuda.h").is_file() and ( - include_path / "nvrtc.h" - ).is_file(): - self.cuda_include_path = str(include_path) - return self.cuda_include_path - - # Check standard locations - cuda_version_str = self.get_cuda_version(out_type="string") - possible_paths = [ - Path("/usr/local/cuda"), - Path(f"/usr/local/cuda-{cuda_version_str}") if cuda_version_str else None, - Path("/opt/cuda"), - ] - # Filter out Nones (if cuda_version_str is None) - possible_paths = [p for p in possible_paths if p is not None] - - for base_path in possible_paths: - include_path = base_path / "include" - if (include_path / "cuda.h").is_file() and ( - include_path / "nvrtc.h" - ).is_file(): - self.cuda_include_path = str(include_path) - return self.cuda_include_path - - # If not found in any known location, try the compiler approach: - cuda_h_path = self.get_include_file_abspath("cuda.h") - nvrtc_h_path = self.get_include_file_abspath("nvrtc.h") - if cuda_h_path and nvrtc_h_path: - if os.path.dirname(cuda_h_path) == os.path.dirname(nvrtc_h_path): - self.cuda_include_path = os.path.dirname(cuda_h_path) - return self.cuda_include_path - - # If still not found, issue a warning - KeOps_Warning( - "CUDA include path not found. Please set the CUDA_PATH or CUDA_HOME environment variable." - ) - self.cuda_include_path = None - return self.cuda_include_path - - def get_include_file_abspath(self, filename): - return get_include_file_abspath(filename, self.cxx_compiler) + return self._cuda_include_path + def print_cuda_include_path(self): + print(f"CUDA Include Path: {":".join(self.get_cuda_include_path()) or not_found_str}") + + # NVRTC Flags def set_nvrtc_flags(self): """Set the NVRTC flags for CUDA compilation.""" + # TODO: redondant with CppConfig compile options, should be refactored to avoid duplication # Ensure that compile_options is set (inherited from ConfigNew) compile_options = " -shared -fPIC -O3 -std=c++11" - # Ensure that libcuda_folder and libnvrtc_folder are set - libcuda_folder = self.libcuda_folder - libnvrtc_folder = self.libnvrtc_folder - # Set the NVRTC flags - self.nvrtc_flags = ( + self._nvrtc_flags = ( compile_options - + f" -fpermissive -L{libcuda_folder} -L{libnvrtc_folder} -lcuda -lnvrtc" + + f" -fpermissive -L{self.get_libcuda_folder()} -L{self.get_libnvrtc_folder()} -lcuda -lnvrtc" ) def get_nvrtc_flags(self): """Get the NVRTC flags for CUDA compilation.""" - return self.nvrtc_flags + return self._nvrtc_flags def print_nvrtc_flags(self): """Print the NVRTC flags for CUDA compilation.""" - print(f"NVRTC Flags: {self.nvrtc_flags}") - + print(f"NVRTC Flags: {self.get_nvrtc_flags()}") + + # GPU compile flags + def set_gpu_compile_flags(self): + """Set GPU compile flags based on detected GPU properties.""" + # This is done in get_gpu_props since it relies on the GPU properties which are only available after checking CUDA availability. + pass + + def get_gpu_compile_flags(self): + """Get GPU compile flags based on detected GPU properties.""" + return self._gpu_compile_flags + + def print_gpu_compile_flags(self): + print(f"GPU Compile Flags: {self.get_gpu_compile_flags() or not_found_str}") + + # GPU Properties def get_gpu_props(self): """ Getting GPU properties and related attributes. @@ -372,11 +439,15 @@ def get_gpu_props(self): if not self._use_cuda: # Already determined that CUDA is unavailable self.n_gpus = 0 - self.gpu_compile_flags = "" - return (self.n_gpus, self.gpu_compile_flags) + self._gpu_compile_flags = "" + return (self.n_gpus, self._gpu_compile_flags) # Attempt to load the CUDA driver library - success, libcuda_path, err_msg = self._try_load_library("cuda") + libcuda_path = self._libcuda_info["library"] if self._libcuda_info else None + success, err_msg = self._try_load_library( + libcuda_path, + mode=RTLD_GLOBAL, + ) if not success: # Something is off at driver level => revert to CPU KeOps_Warning( @@ -385,9 +456,9 @@ def get_gpu_props(self): + " Switching to CPU only." ) self.n_gpus = 0 - self.gpu_compile_flags = "" + self._gpu_compile_flags = "" self._use_cuda = False - return (self.n_gpus, self.gpu_compile_flags) + return (self.n_gpus, self._gpu_compile_flags) # We have a handle, let's proceed libcuda = ctypes.CDLL(libcuda_path) @@ -397,9 +468,9 @@ def get_gpu_props(self): "CUDA was detected, but driver API could not be initialized. Switching to CPU only." ) self.n_gpus = 0 - self.gpu_compile_flags = "" + self._gpu_compile_flags = "" self._use_cuda = False - return (self.n_gpus, self.gpu_compile_flags) + return (self.n_gpus, self._gpu_compile_flags) # Get GPU count nGpus = ctypes.c_int() @@ -410,14 +481,14 @@ def get_gpu_props(self): "Switching to CPU only." ) self.n_gpus = 0 - self.gpu_compile_flags = "" + self._gpu_compile_flags = "" self._use_cuda = False - return (self.n_gpus, self.gpu_compile_flags) + return (self.n_gpus, self._gpu_compile_flags) self.n_gpus = nGpus.value if self.n_gpus == 0: - self.gpu_compile_flags = "" - return (self.n_gpus, self.gpu_compile_flags) + self._gpu_compile_flags = "" + return (self.n_gpus, self._gpu_compile_flags) # Query each GPU for properties MaxThreadsPerBlock = [0] * self.n_gpus @@ -436,9 +507,9 @@ def safe_call(dev_idx, result_code): device = ctypes.c_int() if not safe_call(d, libcuda.cuDeviceGet(ctypes.byref(device), d)): self.n_gpus = 0 - self.gpu_compile_flags = "" + self._gpu_compile_flags = "" self._use_cuda = False - return (self.n_gpus, self.gpu_compile_flags) + return (self.n_gpus, self._gpu_compile_flags) output = ctypes.c_int() if not safe_call( @@ -450,9 +521,9 @@ def safe_call(dev_idx, result_code): ), ): self.n_gpus = 0 - self.gpu_compile_flags = "" + self._gpu_compile_flags = "" self._use_cuda = False - return (self.n_gpus, self.gpu_compile_flags) + return (self.n_gpus, self._gpu_compile_flags) MaxThreadsPerBlock[d] = output.value if not safe_call( @@ -464,55 +535,45 @@ def safe_call(dev_idx, result_code): ), ): self.n_gpus = 0 - self.gpu_compile_flags = "" + self._gpu_compile_flags = "" self._use_cuda = False - return (self.n_gpus, self.gpu_compile_flags) + return (self.n_gpus, self._gpu_compile_flags) SharedMemPerBlock[d] = output.value # Build compile flags string - self.gpu_compile_flags = f"-DMAXIDGPU={self.n_gpus - 1} " + self._gpu_compile_flags = f"-DMAXIDGPU={self.n_gpus - 1} " for d in range(self.n_gpus): - self.gpu_compile_flags += ( + self._gpu_compile_flags += ( f"-DMAXTHREADSPERBLOCK{d}={MaxThreadsPerBlock[d]} " ) - self.gpu_compile_flags += f"-DSHAREDMEMPERBLOCK{d}={SharedMemPerBlock[d]} " + self._gpu_compile_flags += f"-DSHAREDMEMPERBLOCK{d}={SharedMemPerBlock[d]} " - return self.n_gpus, self.gpu_compile_flags + return self.n_gpus, self._gpu_compile_flags - def print_all(self): - """ - Print all CUDA-related configuration and system health status. - """ + def print_cuda(self): + """Print all CUDA-related configuration""" # CUDA Support - cuda_status = CHECK_MARK if self.get_use_cuda() else CROSS_MARK - print(f"\nCUDA Support") - print("-" * 60) + print("=" * 60) + print(f"CUDA Support") + print("=" * 60) + self.print_use_cuda() if self.get_use_cuda(): - print(f"Libcuda Path: {self.libcuda_folder}") - print(f"Libnvrtc Path: {self.libnvrtc_folder}") - print(f"CUDA Version: {self.cuda_version}") print(f"Number of GPUs: {self.n_gpus}") - print(f"GPU Compile Flags: {self.gpu_compile_flags}") - # CUDA Include Path - cuda_include_path = self.cuda_include_path - print(f"CUDA Include Path: {cuda_include_path or 'Not Found'}") + self.print_cuda_version() + self.print_libcuda_folder() + self.print_libnvrtc_folder() + self.print_cuda_include_path() + self.print_nvrtc_flags() + self.print_gpu_compile_flags() # Print relevant environment variables. - print("\nRelevant Environment Variables:") - env_vars = [ - "CUDA_VISIBLE_DEVICES", - "CUDA_PATH", - ] - for var in env_vars: - value = os.environ.get(var, None) - if value: - print(f"{var} = {value}") - else: - print(f"{var} is not set") + print_envs(self.cuda_env_vars) if __name__ == "__main__": - cudastuff = CUDAConfig() - cudastuff.print_all() + cuda_config = CUDAConfig() + cuda_config.print_platform() + cuda_config.print_cpp() + cuda_config.print_cuda() From 6bfa32a642f409fb252734032d432ab319e92496 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Fri, 1 May 2026 12:35:57 +0200 Subject: [PATCH 17/98] Refactor pass test... --- keopscore/keopscore/__init__.py | 40 +- keopscore/keopscore/binders/LinkCompile.py | 18 +- .../binders/nvrtc/Gpu_link_compile.py | 60 +- keopscore/keopscore/config/Cuda.py | 592 ++++++++++++++++++ .../config/{CppConfig.py => CxxCompiler.py} | 150 +++-- keopscore/keopscore/config/Debug.py | 32 + keopscore/keopscore/config/KeOpsPath.py | 231 +++++++ .../keopscore/config/{openmp.py => OpenMP.py} | 154 +++-- keopscore/keopscore/config/Platform.py | 67 +- keopscore/keopscore/config/__init__.py | 95 ++- keopscore/keopscore/config/_shared.py | 146 +---- keopscore/keopscore/config/base_config.py | 440 ------------- keopscore/keopscore/config/chunks.py | 14 +- keopscore/keopscore/config/cuda.py | 579 ----------------- keopscore/keopscore/formulas/GetReduction.py | 17 +- .../LinearOperators/AdjointOperator.py | 2 +- .../LinearOperators/LinearOperator.py | 2 +- .../LinearOperators/SumLinOperator.py | 2 +- .../formulas/LinearOperators/TraceOperator.py | 2 +- .../formulas/LinearOperators/__init__.py | 10 +- keopscore/keopscore/formulas/Operation.py | 10 +- .../formulas/VectorizedComplexScalarOp.py | 2 +- .../keopscore/formulas/VectorizedScalarOp.py | 2 +- .../keopscore/formulas/autodiff/Divergence.py | 2 +- keopscore/keopscore/formulas/autodiff/Grad.py | 2 +- .../keopscore/formulas/autodiff/Laplacian.py | 2 +- .../keopscore/formulas/autodiff/__init__.py | 17 +- .../formulas/complex/ComplexExp1j.py | 2 +- .../keopscore/formulas/complex/ComplexImag.py | 2 +- .../keopscore/formulas/complex/ComplexReal.py | 2 +- .../formulas/complex/ComplexRealScal.py | 2 +- .../formulas/complex/ComplexSquareAbs.py | 2 +- .../keopscore/formulas/complex/ComplexSum.py | 2 +- .../keopscore/formulas/complex/ComplexSumT.py | 2 +- .../keopscore/formulas/complex/__init__.py | 22 + .../formulas/factorization/Factorize.py | 8 +- .../formulas/factorization/__init__.py | 7 + keopscore/keopscore/formulas/maths/ArgMax.py | 2 +- keopscore/keopscore/formulas/maths/ArgMin.py | 2 +- keopscore/keopscore/formulas/maths/BSpline.py | 2 +- keopscore/keopscore/formulas/maths/Concat.py | 2 +- keopscore/keopscore/formulas/maths/Divide.py | 2 +- keopscore/keopscore/formulas/maths/Elem.py | 2 +- keopscore/keopscore/formulas/maths/ElemT.py | 2 +- keopscore/keopscore/formulas/maths/Extract.py | 2 +- .../keopscore/formulas/maths/ExtractT.py | 2 +- .../keopscore/formulas/maths/MatVecMult.py | 2 +- keopscore/keopscore/formulas/maths/Max.py | 2 +- keopscore/keopscore/formulas/maths/Min.py | 2 +- keopscore/keopscore/formulas/maths/Mult.py | 2 +- keopscore/keopscore/formulas/maths/OneHot.py | 2 +- .../keopscore/formulas/maths/Scalprod.py | 2 +- .../formulas/maths/SoftDTW_SqDist.py | 2 +- keopscore/keopscore/formulas/maths/SumT.py | 2 +- .../keopscore/formulas/maths/TensorProd.py | 2 +- .../keopscore/formulas/maths/VecMatMult.py | 2 +- .../keopscore/formulas/maths/__init__.py | 84 +-- .../reductions/KMin_ArgKMin_Reduction.py | 2 +- .../reductions/Max_ArgMax_Reduction_Base.py | 2 +- .../formulas/reductions/Max_Reduction.py | 2 +- .../Max_SumShiftExpWeight_Reduction.py | 5 +- .../reductions/Min_ArgMin_Reduction_Base.py | 2 +- .../formulas/reductions/Min_Reduction.py | 2 +- .../keopscore/formulas/reductions/__init__.py | 24 +- .../keopscore/formulas/variables/__init__.py | 9 + keopscore/keopscore/get_keops_dll.py | 15 +- keopscore/keopscore/mapreduce/MapReduce.py | 9 +- keopscore/keopscore/mapreduce/__init__.py | 8 +- .../keopscore/mapreduce/cpu/CpuAssignZero.py | 8 +- keopscore/keopscore/mapreduce/cpu/CpuReduc.py | 8 +- .../mapreduce/cpu/CpuReduc_ranges.py | 8 +- keopscore/keopscore/mapreduce/cpu/__init__.py | 9 +- .../mapreduce/gpu/GpuReduc1D_chunks.py | 10 +- .../mapreduce/gpu/GpuReduc1D_finalchunks.py | 10 +- .../mapreduce/gpu/GpuReduc1D_ranges_chunks.py | 10 +- .../gpu/GpuReduc1D_ranges_finalchunks.py | 16 +- .../keopscore/mapreduce/gpu/GpuReduc2D.py | 2 +- keopscore/keopscore/mapreduce/gpu/__init__.py | 19 +- keopscore/keopscore/sandbox/do_clean_keops.py | 2 +- keopscore/keopscore/sandbox/formula.py | 5 +- keopscore/keopscore/sandbox/laplacian.py | 4 +- keopscore/keopscore/test/test_op.py | 2 +- keopscore/keopscore/utils/Cache.py | 12 +- keopscore/keopscore/utils/TestFormula.py | 2 +- keopscore/keopscore/utils/TestOperation.py | 2 +- keopscore/keopscore/utils/Tree.py | 2 +- keopscore/keopscore/utils/code_gen_utils.py | 58 +- keopscore/keopscore/utils/file_utils.py | 42 ++ keopscore/keopscore/utils/gpu_utils.py | 228 +------ keopscore/keopscore/utils/math_functions.py | 2 +- keopscore/keopscore/utils/messages.py | 40 ++ keopscore/keopscore/utils/misc_utils.py | 166 ----- keopscore/keopscore/utils/path_utils.py | 98 +++ keopscore/keopscore/utils/system_utils.py | 72 +++ pykeops/pykeops/__init__.py | 37 +- pykeops/pykeops/common/get_options.py | 10 +- pykeops/pykeops/common/gpu_utils.py | 5 +- .../pykeops/common/keops_io/LoadKeOps_cpp.py | 18 +- .../common/keops_io/LoadKeOps_nvrtc.py | 20 +- pykeops/pykeops/common/keops_io/__init__.py | 2 +- pykeops/pykeops/common/lazy_tensor.py | 2 +- pykeops/pykeops/common/operations.py | 2 +- pykeops/pykeops/config.py | 28 +- pykeops/pykeops/numpy/utils.py | 4 +- .../test_soft_dtw_kernel_dissmatrix.py | 4 +- .../sandbox/test_soft_dtw_kernel_v0.py | 2 +- .../sandbox/test_torch_func_hessian.py | 5 +- pykeops/pykeops/test/test_torch_func.py | 5 +- 108 files changed, 1826 insertions(+), 2098 deletions(-) create mode 100644 keopscore/keopscore/config/Cuda.py rename keopscore/keopscore/config/{CppConfig.py => CxxCompiler.py} (50%) create mode 100644 keopscore/keopscore/config/Debug.py create mode 100644 keopscore/keopscore/config/KeOpsPath.py rename keopscore/keopscore/config/{openmp.py => OpenMP.py} (58%) delete mode 100644 keopscore/keopscore/config/base_config.py delete mode 100644 keopscore/keopscore/config/cuda.py create mode 100644 keopscore/keopscore/utils/file_utils.py create mode 100644 keopscore/keopscore/utils/messages.py delete mode 100644 keopscore/keopscore/utils/misc_utils.py create mode 100644 keopscore/keopscore/utils/path_utils.py create mode 100644 keopscore/keopscore/utils/system_utils.py diff --git a/keopscore/keopscore/__init__.py b/keopscore/keopscore/__init__.py index e7beb752e..8ede817db 100644 --- a/keopscore/keopscore/__init__.py +++ b/keopscore/keopscore/__init__.py @@ -1,47 +1,21 @@ import os -from os import path -########################################################### -# Verbosity level -verbose = True -if os.getenv("KEOPS_VERBOSE") == "0": - verbose = False - -here = path.abspath(path.dirname(__file__)) +# Version +here = os.path.abspath(os.path.dirname(__file__)) with open(os.path.join(here, "keops_version"), encoding="utf-8") as v: __version__ = v.read().rstrip() -from keopscore.config import * -from keopscore.utils.code_gen_utils import clean_keops, check_health -from keopscore.utils.misc_utils import CHECK_MARK, CROSS_MARK - -set_build_folder = config.set_different_build_folder - -# flags for debugging : -# prints information about atomic operations during code building -debug_ops = False -# adds C++ code for printing all input and output values -# for all atomic operations during computations -debug_ops_at_exec = False - -# flag for automatic factorization : apply automatic factorization for all formulas before reduction. -auto_factorize = False +# Config +import keopscore.config # Initialize CUDA libraries if CUDA is used -if cuda_config.get_use_cuda(): +if keopscore.config.cuda.get_use_cuda(): # Initialize CUDA libraries if necessary - cuda_config._cuda_libraries_available() from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile from keopscore.binders.nvrtc.Gpu_link_compile import jit_compile_dll if not os.path.exists(jit_compile_dll()): Gpu_link_compile.compile_jit_compile_dll() - -# Retrieve the current build folder -build_folder = config.get_build_folder() - -# Retrieve details about the current CUDA configuration -show_gpu_config = cuda_config -show_cuda_status = cuda_config.get_use_cuda() -cuda_block_size = cuda_config.get_cuda_block_size() +# expose to the user +set_build_folder = keopscore.config.path.set_different_build_folder diff --git a/keopscore/keopscore/binders/LinkCompile.py b/keopscore/keopscore/binders/LinkCompile.py index aae5d70ea..b452e6a48 100644 --- a/keopscore/keopscore/binders/LinkCompile.py +++ b/keopscore/keopscore/binders/LinkCompile.py @@ -1,11 +1,13 @@ import os -from keopscore.config import config -from keopscore.utils.code_gen_utils import get_hash_name -from keopscore.utils.misc_utils import KeOps_Error, KeOps_Message -cpp_flags = config.get_cpp_flags() -get_build_folder = config.get_build_folder +import keopscore.config +from keopscore.utils.code_gen_utils import get_hash_name +from keopscore.utils.messages import KeOps_Error, KeOps_Message +cpp_flag = keopscore.config.cxx.get_compile_options() +cpp_flag += keopscore.config.cxx.get_linking_options() +cpp_flag += keopscore.config.cuda.get_nvrtc_flags() +cpp_flag += keopscore.config.cuda.get_include_options() class LinkCompile: """ @@ -31,17 +33,17 @@ def __init__(self): self.use_half, self.use_fast_math, self.device_id, - cpp_flags, + cpp_flag, # TODO: check that get_envs is sufficient... ) # info_file is the name of the file that will contain some meta-information required by the bindings, e.g. 7b9a611f7e.nfo self.info_file = os.path.join( - get_build_folder(), self.gencode_filename + ".nfo" + keopscore.config.path.get_build_folder(), self.gencode_filename + ".nfo" ) # gencode_file is the name of the source file to be created and then compiled, e.g. 7b9a611f7e.cpp or 7b9a611f7e.cu self.gencode_file = os.path.join( - get_build_folder(), + keopscore.config.path.get_build_folder(), self.gencode_filename + "." + self.source_code_extension, ) diff --git a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py index ccd6eb00d..13e447ea8 100644 --- a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py +++ b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py @@ -1,53 +1,37 @@ +import ctypes import os -from ctypes import create_string_buffer, CDLL, c_int -from os import RTLD_LAZY import sysconfig -from os.path import join -import keopscore +import keopscore.config from keopscore.binders.LinkCompile import LinkCompile -from keopscore.config import * - -from keopscore.utils.misc_utils import KeOps_Error, KeOps_Message, KeOps_OS_Run from keopscore.utils.gpu_utils import custom_cuda_include_fp16_path +from keopscore.utils.messages import KeOps_Error, KeOps_Message +from keopscore.utils.system_utils import KeOps_OS_Run + -cuda_version = cuda_config.get_cuda_version() -jit_binary = config.get_jit_binary() -cxx_compiler = config.get_cxx_compiler() -nvrtc_flags = cuda_config.get_nvrtc_flags() -cuda_available = cuda_config.get_use_cuda() -build_folder = config.get_build_folder() -get_gpu_props = cuda_config.get_gpu_props() - -nvrtc_include = " -I" + config.get_bindings_source_dir() -cuda_include_path = cuda_config.get_cuda_include_path() -if cuda_include_path: - nvrtc_include += " -I" + cuda_include_path - -jit_source_file = join( - config.get_base_dir_path(), "binders", "nvrtc", "keops_nvrtc.cpp" +jit_source_file = os.path.join( + keopscore.config.path.get_base_dir_path(), "binders", "nvrtc", "keops_nvrtc.cpp" ) jit_compile_src = os.path.join( - os.path.abspath(os.path.dirname(__file__)), "nvrtc_jit.cpp" + keopscore.config.path.get_base_dir_path(), "binders", "nvrtc", "nvrtc_jit.cpp" ) def jit_compile_dll(): return os.path.join( - build_folder, + keopscore.config.path.get_build_folder(), "nvrtc_jit" + sysconfig.get_config_var("SHLIB_SUFFIX"), ) class Gpu_link_compile(LinkCompile): source_code_extension = "cu" - low_level_code_prefix = "cubin_" if cuda_version >= 11010 else "ptx_" - ngpu, gpu_props_compile_flags = get_gpu_props + low_level_code_prefix = "cubin_" if keopscore.config.cuda.get_cuda_version() >= 11010 else "ptx_" def __init__(self): # checking that the system has a Gpu : - if not (cuda_available and Gpu_link_compile.ngpu): + if not (keopscore.config.cuda.get_use_cuda() and keopscore.config.cuda.get_n_gpus()): KeOps_Error( "Trying to compile cuda code... but we detected that the system has no properly configured cuda lib." ) @@ -57,13 +41,13 @@ def __init__(self): # low_level_code_file is filename of low level code (PTX for Cuda) or binary (CUBIN for Cuda) # generated by the JIT compiler, e.g. ptx_7b9a611f7e self.low_level_code_file = os.path.join( - build_folder, + keopscore.config.path.get_build_folder(), self.low_level_code_prefix + self.gencode_filename, ).encode("utf-8") - self.my_c_dll = CDLL(jit_compile_dll(), mode=RTLD_LAZY) + self.my_c_dll = ctypes.CDLL(jit_compile_dll(), mode=os.RTLD_LAZY) # actual dll to be called is the jit binary, TODO: check if this is relevent - self.true_dllname = jit_binary + self.true_dllname = keopscore.config.path.get_jit_binary() # file to check for existence to detect compilation is needed self.file_to_check = self.low_level_code_file @@ -76,12 +60,12 @@ def generate_code(self): # we execute the main dll, passing the code as argument, and the name of the low level code file to save the assembly instructions res = self.my_c_dll.Compile( - create_string_buffer(self.low_level_code_file), - create_string_buffer(self.code.encode("utf-8")), - c_int(self.use_half), - c_int(self.use_fast_math), - c_int(self.device_id), - create_string_buffer( + ctypes.create_string_buffer(self.low_level_code_file), + ctypes.create_string_buffer(self.code.encode("utf-8")), + ctypes.c_int(self.use_half), + ctypes.c_int(self.use_fast_math), + ctypes.c_int(self.device_id), + ctypes.create_string_buffer( (custom_cuda_include_fp16_path() + os.path.sep).encode("utf-8") ), ) @@ -95,7 +79,7 @@ def generate_code(self): @staticmethod def get_compile_command( - sourcename=jit_source_file, dllname=jit_binary, extra_flags="" + sourcename=jit_source_file, dllname=keopscore.config.path.get_jit_binary(), extra_flags="" ): # This is about the main KeOps binary (dll) that will be used to JIT compile all formulas. # If the dll is not present, it compiles it from source, except if check_compile is False. @@ -111,7 +95,7 @@ def get_compile_command( else '\\"compute\\"' ) target_type_define = f"-DnvrtcGetTARGET={nvrtcGetTARGET} -DnvrtcGetTARGETSize={nvrtcGetTARGETSize} -DARCHTAG={arch_tag}" - return f"{cxx_compiler} {nvrtc_flags} {extra_flags} {target_type_define} {nvrtc_include} {Gpu_link_compile.gpu_props_compile_flags} {sourcename} -o {dllname}" + return f"{keopscore.config.cxx.get_cxx_compiler()} {keopscore.config.cuda.get_preprocessing_options()} {target_type_define} {keopscore.config.cxx.get_compile_options()} {keopscore.config.cuda.get_nvrtc_flags()} {keopscore.config.cuda.get_include_options()} {keopscore.config.path.get_include_options()} {extra_flags} {sourcename} {keopscore.config.cxx.get_linking_options()} -o {dllname}" @staticmethod def compile_jit_compile_dll(): diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py new file mode 100644 index 000000000..9664f9741 --- /dev/null +++ b/keopscore/keopscore/config/Cuda.py @@ -0,0 +1,592 @@ +import ctypes +import os + +from keopscore.config._shared import print_envs, not_found_str, enabled_dict +from keopscore.utils.messages import KeOps_Warning +from keopscore.utils.path_utils import ( + _first_matching_file, + _ordered_search_roots, + _path_candidates, +) +from keopscore.utils.system_utils import _find_library_by_names + + +class CudaConfig: + """ + Class for CUDA detection and configuration. + """ + + # CUDA constants + CUDA_SUCCESS = 0 + CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK = 1 + CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK = 8 + CUDA_BLOCK_SIZE = 192 + + # Cuda detection variables + _use_cuda = None + _specific_gpus = None + + _cuda_include_path = None + _nvrtc_flags = None + _cuda_version = None + + _n_gpus = 0 + _MaxThreadsPerBlock = [] + _SharedMemPerBlock = [] + + _preprocessing_options = "" + _include_options = "" + _cuda_block_size = None + + # ------------------------ # + # Search location # + # ------------------------ # + + cuda_env_vars = [ + "CUDA_VISIBLE_DEVICES", + "CUDA_PATH", + "CUDA_HOME", + "CUDA_ROOT", + "CUDA_TOOLKIT_ROOT_DIR", + ] + + pip_suffixes = ( + "nvidia/cuda_runtime", + "nvidia/cuda_nvrtc", + ) + + system_suffixes = ( + os.path.join(os.path.sep, "usr", "local", "cuda"), + os.path.join(os.path.sep, "usr", "local"), + os.path.join(os.path.sep, "opt", "cuda"), + os.path.join(os.path.sep, "usr"), + os.path.join(os.path.sep, "lib"), + ) + + library_suffixes = ( + "lib64", + "lib", + os.path.join("lib", "x86_64-linux-gnu"), + ) + + include_suffixes = ( + "include", + os.path.join("targets", "x86_64-linux", "include"), + os.path.join("targets", "sbsa-linux", "include"), + os.path.join("targets", "aarch64-linux", "include"), + ) + + # ------------------------- # + # Library info dicts # + # ------------------------- # + + _libcuda_info = { + "name": "cuda", + "lib_basename_candidate": ["libcuda.so.*", "libcuda.dylib", "cuda.lib"], + "header_basename": "cuda.h", + "library": None, # to be filled later + "header": None, # to be filled later + "ctype_handle": None, # to be filled later + } + _libnvrtc_info = { + "name": "nvrtc", + "lib_basename_candidate": ["libnvrtc.so.*", "libnvrtc.dylib", "nvrtc.lib"], + "header_basename": "nvrtc.h", + "library": None, # to be filled later + "header": None, # to be filled later + "ctype_handle": None, # to be filled later + } + _cudart_info = { + "name": "cudart", + "lib_basename_candidate": ["libcudart.so.*", "libcudart.dylib", "cudart.lib"], + "header_basename": None, + "library": None, # to be filled later + "header": None, # not needed + "ctype_handle": None, # to be filled later + } + + def __init__(self): + + self.set_specific_gpus() + + super().__init__() + + self.set_use_cuda() + + # If cuda is enabled, then we finalize the rest of the config + if self.get_use_cuda(): + self.set_cuda_version() + self.set_cuda_include_path() + self.set_nvrtc_flags() + self.set_cuda_block_size() + self.set_preprocessing_options() + + def find_install_path(self, lib_dict_info, warn=None): + """ + Locate a cuda and headers using an explicit ordered search. + + Arguments: + lib_dict_info (dict): A dictionary containing at least the keys 'name', 'lib_basename_candidate', and 'header_basename' for the library to find. This allows the function to be used for finding libcuda, libcudart, or nvrtc by passing the appropriate info dict. + + Returns: + result (dict): a copy of lib_dict_info completed with the ``library`` and ``header`` keys containing the absolute paths to the library file and include directory, or None if not found. + """ + result = ( + lib_dict_info.copy() + ) # Start with the provided info, which may contain names and file patterns + + candidate_roots = _ordered_search_roots( + env_vars=self.cuda_env_vars, + pip_suffixes=self.pip_suffixes, + conda_root="CONDA_PREFIX", + system_roots=self.system_suffixes, + ) + + # ------------------------ # + # Search for library file # + # ------------------------ # + + # First try to find the library file using the candidate roots and library suffixes + + result["library"] = _first_matching_file( + _path_candidates(candidate_roots, self.library_suffixes), + result["lib_basename_candidate"] + ) + if result["library"] is None: + result["library"] = _find_library_by_names((result["name"],)) + + if result["library"] is None and warn: + KeOps_Warning(f"lib{result['name']} not found.") + + # ------------------------ # + # Search for header files # + # ------------------------ # + + result["header"] = _first_matching_file( + _path_candidates(candidate_roots, self.include_suffixes), + (result["header_basename"],), + ) + + if ( + result["header"] is None + and warn + and result["header_basename"] is not None + ): + KeOps_Warning(f"{result['name']} header files not found.") + + return result + + def _find_and_load_libcuda(self): + """Locate, load, and initialize the CUDA driver library.""" + self._libcuda_info = self.find_install_path(self._libcuda_info, warn=True) + libcuda_path = self._libcuda_info["library"] + if not libcuda_path: + return ( + False, + "libcuda not found. Make sure the CUDA driver is installed and accessible. Switching to CPU only.", + ) + + try: + libcuda = ctypes.CDLL(libcuda_path, mode=ctypes.RTLD_GLOBAL) + except OSError as e: + return ( + False, + f"Failed to load library '{libcuda_path}': {e}", + ) + + if libcuda.cuInit(0) != self.CUDA_SUCCESS: + return ( + False, + "libcuda was detected, but driver API could not be initialized. Rebooting the system may help. Switching to CPU only.", + ) + + # If we successfully loaded libcuda and initialized it, store the handle in the config for potential future use + self._libcuda_info["ctype_handle"] = libcuda + + nGpus = ctypes.c_int() + if libcuda.cuDeviceGetCount(ctypes.byref(nGpus)) != self.CUDA_SUCCESS or nGpus.value == 0: + return ( + False, + "libcuda was detected and driver API was initialized, but no working GPU found. Switching to CPU only.", + ) + + + self._MaxThreadsPerBlock = [0] * nGpus.value + self._SharedMemPerBlock = [0] * nGpus.value + + for d in range(nGpus.value): + self._MaxThreadsPerBlock[d], self._SharedMemPerBlock[d], err_msg = ( + self._get_device_attributes(libcuda, d) + ) + if err_msg: + return ( + False, + f"libcuda was detected and driver API was initialized, but " + + err_msg + + " Switching to CPU only.", + ) + + self._n_gpus = nGpus.value + + return True, "" + + def _find_and_load_libnvrtc(self): + """Locate and load the NVRTC runtime compilation library.""" + self._libnvrtc_info = self.find_install_path(self._libnvrtc_info, warn=True) + libnvrtc_path = self._libnvrtc_info["library"] + if not libnvrtc_path: + return ( + False, + "libnvrtc not found. Make sure the CUDA toolkit is installed and accessible. Switching to CPU only.", + ) + + try: + libnvrtc_handle = ctypes.CDLL(libnvrtc_path, mode=ctypes.RTLD_GLOBAL) + except OSError as e: + return ( + False, + f"Failed to load library '{os.path.basename(libnvrtc_path)}': {e}", + ) + + # If we successfully loaded libnvrtc, store the handle in the config + self._libnvrtc_info["ctype_handle"] = libnvrtc_handle + + return True, "" + + def _find_and_load_cudart(self): + """ + Attempt to find the CUDA version by loading the CUDA runtime library and querying its version. + Returns: + success (bool): True if the version was successfully determined, False otherwise. + error_msg (str): Contains error details if success==False, else "". + + """ + + self._cudart_info = self.find_install_path(self._cudart_info, warn=False) + libcudart_path = self._cudart_info["library"] + if not libcudart_path: + return ( + False, + "libcudart not found. Make sure the CUDA toolkit is installed and accessible. Switching to CPU only.", + ) + + try: + libcudart_handle = ctypes.CDLL(libcudart_path) + except OSError as e: + return ( + False, + f"Failed to load '{os.path.basename(libcudart_path)}': {e}", + ) + + cuda_version = ctypes.c_int() + if libcudart_handle.cudaRuntimeGetVersion(ctypes.byref(cuda_version)) != self.CUDA_SUCCESS: + return ( + False, + "libcudart was found and loaded, but failed to get CUDA runtime version. Switching to CPU only.", + ) + + # If we successfully loaded libcudart, store the handle in the config + self._cudart_info["ctype_handle"] = libcudart_handle + self._cuda_version = int(cuda_version.value) + + return True, "" + + def _cuda_libraries_available(self): + """ + Check if libcuda (driver); libcudart (cudatoolkit) and nvrtc (cudatoolkit) libraries are available. + This is where ```_cuda_include_path``` is set. + + Returns: + True if all three are loadable, False otherwise. + This is also where we handle one single warning if needed. + """ + + # Libcuda (driver) loaded globally so it is available to KeOps shared objects. + success_cuda, err_cuda = self._find_and_load_libcuda() + if not success_cuda: + KeOps_Warning(f"{err_cuda}. Switching to CPU only.") + return False + + # libnvrtc as well, since it's required for the runtime compilation of CUDA code. + success_nvrtc, err_nvrtc = self._find_and_load_libnvrtc() + if not success_nvrtc: + KeOps_Warning(f"{err_nvrtc}. Switching to CPU only.") + return False + + # Finally check cudart + success_cudart, err_cudart = self._find_and_load_cudart() + if not success_cudart: + KeOps_Warning(f"{err_cudart}. Switching to CPU only.") + return False + + return True + + # CUDA Support + def set_use_cuda(self): + """Determine and set whether to use CUDA.""" + self._use_cuda = self._cuda_libraries_available() + + def get_use_cuda(self): + return self._use_cuda + + def print_use_cuda(self): + print(f"CUDA Support: {enabled_dict[self.get_use_cuda() or False]}") + + # CUDA Block Size + def set_cuda_block_size(self): + """Sets default cuda block size.""" + self._cuda_block_size = self.CUDA_BLOCK_SIZE + + def get_cuda_block_size(self): + return self._cuda_block_size + + def print_cuda_block_size(self): + print(f"CUDA Block Size: {self.get_cuda_block_size()}") + + # Specific GPUs + def set_specific_gpus(self): + """Set specific GPUs from CUDA_VISIBLE_DEVICES.""" + if os.getenv("CUDA_VISIBLE_DEVICES"): + self._specific_gpus = os.getenv("CUDA_VISIBLE_DEVICES").replace(",", "_") + + def get_specific_gpus(self): + """Get the specific GPUs.""" + return self._specific_gpus + + # Number of GPUs + def set_n_gpus(self): + """Set the number of GPUs detected. This is done in _find_and_load_libcuda.""" + pass + + def get_n_gpus(self): + """Get the number of GPUs detected.""" + return self._n_gpus + + def print_n_gpus(self): + """Print the number of GPUs detected.""" + print(f"Number of GPUs Detected: {self.get_n_gpus()}") + + # Libcuda folder + def set_libcuda_folder(self): + """ + Is set in _cuda_libraries_available. + """ + pass + + def get_libcuda_folder(self): + return self._libcuda_info["library"] and os.path.dirname( + self._libcuda_info["library"] + ) + + def print_libcuda_folder(self): + print(f"Libcuda Folder: {self.get_libcuda_folder() or not_found_str}") + + # Libnvrtc folder + def set_libnvrtc_folder(self): + """ + Return nothing if not using cuda + self.libnvrtc_folder is already set in _cuda_libraries_available. + """ + pass + + def get_libnvrtc_folder(self): + return self._libnvrtc_info["library"] and os.path.dirname( + self._libnvrtc_info["library"] + ) + + def print_libnvrtc_folder(self): + print(f"Libnvrtc Folder: {self.get_libnvrtc_folder() or not_found_str}") + + # CUDA Version + def set_cuda_version(self, warn=True): + """Set the CUDA version by querying the CUDA runtime library. Done in _find_and_load_cudart.""" + pass + + def get_cuda_version(self, out_type="single_value"): + + major = self._cuda_version // 1000 + minor = (self._cuda_version % 1000) // 10 + + if out_type == "major,minor": + return major, minor + elif out_type == "string": + return f"{major}.{minor}" + else: + return self._cuda_version + + def print_cuda_version(self): + str = f"CUDA Version: {self.get_cuda_version(out_type="string") if self.get_cuda_version() else not_found_str}" + print(str) + + # CUDA Include Path + def set_cuda_include_path(self): + """Set the CUDA include path by searching for cuda.h and nvrtc.h.""" + # This is done in find_cuda_install since it relies on the cuda installation info which is only available after checking library availability. + if not self.get_use_cuda(): + self._cuda_include_path = None + self._cuda_include_path = list( + set( + [ + os.path.dirname(self._libnvrtc_info["header"]), + os.path.dirname(self._libcuda_info["header"]), + ] + ) + ) + + def get_cuda_include_path(self): + """ + Attempt to find CUDA headers (cuda.h, nvrtc.h) using an explicit + ordered search over environment variables and standard locations. + """ + return self._cuda_include_path + + def print_cuda_include_path(self): + print( + f"CUDA Include Path: {":".join(self.get_cuda_include_path()) or not_found_str}" + ) + + # NVRTC Flags + def set_nvrtc_flags(self): + """Set the NVRTC flags for CUDA compilation.""" + + self._nvrtc_flags = f" -fpermissive -L{self.get_libcuda_folder()} -L{self.get_libnvrtc_folder()} -lcuda -lnvrtc" + + def get_nvrtc_flags(self): + """Get the NVRTC flags for CUDA compilation.""" + return self._nvrtc_flags + + def print_nvrtc_flags(self): + """Print the NVRTC flags for CUDA compilation.""" + print(f"NVRTC Flags: {self.get_nvrtc_flags()}") + + # GPU compile flags + def set_preprocessing_options(self): + """Set GPU preprocessing compile flags based on detected GPU properties.""" + if self.get_n_gpus() == 0: + return + + self.add_to_preprocessing_options(f"-DMAXIDGPU={self.get_n_gpus() - 1}") + for d in range(self.get_n_gpus()): + self.add_to_preprocessing_options( + f"-DMAXTHREADSPERBLOCK{d}={self.get_MaxThreadsPerBlock()[d]}" + ) + self.add_to_preprocessing_options( + f"-DSHAREDMEMPERBLOCK{d}={self.get_SharedMemPerBlock()[d]}" + ) + + def add_to_preprocessing_options(self, flags): + self._preprocessing_options += " " + flags + + def get_preprocessing_options(self): + """Get GPU preprocessing compile flags based on detected GPU properties.""" + return self._preprocessing_options + + def print_preprocessing_options(self): + print(f"GPU Preprocessing Options: {self.get_preprocessing_options() or not_found_str}") + + def set_include_options(self): + self._include_options += "".join(f" -I{p}" for p in list(set(self.get_cuda_include_path()))) + + def get_include_options(self): + return self._include_options + + def print_include_options(self): + print(f"GPU Include Options: {self.get_include_options() or not_found_str}") + + # Helper functions for CUDA attribute detection + def _get_device_attributes(self, libcuda, device_index): + """Return CUDA attributes needed to build compile flags for one device.""" + + device = ctypes.c_int() + if libcuda.cuDeviceGet(ctypes.byref(device), device_index) != self.CUDA_SUCCESS: + return ( + 0, + 0, + "failed to get device handle for device index {device_index}.", + ) + + output = ctypes.c_int() + if ( + libcuda.cuDeviceGetAttribute( + ctypes.byref(output), + self.CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK, + device, + ) + != self.CUDA_SUCCESS + ): + return ( + 0, + 0, + "failed to get max threads per block for device index {device_index}.", + ) + max_threads_per_block = output.value + + if ( + libcuda.cuDeviceGetAttribute( + ctypes.byref(output), + self.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK, + device, + ) + != self.CUDA_SUCCESS + ): + return ( + 0, + 0, + "failed to get shared memory per block for device index {device_index}.", + ) + shared_mem_per_block = output.value + + return max_threads_per_block, shared_mem_per_block, "" + + def set_MaxThreadsPerBlock(self): + pass + + def get_MaxThreadsPerBlock(self): + return self._MaxThreadsPerBlock + + def set_SharedMemPerBlock(self): + pass + + def get_SharedMemPerBlock(self): + return self._SharedMemPerBlock + + def print_all(self): + """Print all CUDA-related configuration""" + + # CUDA Support + print("=" * 60) + print(f"CUDA Support") + print("=" * 60) + + self.print_use_cuda() + if self.get_use_cuda(): + self.print_n_gpus() + self.print_cuda_version() + self.print_libcuda_folder() + self.print_libnvrtc_folder() + self.print_cuda_include_path() + self.print_nvrtc_flags() + + self.print_preprocessing_options() + self.print_include_options() + + # Print relevant environment variables. + print_envs(self.cuda_env_vars) + + +if __name__ == "__main__": + from keopscore.config.Platform import PlatformConfig + from keopscore.config.CxxCompiler import CxxCompilerConfig + from keopscore.config.OpenMP import OpenMPConfig + + platform_info = PlatformConfig() + platform_info.print_all() + + cxx_compiler_info = CxxCompilerConfig(platform_info) + cxx_compiler_info.print_all() + + openmp_info = OpenMPConfig(platform_info, cxx_compiler_info) + openmp_info.print_all() + + cuda_info = CudaConfig() + cuda_info.print_all() diff --git a/keopscore/keopscore/config/CppConfig.py b/keopscore/keopscore/config/CxxCompiler.py similarity index 50% rename from keopscore/keopscore/config/CppConfig.py rename to keopscore/keopscore/config/CxxCompiler.py index 1033d184f..3c2a7867a 100644 --- a/keopscore/keopscore/config/CppConfig.py +++ b/keopscore/keopscore/config/CxxCompiler.py @@ -1,27 +1,35 @@ -import os, shutil +import os +import shutil -from keopscore.utils.misc_utils import KeOps_Warning, KeOps_OS_Run -from keopscore.config.Platform import Platform from keopscore.config._shared import print_envs, not_found_str +from keopscore.utils.messages import KeOps_Error +from keopscore.utils.system_utils import KeOps_OS_Run -class CppConfig(Platform): +class CxxCompilerConfig: """ Common platform and C++ compiler configuration. """ - + _cxx_compiler = None - _cpp_env_flags = None - _compile_options = None + _cxx_env_flags = None + _compile_options = "" + _linking_options = "" + _disable_pragma_unrolls = None - cpp_envs = ["CXX", "CXXFLAGS"] - + _use_Apple_clang = False + + cxx_envs = ["CXX", "CXXFLAGS"] + + def __init__(self, platform): + + # Platform instance get info on system + self.platform = platform - def __init__(self): - super().__init__() self.set_cxx_compiler() - self.set_cpp_env_flags() + self.set_cxx_env_flags() self.set_compile_options() + self.set_linking_options() self.set_disable_pragma_unrolls() # C++ Compiler Detection @@ -44,39 +52,54 @@ def detect_cxx_compiler(self): env_cxx = os.getenv("CXX") if env_cxx: preferred_compilers.append(env_cxx) - + # Common alias for the default C++ compiler - preferred_compilers.append("c++") - - if self.get_platform() == "Darwin": + preferred_compilers.append("c++") + + if self.platform.get_platform() == "Darwin": # On macOS, prioritize clang++ over g++ preferred_compilers.extend(("clang++", "g++")) else: preferred_compilers.append(["g++", "clang++"]) - + # Return the first available compiler from the preferred list. for compiler in preferred_compilers: cxx_compiler_path = shutil.which(compiler) if cxx_compiler_path: return cxx_compiler_path else: - KeOps_Warning("No C++ compiler found. Define CXX environment variable or install g++.") + KeOps_Error( + "No C++ compiler found. Define CXX environment variable or install g++." + ) return None def get_cxx_compiler_version(self): """Detect if using Apple Clang.""" - compiler_info = KeOps_OS_Run(f"{self.get_cxx_compiler()} --version").stdout.decode("utf-8").splitlines()[0].strip() + compiler_info = ( + KeOps_OS_Run(f"{self.get_cxx_compiler()} --version") + .stdout.decode("utf-8") + .splitlines()[0] + .strip() + ) return compiler_info - + def print_cxx_compiler_version(self): """Print the C++ compiler version information.""" - print(f"C++ Compiler Version: {self.get_cxx_compiler_version() or not_found_str}") - - @property - def use_Apple_clang(self): + print( + f"C++ Compiler Version: {self.get_cxx_compiler_version() or not_found_str}" + ) + + def set_use_Apple_clang(self): """Detect if using Apple Clang.""" - return "Apple clang" in self.get_cxx_compiler_version() if self.get_cxx_compiler_version() else False + self._use_Apple_clang = ( + "Apple clang" in self.get_cxx_compiler_version() + if self.get_cxx_compiler_version() + else False + ) + + def get_use_Apple_clang(self): + return self._use_Apple_clang # Disable Pragma Unrolls def set_disable_pragma_unrolls(self): @@ -95,15 +118,23 @@ def print_disable_pragma_unrolls(self): # Compile Options def set_compile_options(self): """Set the compile options.""" - self._compile_options = " -shared -fPIC -O3 -std=c++11 -flto=auto" - - # Specific Options for Apple Clang (maybe unnecessary in recent MacOs versions) - if self.get_platform() == "Darwin" and self.use_Apple_clang: - self.cpp_flags += " -undefined dynamic_lookup" - - # ... and Silicon chips - if self.get_platform() == "Darwin" and self.get_machine() in ["arm64", "arm64e"]: - self.cpp_flags += " -arch arm64" + + self._compile_options = self.get_cxx_env_flags() + + # compilation / behavior + self.add_to_compile_option("-std=c++11 -O3 -flto=auto") + # code model + self.add_to_compile_option("-fPIC") + + # architecture (macOS ARM) + if self.platform.get_platform() == "Darwin" and self.platform.get_machine() in [ + "arm64", + "arm64e", + ]: + self.add_to_compile_option("-arch arm64") + + def add_to_compile_option(self, flags): + self._compile_options += " " + flags def get_compile_options(self): """Get the compile options.""" @@ -113,17 +144,44 @@ def print_compile_options(self): """Print the compile options.""" print(f"Compile Options: {self.get_compile_options()}") + # Linking options + def set_linking_options(self): + """Set the linking options.""" + # linking / output + self.add_to_linking_options("-shared") + + # linker behavior (macOS) + if self.platform.get_platform() == "Darwin" and self.get_use_Apple_clang(): + self.add_to_linking_options("-undefined dynamic_lookup") + + # architecture (macOS ARM) + if self.platform.get_platform() == "Darwin" and self.platform.get_machine() in [ + "arm64", + "arm64e", + ]: + self.add_to_linking_options("-arch arm64") + + def add_to_linking_options(self, flags): + self._linking_options += " " + flags + + def get_linking_options(self): + """Get the linking options.""" + return self._linking_options + + def print_linking_options(self): + print(f"Linking Options: {self.get_linking_options()}") + # C++ Environment Flags. Unused yet... - def set_cpp_env_flags(self): + def set_cxx_env_flags(self): """Recover the C++ environment flags.""" - self._cpp_env_flags = os.getenv("CXXFLAGS") if "CXXFLAGS" in os.environ else "" + self._cxx_env_flags = os.getenv("CXXFLAGS") if "CXXFLAGS" in os.environ else "" - def get_cpp_env_flags(self): + def get_cxx_env_flags(self): """Get the C++ environment flags.""" - return self._cpp_env_flags + return self._cxx_env_flags # Comprehensive Platform and Compiler Information - def print_cpp(self): + def print_all(self): """Print all platform and C++ compiler information.""" print("=" * 60) @@ -133,13 +191,17 @@ def print_cpp(self): self.print_cxx_compiler() self.print_cxx_compiler_version() self.print_compile_options() - self.print_disable_pragma_unrolls() + self.print_linking_options() # Print relevant environment variables. - print_envs(self.cpp_envs) + print_envs(self.cxx_envs) if __name__ == "__main__": - cpp_config = CppConfig() - cpp_config.print_platform() - cpp_config.print_cpp() \ No newline at end of file + from keopscore.config.Platform import PlatformConfig + + platform_info = PlatformConfig() + platform_info.print_all() + + cxx_compiler_info = CxxCompilerConfig(platform_info) + cxx_compiler_info.print_all() diff --git a/keopscore/keopscore/config/Debug.py b/keopscore/keopscore/config/Debug.py new file mode 100644 index 000000000..ee2610713 --- /dev/null +++ b/keopscore/keopscore/config/Debug.py @@ -0,0 +1,32 @@ +import os + +class DebugConfig: + + # prints information about atomic operations during code building + _debug_ops = False + # adds C++ code for printing all input and output values for all atomic operations during computations + _debug_ops_at_exec = False + + # Verbosity level (default is False unless KEOPS_VERBOSE define and not 0) + _verbose = os.getenv("KEOPS_VERBOSE") != "0" + + def __init__(self): + pass + + def set_debug_ops(self, debug_ops): + self._debug_ops = debug_ops + + def get_debug_ops(self): + return self._debug_ops + + def set_debug_ops_at_exec(self, debug_ops_at_exec): + self._debug_ops_at_exec = debug_ops_at_exec + + def get_debug_ops_at_exec(self): + return self._debug_ops_at_exec + + def get_verbose(self): + return self._verbose + + def set_verbose(self, verbose): + self._verbose = verbose \ No newline at end of file diff --git a/keopscore/keopscore/config/KeOpsPath.py b/keopscore/keopscore/config/KeOpsPath.py new file mode 100644 index 000000000..951e32e44 --- /dev/null +++ b/keopscore/keopscore/config/KeOpsPath.py @@ -0,0 +1,231 @@ +import os +import sys + +import keopscore +from keopscore.config._shared import not_found_str, ensure_directory, print_envs + + +class KeOpsPathConfig: + """ + Class to manage the path to the KeOps library. + """ + + _base_dir_path = None + _keops_cache_folder = None + _default_build_folder_name = None + _default_build_path = None + _build_folder = None + + _jit_binary = None + _include_options = "" + + path_env_vars = ( + "KEOPS_CACHE_FOLDER", + ) + + def __init__(self, platform, cuda): + + self.platform = platform + self.cuda = cuda + + # Initialize common configuration settings + self.set_base_dir_path() + self.set_keops_cache_folder() + + self.set_default_build_folder_name() + self.set_default_build_path() # should be done at the end.. + self.set_jit_binary() + self.set_include_options() + + # Base Directory Path + def set_base_dir_path(self): + """Set the base directory path.""" + self._base_dir_path = os.path.abspath( + os.path.join(os.path.dirname(os.path.realpath(__file__)), "..") + ) + + def get_base_dir_path(self): + """Get the base directory path.""" + return self._base_dir_path + + def print_base_dir_path(self): + """Print the base directory path.""" + print(f"Base Directory Path: {self.get_base_dir_path() or not_found_str}") + + # KeOps Cache Folder + def set_keops_cache_folder(self): + """Set the KeOps cache folder.""" + cache_folder = os.getenv("KEOPS_CACHE_FOLDER") + if cache_folder is None: + cache_folder = os.path.join( + os.path.expanduser("~"), ".cache", f"keops{keopscore.__version__}" + ) + os.makedirs(cache_folder, exist_ok=True) + self._keops_cache_folder = cache_folder + + def get_keops_cache_folder(self): + """Get the KeOps cache folder.""" + return self._keops_cache_folder + + def print_keops_cache_folder(self): + """Print the KeOps cache folder.""" + print(f"KeOps Cache Folder: {self.get_keops_cache_folder() or not_found_str}") + + # Build Folder Management + def set_default_build_folder_name(self): + """Set the default build folder name.""" + name_parts = [ + "_".join(self.platform.get_uname()[:3]), + f"python{self.platform.get_python_version()}", + ] + if self.cuda.get_use_cuda(): + name_parts.append(f"CUDA{self.cuda.get_cuda_version()}") + specific_gpus = self.cuda.get_specific_gpus() + if specific_gpus: + name_parts.append(f"VISIBLE_DEVICES{specific_gpus}") + + self._default_build_folder_name = "_".join(name_parts) + + def get_default_build_folder_name(self): + """Return the platform-specific default build folder suffix.""" + return self._default_build_folder_name + + def print_default_build_folder_name(self): + """Print the default build folder name.""" + print(f"Default Build Folder Name: {self.get_default_build_folder_name()}") + + def set_different_build_folder( + self, path=None, read_save_file=False, write_save_file=True, reset_all=True + ): + """ + Set or update the build folder path for KeOps. + + Parameters: + - path: The new build folder path. If None, it will be determined based on saved settings or defaults. + - read_save_file: If True, read the build folder path from a save file if path is not provided. + - write_save_file: If True, write the new build folder path to the save file. + - reset_all: If True, reset all cached formulas and recompile necessary components. + """ + # If path is not given, we either read the save file or use the default build path + save_file = os.path.join(self.get_keops_cache_folder(), "build_folder_location.txt") + if not path: + if read_save_file and os.path.isfile(save_file): + with open(save_file, "r") as f: + path = f.read() + else: + path = self.get_default_build_path() + + # Create the folder if not yet done + os.makedirs(path, exist_ok=True) + + # Remove the old build path from sys.path if it's there + if self._build_folder and self._build_folder in sys.path: + sys.path.remove(self._build_folder) + # Update _build_folder to the new path + self._build_folder = path + # Add the new build path to sys.path + if self._build_folder not in sys.path: + sys.path.append(self._build_folder) + + # Saving the location of the build path in a file + if write_save_file: + with open(save_file, "w") as f: + f.write(path) + + # Reset all cached formulas if needed + if reset_all: + # Reset cached formulas + keopscore.get_keops_dll.get_keops_dll.reset( + new_save_folder=self._build_folder + ) + # Handle CUDA-specific recompilation if CUDA is used + if hasattr(self, "_use_cuda") and self._use_cuda: + from keopscore.binders.nvrtc.Gpu_link_compile import ( + Gpu_link_compile, + jit_compile_dll, + ) + + if not os.path.exists(jit_compile_dll()): + Gpu_link_compile.compile_jit_compile_dll() + + def get_build_folder(self): + return self._build_folder + + # Default Build Path + def set_default_build_path(self): + """Set the default build path.""" + self._default_build_path = ensure_directory( + os.path.join(self.get_keops_cache_folder(), self.get_default_build_folder_name()), + add_to_syspath=True, + ) + # Initialize _build_path : TODO : check ?/ + self._build_folder = self._default_build_path + + def get_default_build_path(self): + """Get the default build path.""" + return self._default_build_path + + def print_default_build_path(self): + """Print the default build path.""" + print(f"Default Build Path: {self.get_default_build_path() or not_found_str}") + + # JIT Binary Path + def set_jit_binary(self): + """Set the path to the JIT binary.""" + self._jit_binary = self.get_default_build_path() + + def get_jit_binary(self): + """Get the path to the JIT binary.""" + return self._jit_binary + + def print_jit_binary(self): + """Print the path to the JIT binary.""" + print(f"JIT Binary Path: {self.get_jit_binary() or not_found_str}") + + # include options management + def set_include_options(self): + """Set the include options for compilation.""" + self._include_options = f" -I{self.get_base_dir_path()}" + + def get_include_options(self): + """Get the include options for compilation.""" + return self._include_options + + # Comprehensive Path Information + def print_all(self): + """Print all path-related information.""" + + # Base Configuration + print("=" * 60) + print(f"KeOps Paths Configuration") + print("=" * 60) + + self.print_base_dir_path() + self.print_keops_cache_folder() + self.print_default_build_folder_name() + self.print_jit_binary() + + # Print relevant environment variables. + print_envs(self.path_env_vars) + + +if __name__ == "__main__": + from keopscore.config.Platform import PlatformConfig + from keopscore.config.CxxCompiler import CxxCompilerConfig + from keopscore.config.OpenMP import OpenMPConfig + from keopscore.config.Cuda import CudaConfig + + platform_info = PlatformConfig() + platform_info.print_all() + + cxx_compiler_info = CxxCompilerConfig(platform_info) + cxx_compiler_info.print_all() + + openmp_info = OpenMPConfig(platform_info, cxx_compiler_info) + openmp_info.print_all() + + cuda_info = CudaConfig() + cuda_info.print_all() + + keops_info = KeOpsPathConfig(platform_info, cuda_info) + keops_info.print_all() \ No newline at end of file diff --git a/keopscore/keopscore/config/openmp.py b/keopscore/keopscore/config/OpenMP.py similarity index 58% rename from keopscore/keopscore/config/openmp.py rename to keopscore/keopscore/config/OpenMP.py index 69492decf..9c13e4b9d 100644 --- a/keopscore/keopscore/config/openmp.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -2,21 +2,17 @@ import subprocess import tempfile -from keopscore.config.CppConfig import CppConfig -from keopscore.config._shared import ( - _find_library_by_names, - _first_existing_dir_with_files, - _first_existing_file, +from keopscore.config._shared import print_envs, enabled_dict, not_found_str +from keopscore.utils.messages import KeOps_Warning +from keopscore.utils.path_utils import ( + _first_matching_file, _ordered_search_roots, _path_candidates, - print_envs, - enabled_dict, - not_found_str, ) -from keopscore.utils.misc_utils import KeOps_Warning, get_brew_prefix +from keopscore.utils.system_utils import _find_library_by_names, get_include_file_abspath -class OpenMPConfig(CppConfig): +class OpenMPConfig: """ Class for OpenMP detection and configuration. """ @@ -24,7 +20,13 @@ class OpenMPConfig(CppConfig): _use_OpenMP = None _openmp_lib_name = None _openmp_lib_include_dir = None - _openmp_header_name = "omp.h" + + _openmp_compile_options = "" + _openmp_include_options = "" + _openmp_linking_options = "" + + _openmp_header_basename = "omp.h" + openmp_env_vars = ( "OMP_PATH", "LIBOMP_PATH", @@ -32,24 +34,56 @@ class OpenMPConfig(CppConfig): "OpenMP_ROOT_DIR", ) - openmp_dir_prefixes = ( - get_brew_prefix(), + openmp_system_suffixes = [ + # self.get_brew_prefix() added later on, os.path.join(os.path.sep, "usr", "local", "opt", "libomp"), os.path.join(os.path.sep, "opt", "local"), + os.path.join(os.path.sep, "usr"), + ] + + openmp_basename_candidate = ( + "libomp.dylib", + "libgomp.dylib", + "libomp.so", + ) + + openmp_library_suffixes = ( + "lib", + "lib64", + os.path.join("opt", "libomp", "lib"), ) - def __init__(self): - super().__init__() + _openmp_include_sufixes = ( + "include", + os.path.join("opt", "libomp", "include"), + ) + + def __init__(self, platform, cxx_compiler): + + self.platform = platform + self.cxx_compiler = cxx_compiler + + self.openmp_system_suffixes += [self.platform.get_brew_prefix(),] if self.platform.get_brew_prefix() else [] self.set_openmplib_path() + + self.set_openmp_compile_options() + self.set_openmp_include_options() + self.set_openmp_linking_options() + self.set_use_OpenMP() + if self.get_use_OpenMP(): + self.cxx_compiler.add_to_compile_option(self.get_openmp_compile_options()) + self.cxx_compiler.add_to_include_options(self.get_openmp_include_options()) + self.cxx_compiler.add_to_linking_options(self.get_openmp_linking_options()) + # OpenMP library path def set_openmplib_path(self): """try to locate OpenMP libraries""" openmp_install = self.find_openmp_install() openmp_lib = openmp_install["library"] if openmp_lib: - self._openmp_lib_include_dir = openmp_install["include_dir"] + self._openmp_lib_include_dir = os.path.dirname(openmp_install["header"]) self._openmp_lib_lib_dir = os.path.dirname(openmp_lib) self._openmp_lib_name = openmp_lib else: @@ -74,46 +108,31 @@ def find_openmp_install(self): """ Locate OpenMP runtime files without assuming a specific package manager. - Returns a dict with optional ``library`` and ``include_dir`` entries. + Returns a dict with optional ``library`` and ``header`` entries. """ - result = {"library": None, "include_dir": None} + result = {"library": None, "header": None} # First try to find OpenMP library using standard names via ctypes. result["library"] = _find_library_by_names(("gomp", "omp")) - p = subprocess.run( - [ - self.get_cxx_compiler(), - f"-print-file-name=include/{self._openmp_header_name}", - ], - capture_output=True, - text=True, - ).stdout.strip() - result["include_dir"] = p if os.path.isabs(p) and os.path.exists(p) else None + result["header"] = get_include_file_abspath(self._openmp_header_basename, self.cxx_compiler.get_cxx_compiler()) # If that fails, search for OpenMP headers and libraries in common locations. if not result["library"]: candidate_roots = _ordered_search_roots( env_vars=self.openmp_env_vars, conda_root="CONDA_PREFIX", - system_roots=self.openmp_dir_prefixes, + system_roots=self.openmp_system_suffixes, ) - result["library"] = _first_existing_file( - _path_candidates( - candidate_roots, - ( - "lib/libomp.dylib", - "lib/libgomp.dylib", - "lib/libomp.so", - "opt/libomp/lib/libomp.dylib", - ), - ) + result["library"] = _first_matching_file( + _path_candidates(candidate_roots, self.openmp_library_suffixes), + self.openmp_basename_candidate, ) # Finally, search for OpenMP headers in common locations. - result["include_dir"] = _first_existing_dir_with_files( - _path_candidates(candidate_roots, ("include", "opt/libomp/include")), - (self._openmp_header_name,), + result["header"] = _first_matching_file( + _path_candidates(candidate_roots, self._openmp_include_sufixes), + (self._openmp_header_basename,), ) return result @@ -168,7 +187,7 @@ def print_use_OpenMP(self): def check_compiler_for_openmp(self): """Attempt to compile a simple OpenMP program to check if the compiler supports OpenMP.""" - if not self.get_cxx_compiler(): + if not self.cxx_compiler.get_cxx_compiler(): KeOps_Warning("No C++ compiler available to check for OpenMP support.") return False @@ -185,12 +204,13 @@ def check_compiler_for_openmp(self): test_file = f.name compile_command = [ - self.get_cxx_compiler(), + self.cxx_compiler.get_cxx_compiler(), test_file, - "-fopenmp", + self.get_openmp_include_options(), + self.get_openmp_compile_options(), ] - if self.get_openmp_lib_dir(): - compile_command.append(f"-L{self.get_openmp_lib_dir()}") + if self.get_openmp_linking_options(): + compile_command.append(self.get_openmp_linking_options()) compile_command.extend(["-o", f"{test_file}.out"]) try: @@ -203,8 +223,33 @@ def check_compiler_for_openmp(self): os.remove(test_file) return False + # C++ Compiler Options + def set_openmp_compile_options(self): + # Add special fix for openMP prgama and Apple Clang. Order matters. + if self.cxx_compiler.get_use_Apple_clang() and self.platform.get_brew_prefix(): + self._openmp_compile_options += "-Xpreprocessor " + + self._openmp_compile_options += "-fopenmp" + + def get_openmp_compile_options(self): + return self._openmp_compile_options + + # C++ Compiler Include Options + def set_openmp_include_options(self): + self._openmp_include_options = f'-I{self.get_openmp_include_dir()}' + + def get_openmp_include_options(self): + return self._openmp_include_options + + # C++ linking Options + def set_openmp_linking_options(self): + self._openmp_linking_options += f'-L{self.get_openmp_lib_dir()}' + + def get_openmp_linking_options(self): + return self._openmp_linking_options + # OpenMP configuration printing - def print_openmp(self): + def print_all(self): """ Print all OpenMP-related configuration and system health status. """ @@ -222,7 +267,14 @@ def print_openmp(self): if __name__ == "__main__": - openmp_config = OpenMPConfig() - openmp_config.print_platform() - openmp_config.print_cpp() - openmp_config.print_openmp() + from keopscore.config.Platform import PlatformConfig + from keopscore.config.CxxCompiler import CxxCompilerConfig + + platform_info = PlatformConfig() + platform_info.print_all() + + cxx_compiler_info = CxxCompilerConfig(platform_info) + cxx_compiler_info.print_all() + + openmp_info = OpenMPConfig(platform_info, cxx_compiler_info) + openmp_info.print_all() diff --git a/keopscore/keopscore/config/Platform.py b/keopscore/keopscore/config/Platform.py index 35687a4ce..a9a63bb04 100644 --- a/keopscore/keopscore/config/Platform.py +++ b/keopscore/keopscore/config/Platform.py @@ -2,13 +2,15 @@ import platform import sys -from keopscore.config._shared import print_envs +from ._shared import print_envs +from keopscore.utils.system_utils import KeOps_OS_Run -class Platform: +class PlatformConfig: """ Class for detecting the operating system, Python version, and environment type. """ + _os = None _platform = None _machine = None @@ -16,14 +18,16 @@ class Platform: _python_version = None _python_executable = None _env_type = None + _brew_prefix = None + platform_envs = [ - "PYTHONPATH", - "PATH", - "VIRTUAL_ENV", - "CONDA_DEFAULT_ENV", - "CONDA_PREFIX", - ] - + "PYTHONPATH", + "PATH", + "VIRTUAL_ENV", + "CONDA_DEFAULT_ENV", + "CONDA_PREFIX", + ] + def __init__(self): self.set_os() self.set_platform() @@ -56,33 +60,33 @@ def detect_os(): return f"{platform.system()} {name} {version}" except FileNotFoundError: return "Linux (distribution info not found)" - + return platform.system() + " " + platform.version() # Platform detection (Darwin, Windows, Linux, etc.) def set_platform(self): self._platform = platform.system() - + def get_platform(self): return self._platform # Machine architecture detection (x86_64, arm64, etc.) def set_machine(self): self._machine = platform.machine() - + def get_machine(self): return self._machine - + def print_machine(self): print(f"Machine Architecture: {self.get_machine()}") # uname detection def set_uname(self): self._uname = platform.uname() - + def get_uname(self): return self._uname - + # Python Version Detection def set_python_version(self): """Set the Python version.""" @@ -98,13 +102,13 @@ def print_python_version(self): def set_python_executable(self): """Set the Python executable path.""" self._python_executable = sys.executable - + def get_python_executable(self): return self._python_executable - + def print_python_executable(self): print(f"Python Executable: {self.get_python_executable()}") - + # Environment Type Detection def set_env_type(self): """Set the environment type (conda, virtualenv, or system).""" @@ -115,10 +119,23 @@ def get_env_type(self): def print_env_type(self): print(f"Environment Type: {self.get_env_type()} {sys.prefix}", end="") - if self.get_env_type() != 'system': - print(f' (base at {sys.base_prefix})', end="") + if self.get_env_type() != "system": + print(f" (base at {sys.base_prefix})", end="") print() + # Brew package system + def set_brew_prefix(self): + """Get Homebrew prefix path using KeOps_OS_Run""" + if self.get_platform() != "Darwin": + return + + out = KeOps_OS_Run(f"brew --prefix", print_warning=False) + self._brew_prefix = out.stdout.decode("utf-8").strip() if out.stderr != b"" else None + + def get_brew_prefix(self): + """Get Homebrew prefix path using KeOps_OS_Run""" + return self._brew_prefix + @staticmethod def detect_env_type(): """Return whether Python runs in conda, virtualenv, or the system env.""" @@ -129,9 +146,9 @@ def detect_env_type(): ): return "virtualenv" return "system" - + # Comprehensive Platform Information - def print_platform(self): + def print_all(self): """ Print all platform-related information. """ @@ -147,8 +164,8 @@ def print_platform(self): # Print relevant environment variables. print_envs(self.platform_envs) - + if __name__ == "__main__": - platform_info = Platform() - platform_info.print_platform() \ No newline at end of file + platform_info = PlatformConfig() + platform_info.print_all() diff --git a/keopscore/keopscore/config/__init__.py b/keopscore/keopscore/config/__init__.py index fc31a0b9f..ea0efe90f 100644 --- a/keopscore/keopscore/config/__init__.py +++ b/keopscore/keopscore/config/__init__.py @@ -1,47 +1,76 @@ +import os +import shutil + # Import the configuration classes -from .base_config import Config -from .cuda import CUDAConfig -from .openmp import OpenMPConfig -from .Platform import DetectPlatform +from .Cuda import CudaConfig +from .CxxCompiler import CxxCompilerConfig +from .Debug import DebugConfig +from .KeOpsPath import KeOpsPathConfig +from .OpenMP import OpenMPConfig +from .Platform import PlatformConfig +from keopscore.utils.messages import KeOps_Message -# Instantiate the configurations -config = Config() -platform_detector = DetectPlatform() -cuda_config = CUDAConfig() -openmp_config = OpenMPConfig() +# Instantiate the configurations once at import time to preserve the existing API. +debug = DebugConfig() +platform = PlatformConfig() +cxx = CxxCompilerConfig(platform) +openmp = OpenMPConfig(platform, cxx) +cuda = CudaConfig() +path = KeOpsPathConfig(platform, cuda) -__all__ = [ - "config", - "platform_detector", - "cuda_config", - "openmp_config", - "get_config", - "get_platform_config", - "get_cuda_config", - "get_openmp_config", -] +# flag for automatic factorization : apply automatic factorization for all formulas before reduction. +auto_factorize = False -# Lazy initializers -_instances = {} +def check_health(infos="all"): + """ + Check the health of the specified configuration. + """ -def get_instance(key, factory): - if key not in _instances: - _instances[key] = factory - return _instances[key] + if infos == "all" or infos == "platform": + platform.print_all() + if infos == "all" or infos == "cxx": + cxx.print_all() + if infos == "all" or infos == "openmp": + openmp.print_all() + if infos == "all" or infos == "cuda": + cuda.print_all() + if infos == "all" or infos == "path": + path.print_all() -def get_config(): - return get_instance("config", config) +def clean_keops(recompile_jit_binary=True, verbose=True): + build_path = path.get_build_folder() + use_cuda = cuda.get_use_cuda() + jit_binary = path.get_jit_binary() if use_cuda else None + if build_path and os.path.isdir(build_path): + for entry in os.scandir(build_path): + if recompile_jit_binary or entry.path != jit_binary: + if entry.is_dir(follow_symlinks=False): + shutil.rmtree(entry.path) + else: + os.remove(entry.path) -def get_cuda_config(): - return get_instance("cuda_config", cuda_config) + if verbose: + KeOps_Message(f"{build_path} has been cleaned.") + from keopscore.get_keops_dll import get_keops_dll -def get_openmp_config(): - return get_instance("openmp_config", openmp_config) + get_keops_dll.reset() + if use_cuda and recompile_jit_binary: + from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile + Gpu_link_compile.compile_jit_compile_dll() -def get_platform_config(): - return get_instance("platform_detector", platform_config) +__all__ = [ + "platform", + "cxx", + "openmp", + "cuda", + "path", + "debug", + "auto_factorize", + "check_health", + "clean_keops", +] diff --git a/keopscore/keopscore/config/_shared.py b/keopscore/keopscore/config/_shared.py index 8283040ca..6476d52d3 100644 --- a/keopscore/keopscore/config/_shared.py +++ b/keopscore/keopscore/config/_shared.py @@ -1,139 +1,21 @@ import os -import site import sys -import sysconfig -from pathlib import Path -from keopscore.utils.misc_utils import CHECK_MARK, CROSS_MARK, find_library_abspath +from keopscore.utils.path_utils import ( + _env_roots, + _ordered_search_roots, + _path_candidates, + _python_package_roots, + _unique_paths, +) +from keopscore.utils.system_utils import _find_library_by_names + +CHECK_MARK = "✅" +CROSS_MARK = "❌" not_found_str = f"Not Found. {CROSS_MARK}" enabled_dict = {True: f"Enabled {CHECK_MARK}", False: f"Disabled {CROSS_MARK}"} -def _unique_paths(paths): - """Return paths in first-seen order with duplicates and None values removed.""" - unique = [] - seen = set() - for path in paths: - if path is None: - continue - path = Path(path) - key = str(path) - if key not in seen: - seen.add(key) - unique.append(path) - return unique - - -def _env_roots(env_vars): - """Return existing environment variable values as normalized Path objects.""" - return _unique_paths(os.getenv(env_var) for env_var in env_vars) - - -def _path_candidates(roots, suffixes): - """Expand root directories with relative suffixes while preserving order.""" - candidates = [] - for root in _unique_paths(roots): - for suffix in suffixes: - candidates.append(root / suffix if suffix else root) - return _unique_paths(candidates) - - -def _first_existing_file(paths): - """Return the first path that points to an existing file.""" - for path in _unique_paths(paths): - if path.is_file(): - return str(path) - return None - - -def _first_matching_file(directories, patterns): - """Return the first file matching glob patterns inside ordered directories.""" - for directory in _unique_paths(directories): - if not directory.is_dir(): - continue - for pattern in patterns: - matches = sorted(directory.glob(pattern)) - for match in matches: - if match.is_file(): - return str(match) - return None - - -def _first_existing_dir_with_files(directories, required_filenames): - """Return the first directory containing all required filenames.""" - for directory in _unique_paths(directories): - if all((directory / filename).is_file() for filename in required_filenames): - return str(directory) - return None - - -def _find_library_by_names(library_names): - """Return the first library path resolved by ctypes for known library names.""" - from ctypes.util import find_library - - # ctypes delegates to the platform loader, so this is the final fallback. - for library_name in library_names: - library_path = find_library(library_name) - if library_path: - return find_library_abspath(library_name) - """ - library_path = find_library(library_name) - - if library_path: - full_path = KeOps_OS_Run("ldconfig -p").stdout.decode("utf-8") - - #print(f"Checking for OpenMP library '{full_path}' in ldconfig output...") - for line in full_path.splitlines(): - if library_path in line: - library_full_path = line.split("=>", 1)[1].strip() - return library_path, os.path.dirname(library_full_path) - """ - - return None - - -def _python_package_roots(): - """Return Python package roots where pip-installed wheels may live.""" - package_roots = [] - getters = [ - lambda: getattr(site, "getsitepackages", lambda: [])(), - lambda: [site.getusersitepackages()], - lambda: [sysconfig.get_path("purelib")], - lambda: [sysconfig.get_path("platlib")], - lambda: sys.path, - ] - for getter in getters: - try: - # Some site helpers are unavailable in embedded or non-standard Python builds. - package_roots.extend(getter() or []) - except Exception: - continue - return _unique_paths(package_roots) - - -def _ordered_search_roots(env_vars=(), pip_suffixes=(), conda_root=None, system_roots=()): - """Return roots ordered by explicit env vars, pip, conda, then system paths. - - :argument - - env_vars (list): environment variable names to search for paths - - pip_suffixes (list): suffixes to append to pip-installed wheel paths - - conda_root (str): environment variable name for conda root - - system_roots (list): explicit system paths to search - """ - roots = [] - - if env_vars: - roots.extend(_env_roots(env_vars)) - - if pip_suffixes: - # Pip wheels such as nvidia-cuda-runtime expose libraries under site-packages. - roots.extend(_path_candidates(_python_package_roots(), pip_suffixes)) - - if conda_root: - roots.extend(_env_roots((conda_root,))) - roots.extend(system_roots) - return _unique_paths(roots) - def ensure_directory(path, add_to_syspath=False): """Create a directory and optionally register it in sys.path.""" @@ -145,6 +27,10 @@ def ensure_directory(path, add_to_syspath=False): def print_envs(env_vars): """Print the values of specified environment variables.""" + + if not env_vars: + return + print("\nRelevant Environment Variables") print("-" * 60) for var in env_vars: @@ -152,4 +38,4 @@ def print_envs(env_vars): if value: print(f"{var} = {value}") else: - print(f"{var} is not set") \ No newline at end of file + print(f"{var} is not set") diff --git a/keopscore/keopscore/config/base_config.py b/keopscore/keopscore/config/base_config.py deleted file mode 100644 index 3519bc5d2..000000000 --- a/keopscore/keopscore/config/base_config.py +++ /dev/null @@ -1,440 +0,0 @@ -import os -from os.path import join -import platform -import sys -import shutil -from pathlib import Path -import keopscore -from keopscore.utils.misc_utils import KeOps_Warning, KeOps_OS_Run, get_brew_prefix -from keopscore.utils.misc_utils import CHECK_MARK, CROSS_MARK - - -class Config: - """ - Base configuration class for the KeOps library. - This class contains common attributes and methods shared by other configuration classes. - """ - - # Common attributes - base_dir_path = None - bindings_source_dir = None - keops_cache_folder = None - default_build_folder_name = None - default_build_path = None - jit_binary = None - cxx_compiler = None - cpp_env_flags = None - compile_options = None - cpp_flags = None - disable_pragma_unrolls = None - os = None - _build_folder = None - - def __init__(self): - - # Initialize common configuration settings - self.set_base_dir_path() - self.set_bindings_source_dir() - self.set_keops_cache_folder() - self.set_default_build_folder_name() - self.set_default_build_path() - self.set_jit_binary() - self.set_cxx_compiler() - self.set_cpp_env_flags() - self.set_compile_options() - self.set_cpp_flags() - self.set_disable_pragma_unrolls() - self.set_os() - - # Setters, getters, and print methods for common attributes - - def set_base_dir_path(self): - """Set the base directory path.""" - self.base_dir_path = os.path.abspath( - join(os.path.dirname(os.path.realpath(__file__)), "..") - ) - - def get_base_dir_path(self): - """Get the base directory path.""" - return self.base_dir_path - - def print_base_dir_path(self): - """Print the base directory path.""" - print(f"Base Directory Path: {self.base_dir_path}") - - def set_bindings_source_dir(self): - """Set the bindings source directory.""" - self.bindings_source_dir = self.get_base_dir_path() - - def get_bindings_source_dir(self): - """Get the bindings source directory.""" - return self.bindings_source_dir - - def print_bindings_source_dir(self): - """Print the bindings source directory.""" - print(f"Bindings Source Directory: {self.bindings_source_dir}") - - def set_keops_cache_folder(self): - """Set the KeOps cache folder.""" - self.keops_cache_folder = os.getenv("KEOPS_CACHE_FOLDER") - if self.keops_cache_folder is None: - self.keops_cache_folder = join( - os.path.expanduser("~"), ".cache", f"keops{keopscore.__version__}" - ) - # Ensure the cache folder exists - os.makedirs(self.keops_cache_folder, exist_ok=True) - - def get_keops_cache_folder(self): - """Get the KeOps cache folder.""" - return self.keops_cache_folder - - def print_keops_cache_folder(self): - """Print the KeOps cache folder.""" - print(f"KeOps Cache Folder: {self.keops_cache_folder}") - - def set_default_build_folder_name(self): - """Set the default build folder name.""" - uname = platform.uname() - self.default_build_folder_name = ( - "_".join(uname[:3]) + f"_p{sys.version.split(' ')[0]}" - ) - - def get_default_build_folder_name(self): - """Get the default build folder name.""" - return self.default_build_folder_name - - def print_default_build_folder_name(self): - """Print the default build folder name.""" - print(f"Default Build Folder Name: {self.default_build_folder_name}") - - def set_default_build_path(self): - """Set the default build path.""" - self.default_build_path = join( - self.keops_cache_folder, self.default_build_folder_name - ) - # Ensure the build path exists - os.makedirs(self.default_build_path, exist_ok=True) - # Add the build path to sys.path - if self.default_build_path not in sys.path: - sys.path.append(self.default_build_path) - # Initialize _build_path - self._build_folder = self.default_build_path - - def get_default_build_path(self): - """Get the default build path.""" - return self.default_build_path - - def print_default_build_path(self): - """Print the default build path.""" - print(f"Default Build Path: {self.default_build_path}") - - def set_jit_binary(self): - """Set the path to the JIT binary.""" - self.jit_binary = join(self.default_build_path) - - def get_jit_binary(self): - """Get the path to the JIT binary.""" - return self.jit_binary - - def print_jit_binary(self): - """Print the path to the JIT binary.""" - print(f"JIT Binary Path: {self.jit_binary}") - - def set_os(self): - """Set the operating system.""" - if platform.system() == "Linux": - try: - with open("/etc/os-release") as f: - info = dict(line.strip().split("=", 1) for line in f if "=" in line) - name = info.get("NAME", "Linux").strip('"') - version = info.get("VERSION_ID", "").strip('"') - self.os = f"{platform.system()} {name} {version}" - except FileNotFoundError: - self.os = "Linux (distribution info not found)" - else: - self.os = platform.system() + " " + platform.version() - - def get_os(self): - return self.os - - def print_os(self): - print(f"Operating System: {self.os}") - - def set_cxx_compiler(self): - """Set the C++ compiler.""" - env_cxx = os.getenv("CXX") - if env_cxx and shutil.which(env_cxx): - self.cxx_compiler = env_cxx - else: - # On macOS, try clang++ first, as it's the canonical C++ compiler driver for Apple Clang - if platform.system() == "Darwin": - if shutil.which("clang++"): - self.cxx_compiler = "clang++" - elif shutil.which("g++"): - self.cxx_compiler = "g++" - else: - self.cxx_compiler = None - KeOps_Warning( - "No suitable C++ compiler (clang++ or g++) found on macOS." - ) - else: - # On Linux or other systems, fall back to g++ if available - if shutil.which("g++"): - self.cxx_compiler = "g++" - else: - self.cxx_compiler = None - KeOps_Warning( - "No C++ compiler found. Define CXX environment variable or install g++." - ) - - def get_cxx_compiler(self): - """Get the C++ compiler.""" - return self.cxx_compiler - - def print_cxx_compiler(self): - """Print the C++ compiler.""" - print(f"C++ Compiler: {self.cxx_compiler}") - - def set_cpp_env_flags(self): - """Set the C++ environment flags.""" - self.cpp_env_flags = os.getenv("CXXFLAGS") if "CXXFLAGS" in os.environ else "" - - def get_cpp_env_flags(self): - """Get the C++ environment flags.""" - return self.cpp_env_flags - - def print_cpp_env_flags(self): - """Print the C++ environment flags.""" - print(f"C++ Environment Flags (CXXFLAGS): {self.cpp_env_flags}") - - def set_compile_options(self): - """Set the compile options.""" - self.compile_options = " -shared -fPIC -O3 -std=c++11" - - def get_compile_options(self): - """Get the compile options.""" - return self.compile_options - - def print_compile_options(self): - """Print the compile options.""" - print(f"Compile Options: {self.compile_options}") - - def get_use_Apple_clang(self): - """Detect if using Apple Clang.""" - is_apple_clang = False - if platform.system() == "Darwin": - compiler_info = KeOps_OS_Run(f"c++ --version").stdout.decode("utf-8") - # Check if 'Apple clang' appears in the output - is_apple_clang = "Apple clang" in compiler_info - return is_apple_clang - - def set_cpp_flags(self): - """Set the C++ compiler flags.""" - self.cpp_flags = f"{self.cpp_env_flags} {self.compile_options}" - self.cpp_flags += f" -I{self.base_dir_path}/include" - self.cpp_flags += f" -I{self.bindings_source_dir}" - - # Add OpenMP flags based on compiler - if self.get_use_Apple_clang(): - # For Apple Clang, you need to specify OpenMP library location - brew_prefix = get_brew_prefix() - if brew_prefix is not None: - self.cpp_flags += f" -Xpreprocessor -fopenmp" - self.cpp_flags += f" -I{brew_prefix}/opt/libomp/include" - self.cpp_flags += f" -L{brew_prefix}/opt/libomp/lib" - else: - # For GCC and other compilers - self.cpp_flags += " -fopenmp" - - # Specific check for Apple Silicon chips - if platform.system() == "Darwin" and platform.machine() in ["arm64", "arm64e"]: - self.cpp_flags += " -arch arm64" - - if platform.system() == "Darwin": - self.cpp_flags += " -undefined dynamic_lookup" - self.cpp_flags += " -flto" - - self.cpp_flags += " -flto=auto" - - def get_cpp_flags(self): - """Get the C++ compiler flags.""" - return self.cpp_flags - - def print_cpp_flags(self): - """Print the C++ compiler flags.""" - print(f"C++ Compiler Flags: {self.cpp_flags}") - - def set_disable_pragma_unrolls(self): - """Set the flag for disabling pragma unrolls.""" - self.disable_pragma_unrolls = True - - def get_disable_pragma_unrolls(self): - """Get the flag for disabling pragma unrolls.""" - return self.disable_pragma_unrolls - - def print_disable_pragma_unrolls(self): - """Print the flag for disabling pragma unrolls.""" - status = "Enabled" if self.disable_pragma_unrolls else "Disabled" - print(f"Disable Pragma Unrolls: {status}") - - def set_different_build_folder( - self, path=None, read_save_file=False, write_save_file=True, reset_all=True - ): - """ - Set or update the build folder path for KeOps. - - Parameters: - - path: The new build folder path. If None, it will be determined based on saved settings or defaults. - - read_save_file: If True, read the build folder path from a save file if path is not provided. - - write_save_file: If True, write the new build folder path to the save file. - - reset_all: If True, reset all cached formulas and recompile necessary components. - """ - # If path is not given, we either read the save file or use the default build path - save_file = join(self.keops_cache_folder, "build_folder_location.txt") - if not path: - if read_save_file and os.path.isfile(save_file): - with open(save_file, "r") as f: - path = f.read() - else: - path = self.default_build_path - - # Create the folder if not yet done - os.makedirs(path, exist_ok=True) - - # Remove the old build path from sys.path if it's there - if self._build_folder and self._build_folder in sys.path: - sys.path.remove(self._build_folder) - # Update _build_folder to the new path - self._build_folder = path - # Add the new build path to sys.path - if self._build_folder not in sys.path: - sys.path.append(self._build_folder) - - # Saving the location of the build path in a file - if write_save_file: - with open(save_file, "w") as f: - f.write(path) - - # Reset all cached formulas if needed - if reset_all: - # Reset cached formulas - keopscore.get_keops_dll.get_keops_dll.reset( - new_save_folder=self._build_folder - ) - # Handle CUDA-specific recompilation if CUDA is used - if hasattr(self, "_use_cuda") and self._use_cuda: - from keopscore.binders.nvrtc.Gpu_link_compile import ( - Gpu_link_compile, - jit_compile_dll, - ) - - if not os.path.exists(jit_compile_dll()): - Gpu_link_compile.compile_jit_compile_dll() - - def get_build_folder(self): - return self._build_folder - - # Environment variables printing method - def print_environment_variables(self): - """Print relevant environment variables.""" - print("\nRelevant Environment Variables:") - env_vars = [ - "KEOPS_CACHE_FOLDER", - "CXX", - "CXXFLAGS", - ] - for var in env_vars: - value = os.environ.get(var, None) - if value: - print(f"{var} = {value}") - else: - print(f"{var} is not set") - - def print_all(self): - """ - Print all base configuration - """ - - # Base Configuration - print(f"\nBase Configuration") - print("-" * 60) - - # Base Directory Path - base_dir_path = self.get_base_dir_path() - base_dir_status = ( - CHECK_MARK - if base_dir_path and os.path.exists(base_dir_path) - else CROSS_MARK - ) - print(f"Base Directory Path: {base_dir_path or 'Not Found'} {base_dir_status}") - - # Bindings Source Directory - bindings_source_dir = self.get_bindings_source_dir() - bindings_source_dir_status = ( - CHECK_MARK - if bindings_source_dir and os.path.exists(bindings_source_dir) - else CROSS_MARK - ) - print( - f"Bindings Source Directory: {bindings_source_dir or 'Not Found'} {bindings_source_dir_status}" - ) - - # KeOps Cache Folder - keops_cache_folder = self.get_keops_cache_folder() - keops_cache_folder_status = ( - CHECK_MARK - if keops_cache_folder and os.path.exists(keops_cache_folder) - else CROSS_MARK - ) - print( - f"KeOps Cache Folder: {keops_cache_folder or 'Not Found'} {keops_cache_folder_status}" - ) - - # Default Build Folder Name - default_build_folder_name = self.get_default_build_folder_name() - print(f"Default Build Folder Name: {default_build_folder_name}") - - # Default Build Path - default_build_path = self.get_default_build_path() - default_build_path_status = ( - CHECK_MARK - if default_build_path and os.path.exists(default_build_path) - else CROSS_MARK - ) - print( - f"Default Build Path: {default_build_path or 'Not Found'} {default_build_path_status}" - ) - - # JIT Binary Path - jit_binary = self.get_jit_binary() - jit_binary_status = ( - CHECK_MARK if jit_binary and os.path.exists(jit_binary) else CROSS_MARK - ) - print(f"JIT Binary Path: {jit_binary or 'Not Found'} {jit_binary_status}") - - # Disable Pragma Unrolls - disable_pragma_unrolls = self.get_disable_pragma_unrolls() - disable_status = CHECK_MARK if disable_pragma_unrolls else CROSS_MARK - status_text = "Enabled" if disable_pragma_unrolls else "Disabled" - print(f"Disable Pragma Unrolls: {status_text} {disable_status}") - - # Compile Options - compile_options = self.get_compile_options() - print(f"Compile Options: {compile_options}") - - # C++ Compiler Flags - cpp_flags = self.get_cpp_flags() - print(f"C++ Compiler Flags: {cpp_flags}") - - # Print relevant environment variables. - print("\nRelevant Environment Variables:") - env_vars = [ - "KEOPS_CACHE_FOLDER", - "CXXFLAGS", - ] - for var in env_vars: - value = os.environ.get(var, None) - if value: - print(f"{var} = {value}") - else: - print(f"{var} is not set") diff --git a/keopscore/keopscore/config/chunks.py b/keopscore/keopscore/config/chunks.py index 4efb79c97..531e52676 100644 --- a/keopscore/keopscore/config/chunks.py +++ b/keopscore/keopscore/config/chunks.py @@ -1,7 +1,13 @@ # special computation scheme for dim>100 -enable_chunk = True +dimchunk = 64 +dim_treshold_chunk = 146 +specdims_use_chunk = [99, 100, 102, 120, 133, 138, 139, 140, 141, 142] + +dimfinalchunk = 64 + +enable_chunk = True def get_enable_chunk(): global enable_chunk @@ -16,10 +22,6 @@ def set_enable_chunk(val): enable_chunk = False -dimchunk = 64 -dim_treshold_chunk = 146 -specdims_use_chunk = [99, 100, 102, 120, 133, 138, 139, 140, 141, 142] - # special mode for formula of the type sum_j k(x_i,y_j)*b_j with high dimensional b_j enable_final_chunk = True @@ -32,8 +34,6 @@ def set_enable_finalchunk(val): enable_final_chunk = False -dimfinalchunk = 64 - def get_dimfinalchunk(): global dimfinalchunk diff --git a/keopscore/keopscore/config/cuda.py b/keopscore/keopscore/config/cuda.py deleted file mode 100644 index 4252ad78d..000000000 --- a/keopscore/keopscore/config/cuda.py +++ /dev/null @@ -1,579 +0,0 @@ -import ctypes -import os -from ctypes import ( - DEFAULT_MODE, - c_int, - CDLL, - byref, - RTLD_GLOBAL, -) - -from keopscore.config.CppConfig import CppConfig -from keopscore.config._shared import ( - _find_library_by_names, - _first_existing_dir_with_files, - _first_matching_file, - _ordered_search_roots, - _path_candidates, - print_envs, - not_found_str, -) -from keopscore.utils.misc_utils import ( - CHECK_MARK, - CROSS_MARK, -) -from keopscore.utils.misc_utils import KeOps_Warning - - -class CUDAConfig(CppConfig): - """ - Class for CUDA detection and configuration. - """ - - # CUDA constants - CUDA_SUCCESS = 0 - CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK = 1 - CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK = 8 - - # Cuda detection variables - _use_cuda = None - _specific_gpus = None - - _cuda_include_path = None - _nvrtc_flags = None - _cuda_version = None - n_gpus = 0 - _gpu_compile_flags = "" - cuda_message = "" - cuda_block_size = None - cuda_install_info = None - - # ------------------------ # - # Search location # - # ------------------------ # - - cuda_env_vars = [ - "CUDA_VISIBLE_DEVICES", - "CUDA_PATH", - "CUDA_HOME", - "CUDA_ROOT", - "CUDA_TOOLKIT_ROOT_DIR", - ] - - pip_suffixes = ( - "nvidia/cuda_runtime", - "nvidia/cuda_nvrtc", - ) - - system_suffixes = ( - os.path.join(os.path.sep, "usr", "local", "cuda"), - os.path.join(os.path.sep, "usr", "local"), - os.path.join(os.path.sep, "opt", "cuda"), - os.path.join(os.path.sep, "usr"), - os.path.join(os.path.sep, "lib"), - ) - - library_suffixes = ( - "lib64", - "lib", - "lib/x86_64-linux-gnu/", - ) - - include_suffixes = ( - "include", - "targets/x86_64-linux/include", - "targets/sbsa-linux/include", - "targets/aarch64-linux/include", - ) - - # ------------------------- # - # Library info dicts # - # ------------------------- # - - _libcuda_info = { - "name": "cuda", - "lib_file_name_candidate": ["libcuda.so.*", "libcuda.dylib", "cuda.lib"], - "header_file_name": "cuda.h", - "library": None, # to be filled later - "include_dir": None, # to be filled later - } - _libnvrtc_info = { - "name": "nvrtc", - "lib_file_name_candidate": ["libnvrtc.so.*", "libnvrtc.dylib", "nvrtc.lib"], - "header_file_name": "nvrtc.h", - "library": None, # to be filled later - "include_dir": None, # to be filled later - } - _cudart_info = { - "name": "cudart", - "lib_file_name_candidate": ["libcudart.so.*", "libcudart.dylib", "cudart.lib"], - "header_file_name": None, - "library": None, # to be filled later - "include_dir": None, # not needed - } - - def __init__(self): - - self.set_specific_gpus() - - super().__init__() - - self.set_use_cuda() - - # If cuda is enabled, then we finalize the rest of the config - if self.get_use_cuda(): - self.get_cuda_version() - self.get_cuda_include_path() - self.set_nvrtc_flags() - self.set_cuda_block_size() - - def find_cuda_install(self, lib_dict_info, warn=None): - """ - Locate a cuda and headers using an explicit ordered search. - - Arguments: - lib_dict_info (dict): A dictionary containing at least the keys 'name', 'lib_file_name_candidate', and 'header_file_name' for the library to find. This allows the function to be used for finding libcuda, libcudart, or nvrtc by passing the appropriate info dict. - - Returns: - result (dict): a copy of lib_dict_info completed with the ``library`` and ``include_dir`` keys containing the absolute paths to the library file and include directory, or None if not found. - """ - result = ( - lib_dict_info.copy() - ) # Start with the provided info, which may contain names and file patterns - - candidate_roots = _ordered_search_roots( - env_vars=self.cuda_env_vars, - pip_suffixes=self.pip_suffixes, - conda_root="CONDA_PREFIX", - system_roots=self.system_suffixes, - ) - - # ------------------------ # - # Search for library file # - # ------------------------ # - - # First try to find the library file using the candidate roots and library suffixes - candidate_library_dirs = _path_candidates(candidate_roots, self.library_suffixes) - result["library"] = _first_matching_file( - candidate_library_dirs, result["lib_file_name_candidate"] - ) - if result["library"] is None: - result["library"] = _find_library_by_names((result["name"],)) - - if result["library"] is None and warn: - KeOps_Warning(f"lib{result['name']} not found.") - - # ------------------------ # - # Search for header files # - # ------------------------ # - - result["include_dir"] = _first_existing_dir_with_files( - _path_candidates(candidate_roots, self.include_suffixes), - (result["header_file_name"],), - ) - - if ( - result["include_dir"] is None - and warn - and result["header_file_name"] is not None - ): - KeOps_Warning(f"{result['name']} header files not found.") - - return result - - def _try_load_library(self, lib_full_path_path, mode=DEFAULT_MODE): - """ - Attempt to load a shared library. - Returns: - success (bool): True if the library was found and loaded. - error_msg (str): Contains error details if success==False, else "". - """ - if not lib_full_path_path: - return False, "Library path not found" - try: - CDLL(lib_full_path_path, mode=mode) - except OSError as e: - return ( - False, - f"Failed to load library '{os.path.basename(lib_full_path_path)}': {e}", - ) - - return True, "" - - def _try_find_cuda_version(self, libcudart_full_path_path): - """ - Attempt to find the CUDA version by loading the CUDA runtime library and querying its version. - Returns: - success (bool): True if the version was successfully determined, False otherwise. - error_msg (str): Contains error details if success==False, else "". - - """ - try: - libcudart = ctypes.CDLL(libcudart_full_path_path) - cuda_version = ctypes.c_int() - libcudart.cudaRuntimeGetVersion(ctypes.byref(cuda_version)) - self._cuda_version = int(cuda_version.value) - return True, "" - - except OSError as e: - self._cuda_version = None - return ( - False, - f"Failed to load '{os.path.basename(libcudart_full_path_path)}': {e}", - ) - - def _cuda_libraries_available(self): - """ - Check if libcuda (driver); libcudart (cudatoolkit) and nvrtc (cudatoolkit) libraries are available. - This is where ```_cuda_include_path``` is set. - - Returns: - True if all three are loadable, False otherwise. - This is also where we handle one single warning if needed. - """ - - # Libcuda (driver) loaded globally so they are available to KeOps shared objects. - self._libcuda_info = self.find_cuda_install(self._libcuda_info, warn=True) - success_cuda = False - err_cuda = "libcuda not found" - if self._libcuda_info["library"] is not None: - success_cuda, err_cuda = self._try_load_library( - self._libcuda_info["library"], mode=RTLD_GLOBAL - ) - if not success_cuda: - KeOps_Warning(f"{err_cuda}. Switching to CPU only.") - return False - - # libnvrtc as well, since it's required for the runtime compilation of CUDA code. - self._libnvrtc_info = self.find_cuda_install(self._libnvrtc_info, warn=True) - success_nvrtc = False - err_nvrtc = "libnvrtc not found" - if self._libnvrtc_info["library"] is not None: - success_nvrtc, err_nvrtc = self._try_load_library( - self._libnvrtc_info["library"], mode=RTLD_GLOBAL - ) - if not success_nvrtc: - KeOps_Warning(f"{err_nvrtc}. Switching to CPU only.") - return False - - # Populate the cuda_install_info which is needed for include path and cuda version detection - self._cuda_include_path = list(set([ - self._libnvrtc_info["include_dir"], - self._libcuda_info["include_dir"], - ])) - - # Finally check cudart - self._cudart_info = self.find_cuda_install(self._cudart_info, warn=False) - success_cudart = False - err_cudart = "libcudart not found" - if self._cudart_info["library"] is not None: - success_cudart, err_cudart = self._try_find_cuda_version( - self._cudart_info["library"] - ) - if not success_cudart: - KeOps_Warning(f"{err_cudart}. Switching to CPU only.") - return False - - return True - - # CUDA Support - def set_use_cuda(self): - """Determine and set whether to use CUDA.""" - self._use_cuda = self._cuda_libraries_available() - - self.get_gpu_props() - if self.n_gpus == 0 and self._use_cuda: - self._use_cuda = False - self.cuda_message = "CUDA libraries detected, but no GPUs found on this system; Switching to CPU only." - KeOps_Warning(self.cuda_message) - - def get_use_cuda(self): - return self._use_cuda - - def print_use_cuda(self): - status = f"Enabled {CHECK_MARK}" if self._use_cuda else f"Disabled {CROSS_MARK}" - print(f"CUDA Support: {status}") - - # CUDA Block Size - def set_cuda_block_size(self, cuda_block_size=192): - """Sets default cuda block size.""" - self.cuda_block_size = cuda_block_size - - def get_cuda_block_size(self): - return self.cuda_block_size - - def print_cuda_block_size(self): - print(f"CUDA Block Size: {self.cuda_block_size}") - - # Specific GPUs - def set_specific_gpus(self): - """Set specific GPUs from CUDA_VISIBLE_DEVICES.""" - if os.getenv("CUDA_VISIBLE_DEVICES"): - self._specific_gpus = os.getenv("CUDA_VISIBLE_DEVICES") - # Modify the build folder name to include GPU specifics - gpu_suffix = self._specific_gpus.replace(",", "_") - self.set_default_build_folder_name( - appended_name=f"CUDA_VISIBLE_DEVICES_{gpu_suffix}" - ) - - def get_specific_gpus(self): - """Get the specific GPUs.""" - return self._specific_gpus - - def print_specific_gpus(self): - """Print the specific GPUs.""" - print( - f"Specific GPUs (CUDA_VISIBLE_DEVICES): {self.get_specific_gpus() or "Not set"}" - ) - - # Libcuda folder - def set_libcuda_folder(self): - """ - Is set in _cuda_libraries_available. - """ - pass - - def get_libcuda_folder(self): - return self._libcuda_info["library"] and os.path.dirname( - self._libcuda_info["library"] - ) - - def print_libcuda_folder(self): - print(f"Libcuda Folder: {self.get_libcuda_folder() or not_found_str}") - - # Libnvrtc folder - def set_libnvrtc_folder(self): - """ - Return nothing if not using cuda - self.libnvrtc_folder is already set in _cuda_libraries_available. - """ - pass - - def get_libnvrtc_folder(self): - return self._libnvrtc_info["library"] and os.path.dirname( - self._libnvrtc_info["library"] - ) - - def print_libnvrtc_folder(self): - print(f"Libnvrtc Folder: {self.get_libnvrtc_folder() or not_found_str}") - - # CUDA Version - def set_cuda_version(self, warn=True): - """Set the CUDA version by querying the CUDA runtime library.""" - self._cuda_version = self.find_cuda_version(warn=warn) - - def get_cuda_version(self, out_type="single_value"): - - major = self._cuda_version // 1000 - minor = (self._cuda_version % 1000) // 10 - - if out_type == "major,minor": - return major, minor - elif out_type == "string": - return f"{major}.{minor}" - else: - return self._cuda_version - - def print_cuda_version(self): - str = f"CUDA Version: {self.get_cuda_version(out_type="string") if self.get_cuda_version() else not_found_str}" - print(str) - - # CUDA Include Path - def set_cuda_include_path(self): - """Set the CUDA include path by searching for cuda.h and nvrtc.h.""" - # This is done in find_cuda_install since it relies on the cuda installation info which is only available after checking library availability. - pass - - def get_cuda_include_path(self): - """ - Attempt to find CUDA headers (cuda.h, nvrtc.h) using an explicit - ordered search over environment variables and standard locations. - """ - if not self._use_cuda: - return None - - return self._cuda_include_path - - def print_cuda_include_path(self): - print(f"CUDA Include Path: {":".join(self.get_cuda_include_path()) or not_found_str}") - - # NVRTC Flags - def set_nvrtc_flags(self): - """Set the NVRTC flags for CUDA compilation.""" - # TODO: redondant with CppConfig compile options, should be refactored to avoid duplication - # Ensure that compile_options is set (inherited from ConfigNew) - compile_options = " -shared -fPIC -O3 -std=c++11" - - # Set the NVRTC flags - self._nvrtc_flags = ( - compile_options - + f" -fpermissive -L{self.get_libcuda_folder()} -L{self.get_libnvrtc_folder()} -lcuda -lnvrtc" - ) - - def get_nvrtc_flags(self): - """Get the NVRTC flags for CUDA compilation.""" - return self._nvrtc_flags - - def print_nvrtc_flags(self): - """Print the NVRTC flags for CUDA compilation.""" - print(f"NVRTC Flags: {self.get_nvrtc_flags()}") - - # GPU compile flags - def set_gpu_compile_flags(self): - """Set GPU compile flags based on detected GPU properties.""" - # This is done in get_gpu_props since it relies on the GPU properties which are only available after checking CUDA availability. - pass - - def get_gpu_compile_flags(self): - """Get GPU compile flags based on detected GPU properties.""" - return self._gpu_compile_flags - - def print_gpu_compile_flags(self): - print(f"GPU Compile Flags: {self.get_gpu_compile_flags() or not_found_str}") - - # GPU Properties - def get_gpu_props(self): - """ - Getting GPU properties and related attributes. - """ - if not self._use_cuda: - # Already determined that CUDA is unavailable - self.n_gpus = 0 - self._gpu_compile_flags = "" - return (self.n_gpus, self._gpu_compile_flags) - - # Attempt to load the CUDA driver library - libcuda_path = self._libcuda_info["library"] if self._libcuda_info else None - success, err_msg = self._try_load_library( - libcuda_path, - mode=RTLD_GLOBAL, - ) - if not success: - # Something is off at driver level => revert to CPU - KeOps_Warning( - "cuda library not fully accessible. " - + err_msg - + " Switching to CPU only." - ) - self.n_gpus = 0 - self._gpu_compile_flags = "" - self._use_cuda = False - return (self.n_gpus, self._gpu_compile_flags) - - # We have a handle, let's proceed - libcuda = ctypes.CDLL(libcuda_path) - result = libcuda.cuInit(0) - if result != self.CUDA_SUCCESS: - KeOps_Warning( - "CUDA was detected, but driver API could not be initialized. Switching to CPU only." - ) - self.n_gpus = 0 - self._gpu_compile_flags = "" - self._use_cuda = False - return (self.n_gpus, self._gpu_compile_flags) - - # Get GPU count - nGpus = ctypes.c_int() - result = libcuda.cuDeviceGetCount(ctypes.byref(nGpus)) - if result != self.CUDA_SUCCESS: - KeOps_Warning( - "CUDA was detected and driver API was initialized, but no working GPU found. " - "Switching to CPU only." - ) - self.n_gpus = 0 - self._gpu_compile_flags = "" - self._use_cuda = False - return (self.n_gpus, self._gpu_compile_flags) - - self.n_gpus = nGpus.value - if self.n_gpus == 0: - self._gpu_compile_flags = "" - return (self.n_gpus, self._gpu_compile_flags) - - # Query each GPU for properties - MaxThreadsPerBlock = [0] * self.n_gpus - SharedMemPerBlock = [0] * self.n_gpus - - def safe_call(dev_idx, result_code): - if result_code != self.CUDA_SUCCESS: - KeOps_Warning( - f"Error detecting properties for GPU device {dev_idx}. " - "Switching to CPU only." - ) - return False - return True - - for d in range(self.n_gpus): - device = ctypes.c_int() - if not safe_call(d, libcuda.cuDeviceGet(ctypes.byref(device), d)): - self.n_gpus = 0 - self._gpu_compile_flags = "" - self._use_cuda = False - return (self.n_gpus, self._gpu_compile_flags) - - output = ctypes.c_int() - if not safe_call( - d, - libcuda.cuDeviceGetAttribute( - byref(output), - self.CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK, - device, - ), - ): - self.n_gpus = 0 - self._gpu_compile_flags = "" - self._use_cuda = False - return (self.n_gpus, self._gpu_compile_flags) - MaxThreadsPerBlock[d] = output.value - - if not safe_call( - d, - libcuda.cuDeviceGetAttribute( - byref(output), - self.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK, - device, - ), - ): - self.n_gpus = 0 - self._gpu_compile_flags = "" - self._use_cuda = False - return (self.n_gpus, self._gpu_compile_flags) - SharedMemPerBlock[d] = output.value - - # Build compile flags string - self._gpu_compile_flags = f"-DMAXIDGPU={self.n_gpus - 1} " - for d in range(self.n_gpus): - self._gpu_compile_flags += ( - f"-DMAXTHREADSPERBLOCK{d}={MaxThreadsPerBlock[d]} " - ) - self._gpu_compile_flags += f"-DSHAREDMEMPERBLOCK{d}={SharedMemPerBlock[d]} " - - return self.n_gpus, self._gpu_compile_flags - - def print_cuda(self): - """Print all CUDA-related configuration""" - - # CUDA Support - print("=" * 60) - print(f"CUDA Support") - print("=" * 60) - - self.print_use_cuda() - if self.get_use_cuda(): - print(f"Number of GPUs: {self.n_gpus}") - self.print_cuda_version() - self.print_libcuda_folder() - self.print_libnvrtc_folder() - self.print_cuda_include_path() - self.print_nvrtc_flags() - self.print_gpu_compile_flags() - - # Print relevant environment variables. - print_envs(self.cuda_env_vars) - - -if __name__ == "__main__": - cuda_config = CUDAConfig() - cuda_config.print_platform() - cuda_config.print_cpp() - cuda_config.print_cuda() diff --git a/keopscore/keopscore/formulas/GetReduction.py b/keopscore/keopscore/formulas/GetReduction.py index 60afe3e5c..b1a3ad88d 100644 --- a/keopscore/keopscore/formulas/GetReduction.py +++ b/keopscore/keopscore/formulas/GetReduction.py @@ -1,16 +1,9 @@ import ast, inspect -import keopscore -import keopscore.formulas -from keopscore.utils.misc_utils import KeOps_Print +import keopscore.config +from keopscore.utils.messages import KeOps_Print from keopscore.utils.code_gen_utils import get_hash_name -from keopscore.formulas.reductions import * -from keopscore.formulas.maths import * -from keopscore.formulas.complex import * -from keopscore.formulas.variables import * -from keopscore.formulas.autodiff import * -from keopscore.formulas.LinearOperators import * -from keopscore.formulas.factorization import * +from keopscore.formulas import * class GetReduction: @@ -18,7 +11,7 @@ class GetReduction: def __new__(self, red_formula_string, aliases=[]): string_id_hash = get_hash_name( - red_formula_string, aliases, keopscore.auto_factorize + red_formula_string, aliases, keopscore.config.auto_factorize ) if string_id_hash in GetReduction.library: return GetReduction.library[string_id_hash] @@ -31,7 +24,7 @@ def __new__(self, red_formula_string, aliases=[]): varname, var = alias.split("=") aliases_dict[varname] = eval(var) reduction = eval(red_formula_string, globals(), aliases_dict) - if keopscore.auto_factorize: + if keopscore.config.auto_factorize: formula = reduction.children[0] new_formula = AutoFactorize(formula) reduction.children[0] = new_formula diff --git a/keopscore/keopscore/formulas/LinearOperators/AdjointOperator.py b/keopscore/keopscore/formulas/LinearOperators/AdjointOperator.py index f6b980fac..ff1ab0544 100644 --- a/keopscore/keopscore/formulas/LinearOperators/AdjointOperator.py +++ b/keopscore/keopscore/formulas/LinearOperators/AdjointOperator.py @@ -2,7 +2,7 @@ from keopscore.formulas import Var from keopscore.utils.code_gen_utils import GetInds -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.formulas.maths.Minus import Minus_Impl from keopscore.formulas.maths.Add import Add_Impl from keopscore.formulas.maths.Subtract import Subtract_Impl diff --git a/keopscore/keopscore/formulas/LinearOperators/LinearOperator.py b/keopscore/keopscore/formulas/LinearOperators/LinearOperator.py index 503993799..df453d94e 100644 --- a/keopscore/keopscore/formulas/LinearOperators/LinearOperator.py +++ b/keopscore/keopscore/formulas/LinearOperators/LinearOperator.py @@ -1,6 +1,6 @@ from keopscore.formulas import Var from keopscore.utils.code_gen_utils import GetInds -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.formulas.maths.Minus import Minus_Impl from keopscore.formulas.maths.Add import Add_Impl from keopscore.formulas.maths.Subtract import Subtract_Impl diff --git a/keopscore/keopscore/formulas/LinearOperators/SumLinOperator.py b/keopscore/keopscore/formulas/LinearOperators/SumLinOperator.py index 337db5301..55f8bcd2e 100644 --- a/keopscore/keopscore/formulas/LinearOperators/SumLinOperator.py +++ b/keopscore/keopscore/formulas/LinearOperators/SumLinOperator.py @@ -2,7 +2,7 @@ from keopscore.formulas import Var from keopscore.utils.code_gen_utils import GetInds -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.formulas.maths.Minus import Minus_Impl from keopscore.formulas.maths.Add import Add_Impl from keopscore.formulas.maths.Subtract import Subtract_Impl diff --git a/keopscore/keopscore/formulas/LinearOperators/TraceOperator.py b/keopscore/keopscore/formulas/LinearOperators/TraceOperator.py index ee659857b..2eff1ffce 100644 --- a/keopscore/keopscore/formulas/LinearOperators/TraceOperator.py +++ b/keopscore/keopscore/formulas/LinearOperators/TraceOperator.py @@ -3,7 +3,7 @@ from keopscore.formulas import Var from keopscore.utils.code_gen_utils import GetInds -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.formulas.maths.Minus import Minus_Impl from keopscore.formulas.maths.Add import Add_Impl from keopscore.formulas.maths.Subtract import Subtract_Impl diff --git a/keopscore/keopscore/formulas/LinearOperators/__init__.py b/keopscore/keopscore/formulas/LinearOperators/__init__.py index 8e2515c48..b68ec01ea 100644 --- a/keopscore/keopscore/formulas/LinearOperators/__init__.py +++ b/keopscore/keopscore/formulas/LinearOperators/__init__.py @@ -1,3 +1,11 @@ -from .TraceOperator import TraceOperator from .AdjointOperator import AdjointOperator from .SumLinOperator import SumLinOperator +from .TraceOperator import TraceOperator + +_exports = [ + TraceOperator, + AdjointOperator, + SumLinOperator, +] + +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/formulas/Operation.py b/keopscore/keopscore/formulas/Operation.py index 62019af5c..8f71ea066 100644 --- a/keopscore/keopscore/formulas/Operation.py +++ b/keopscore/keopscore/formulas/Operation.py @@ -1,7 +1,7 @@ from keopscore.utils.code_gen_utils import new_c_varname, c_array from keopscore.utils.Tree import Tree import keopscore -from keopscore.utils.misc_utils import KeOps_Error, KeOps_Print +from keopscore.utils.messages import KeOps_Error, KeOps_Print ################### ## Base class @@ -71,14 +71,14 @@ def __call__(self, out, table): from keopscore.formulas.variables.Var import Var string = f"\n{{\n// Starting code block for {self.__repr__()}.\n\n" - if keopscore.debug_ops: + if keopscore.config.debug.get_debug_ops(): KeOps_Print(f"Building code block for {self.__repr__()}") KeOps_Print("out=", out) KeOps_Print("dim of out : ", out.dim) KeOps_Print("table=", table) for v in table: KeOps_Print(f"dim of {v} : ", v.dim) - if keopscore.debug_ops_at_exec: + if keopscore.config.debug.get_debug_ops_at_exec(): string += f'printf("\\n\\nComputing {self.__repr__()} :\\n");\n' args = [] # Evaluation of the child operations @@ -104,12 +104,12 @@ def __call__(self, out, table): string += self.Op(out, table, *args) # some debugging helper : - if keopscore.debug_ops_at_exec: + if keopscore.config.debug.get_debug_ops_at_exec(): for arg in args: string += arg.c_print string += out.c_print string += f'printf("\\n\\n");\n' - if keopscore.debug_ops: + if keopscore.config.debug.get_debug_ops(): KeOps_Print(f"Finished building code block for {self.__repr__()}") string += f"\n\n// Finished code block for {self.__repr__()}.\n}}\n\n" diff --git a/keopscore/keopscore/formulas/VectorizedComplexScalarOp.py b/keopscore/keopscore/formulas/VectorizedComplexScalarOp.py index 07207d331..b221d5ce5 100644 --- a/keopscore/keopscore/formulas/VectorizedComplexScalarOp.py +++ b/keopscore/keopscore/formulas/VectorizedComplexScalarOp.py @@ -1,6 +1,6 @@ from keopscore.utils.code_gen_utils import ComplexVectApply from keopscore.formulas.Operation import Operation -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error class VectorizedComplexScalarOp(Operation): diff --git a/keopscore/keopscore/formulas/VectorizedScalarOp.py b/keopscore/keopscore/formulas/VectorizedScalarOp.py index efee4d041..0750a9d47 100644 --- a/keopscore/keopscore/formulas/VectorizedScalarOp.py +++ b/keopscore/keopscore/formulas/VectorizedScalarOp.py @@ -1,6 +1,6 @@ from keopscore.utils.code_gen_utils import VectApply from keopscore.formulas.Operation import Operation -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error class VectorizedScalarOp(Operation): diff --git a/keopscore/keopscore/formulas/autodiff/Divergence.py b/keopscore/keopscore/formulas/autodiff/Divergence.py index 18c69affc..f9bcadb34 100644 --- a/keopscore/keopscore/formulas/autodiff/Divergence.py +++ b/keopscore/keopscore/formulas/autodiff/Divergence.py @@ -2,7 +2,7 @@ from keopscore.formulas.LinearOperators import TraceOperator from keopscore.utils.code_gen_utils import GetInds from keopscore.formulas.variables.IntCst import IntCst -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////// # /// DIVERGENCE OPERATOR : Divergence< F, V, U > //// diff --git a/keopscore/keopscore/formulas/autodiff/Grad.py b/keopscore/keopscore/formulas/autodiff/Grad.py index ae8280058..2603b8f19 100644 --- a/keopscore/keopscore/formulas/autodiff/Grad.py +++ b/keopscore/keopscore/formulas/autodiff/Grad.py @@ -1,6 +1,6 @@ from keopscore.formulas import Var from keopscore.utils.code_gen_utils import GetInds -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////// # /// GRADIENT OPERATOR : Grad< F, V, Gradin > //// diff --git a/keopscore/keopscore/formulas/autodiff/Laplacian.py b/keopscore/keopscore/formulas/autodiff/Laplacian.py index cf217c048..752748c46 100644 --- a/keopscore/keopscore/formulas/autodiff/Laplacian.py +++ b/keopscore/keopscore/formulas/autodiff/Laplacian.py @@ -3,7 +3,7 @@ from keopscore.formulas.autodiff.Divergence import Divergence from keopscore.utils.code_gen_utils import GetInds from keopscore.formulas.variables.IntCst import IntCst -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////// # /// LAPLACIAN OPERATOR : Laplacian< F, V, U > //// diff --git a/keopscore/keopscore/formulas/autodiff/__init__.py b/keopscore/keopscore/formulas/autodiff/__init__.py index f29bf31cf..6452b9a96 100644 --- a/keopscore/keopscore/formulas/autodiff/__init__.py +++ b/keopscore/keopscore/formulas/autodiff/__init__.py @@ -1,6 +1,17 @@ -from .Grad import Grad -from .Grad_WithSavedForward import Grad_WithSavedForward from .Diff import Diff from .Diff_WithSavedForward import Diff_WithSavedForward -from .Laplacian import Laplacian from .Divergence import Divergence +from .Grad import Grad +from .Grad_WithSavedForward import Grad_WithSavedForward +from .Laplacian import Laplacian + +_exports = [ + Grad, + Grad_WithSavedForward, + Diff, + Diff_WithSavedForward, + Laplacian, + Divergence, +] + +__all__ = [cls.__name__ for cls in _exports] \ No newline at end of file diff --git a/keopscore/keopscore/formulas/complex/ComplexExp1j.py b/keopscore/keopscore/formulas/complex/ComplexExp1j.py index 395a71ae8..68ec03823 100644 --- a/keopscore/keopscore/formulas/complex/ComplexExp1j.py +++ b/keopscore/keopscore/formulas/complex/ComplexExp1j.py @@ -5,7 +5,7 @@ from keopscore.formulas.complex.Imag2Complex import Imag2Complex from keopscore.formulas.maths.Cos import Cos from keopscore.formulas.maths.Sin import Sin -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// ComplexExp1j //// diff --git a/keopscore/keopscore/formulas/complex/ComplexImag.py b/keopscore/keopscore/formulas/complex/ComplexImag.py index 8c1e3a8e2..3309a29ee 100644 --- a/keopscore/keopscore/formulas/complex/ComplexImag.py +++ b/keopscore/keopscore/formulas/complex/ComplexImag.py @@ -1,7 +1,7 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import c_for_loop from keopscore.formulas.complex.Imag2Complex import Imag2Complex -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// ComplexImag //// diff --git a/keopscore/keopscore/formulas/complex/ComplexReal.py b/keopscore/keopscore/formulas/complex/ComplexReal.py index d2ffbb1db..c08e73d14 100644 --- a/keopscore/keopscore/formulas/complex/ComplexReal.py +++ b/keopscore/keopscore/formulas/complex/ComplexReal.py @@ -1,7 +1,7 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import c_for_loop from keopscore.formulas.complex.Real2Complex import Real2Complex -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// ComplexReal //// diff --git a/keopscore/keopscore/formulas/complex/ComplexRealScal.py b/keopscore/keopscore/formulas/complex/ComplexRealScal.py index e283ef00f..dd2b33384 100644 --- a/keopscore/keopscore/formulas/complex/ComplexRealScal.py +++ b/keopscore/keopscore/formulas/complex/ComplexRealScal.py @@ -2,7 +2,7 @@ from keopscore.utils.code_gen_utils import c_for_loop from keopscore.formulas.complex.Real2Complex import Real2Complex from keopscore.formulas.complex.ComplexMult import ComplexMult -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// ComplexRealScal //// diff --git a/keopscore/keopscore/formulas/complex/ComplexSquareAbs.py b/keopscore/keopscore/formulas/complex/ComplexSquareAbs.py index a95a5626f..184448a77 100644 --- a/keopscore/keopscore/formulas/complex/ComplexSquareAbs.py +++ b/keopscore/keopscore/formulas/complex/ComplexSquareAbs.py @@ -3,7 +3,7 @@ from keopscore.formulas.complex.ComplexReal import ComplexReal from keopscore.formulas.complex.ComplexMult import ComplexMult from keopscore.formulas.complex.Conj import Conj -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// ComplexSquareAbs //// diff --git a/keopscore/keopscore/formulas/complex/ComplexSum.py b/keopscore/keopscore/formulas/complex/ComplexSum.py index c244b952c..0fa3c8ab4 100644 --- a/keopscore/keopscore/formulas/complex/ComplexSum.py +++ b/keopscore/keopscore/formulas/complex/ComplexSum.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import c_zero_float, c_for_loop -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// ComplexSum //// diff --git a/keopscore/keopscore/formulas/complex/ComplexSumT.py b/keopscore/keopscore/formulas/complex/ComplexSumT.py index 634f6b283..d2ec2797b 100644 --- a/keopscore/keopscore/formulas/complex/ComplexSumT.py +++ b/keopscore/keopscore/formulas/complex/ComplexSumT.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import c_for_loop -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// adjoint of ComplexSum //// diff --git a/keopscore/keopscore/formulas/complex/__init__.py b/keopscore/keopscore/formulas/complex/__init__.py index 0008a4202..b0bf68c61 100644 --- a/keopscore/keopscore/formulas/complex/__init__.py +++ b/keopscore/keopscore/formulas/complex/__init__.py @@ -15,3 +15,25 @@ from .Conj import Conj from .Imag2Complex import Imag2Complex from .Real2Complex import Real2Complex + +_exports = [ + ComplexAbs, + ComplexAdd, + ComplexAngle, + ComplexDivide, + ComplexExp, + ComplexExp1j, + ComplexImag, + ComplexMult, + ComplexReal, + ComplexRealScal, + ComplexSquareAbs, + ComplexSubtract, + ComplexSum, + ComplexSumT, + Conj, + Imag2Complex, + Real2Complex, +] + +__all__ = [cls.__name__ for cls in _exports] \ No newline at end of file diff --git a/keopscore/keopscore/formulas/factorization/Factorize.py b/keopscore/keopscore/formulas/factorization/Factorize.py index cf9f4907c..e92d4d147 100644 --- a/keopscore/keopscore/formulas/factorization/Factorize.py +++ b/keopscore/keopscore/formulas/factorization/Factorize.py @@ -4,7 +4,7 @@ from keopscore.formulas.Operation import Operation import keopscore from keopscore.utils.code_gen_utils import GetInds -from keopscore.utils.misc_utils import KeOps_Print +from keopscore.utils.messages import KeOps_Print class Factorize_Impl(Operation): @@ -32,14 +32,14 @@ def __call__(self, out, table): from keopscore.formulas.variables.Var import Var string = f"\n{{\n// Starting code block for {self.__repr__()}.\n\n" - if keopscore.debug_ops: + if keopscore.config.debug.get_debug_ops(): KeOps_Print(f"Building code block for {self.__repr__()}") KeOps_Print("out=", out) KeOps_Print("dim of out : ", out.dim) KeOps_Print("table=", table) for v in table: KeOps_Print(f"dim of {v} : ", v.dim) - if keopscore.debug_ops_at_exec: + if keopscore.config.debug.get_debug_ops_at_exec(): string += f'printf("\\n\\nComputing {self.__repr__()} :\\n");\n' f, g = self.children @@ -68,7 +68,7 @@ def __call__(self, out, table): # Evaluation of f string += newf(out, table) - if keopscore.debug_ops: + if keopscore.config.debug.get_debug_ops(): KeOps_Print(f"Finished building code block for {self.__repr__()}") string += f"\n\n// Finished code block for {self.__repr__()}.\n}}\n\n" diff --git a/keopscore/keopscore/formulas/factorization/__init__.py b/keopscore/keopscore/formulas/factorization/__init__.py index 970407e18..8668dab65 100644 --- a/keopscore/keopscore/formulas/factorization/__init__.py +++ b/keopscore/keopscore/formulas/factorization/__init__.py @@ -1 +1,8 @@ from .Factorize import Factorize, AutoFactorize + +_exports = [ + Factorize, + AutoFactorize, +] + +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/formulas/maths/ArgMax.py b/keopscore/keopscore/formulas/maths/ArgMax.py index bf42c1237..dc10f9f56 100644 --- a/keopscore/keopscore/formulas/maths/ArgMax.py +++ b/keopscore/keopscore/formulas/maths/ArgMax.py @@ -7,7 +7,7 @@ value, c_variable, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ############################ ###### ArgMax ##### diff --git a/keopscore/keopscore/formulas/maths/ArgMin.py b/keopscore/keopscore/formulas/maths/ArgMin.py index d8fcba980..0a798ab3a 100644 --- a/keopscore/keopscore/formulas/maths/ArgMin.py +++ b/keopscore/keopscore/formulas/maths/ArgMin.py @@ -7,7 +7,7 @@ value, c_variable, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ############################ ###### ArgMin ##### diff --git a/keopscore/keopscore/formulas/maths/BSpline.py b/keopscore/keopscore/formulas/maths/BSpline.py index 44f231cb9..929344598 100644 --- a/keopscore/keopscore/formulas/maths/BSpline.py +++ b/keopscore/keopscore/formulas/maths/BSpline.py @@ -2,7 +2,7 @@ from keopscore.formulas.maths.Extract import Extract from keopscore.utils.code_gen_utils import c_variable, c_for_loop, c_zero_float from keopscore.utils.code_gen_utils import c_array, VectCopy -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ////////////////////////////////////////////////////////////// # //// BSPLINE VECTOR : BSPLINE //// diff --git a/keopscore/keopscore/formulas/maths/Concat.py b/keopscore/keopscore/formulas/maths/Concat.py index a104733bf..4dafd03ad 100644 --- a/keopscore/keopscore/formulas/maths/Concat.py +++ b/keopscore/keopscore/formulas/maths/Concat.py @@ -1,7 +1,7 @@ from keopscore.formulas.Operation import Operation from keopscore.formulas.maths.Extract import Extract from keopscore.utils.code_gen_utils import VectCopy -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ############################ ###### Concat ##### diff --git a/keopscore/keopscore/formulas/maths/Divide.py b/keopscore/keopscore/formulas/maths/Divide.py index cd62f304e..9b61f0de8 100644 --- a/keopscore/keopscore/formulas/maths/Divide.py +++ b/keopscore/keopscore/formulas/maths/Divide.py @@ -6,7 +6,7 @@ from keopscore.formulas.variables.Zero import Zero from keopscore.formulas.variables.IntCst import IntCst, IntCst_Impl from keopscore.formulas.variables.RatCst import RatCst, RatCst_Impl -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ########################## ###### Divide ##### diff --git a/keopscore/keopscore/formulas/maths/Elem.py b/keopscore/keopscore/formulas/maths/Elem.py index e802e2dfb..d061b4048 100644 --- a/keopscore/keopscore/formulas/maths/Elem.py +++ b/keopscore/keopscore/formulas/maths/Elem.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import value -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ############################ ###### ELEMENT EXTRACTION : Elem(f,m) (aka get_item) ##### diff --git a/keopscore/keopscore/formulas/maths/ElemT.py b/keopscore/keopscore/formulas/maths/ElemT.py index 652ad4c79..aaa1fa88e 100644 --- a/keopscore/keopscore/formulas/maths/ElemT.py +++ b/keopscore/keopscore/formulas/maths/ElemT.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import value, c_zero_float, c_for_loop -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ############################ ###### ELEMENT "INJECTION" : ElemT(f,n,m) diff --git a/keopscore/keopscore/formulas/maths/Extract.py b/keopscore/keopscore/formulas/maths/Extract.py index 62bc63761..9d2cdd8ae 100644 --- a/keopscore/keopscore/formulas/maths/Extract.py +++ b/keopscore/keopscore/formulas/maths/Extract.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import c_array, VectCopy -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ////////////////////////////////////////////////////////////// # //// VECTOR EXTRACTION : Extract //// diff --git a/keopscore/keopscore/formulas/maths/ExtractT.py b/keopscore/keopscore/formulas/maths/ExtractT.py index 9263a2b8c..11f89a1b1 100644 --- a/keopscore/keopscore/formulas/maths/ExtractT.py +++ b/keopscore/keopscore/formulas/maths/ExtractT.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import c_zero_float, VectCopy -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ////////////////////////////////////////////////////////////// # //// VECTOR "INJECTION" : ExtractT //// diff --git a/keopscore/keopscore/formulas/maths/MatVecMult.py b/keopscore/keopscore/formulas/maths/MatVecMult.py index 0369384de..927400f9b 100644 --- a/keopscore/keopscore/formulas/maths/MatVecMult.py +++ b/keopscore/keopscore/formulas/maths/MatVecMult.py @@ -4,7 +4,7 @@ c_for_loop, c_zero_float, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// Matrix-vector product A x b //// diff --git a/keopscore/keopscore/formulas/maths/Max.py b/keopscore/keopscore/formulas/maths/Max.py index cdbf13f3e..d780b33d4 100644 --- a/keopscore/keopscore/formulas/maths/Max.py +++ b/keopscore/keopscore/formulas/maths/Max.py @@ -2,7 +2,7 @@ from keopscore.formulas.maths.ArgMax import ArgMax from keopscore.formulas.maths.OneHot import OneHot from keopscore.utils.code_gen_utils import c_for_loop, c_if, value -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ############################ ###### Max ##### diff --git a/keopscore/keopscore/formulas/maths/Min.py b/keopscore/keopscore/formulas/maths/Min.py index 6995c2278..1956b80d1 100644 --- a/keopscore/keopscore/formulas/maths/Min.py +++ b/keopscore/keopscore/formulas/maths/Min.py @@ -2,7 +2,7 @@ from keopscore.formulas.maths.ArgMin import ArgMin from keopscore.formulas.maths.OneHot import OneHot from keopscore.utils.code_gen_utils import c_for_loop, c_if, value -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ############################ ###### Min ##### diff --git a/keopscore/keopscore/formulas/maths/Mult.py b/keopscore/keopscore/formulas/maths/Mult.py index c8f09c6a5..870f8160d 100644 --- a/keopscore/keopscore/formulas/maths/Mult.py +++ b/keopscore/keopscore/formulas/maths/Mult.py @@ -8,7 +8,7 @@ from keopscore.formulas.variables.IntCst import IntCst_Impl, IntCst from keopscore.formulas.variables.RatCst import RatCst_Impl, RatCst from keopscore.formulas.maths.SumT import SumT, SumT_Impl -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ########################## ###### Mult ##### diff --git a/keopscore/keopscore/formulas/maths/OneHot.py b/keopscore/keopscore/formulas/maths/OneHot.py index 8cf93fb7c..7261ae005 100644 --- a/keopscore/keopscore/formulas/maths/OneHot.py +++ b/keopscore/keopscore/formulas/maths/OneHot.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.utils.code_gen_utils import c_zero_float -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ////////////////////////////////////////////////////////////// # //// ONE-HOT REPRESENTATION : OneHot //// diff --git a/keopscore/keopscore/formulas/maths/Scalprod.py b/keopscore/keopscore/formulas/maths/Scalprod.py index 4da68f8b2..5d72d4fd6 100644 --- a/keopscore/keopscore/formulas/maths/Scalprod.py +++ b/keopscore/keopscore/formulas/maths/Scalprod.py @@ -5,7 +5,7 @@ VectApply, ) from keopscore.utils.math_functions import keops_fma -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ########################## ##### Scalprod #### diff --git a/keopscore/keopscore/formulas/maths/SoftDTW_SqDist.py b/keopscore/keopscore/formulas/maths/SoftDTW_SqDist.py index 6aff05757..9387d9f2f 100644 --- a/keopscore/keopscore/formulas/maths/SoftDTW_SqDist.py +++ b/keopscore/keopscore/formulas/maths/SoftDTW_SqDist.py @@ -3,7 +3,7 @@ #################################################################### from keopscore.formulas.Operation import Operation -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.utils.code_gen_utils import ( c_variable, pointer, diff --git a/keopscore/keopscore/formulas/maths/SumT.py b/keopscore/keopscore/formulas/maths/SumT.py index 786c44536..795b18b32 100644 --- a/keopscore/keopscore/formulas/maths/SumT.py +++ b/keopscore/keopscore/formulas/maths/SumT.py @@ -1,7 +1,7 @@ from keopscore.formulas.Operation import Operation from keopscore.formulas.variables.Zero import Zero from keopscore.utils.code_gen_utils import value -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error ########################## ###### SumT ##### diff --git a/keopscore/keopscore/formulas/maths/TensorProd.py b/keopscore/keopscore/formulas/maths/TensorProd.py index 44f5e72eb..2320a18ff 100644 --- a/keopscore/keopscore/formulas/maths/TensorProd.py +++ b/keopscore/keopscore/formulas/maths/TensorProd.py @@ -3,7 +3,7 @@ c_variable, c_for_loop, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error #################################### ###### Tensor product ##### diff --git a/keopscore/keopscore/formulas/maths/VecMatMult.py b/keopscore/keopscore/formulas/maths/VecMatMult.py index 93b400529..369f0baa3 100644 --- a/keopscore/keopscore/formulas/maths/VecMatMult.py +++ b/keopscore/keopscore/formulas/maths/VecMatMult.py @@ -4,7 +4,7 @@ c_for_loop, c_zero_float, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error # ///////////////////////////////////////////////////////////////////////// # //// Vector-matrix product b x A //// diff --git a/keopscore/keopscore/formulas/maths/__init__.py b/keopscore/keopscore/formulas/maths/__init__.py index c036915f3..5914ca5e4 100644 --- a/keopscore/keopscore/formulas/maths/__init__.py +++ b/keopscore/keopscore/formulas/maths/__init__.py @@ -66,72 +66,20 @@ from .WeightedSqNorm import WeightedSqNorm from .XLogX import XLogX -__all__ = [ - "Abs", - "Acos", - "Add", - "ArgMax", - "ArgMin", - "Asin", - "Atan", - "Atan2", - "BSpline", - "Clamp", - "ClampInt", - "Concat", - "Cos", - "DiffClampInt", - "Divide", - "Elem", - "ElemT", - "Equal", - "Exp", - "Extract", - "ExtractT", - "GradMatrix", - "Floor", - "IfElse", - "IntInv", - "Inv", - "Kron", - "LessOrEqual", - "LessThan", - "Log", - "MatVecMult", - "Max", - "Min", - "Minus", - "Mod", - "Mult", - "Norm2", - "Normalize", - "NotEqual", - "OneHot", - "Pow", - "Powf", - "ReLU", - "Round", - "Rsqrt", - "Scalprod", - "Sign", - "Sin", - "SinXDivX", - "SoftDTW_SqDist", - "SqDist", - "SqNorm2", - "SqNormDiag", - "SqNormIso", - "Sqrt", - "Square", - "Step", - "Subtract", - "Sum", - "SumT", - "SymSqNorm", - "TensorDot", - "TensorProd", - "VecMatMult", - "WeightedSqDist", - "WeightedSqNorm", - "XLogX", +_exports = [ + Abs, Acos, Add, ArgMax, ArgMin, Asin, Atan, Atan2, BSpline, + Clamp, ClampInt, Concat, Cos, DiffClampInt, Divide, + Elem, ElemT, Equal, Exp, Extract, ExtractT, Floor, + GradMatrix, IfElse, IntInv, Inv, Kron, + LessOrEqual, LessThan, Log, MatVecMult, Max, Min, Minus, Mod, Mult, + Norm2, Normalize, NotEqual, OneHot, + Pow, Powf, ReLU, Round, Rsqrt, Scalprod, Sign, + Sin, SinXDivX, SoftDTW_SqDist, + SqDist, SqNorm2, SqNormDiag, SqNormIso, + Sqrt, Square, Step, Subtract, Sum, SumT, + SymSqNorm, TensorDot, TensorProd, + VecMatMult, WeightedSqDist, WeightedSqNorm, + XLogX, ] + +__all__ = [cls.__name__ for cls in _exports] \ No newline at end of file diff --git a/keopscore/keopscore/formulas/reductions/KMin_ArgKMin_Reduction.py b/keopscore/keopscore/formulas/reductions/KMin_ArgKMin_Reduction.py index cb19df7a1..a4d5d38d8 100644 --- a/keopscore/keopscore/formulas/reductions/KMin_ArgKMin_Reduction.py +++ b/keopscore/keopscore/formulas/reductions/KMin_ArgKMin_Reduction.py @@ -10,7 +10,7 @@ use_pragma_unroll, ) from keopscore.formulas.reductions.Reduction import Reduction -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error class KMin_ArgKMin_Reduction(Reduction): diff --git a/keopscore/keopscore/formulas/reductions/Max_ArgMax_Reduction_Base.py b/keopscore/keopscore/formulas/reductions/Max_ArgMax_Reduction_Base.py index 6963e2195..530489a64 100644 --- a/keopscore/keopscore/formulas/reductions/Max_ArgMax_Reduction_Base.py +++ b/keopscore/keopscore/formulas/reductions/Max_ArgMax_Reduction_Base.py @@ -5,7 +5,7 @@ c_if, ) from keopscore.formulas.reductions.Reduction import Reduction -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.utils.code_gen_utils import c_variable diff --git a/keopscore/keopscore/formulas/reductions/Max_Reduction.py b/keopscore/keopscore/formulas/reductions/Max_Reduction.py index ca0837bd8..91a439ad2 100644 --- a/keopscore/keopscore/formulas/reductions/Max_Reduction.py +++ b/keopscore/keopscore/formulas/reductions/Max_Reduction.py @@ -1,6 +1,6 @@ from keopscore.utils.code_gen_utils import neg_infinity, c_if from keopscore.formulas.reductions.Reduction import Reduction -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error class Max_Reduction(Reduction): diff --git a/keopscore/keopscore/formulas/reductions/Max_SumShiftExpWeight_Reduction.py b/keopscore/keopscore/formulas/reductions/Max_SumShiftExpWeight_Reduction.py index 364e19e4f..c3760c446 100644 --- a/keopscore/keopscore/formulas/reductions/Max_SumShiftExpWeight_Reduction.py +++ b/keopscore/keopscore/formulas/reductions/Max_SumShiftExpWeight_Reduction.py @@ -10,7 +10,7 @@ c_for_loop, ) from keopscore.utils.math_functions import keops_exp -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error class Max_SumShiftExpWeight_Reduction(Reduction): @@ -120,4 +120,5 @@ def Diff(self, v, diffin, MS): ) -Max_SumShiftExp_Reduction = Max_SumShiftExpWeight_Reduction +def Max_SumShiftExp_Reduction(*args, **kwargs): + return Max_SumShiftExpWeight_Reduction(*args, **kwargs) diff --git a/keopscore/keopscore/formulas/reductions/Min_ArgMin_Reduction_Base.py b/keopscore/keopscore/formulas/reductions/Min_ArgMin_Reduction_Base.py index aad023907..f5e66858a 100644 --- a/keopscore/keopscore/formulas/reductions/Min_ArgMin_Reduction_Base.py +++ b/keopscore/keopscore/formulas/reductions/Min_ArgMin_Reduction_Base.py @@ -5,7 +5,7 @@ c_if, ) from keopscore.formulas.reductions.Reduction import Reduction -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.utils.code_gen_utils import c_variable diff --git a/keopscore/keopscore/formulas/reductions/Min_Reduction.py b/keopscore/keopscore/formulas/reductions/Min_Reduction.py index 33e014f6d..f8b3000aa 100644 --- a/keopscore/keopscore/formulas/reductions/Min_Reduction.py +++ b/keopscore/keopscore/formulas/reductions/Min_Reduction.py @@ -1,6 +1,6 @@ from keopscore.utils.code_gen_utils import infinity, c_if from keopscore.formulas.reductions.Reduction import Reduction -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error class Min_Reduction(Reduction): diff --git a/keopscore/keopscore/formulas/reductions/__init__.py b/keopscore/keopscore/formulas/reductions/__init__.py index 956363d88..4c057081d 100644 --- a/keopscore/keopscore/formulas/reductions/__init__.py +++ b/keopscore/keopscore/formulas/reductions/__init__.py @@ -14,5 +14,27 @@ from .Min_Reduction import Min_Reduction from .Sum_Reduction import Sum_Reduction from .Zero_Reduction import Zero_Reduction +from .sum_schemes import block_sum, kahan_scheme, direct_sum -from .sum_schemes import * +_exports = [ + Reduction, + ArgKMin_Reduction, + ArgMax_Reduction, + ArgMin_Reduction, + KMin_ArgKMin_Reduction, + KMin_Reduction, + Max_ArgMax_Reduction, + Max_Reduction, + Max_SumShiftExpWeight_Reduction, + Max_SumShiftExp_Reduction, + Min_ArgMin_Reduction, + Min_Reduction, + Sum_Reduction, + Zero_Reduction, + block_sum, + kahan_scheme, + direct_sum +] + + +__all__ = [cls.__name__ for cls in _exports] \ No newline at end of file diff --git a/keopscore/keopscore/formulas/variables/__init__.py b/keopscore/keopscore/formulas/variables/__init__.py index cc9af2fb4..d39491a44 100644 --- a/keopscore/keopscore/formulas/variables/__init__.py +++ b/keopscore/keopscore/formulas/variables/__init__.py @@ -2,3 +2,12 @@ from .RatCst import RatCst from .Var import Var from .Zero import Zero + +_exports = [ + IntCst, + RatCst, + Var, + Zero, +] + +__all__ = [cls.__name__ for cls in _exports] \ No newline at end of file diff --git a/keopscore/keopscore/get_keops_dll.py b/keopscore/keopscore/get_keops_dll.py index 6d655609a..91a4fd10c 100644 --- a/keopscore/keopscore/get_keops_dll.py +++ b/keopscore/keopscore/get_keops_dll.py @@ -52,11 +52,8 @@ import inspect import sys -import keopscore -from keopscore.config import * - +import keopscore.config import keopscore.mapreduce -from keopscore import cuda_block_size from keopscore.config.chunks import ( get_enable_chunk, set_enable_chunk, @@ -69,7 +66,7 @@ from keopscore.formulas.GetReduction import GetReduction from keopscore.formulas.variables.Zero import Zero from keopscore.utils.Cache import Cache -from keopscore.utils.misc_utils import KeOps_Error, KeOps_Print +from keopscore.utils.messages import KeOps_Error, KeOps_Print # Get every classes in mapreduce map_reduce = dict(inspect.getmembers(keopscore.mapreduce, inspect.isclass)) @@ -87,7 +84,7 @@ def get_keops_dll_impl( # detecting the need for special chunked computation modes : use_chunk_mode = 0 if "Gpu" in map_reduce_id: - if not cuda_config.get_use_cuda(): + if not keopscore.config.cuda.get_use_cuda(): KeOps_Error( "You selected a Gpu reduce scheme but KeOps is in Cpu only mode." ) @@ -115,7 +112,7 @@ def get_keops_dll_impl( rf = map_reduce_obj.red_formula - if keopscore.debug_ops: + if keopscore.config.debug.get_debug_ops(): KeOps_Print("In get_keops_dll, formula is :", rf) KeOps_Print("formula.__repr__() is : ", rf.__repr__()) rf.make_dot() @@ -144,7 +141,7 @@ def get_keops_dll_impl( tagZero, res["use_half"], res["use_fast_math"], - cuda_block_size, + keopscore.config.cuda.get_cuda_block_size(), use_chunk_mode, tag1D2D, res["dimred"], @@ -162,7 +159,7 @@ def get_keops_dll_impl( get_keops_dll = Cache( get_keops_dll_impl, use_cache_file=True, - save_folder=config.get_build_folder(), + save_folder=keopscore.config.path.get_build_folder(), ) diff --git a/keopscore/keopscore/mapreduce/MapReduce.py b/keopscore/keopscore/mapreduce/MapReduce.py index efd9fc80c..8fa17aeae 100644 --- a/keopscore/keopscore/mapreduce/MapReduce.py +++ b/keopscore/keopscore/mapreduce/MapReduce.py @@ -1,6 +1,13 @@ from keopscore.formulas.reductions import * from keopscore.formulas.GetReduction import GetReduction -from keopscore.utils.code_gen_utils import Var_loader, new_c_varname, pointer, c_include +from keopscore.utils.code_gen_utils import ( + Var_loader, + new_c_varname, + pointer, + c_include, + c_variable, + c_array, +) class MapReduce: diff --git a/keopscore/keopscore/mapreduce/__init__.py b/keopscore/keopscore/mapreduce/__init__.py index 683989a1a..b490efc0d 100644 --- a/keopscore/keopscore/mapreduce/__init__.py +++ b/keopscore/keopscore/mapreduce/__init__.py @@ -1,7 +1,5 @@ -from .cpu import * -from ..config import get_cuda_config - -cuda_config = get_cuda_config() +from keopscore.config import cuda -if cuda_config._use_cuda: +from .cpu import * +if cuda.get_use_cuda(): from .gpu import * diff --git a/keopscore/keopscore/mapreduce/cpu/CpuAssignZero.py b/keopscore/keopscore/mapreduce/cpu/CpuAssignZero.py index de1333ac2..e66103e4b 100644 --- a/keopscore/keopscore/mapreduce/cpu/CpuAssignZero.py +++ b/keopscore/keopscore/mapreduce/cpu/CpuAssignZero.py @@ -1,12 +1,10 @@ -import keopscore +import keopscore.config from keopscore.binders.cpp.Cpu_link_compile import Cpu_link_compile from keopscore.mapreduce.MapReduce import MapReduce from keopscore.utils.code_gen_utils import ( c_include, c_zero_float, ) -import keopscore -from keopscore.config import * class CpuAssignZero(MapReduce, Cpu_link_compile): @@ -26,9 +24,9 @@ def get_code(self): args = self.args headers = ["stdlib.h"] - if keopscore.openmp_config.get_use_OpenMP(): + if keopscore.config.openmp.get_use_OpenMP(): headers.append("omp.h") - if keopscore.debug_ops_at_exec: + if keopscore.config.debug.get_debug_ops_at_exec(): headers.append("iostream") self.headers += c_include(*headers) diff --git a/keopscore/keopscore/mapreduce/cpu/CpuReduc.py b/keopscore/keopscore/mapreduce/cpu/CpuReduc.py index 29e8fb6b1..3b702e720 100644 --- a/keopscore/keopscore/mapreduce/cpu/CpuReduc.py +++ b/keopscore/keopscore/mapreduce/cpu/CpuReduc.py @@ -1,10 +1,8 @@ -import keopscore +import keopscore.config from keopscore.binders.cpp.Cpu_link_compile import Cpu_link_compile from keopscore.mapreduce.cpu.CpuAssignZero import CpuAssignZero from keopscore.mapreduce.MapReduce import MapReduce from keopscore.utils.code_gen_utils import c_include -import keopscore -from keopscore.config import * class CpuReduc(MapReduce, Cpu_link_compile): @@ -34,9 +32,9 @@ def get_code(self): sum_scheme = self.sum_scheme headers = ["cmath", "stdlib.h"] - if keopscore.openmp_config.get_use_OpenMP(): + if keopscore.config.openmp.get_use_OpenMP(): headers.append("omp.h") - if keopscore.debug_ops_at_exec: + if keopscore.config.debug.get_debug_ops_at_exec(): headers.append("iostream") self.headers += c_include(*headers) diff --git a/keopscore/keopscore/mapreduce/cpu/CpuReduc_ranges.py b/keopscore/keopscore/mapreduce/cpu/CpuReduc_ranges.py index 374c9382d..c4dfdade4 100644 --- a/keopscore/keopscore/mapreduce/cpu/CpuReduc_ranges.py +++ b/keopscore/keopscore/mapreduce/cpu/CpuReduc_ranges.py @@ -1,4 +1,4 @@ -import keopscore +import keopscore.config from keopscore.binders.cpp.Cpu_link_compile import Cpu_link_compile from keopscore.mapreduce.cpu.CpuAssignZero import CpuAssignZero @@ -8,8 +8,6 @@ c_array, c_include, ) -import keopscore -from keopscore.config import * class CpuReduc_ranges(MapReduce, Cpu_link_compile): @@ -65,9 +63,9 @@ def get_code(self): jmstarty = c_variable("int", "j-start_y") headers = ["cmath", "stdlib.h"] - if keopscore.openmp_config.get_use_OpenMP: + if keopscore.config.openmp.get_use_OpenMP(): headers.append("omp.h") - if keopscore.debug_ops_at_exec: + if keopscore.config.debug.get_debug_ops_at_exec(): headers.append("iostream") self.headers += c_include(*headers) diff --git a/keopscore/keopscore/mapreduce/cpu/__init__.py b/keopscore/keopscore/mapreduce/cpu/__init__.py index 49e1c5474..adf1ff1ff 100644 --- a/keopscore/keopscore/mapreduce/cpu/__init__.py +++ b/keopscore/keopscore/mapreduce/cpu/__init__.py @@ -1,3 +1,8 @@ -from .CpuReduc_ranges import CpuReduc_ranges -from .CpuReduc import CpuReduc from .CpuAssignZero import CpuAssignZero +from .CpuReduc import CpuReduc +from .CpuReduc_ranges import CpuReduc_ranges + + +_exports = [CpuAssignZero, CpuReduc, CpuReduc_ranges,] + +__all__ = [cls.__name__ for cls in _exports] \ No newline at end of file diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py index d664a3fa3..0d2617e92 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py @@ -1,9 +1,10 @@ -from keopscore import cuda_block_size -from keopscore.config.chunks import dimchunk from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile +from keopscore.config import cuda +from keopscore.config.chunks import dimchunk from keopscore.formulas.reductions.sum_schemes import * -from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero +from keopscore.mapreduce.Chunk_Mode_Constants import Chunk_Mode_Constants from keopscore.mapreduce.MapReduce import MapReduce +from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero from keopscore.utils.code_gen_utils import ( load_vars, load_vars_chunks, @@ -13,7 +14,6 @@ table4, use_pragma_unroll, ) -from keopscore.mapreduce.Chunk_Mode_Constants import Chunk_Mode_Constants def do_chunk_sub( @@ -129,7 +129,7 @@ def __init__(self, *args): self.chk = Chunk_Mode_Constants(self.red_formula) self.dimy = self.chk.dimy self.blocksize_chunks = min( - cuda_block_size, 1024, 49152 // max(1, self.dimy * sizeof(self.dtype)) + cuda.get_cuda_block_size(), 1024, 49152 // max(1, self.dimy * sizeof(self.dtype)) ) def get_code(self): diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py index 9e7ad6c7f..1d2374a70 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py @@ -1,10 +1,10 @@ -from keopscore import cuda_block_size -from keopscore.config.chunks import dimfinalchunk from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile +from keopscore.config import cuda +from keopscore.config.chunks import dimfinalchunk from keopscore.formulas.reductions.Sum_Reduction import Sum_Reduction from keopscore.formulas.reductions.sum_schemes import * -from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero from keopscore.mapreduce.MapReduce import MapReduce +from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero from keopscore.utils.code_gen_utils import ( load_vars, load_vars_chunks, @@ -13,7 +13,7 @@ Var_loader, use_pragma_unroll, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error def do_finalchunk_sub( @@ -118,7 +118,7 @@ def get_code(self): self.dimy = max(dimfinalchunk, dimy) blocksize_chunks = min( - cuda_block_size, 1024, 49152 // max(1, self.dimy * sizeof(self.dtype)) + cuda.get_cuda_block_size(), 1024, 49152 // max(1, self.dimy * sizeof(self.dtype)) ) if not isinstance(sum_scheme, block_sum): diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py index ca3dcb254..5a964ac5b 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py @@ -1,9 +1,10 @@ -from keopscore import cuda_block_size -from keopscore.config.chunks import dimchunk from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile +from keopscore.config import cuda +from keopscore.config.chunks import dimchunk from keopscore.formulas.reductions.sum_schemes import * -from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero +from keopscore.mapreduce.Chunk_Mode_Constants import Chunk_Mode_Constants from keopscore.mapreduce.MapReduce import MapReduce +from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero from keopscore.utils.code_gen_utils import ( load_vars, load_vars_chunks, @@ -15,7 +16,6 @@ Var_loader, use_pragma_unroll, ) -from keopscore.mapreduce.Chunk_Mode_Constants import Chunk_Mode_Constants def do_chunk_sub_ranges( @@ -194,7 +194,7 @@ def __init__(self, *args): self.chk = Chunk_Mode_Constants(self.red_formula) self.dimy = self.chk.dimy self.blocksize_chunks = min( - cuda_block_size, 1024, 49152 // max(1, self.dimy * sizeof(self.dtype)) + cuda.get_cuda_block_size(), 1024, 49152 // max(1, self.dimy * sizeof(self.dtype)) ) def get_code(self): diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py index f968ecd9f..e4d61f01c 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py @@ -1,10 +1,14 @@ -from keopscore import cuda_block_size -from keopscore.config.chunks import dimfinalchunk from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile +from keopscore.config import cuda +from keopscore.config.chunks import dimfinalchunk from keopscore.formulas.reductions.Sum_Reduction import Sum_Reduction -from keopscore.formulas.reductions.sum_schemes import * -from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero +from keopscore.formulas.reductions.sum_schemes import ( + block_sum, + kahan_scheme, + direct_sum, +) from keopscore.mapreduce.MapReduce import MapReduce +from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero from keopscore.utils.code_gen_utils import ( load_vars, load_vars_chunks, @@ -14,7 +18,7 @@ Var_loader, use_pragma_unroll, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error def do_finalchunk_sub_ranges( @@ -141,7 +145,7 @@ def get_code(self): self.dimy = max(dimfinalchunk, dimy) blocksize_chunks = min( - cuda_block_size, 1024, 49152 // max(1, self.dimy * sizeof(self.dtype)) + cuda.get_cuda_block_size(), 1024, 49152 // max(1, self.dimy * sizeof(self.dtype)) ) if not isinstance(sum_scheme, block_sum): diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc2D.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc2D.py index 6abc89fb9..552fa85fc 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc2D.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc2D.py @@ -12,7 +12,7 @@ use_pragma_unroll, c_zero_float, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error class GpuReduc2D(MapReduce, Gpu_link_compile): diff --git a/keopscore/keopscore/mapreduce/gpu/__init__.py b/keopscore/keopscore/mapreduce/gpu/__init__.py index cb53e5618..d74c11a76 100644 --- a/keopscore/keopscore/mapreduce/gpu/__init__.py +++ b/keopscore/keopscore/mapreduce/gpu/__init__.py @@ -1,10 +1,23 @@ from .GpuAssignZero import GpuAssignZero + from .GpuReduc1D import GpuReduc1D from .GpuReduc1D_chunks import GpuReduc1D_chunks from .GpuReduc1D_finalchunks import GpuReduc1D_finalchunks from .GpuReduc1D_ranges import GpuReduc1D_ranges from .GpuReduc1D_ranges_chunks import GpuReduc1D_ranges_chunks -from .GpuReduc1D_ranges_finalchunks import ( - GpuReduc1D_ranges_finalchunks, -) +from .GpuReduc1D_ranges_finalchunks import GpuReduc1D_ranges_finalchunks + from .GpuReduc2D import GpuReduc2D + +_exports = [ + GpuAssignZero, + GpuReduc1D, + GpuReduc1D_chunks, + GpuReduc1D_finalchunks, + GpuReduc1D_ranges, + GpuReduc1D_ranges_chunks, + GpuReduc1D_ranges_finalchunks, + GpuReduc2D +] + +__all__ = [cls.__name__ for cls in _exports] \ No newline at end of file diff --git a/keopscore/keopscore/sandbox/do_clean_keops.py b/keopscore/keopscore/sandbox/do_clean_keops.py index 4fe078e82..70b357ac8 100644 --- a/keopscore/keopscore/sandbox/do_clean_keops.py +++ b/keopscore/keopscore/sandbox/do_clean_keops.py @@ -1,3 +1,3 @@ -from keopscore.utils.code_gen_utils import clean_keops +from keopscore.config import clean_keops clean_keops() diff --git a/keopscore/keopscore/sandbox/formula.py b/keopscore/keopscore/sandbox/formula.py index 03e93c095..8ec57b364 100644 --- a/keopscore/keopscore/sandbox/formula.py +++ b/keopscore/keopscore/sandbox/formula.py @@ -1,10 +1,9 @@ # testing some formulas with keopscore +import keopscore.config from keopscore.formulas import * -import keopscore - -keopscore.auto_factorize = True +keopscore.config.auto_factorize = True print("********************************") print("test 1") diff --git a/keopscore/keopscore/sandbox/laplacian.py b/keopscore/keopscore/sandbox/laplacian.py index 4404eb9d3..7792212e5 100644 --- a/keopscore/keopscore/sandbox/laplacian.py +++ b/keopscore/keopscore/sandbox/laplacian.py @@ -3,8 +3,8 @@ from time import time from keopscore.formulas import * -# import keopscore -# keopscore.debug_ops = True +import keopscore +keopscore.config.debug.set_debug_ops(True) def GaussLapKernel(sigma, D): diff --git a/keopscore/keopscore/test/test_op.py b/keopscore/keopscore/test/test_op.py index 34512a12d..1aed2eb83 100644 --- a/keopscore/keopscore/test/test_op.py +++ b/keopscore/keopscore/test/test_op.py @@ -8,7 +8,7 @@ import keopscore import keopscore.formulas -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from pykeops.torch import Genred # fix seed for reproducibility diff --git a/keopscore/keopscore/utils/Cache.py b/keopscore/keopscore/utils/Cache.py index 0a76ed597..e82552508 100644 --- a/keopscore/keopscore/utils/Cache.py +++ b/keopscore/keopscore/utils/Cache.py @@ -1,14 +1,18 @@ import os import pickle + import keopscore -from keopscore.config import * +import keopscore.config # global configuration parameter to be added for the lookup : # N.B we turn this into a function because the parameters need to be read dynamically. env_param = ( - lambda: keopscore.config.get_cpp_flags() - + " auto_factorize=" - + str(keopscore.auto_factorize) + lambda: keopscore.config.cxx.get_compile_options() + + keopscore.config.cxx.get_linking_options() + + " auto_factorize=" + str(keopscore.config.auto_factorize) + + keopscore.config.cuda.get_include_options() + + keopscore.config.cuda.get_preprocessing_options() + + keopscore.config.cuda.get_nvrtc_flags() ) diff --git a/keopscore/keopscore/utils/TestFormula.py b/keopscore/keopscore/utils/TestFormula.py index 21857d45d..bbff20bb4 100644 --- a/keopscore/keopscore/utils/TestFormula.py +++ b/keopscore/keopscore/utils/TestFormula.py @@ -5,7 +5,7 @@ from pykeops.torch import Genred from keopscore.formulas import * import types -from keopscore.utils.misc_utils import KeOps_Print +from keopscore.utils.messages import KeOps_Print def TestFormula(formula, tol=1e-4, dtype="float32", test_grad=False, randseed=None): diff --git a/keopscore/keopscore/utils/TestOperation.py b/keopscore/keopscore/utils/TestOperation.py index 56c8ee355..b8ce7c7e6 100644 --- a/keopscore/keopscore/utils/TestOperation.py +++ b/keopscore/keopscore/utils/TestOperation.py @@ -5,7 +5,7 @@ from pykeops.torch import Genred from keopscore.formulas import * import types -from keopscore.utils.misc_utils import KeOps_Print +from keopscore.utils.messages import KeOps_Print def TestOperation(op_str, tol=1e-4, dtype="float32", test_grad=True): diff --git a/keopscore/keopscore/utils/Tree.py b/keopscore/keopscore/utils/Tree.py index 9bd3f84f7..ffc608d78 100644 --- a/keopscore/keopscore/utils/Tree.py +++ b/keopscore/keopscore/utils/Tree.py @@ -1,4 +1,4 @@ -from keopscore.utils.misc_utils import KeOps_Print +from keopscore.utils.messages import KeOps_Print class Tree: diff --git a/keopscore/keopscore/utils/code_gen_utils.py b/keopscore/keopscore/utils/code_gen_utils.py index 10647c239..1102e0d4e 100644 --- a/keopscore/keopscore/utils/code_gen_utils.py +++ b/keopscore/keopscore/utils/code_gen_utils.py @@ -1,11 +1,7 @@ -import os from hashlib import sha256 -import keopscore -from keopscore.config import * -from keopscore.utils.misc_utils import KeOps_Error, KeOps_Message - -disable_pragma_unrolls = config.get_disable_pragma_unrolls() +import keopscore.config +from keopscore.utils.messages import KeOps_Error, KeOps_Message def get_hash_name(*args): @@ -216,7 +212,7 @@ def __getitem__(self, other): def use_pragma_unroll(n=64): - if disable_pragma_unrolls: + if keopscore.config.cxx.get_disable_pragma_unrolls(): return "\n" else: if n is None: @@ -827,51 +823,3 @@ def varseq_to_array(vars, vars_ptr_name): """ return string - -def clean_keops(recompile_jit_binary=True, verbose=True): - - build_path = config.get_build_folder() - use_cuda = keopscore.cuda_config.get_use_cuda() - if use_cuda: - jit_binary = config.get_jit_binary() - else: - jit_binary = None - for f in os.scandir(build_path): - if recompile_jit_binary or f.path != jit_binary: - os.remove(f.path) - if verbose: - KeOps_Message(f"{build_path} has been cleaned.") - from keopscore.get_keops_dll import get_keops_dll - - get_keops_dll.reset() - if use_cuda and recompile_jit_binary: - from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile - - Gpu_link_compile.compile_jit_compile_dll() - - -def check_health(config_type="all"): - """ - Check the health of the specified configuration. - - Parameters: - config_type (str): The configuration to check. Options are: - 'cuda', 'openmp', 'platform', 'base', 'all'. - Default is 'all'. - """ - if config_type == "cuda": - cuda_config.print_all() - elif config_type == "openmp": - openmp_config.print_all() - elif config_type == "platform": - platform_detector.print_all() - elif config_type == ("base"): - config.print_all() - elif config_type == "all": - platform_detector.print_all() - config.print_all() - cuda_config.print_all() - openmp_config.print_all() - else: - print(f"Unknown configuration type: '{config_type}'") - print("Please specify one of: 'cuda', 'openmp', 'platform', 'base', 'all'") diff --git a/keopscore/keopscore/utils/file_utils.py b/keopscore/keopscore/utils/file_utils.py new file mode 100644 index 000000000..a352bc37e --- /dev/null +++ b/keopscore/keopscore/utils/file_utils.py @@ -0,0 +1,42 @@ +import os +import re + + +def file_to_string(file_path): + """Read a text file and return its content as a string.""" + with open(file_path, "r", encoding="utf-8") as file: + return file.read() + + +def string_to_file(string, file_path): + """Write a string to a text file.""" + with open(file_path, "w", encoding="utf-8") as file: + file.write(string) + + +def pack_header(filename, origin_folder, target_folder): + """ + Produce a stand-alone C/C++ header by inlining local quoted includes. + """ + code = file_to_string(os.path.join(origin_folder, filename)) + used_headers = [] + while True: + match = re.search('#include *"([^"]*)"', code) + if match is None: + break + header = match.groups()[0] + if header in used_headers: + code_to_insert = "" + else: + if not os.path.exists(os.path.join(origin_folder, header)): + header_found = os.path.basename(header) + elif os.path.exists(os.path.join(origin_folder, header)): + header_found = header + else: + raise FileNotFoundError( + f"Header file {header} not found in {origin_folder}." + ) + code_to_insert = file_to_string(os.path.join(origin_folder, header_found)) + used_headers.append(header) + code = code[: match.start()] + code_to_insert + code[match.end() :] + string_to_file(code, os.path.join(target_folder, filename)) diff --git a/keopscore/keopscore/utils/gpu_utils.py b/keopscore/keopscore/utils/gpu_utils.py index 878e94f4a..6569bf983 100644 --- a/keopscore/keopscore/utils/gpu_utils.py +++ b/keopscore/keopscore/utils/gpu_utils.py @@ -1,101 +1,10 @@ -import ctypes -from ctypes.util import find_library - - -from keopscore.utils.misc_utils import ( - KeOps_Error, - KeOps_Warning, - find_library_abspath, - KeOps_OS_Run, - get_include_file_abspath, -) - -import keopscore -from keopscore.config import * - import os -from os.path import join - -# Some constants taken from cuda.h -CUDA_SUCCESS = 0 -CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK = 1 -CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK = 8 - -# warning : libcuda.so is shipped with nvidia drivers (usually system-wide) -# but nvrtc is shipped with cuda-toolkit-dev (may be user installed) -libcuda_folder = os.path.dirname(find_library_abspath("cuda")) -libnvrtc_folder = os.path.dirname(find_library_abspath("nvrtc")) - - -def get_cuda_include_path(): - # auto detect location of cuda headers - - # First we look at CUDA_PATH env variable if it is set - path = os.getenv("CUDA_PATH") - if path: - path = join(path, "include") - if os.path.isfile(join(path, "cuda.h")) and os.path.isfile( - join(path, "nvrtc.h") - ): - return path - - # if not successfull, we try a few standard locations: - cuda_paths_to_try_start = [] - # if user has installed cuda toolkit via conda in the current env, - # we will find it via CONDA_PREFIX environment variable - path_conda = os.getenv("CONDA_PREFIX") - if path_conda is not None: - cuda_paths_to_try_start.append(path_conda) - cuda_version = get_cuda_version(out_type="string") - cuda_paths_to_try_start += [ - join(os.path.sep, "opt", "cuda"), - join(os.path.sep, "usr", "local", "cuda"), - join(os.path.sep, "usr", "local", f"cuda-{cuda_version}"), - ] - - cuda_paths_to_try_end = [ - "include", - join("targets", "x86_64-linux", "include"), - ] - for path_start in cuda_paths_to_try_start: - for path_end in cuda_paths_to_try_end: - path = join(path_start, path_end) - if os.path.isfile(join(path, "cuda.h")) and os.path.isfile( - join(path, "nvrtc.h") - ): - return path - - # if not successfull, we try to infer location from the libs - cuda_include_path = None - for libpath in libcuda_folder, libnvrtc_folder: - for libtag in "lib", "lib64": - libtag = os.path.sep + libtag + os.path.sep - if libtag in libpath: - includetag = os.path.sep + "include" + os.path.sep - includepath = libpath.replace(libtag, includetag) - if os.path.isfile(join(includepath, "cuda.h")) and os.path.isfile( - join(includepath, "nvrtc.h") - ): - return includepath - - # last try, testing if by any chance the header is already in the default - # include path of gcc - path_cudah = get_include_file_abspath("cuda.h", config.cxx_compiler()) - if path_cudah: - path = os.path.dirname(path_cudah) - if os.path.isfile(join(path, "nvrtc.h")): - return path - - # finally nothing found, so we display a warning asking the user to do something - KeOps_Warning(""" - The location of Cuda header files cuda.h and nvrtc.h could not be detected on your system. - You must determine their location and then define the environment variable CUDA_PATH, - either before launching Python or using os.environ before importing keops. For example - if these files are in /vol/cuda/10.2.89-cudnn7.6.4.38/include you can do : - import os - os.environ['CUDA_PATH'] = '/vol/cuda/10.2.89-cudnn7.6.4.38' - """) +import keopscore.config +from keopscore.utils.file_utils import pack_header +from keopscore.utils.messages import KeOps_Error +from keopscore.utils.path_utils import _first_matching_file +from keopscore.utils.system_utils import get_include_file_abspath def orig_cuda_include_fp16_path(): @@ -103,12 +12,24 @@ def orig_cuda_include_fp16_path(): We look for float 16 cuda headers cuda_fp16.h and cuda_fp16.hpp based on cuda_path locations and return their directory """ - cuda_include_path = cuda_config.get_cuda_include_path() - if cuda_include_path: - return cuda_include_path - cuda_fp16_h_abspath = cuda_config.get_include_file_abspath("cuda_fp16.h") - cuda_fp16_hpp_abspath = cuda_config.get_include_file_abspath("cuda_fp16.hpp") + # First try to find the library file using the cuda includes + cuda_include_path = keopscore.config.cuda.get_cuda_include_path() + cuda_fp16_h_abspath = _first_matching_file( + cuda_include_path, + "cuda_fp16.h", + ) + cuda_fp16_hpp_abspath =_first_matching_file( + cuda_include_path, + "cuda_fp16.hpp", + ) + + if cuda_fp16_h_abspath and cuda_fp16_hpp_abspath: + return os.path.dirname(cuda_fp16_h_abspath) + + # Second try with compiler + cuda_fp16_h_abspath = get_include_file_abspath("cuda_fp16.h") + cuda_fp16_hpp_abspath = get_include_file_abspath("cuda_fp16.hpp") if cuda_fp16_h_abspath and cuda_fp16_hpp_abspath: path = os.path.dirname(cuda_fp16_h_abspath) if path != os.path.dirname(cuda_fp16_hpp_abspath): @@ -129,112 +50,11 @@ def custom_cuda_include_fp16_path(): does not work. Hence we produce a packed stand-alone version of cuda_fp16.h by replacing all #include statements by the corresponding headers contents. """ - from keopscore.utils.misc_utils import pack_header - build_folder = config.get_build_folder() + build_folder = keopscore.config.path.get_build_folder() fp16_header = "cuda_fp16.h" - fp16_header_path = join(build_folder, fp16_header) + fp16_header_path = os.path.join(build_folder, fp16_header) if not os.path.isfile(fp16_header_path): pack_header(fp16_header, orig_cuda_include_fp16_path(), build_folder) return build_folder - -def get_cuda_version(out_type="single_value"): - cuda = ctypes.CDLL(find_library("cudart")) - cuda_version = ctypes.c_int() - cuda.cudaRuntimeGetVersion(ctypes.byref(cuda_version)) - cuda_version = int(cuda_version.value) - if out_type == "single_value": - return cuda_version - cuda_version_major = cuda_version // 1000 - cuda_version_minor = (cuda_version - (1000 * cuda_version_major)) // 10 - if out_type == "major,minor": - return cuda_version_major, cuda_version_minor - elif out_type == "string": - return f"{cuda_version_major}.{cuda_version_minor}" - - -def get_gpu_props(): - """ - Return number of GPU by reading libcuda. - Here we assume the system has cuda support (more precisely that libcuda can be loaded) - Adapted from https://gist.github.com/f0k/0d6431e3faa60bffc788f8b4daa029b1 - credit: Jan Schlüter - """ - cuda = ctypes.CDLL(find_library("cuda")) - - nGpus = ctypes.c_int() - error_str = ctypes.c_char_p() - - result = cuda.cuInit(0) - if result != CUDA_SUCCESS: - # cuda.cuGetErrorString(result, ctypes.byref(error_str)) - # KeOps_Warning("cuInit failed with error code %d: %s" % (result, error_str.value.decode())) - KeOps_Warning( - "cuda was detected, but driver API could not be initialized. Switching to cpu only." - ) - return 0, "" - - result = cuda.cuDeviceGetCount(ctypes.byref(nGpus)) - if result != CUDA_SUCCESS: - # cuda.cuGetErrorString(result, ctypes.byref(error_str)) - # KeOps_Warning("cuDeviceGetCount failed with error code %d: %s" % (result, error_str.value.decode())) - KeOps_Warning( - "cuda was detected, driver API has been initialized, but no working GPU has been found. Switching to cpu only." - ) - return 0, "" - - nGpus = nGpus.value - - def safe_call(d, result): - test = result == CUDA_SUCCESS - if not test: - KeOps_Warning(f""" - cuda was detected, driver API has been initialized, - but there was an error for detecting properties of GPU device nr {d}. - Switching to cpu only. - """) - return test - - test = True - MaxThreadsPerBlock = [0] * (nGpus) - SharedMemPerBlock = [0] * (nGpus) - for d in range(nGpus): - # getting handle to cuda device - device = ctypes.c_int() - result &= safe_call(d, cuda.cuDeviceGet(ctypes.byref(device), ctypes.c_int(d))) - - # getting MaxThreadsPerBlock info for device - output = ctypes.c_int() - result &= safe_call( - d, - cuda.cuDeviceGetAttribute( - ctypes.byref(output), - ctypes.c_int(CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK), - device, - ), - ) - MaxThreadsPerBlock[d] = output.value - - # getting SharedMemPerBlock info for device - result &= safe_call( - d, - cuda.cuDeviceGetAttribute( - ctypes.byref(output), - ctypes.c_int(CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK), - device, - ), - ) - SharedMemPerBlock[d] = output.value - - # Building compile flags in the form "-D..." options for further compilations - # (N.B. the purpose is to avoid the device query at runtime because it would slow down computations) - string_flags = f"-DMAXIDGPU={nGpus-1} " - for d in range(nGpus): - string_flags += f"-DMAXTHREADSPERBLOCK{d}={MaxThreadsPerBlock[d]} " - string_flags += f"-DSHAREDMEMPERBLOCK{d}={SharedMemPerBlock[d]} " - - if test: - return nGpus, string_flags - else: - return 0, 0, "" diff --git a/keopscore/keopscore/utils/math_functions.py b/keopscore/keopscore/utils/math_functions.py index e883d5072..b93aa2b1a 100644 --- a/keopscore/keopscore/utils/math_functions.py +++ b/keopscore/keopscore/utils/math_functions.py @@ -3,7 +3,7 @@ new_c_varname, c_variable, ) -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error import keopscore from keopscore.config import * diff --git a/keopscore/keopscore/utils/messages.py b/keopscore/keopscore/utils/messages.py new file mode 100644 index 000000000..290b0973c --- /dev/null +++ b/keopscore/keopscore/utils/messages.py @@ -0,0 +1,40 @@ +import os +import sys + + +def _keops_verbose(): + config = sys.modules.get("keopscore.config") + debug = getattr(config, "debug", None) + if debug is not None: + return debug.get_verbose() + + keopscore = sys.modules.get("keopscore") + return getattr(keopscore, "verbose", os.getenv("KEOPS_VERBOSE") != "0") + + +def KeOps_Print(*messages, force_print=False, **kwargs): + if _keops_verbose() or force_print: + print(*messages, **kwargs) + + +def KeOps_Message(message, use_tag=True, **kwargs): + if _keops_verbose(): + tag = "[KeOps] " if use_tag else "" + message = tag + message + print(message, **kwargs) + + +def KeOps_Warning(message, newline=False): + if _keops_verbose(): + message = ("\n" if newline else "") + "[KeOps] Warning : " + message + print(message) + + +def KeOps_Error(message, show_line_number=True): + message = "[KeOps] Error : " + message + if show_line_number: + from inspect import currentframe, getframeinfo + + frameinfo = getframeinfo(currentframe().f_back) + message += f" (error at line {frameinfo.lineno} in file {frameinfo.filename})" + raise ValueError(message) diff --git a/keopscore/keopscore/utils/misc_utils.py b/keopscore/keopscore/utils/misc_utils.py deleted file mode 100644 index bced3371d..000000000 --- a/keopscore/keopscore/utils/misc_utils.py +++ /dev/null @@ -1,166 +0,0 @@ -####################################################################### -# . Warnings, Errors, etc. -####################################################################### - -import keopscore -from os.path import join -import re -import platform - - -def KeOps_Print(*messages, force_print=False, **kwargs): - if keopscore.verbose or force_print: - print(*messages, **kwargs) - - -def KeOps_Message(message, use_tag=True, **kwargs): - if keopscore.verbose: - tag = "[KeOps] " if use_tag else "" - message = tag + message - print(message, **kwargs) - - -def KeOps_Warning(message, newline=False): - if keopscore.verbose: - message = ("\n" if newline else "") + "[KeOps] Warning : " + message - print(message) - - -def KeOps_Error(message, show_line_number=True): - message = "[KeOps] Error : " + message - if show_line_number: - from inspect import currentframe, getframeinfo - - frameinfo = getframeinfo(currentframe().f_back) - message += f" (error at line {frameinfo.lineno} in file {frameinfo.filename})" - raise ValueError(message) - - -def KeOps_OS_Run(command, print_warning=True): - import sys - - python_version = sys.version_info - if python_version >= (3, 7): - import subprocess - - out = subprocess.run(command, shell=True, capture_output=True) - if out.stderr != b"" and print_warning: - KeOps_Warning("There were warnings or errors :", newline=True) - KeOps_Print(out.stderr.decode("utf-8")) - elif python_version >= (3, 5): - import subprocess - - out = subprocess.run( - command, - shell=True, - ) - else: - KeOps_Error("Python version >= 3.5 required.", newline=True) - return out - - -def get_include_file_abspath(filename, compiler): - out = KeOps_OS_Run( - f'echo "#include <{filename}>" | {compiler} -M -E -x c++ - | head -n 2' - ) - strings = out.stdout.decode("utf8").split() - abspath = None - for s in strings: - if filename in s: - abspath = s - return abspath - - -def get_brew_prefix(): - """Get Homebrew prefix path using KeOps_OS_Run""" - if platform.system() != "Darwin": - return None - out = KeOps_OS_Run(f"brew --prefix", print_warning=False) - if out.stderr != b"": - return None - return out.stdout.decode("utf-8").strip() - - -def find_library_abspath(lib): - """ - wrapper around ctypes find_library that returns the full path - of the library. - Warning : it also opens the shared library ! - Adapted from - https://stackoverflow.com/questions/35682600/get-absolute-path-of-shared-library-in-python/35683698 - """ - from ctypes import c_int, c_void_p, c_char_p, CDLL, byref, cast, POINTER, Structure - from ctypes.util import find_library - - # linkmap structure, we only need the second entry - class LINKMAP(Structure): - _fields_ = [("l_addr", c_void_p), ("l_name", c_char_p)] - - res = find_library(lib) - if res is None: - return "" - - lib = CDLL(res) - libdl = CDLL(find_library("dl")) - - dlinfo = libdl.dlinfo - dlinfo.argtypes = c_void_p, c_int, c_void_p - dlinfo.restype = c_int - - # gets typecasted later, I dont know how to create a ctypes struct pointer instance - lmptr = c_void_p() - - # 2 equals RTLD_DI_LINKMAP, pass pointer by reference - dlinfo(lib._handle, 2, byref(lmptr)) - - # typecast to a linkmap pointer and retrieve the name. - abspath = cast(lmptr, POINTER(LINKMAP)).contents.l_name - - return abspath.decode("utf-8") - - -def file_to_string(file_path): - """ - reads text file and returns its content as a string - """ - f = open(file_path, "r", encoding="utf-8") - out = f.read() - f.close() - return out - - -def string_to_file(string, file_path): - """ - writes string to file - """ - f = open(file_path, "w", encoding="utf-8") - out = f.write(string) - f.close() - - -def pack_header(filename, origin_folder, target_folder): - """ - Given a C/C++ header file "filename", located in folder "origin_folder", - produces a stand-alone version of the header by recursiveley - replacing all #include "xxx" statements by the content of the corresponding - file. The resulting file is put into "target_folder". - """ - code = file_to_string(join(origin_folder, filename)) - used_headers = [] - while True: - match = re.search('#include *"([^"]*)"', code) - if match is None: - break - header = match.groups()[0] - if header in used_headers: - code_to_insert = "" - else: - code_to_insert = file_to_string(join(origin_folder, header)) - used_headers.append(header) - code = code[: match.start()] + code_to_insert + code[match.end() :] - string_to_file(code, join(target_folder, filename)) - - -# Factoring the check/cross marks used in the config classes prints -CHECK_MARK = "✅" -CROSS_MARK = "❌" diff --git a/keopscore/keopscore/utils/path_utils.py b/keopscore/keopscore/utils/path_utils.py new file mode 100644 index 000000000..ed21fc170 --- /dev/null +++ b/keopscore/keopscore/utils/path_utils.py @@ -0,0 +1,98 @@ +import os +import site +import sys +import sysconfig +from pathlib import Path + + +def _unique_paths(paths): + """Return paths in first-seen order with duplicates and None values removed.""" + unique = [] + seen = set() + for path in paths: + if path is None: + continue + path = Path(path) + key = str(path) + if key not in seen: + seen.add(key) + unique.append(path) + return unique + + +def _env_roots(env_vars): + """Return existing environment variable values as normalized Path objects.""" + return _unique_paths(os.getenv(env_var) for env_var in env_vars) + + +def _path_candidates(roots, suffixes): + """Expand root directories with relative suffixes while preserving order.""" + candidates = [] + for root in _unique_paths(roots): + for suffix in suffixes: + candidates.append(root / suffix if suffix else root) + return _unique_paths(candidates) + + +def _python_package_roots(): + """Return Python package roots where pip-installed wheels may live.""" + package_roots = [] + getters = [ + lambda: getattr(site, "getsitepackages", lambda: [])(), + lambda: [site.getusersitepackages()], + lambda: [sysconfig.get_path("purelib")], + lambda: [sysconfig.get_path("platlib")], + lambda: sys.path, + ] + for getter in getters: + try: + # Some site helpers are unavailable in embedded or non-standard Python builds. + package_roots.extend(getter() or []) + except Exception: + continue + return _unique_paths(package_roots) + + +def _ordered_search_roots( + env_vars=(), pip_suffixes=(), conda_root=None, system_roots=() +): + """Return roots ordered by explicit env vars, pip, conda, then system paths.""" + roots = [] + + if env_vars: + roots.extend(_env_roots(env_vars)) + + if pip_suffixes: + # Pip wheels such as nvidia-cuda-runtime expose libraries under site-packages. + roots.extend(_path_candidates(_python_package_roots(), pip_suffixes)) + + if conda_root: + roots.extend(_env_roots((conda_root,))) + roots.extend(system_roots) + return _unique_paths(roots) + + +def _first_matching_file(directories, patterns): + """Return the first file matching glob patterns inside ordered directories.""" + if not directories or not patterns: + return None + + if isinstance(patterns, str): + patterns = (patterns,) + patterns = [pattern for pattern in patterns if pattern] + if not patterns: + return None + + for directory in _unique_paths(directories): + try: + directory = Path(directory) + if not directory.is_dir(): + continue + for pattern in patterns: + for match in sorted(directory.glob(pattern)): + if match.is_file(): + return str(match) + except Exception: + continue + + return None diff --git a/keopscore/keopscore/utils/system_utils.py b/keopscore/keopscore/utils/system_utils.py new file mode 100644 index 000000000..82392b1bd --- /dev/null +++ b/keopscore/keopscore/utils/system_utils.py @@ -0,0 +1,72 @@ +import os +import subprocess +from ctypes import c_int, c_void_p, c_char_p, CDLL, byref, cast, POINTER, Structure +from ctypes.util import find_library + +from keopscore.utils.messages import KeOps_Print, KeOps_Warning + + +def KeOps_OS_Run(command, print_warning=True): + out = subprocess.run(command, shell=True, capture_output=True) + if out.stderr != b"" and print_warning: + KeOps_Warning("There were warnings or errors while executing: " + command, newline=True) + KeOps_Print(out.stderr.decode("utf-8")) + return out + + +def get_include_file_abspath(filename, compiler): + """Return the full path of the header filename using compiler.""" + + cmd = f'echo "#include <{filename}>" | {compiler} -M -E -x c++ -' + out = KeOps_OS_Run(cmd) + + text = out.stdout.decode("utf8") + text = text.replace("\\\n", " ") + + try: + deps = text.split(":", 1)[1] + except IndexError: + return None + + for path in deps.split(): + if os.path.basename(path) == filename: + return path + + return None + + +def find_library_abspath(lib): + """ + Wrapper around ctypes find_library that returns the full path of the library. + Warning: it also opens the shared library. + """ + + class LINKMAP(Structure): + _fields_ = [("l_addr", c_void_p), ("l_name", c_char_p)] + + res = find_library(lib) + if res is None: + return "" + + lib = CDLL(res) + libdl = CDLL(find_library("dl")) + + dlinfo = libdl.dlinfo + dlinfo.argtypes = c_void_p, c_int, c_void_p + dlinfo.restype = c_int + + lmptr = c_void_p() + dlinfo(lib._handle, 2, byref(lmptr)) + + abspath = cast(lmptr, POINTER(LINKMAP)).contents.l_name + return abspath.decode("utf-8") + + +def _find_library_by_names(library_names): + """Return the first library path resolved by ctypes for known library names.""" + for library_name in library_names: + library_path = find_library(library_name) + if library_path: + return find_library_abspath(library_name) + + return None diff --git a/pykeops/pykeops/__init__.py b/pykeops/pykeops/__init__.py index 9d08d6e12..68910620d 100644 --- a/pykeops/pykeops/__init__.py +++ b/pykeops/pykeops/__init__.py @@ -1,4 +1,5 @@ import os + import keopscore ############################################################## @@ -10,17 +11,12 @@ from . import config as pykeopsconfig -from keopscore import show_cuda_status - -keops_get_build_folder = pykeopsconfig.pykeops_base.get_build_folder -from .config import pykeops_nvrtc_name -from .config import numpy_found, torch_found def set_verbose(val): global verbose verbose = val - keopscore.verbose = val + keopscore.config.debug.set_verbose(val) ########################################################### @@ -37,13 +33,12 @@ def set_verbose(val): default_device_id = 0 # default Gpu device number -if pykeopsconfig.pykeops_cuda.get_use_cuda(): - if not os.path.exists(pykeops_nvrtc_name(type="target")): +if pykeopsconfig.cuda.get_use_cuda(): + if not os.path.exists(pykeopsconfig.pykeops_nvrtc_name(type="target")): from .common.keops_io.LoadKeOps_nvrtc import compile_jit_binary compile_jit_binary() - def clean_pykeops(recompile_jit_binaries=True): r""" This function cleans the KeOps cache and recompiles the JIT binaries if necessary. @@ -53,31 +48,29 @@ def clean_pykeops(recompile_jit_binaries=True): """ import pykeops - keopscore.clean_keops(recompile_jit_binary=recompile_jit_binaries) + keopscore.config.clean_keops(recompile_jit_binary=recompile_jit_binaries) keops_binder = pykeops.common.keops_io.keops_binder for key in keops_binder: keops_binder[key].reset() - if recompile_jit_binaries and pykeopsconfig.pykeops_cuda.get_use_cuda(): + if recompile_jit_binaries and pykeopsconfig.cuda.get_use_cuda(): pykeops.common.keops_io.LoadKeOps_nvrtc.compile_jit_binary() -def check_health(): +def check_health(infos="all"): r""" Runs a complete sanity check of the KeOps installation within your system. This function verifies the setup and configuration of KeOps, including compilation flags, paths, .... Parameters: - config_type (str): The configuration to check. Options are: - 'base', 'cuda', 'openmp', 'platform', 'all'. + infos (str): The configuration to check. Options are: + 'cuda', 'cxx', 'openmp', 'platform', 'path', 'all'. Default is 'all'. Returns: None """ - import pykeops - - keopscore.check_health() + keopscore.config.check_health(infos=infos) def set_build_folder(path=None): @@ -87,20 +80,20 @@ def set_build_folder(path=None): keops_binder = pykeops.common.keops_io.keops_binder for key in keops_binder: keops_binder[key].reset(new_save_folder=get_build_folder()) - if pykeopsconfig.pykeops_cuda.get_use_cuda() and not os.path.exists( - pykeops.config.pykeops_nvrtc_name(type="target") + if pykeopsconfig.cuda.get_use_cuda() and not os.path.exists( + pykeopsconfig.pykeops_nvrtc_name(type="target") ): pykeops.common.keops_io.LoadKeOps_nvrtc.compile_jit_binary() def get_build_folder(): - return keops_get_build_folder() + return pykeopsconfig.get_build_folder() -if numpy_found: +if pykeopsconfig.numpy_found: from .numpy.test_install import test_numpy_bindings -if torch_found: +if pykeopsconfig.torch_found: from .torch.test_install import test_torch_bindings # next line is to ensure that cache file for formulas is loaded at import diff --git a/pykeops/pykeops/common/get_options.py b/pykeops/pykeops/common/get_options.py index 084371304..ce4cb144b 100644 --- a/pykeops/pykeops/common/get_options.py +++ b/pykeops/pykeops/common/get_options.py @@ -1,8 +1,8 @@ import re import numpy as np from collections import OrderedDict -import pykeops -import pykeops.config + +import pykeops.config as pykeopsconfig ############################################################ # define backend @@ -53,7 +53,7 @@ def define_tag_backend(self, backend, variables): # auto : infer everything if backend == "auto": return ( - int(pykeops.config.gpu_available), + int(pykeopsconfig.gpu_available), self._find_grid(), self._find_mem(variables), ) @@ -84,7 +84,7 @@ def define_backend(self, backend, variables): @staticmethod def _find_dev(): - return int(pykeops.config.gpu_available) + return int(pykeopsconfig.gpu_available) @staticmethod def _find_mem(variables): @@ -92,7 +92,7 @@ def _find_mem(variables): [type(var) is np.ndarray for var in variables] ): # Infer if we're working with numpy arrays or torch tensors: MemType = 0 - elif pykeops.config.torch_found: + elif pykeopsconfig.torch_found: import torch if all( diff --git a/pykeops/pykeops/common/gpu_utils.py b/pykeops/pykeops/common/gpu_utils.py index 100995c0a..d5c45638c 100644 --- a/pykeops/pykeops/common/gpu_utils.py +++ b/pykeops/pykeops/common/gpu_utils.py @@ -1,5 +1,4 @@ -from keopscore.utils.gpu_utils import get_gpu_props - +import pykeops.config as pykeopsconfig def get_gpu_number(): - return get_gpu_props()[0] + return pykeopsconfig.cuda.get_n_gpus() diff --git a/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py b/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py index 7b632f3c9..8c2a0b309 100644 --- a/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py +++ b/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py @@ -3,13 +3,11 @@ import pykeops.config as pykeopsconfig -get_build_folder = pykeopsconfig.pykeops_base.get_build_folder - from keopscore.utils.Cache import Cache_partial +from keopscore.utils.system_utils import KeOps_OS_Run + from pykeops.common.keops_io.LoadKeOps import LoadKeOps from pykeops.common.utils import pyKeOps_Message -from keopscore.utils.misc_utils import KeOps_OS_Run -from pykeops.config import pykeops_cpp_name, python_includes class LoadKeOps_cpp_class(LoadKeOps): @@ -17,9 +15,9 @@ def __init__(self, *args, fast_init=False): super().__init__(*args, fast_init=fast_init) def init_phase1(self): - srcname = pykeops_cpp_name(tag=self.params.tag, extension=".cpp") + srcname = pykeopsconfig.pykeops_cpp_name(tag=self.params.tag, extension=".cpp") - dllname = pykeops_cpp_name( + dllname = pykeopsconfig.pykeops_cpp_name( tag=self.params.tag, extension=sysconfig.get_config_var("EXT_SUFFIX") ) @@ -27,7 +25,7 @@ def init_phase1(self): f = open(srcname, "w") f.write(self.get_pybind11_code()) f.close() - compile_command = f"{pykeopsconfig.pykeops_base.get_cxx_compiler()} {pykeopsconfig.pykeops_base.get_cpp_flags()} {python_includes} {srcname} -o {dllname}" + compile_command = f"{pykeopsconfig.cxx.get_cxx_compiler()} {pykeopsconfig.cxx.get_compile_options()} {pykeopsconfig.cxx.get_include_options()} {pykeopsconfig.path.get_include_options()} {pykeopsconfig.python_includes} {srcname} {pykeopsconfig.cxx.get_linking_options()} -o {dllname}" pyKeOps_Message( "Compiling pykeops cpp " + self.params.tag + " module ... ", flush=True, @@ -40,7 +38,7 @@ def init_phase2(self): import importlib mylib = importlib.import_module( - os.path.basename(pykeops_cpp_name(tag=self.params.tag)) + os.path.basename(pykeopsconfig.pykeops_cpp_name(tag=self.params.tag)) ) self.launch_keops_cpu = mylib.launch_pykeops_cpu @@ -183,7 +181,9 @@ def get_pybind11_code(self): LoadKeOps_cpp = Cache_partial( - LoadKeOps_cpp_class, use_cache_file=True, save_folder=get_build_folder() + LoadKeOps_cpp_class, + use_cache_file=True, + save_folder=pykeopsconfig.path.get_build_folder() ) cpp_dtype = { diff --git a/pykeops/pykeops/common/keops_io/LoadKeOps_nvrtc.py b/pykeops/pykeops/common/keops_io/LoadKeOps_nvrtc.py index 95ef25489..d51187637 100644 --- a/pykeops/pykeops/common/keops_io/LoadKeOps_nvrtc.py +++ b/pykeops/pykeops/common/keops_io/LoadKeOps_nvrtc.py @@ -1,14 +1,14 @@ import os import sys -import pykeops +import pykeops.config as pykeopsconfig + from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile from keopscore.utils.Cache import Cache_partial +from keopscore.utils.system_utils import KeOps_OS_Run + from pykeops.common.keops_io.LoadKeOps import LoadKeOps from pykeops.common.utils import pyKeOps_Message -from keopscore.utils.misc_utils import KeOps_OS_Run - -get_build_folder = pykeops.config.pykeops_base.get_build_folder class LoadKeOps_nvrtc_class(LoadKeOps): @@ -18,10 +18,10 @@ def __init__(self, *args, fast_init=False): def init_phase2(self): import importlib - if pykeops.get_build_folder() not in sys.path: + if pykeopsconfig.get_build_folder() not in sys.path: # The build folder is supposed to be in the python path, if not, # we add it - sys.path.append(pykeops.get_build_folder()) + sys.path.append(pykeopsconfig.get_build_folder()) pykeops_nvrtc = importlib.import_module("pykeops_nvrtc") @@ -80,9 +80,9 @@ def compile_jit_binary(): This function compile the main .so entry point to keops_nvrt binder... """ compile_command = Gpu_link_compile.get_compile_command( - extra_flags=pykeops.config.python_includes, - sourcename=pykeops.config.pykeops_nvrtc_name(type="src"), - dllname=pykeops.config.pykeops_nvrtc_name(type="target"), + extra_flags=pykeopsconfig.python_includes, + sourcename=pykeopsconfig.pykeops_nvrtc_name(type="src"), + dllname=pykeopsconfig.pykeops_nvrtc_name(type="target"), ) pyKeOps_Message("Compiling nvrtc binder for python ... ", flush=True, end="") KeOps_OS_Run(compile_command) @@ -92,5 +92,5 @@ def compile_jit_binary(): LoadKeOps_nvrtc = Cache_partial( LoadKeOps_nvrtc_class, use_cache_file=True, - save_folder=get_build_folder(), + save_folder=pykeopsconfig.get_build_folder(), ) diff --git a/pykeops/pykeops/common/keops_io/__init__.py b/pykeops/pykeops/common/keops_io/__init__.py index b8b2eb851..15271bf6e 100644 --- a/pykeops/pykeops/common/keops_io/__init__.py +++ b/pykeops/pykeops/common/keops_io/__init__.py @@ -1,6 +1,6 @@ import pykeops.config as pykeopsconfig -if pykeopsconfig.pykeops_cuda.get_use_cuda(): +if pykeopsconfig.cuda.get_use_cuda(): from . import LoadKeOps_nvrtc, LoadKeOps_cpp keops_binder = { diff --git a/pykeops/pykeops/common/lazy_tensor.py b/pykeops/pykeops/common/lazy_tensor.py index 8c381d881..51356f800 100644 --- a/pykeops/pykeops/common/lazy_tensor.py +++ b/pykeops/pykeops/common/lazy_tensor.py @@ -5,7 +5,7 @@ import numpy as np -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from pykeops.common.utils import check_broadcasting diff --git a/pykeops/pykeops/common/operations.py b/pykeops/pykeops/common/operations.py index 9b51bb9f1..ace359afe 100644 --- a/pykeops/pykeops/common/operations.py +++ b/pykeops/pykeops/common/operations.py @@ -1,6 +1,6 @@ import numpy as np -from keopscore.utils.misc_utils import KeOps_Print, KeOps_Warning +from keopscore.utils.messages import KeOps_Print, KeOps_Warning from pykeops.common.utils import get_tools diff --git a/pykeops/pykeops/config.py b/pykeops/pykeops/config.py index ab14cebea..2c9a44c1e 100644 --- a/pykeops/pykeops/config.py +++ b/pykeops/pykeops/config.py @@ -1,7 +1,7 @@ import importlib.util +import os import sys import sysconfig -from os.path import join, dirname, realpath ############################################################### # Initialize some variables: the values may be redefined later @@ -9,26 +9,26 @@ numpy_found = importlib.util.find_spec("numpy") is not None torch_found = importlib.util.find_spec("torch") is not None -import keopscore -from keopscore.config import * - # Instantiating the keopscore.config main classes for pykeops -pykeops_cuda = cuda_config -pykeops_openmp = openmp_config -pykeops_base = config +import keopscore.config + +cuda = keopscore.config.cuda +path = keopscore.config.path +openmp = keopscore.config.openmp +cxx = keopscore.config.cxx -get_build_folder = pykeops_base.get_build_folder -gpu_available = pykeops_cuda.get_use_cuda() +get_build_folder = path.get_build_folder +gpu_available = cuda.get_use_cuda() def pykeops_nvrtc_name(type="src"): basename = "pykeops_nvrtc" extension = ".cpp" if type == "src" else sysconfig.get_config_var("EXT_SUFFIX") - return join( + return os.path.join( ( - join(dirname(realpath(__file__)), "common", "keops_io") + os.path.join(os.path.dirname(os.path.realpath(__file__)), "common", "keops_io") if type == "src" - else config.get_build_folder() + else get_build_folder() ), basename + extension, ) @@ -36,8 +36,8 @@ def pykeops_nvrtc_name(type="src"): def pykeops_cpp_name(tag="", extension=""): basename = "pykeops_cpp_" - return join( - config.get_build_folder(), + return os.path.join( + get_build_folder(), basename + tag + extension, ) diff --git a/pykeops/pykeops/numpy/utils.py b/pykeops/pykeops/numpy/utils.py index 0322596e1..d54518009 100644 --- a/pykeops/pykeops/numpy/utils.py +++ b/pykeops/pykeops/numpy/utils.py @@ -1,5 +1,5 @@ import numpy as np -import pykeops.config +import pykeops.config as pykeopsconfig from pykeops.numpy import Genred, KernelSolve from pykeops.numpy.cluster import swap_axes as np_swap_axes @@ -236,7 +236,7 @@ def WarmUpGpu(): from pykeops.common.utils import pyKeOps_Message pyKeOps_Message("Warming up the Gpu (numpy bindings) !!!") - if pykeops.config.gpu_available: + if pykeopsconfig.gpu_available: formula = "Exp(-oos2*SqDist(x,y))*b" aliases = [ "x = Vi(1)", # First arg : i-variable, of size 1 diff --git a/pykeops/pykeops/sandbox/test_soft_dtw_kernel_dissmatrix.py b/pykeops/pykeops/sandbox/test_soft_dtw_kernel_dissmatrix.py index 46cb12097..eeab3f0b7 100644 --- a/pykeops/pykeops/sandbox/test_soft_dtw_kernel_dissmatrix.py +++ b/pykeops/pykeops/sandbox/test_soft_dtw_kernel_dissmatrix.py @@ -122,7 +122,7 @@ def diss(i, j): #################################################################### from keopscore.formulas.Operation import Operation -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.utils.code_gen_utils import ( c_variable, pointer, @@ -166,7 +166,7 @@ def DiffT(self, v, gradin): from keopscore.formulas.Operation import Operation from keopscore.formulas.variables.Zero import Zero -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.utils.code_gen_utils import ( c_variable, pointer, diff --git a/pykeops/pykeops/sandbox/test_soft_dtw_kernel_v0.py b/pykeops/pykeops/sandbox/test_soft_dtw_kernel_v0.py index 7e3df470c..2aa8ec433 100644 --- a/pykeops/pykeops/sandbox/test_soft_dtw_kernel_v0.py +++ b/pykeops/pykeops/sandbox/test_soft_dtw_kernel_v0.py @@ -53,7 +53,7 @@ def SoftDTW_torch(x, y, gamma): ################################## from keopscore.formulas.Operation import Operation -from keopscore.utils.misc_utils import KeOps_Error +from keopscore.utils.messages import KeOps_Error from keopscore.utils.code_gen_utils import ( c_variable, pointer, diff --git a/pykeops/pykeops/sandbox/test_torch_func_hessian.py b/pykeops/pykeops/sandbox/test_torch_func_hessian.py index f91422fc5..7e574e2dd 100644 --- a/pykeops/pykeops/sandbox/test_torch_func_hessian.py +++ b/pykeops/pykeops/sandbox/test_torch_func_hessian.py @@ -1,9 +1,8 @@ +import keopscore import torch from pykeops.torch import LazyTensor -import keopscore - -keopscore.auto_factorize = False +keopscore.config.auto_factorize = False def fn_torch(x_i): diff --git a/pykeops/pykeops/test/test_torch_func.py b/pykeops/pykeops/test/test_torch_func.py index a6657c56a..c518cbe1a 100644 --- a/pykeops/pykeops/test/test_torch_func.py +++ b/pykeops/pykeops/test/test_torch_func.py @@ -1,11 +1,10 @@ +import keopscore import torch from pykeops.torch import LazyTensor -import keopscore - torch.manual_seed(0) -keopscore.auto_factorize = False +keopscore.config.auto_factorize = False B1, B2, M, N, D = 5, 4, 10, 20, 2 From 5cb80e005423afa0035c1bc3ac5fae5fa3886523 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Fri, 1 May 2026 13:04:28 +0200 Subject: [PATCH 18/98] fix openmp linking --- keopscore/keopscore/binders/LinkCompile.py | 2 +- keopscore/keopscore/config/OpenMP.py | 51 +++++++++---------- keopscore/keopscore/mapreduce/MapReduce.py | 2 +- .../mapreduce/gpu/GpuReduc1D_chunks.py | 2 + .../mapreduce/gpu/GpuReduc1D_finalchunks.py | 5 +- .../mapreduce/gpu/GpuReduc1D_ranges_chunks.py | 4 +- .../gpu/GpuReduc1D_ranges_finalchunks.py | 9 ++-- .../pykeops/common/keops_io/LoadKeOps_cpp.py | 2 +- 8 files changed, 39 insertions(+), 38 deletions(-) diff --git a/keopscore/keopscore/binders/LinkCompile.py b/keopscore/keopscore/binders/LinkCompile.py index b452e6a48..a94f0c01a 100644 --- a/keopscore/keopscore/binders/LinkCompile.py +++ b/keopscore/keopscore/binders/LinkCompile.py @@ -33,7 +33,7 @@ def __init__(self): self.use_half, self.use_fast_math, self.device_id, - cpp_flag, # TODO: check that get_envs is sufficient... + cpp_flag, ) # info_file is the name of the file that will contain some meta-information required by the bindings, e.g. 7b9a611f7e.nfo diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 9c13e4b9d..c47c02825 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -21,9 +21,9 @@ class OpenMPConfig: _openmp_lib_name = None _openmp_lib_include_dir = None - _openmp_compile_options = "" - _openmp_include_options = "" - _openmp_linking_options = "" + _compile_options = "" + _include_options = "" + _linking_options = "" _openmp_header_basename = "omp.h" @@ -66,17 +66,12 @@ def __init__(self, platform, cxx_compiler): self.openmp_system_suffixes += [self.platform.get_brew_prefix(),] if self.platform.get_brew_prefix() else [] self.set_openmplib_path() - self.set_openmp_compile_options() - self.set_openmp_include_options() - self.set_openmp_linking_options() + self.set_compile_options() + self.set_include_options() + self.set_linking_options() self.set_use_OpenMP() - if self.get_use_OpenMP(): - self.cxx_compiler.add_to_compile_option(self.get_openmp_compile_options()) - self.cxx_compiler.add_to_include_options(self.get_openmp_include_options()) - self.cxx_compiler.add_to_linking_options(self.get_openmp_linking_options()) - # OpenMP library path def set_openmplib_path(self): """try to locate OpenMP libraries""" @@ -206,11 +201,11 @@ def check_compiler_for_openmp(self): compile_command = [ self.cxx_compiler.get_cxx_compiler(), test_file, - self.get_openmp_include_options(), - self.get_openmp_compile_options(), + self.get_include_options(), + self.get_compile_options(), ] - if self.get_openmp_linking_options(): - compile_command.append(self.get_openmp_linking_options()) + if self.get_linking_options(): + compile_command.append(self.get_linking_options()) compile_command.extend(["-o", f"{test_file}.out"]) try: @@ -224,29 +219,29 @@ def check_compiler_for_openmp(self): return False # C++ Compiler Options - def set_openmp_compile_options(self): + def set_compile_options(self): # Add special fix for openMP prgama and Apple Clang. Order matters. if self.cxx_compiler.get_use_Apple_clang() and self.platform.get_brew_prefix(): - self._openmp_compile_options += "-Xpreprocessor " + self._compile_options += "-Xpreprocessor " - self._openmp_compile_options += "-fopenmp" + self._compile_options += "-fopenmp" - def get_openmp_compile_options(self): - return self._openmp_compile_options + def get_compile_options(self): + return self._compile_options # C++ Compiler Include Options - def set_openmp_include_options(self): - self._openmp_include_options = f'-I{self.get_openmp_include_dir()}' + def set_include_options(self): + self._include_options = f'-I{self.get_openmp_include_dir()}' - def get_openmp_include_options(self): - return self._openmp_include_options + def get_include_options(self): + return self._include_options # C++ linking Options - def set_openmp_linking_options(self): - self._openmp_linking_options += f'-L{self.get_openmp_lib_dir()}' + def set_linking_options(self): + self._linking_options += f'-L{self.get_openmp_lib_dir()}' - def get_openmp_linking_options(self): - return self._openmp_linking_options + def get_linking_options(self): + return self._linking_options # OpenMP configuration printing def print_all(self): diff --git a/keopscore/keopscore/mapreduce/MapReduce.py b/keopscore/keopscore/mapreduce/MapReduce.py index 8fa17aeae..d824cb3d7 100644 --- a/keopscore/keopscore/mapreduce/MapReduce.py +++ b/keopscore/keopscore/mapreduce/MapReduce.py @@ -1,4 +1,4 @@ -from keopscore.formulas.reductions import * +from keopscore.formulas.reductions import block_sum, kahan_scheme, direct_sum from keopscore.formulas.GetReduction import GetReduction from keopscore.utils.code_gen_utils import ( Var_loader, diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py index 0d2617e92..db8295012 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py @@ -13,6 +13,8 @@ table, table4, use_pragma_unroll, + c_variable, + c_array, ) diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py index 1d2374a70..da12b3b4c 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py @@ -2,7 +2,7 @@ from keopscore.config import cuda from keopscore.config.chunks import dimfinalchunk from keopscore.formulas.reductions.Sum_Reduction import Sum_Reduction -from keopscore.formulas.reductions.sum_schemes import * +from keopscore.formulas.reductions.sum_schemes import block_sum from keopscore.mapreduce.MapReduce import MapReduce from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero from keopscore.utils.code_gen_utils import ( @@ -12,6 +12,9 @@ pointer, Var_loader, use_pragma_unroll, + c_variable, + c_zero_float, + c_array, ) from keopscore.utils.messages import KeOps_Error diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py index 5a964ac5b..ac4242ba4 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py @@ -1,7 +1,7 @@ from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile from keopscore.config import cuda from keopscore.config.chunks import dimchunk -from keopscore.formulas.reductions.sum_schemes import * +from keopscore.formulas.reductions.sum_schemes import block_sum, kahan_scheme, direct_sum from keopscore.mapreduce.Chunk_Mode_Constants import Chunk_Mode_Constants from keopscore.mapreduce.MapReduce import MapReduce from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero @@ -15,6 +15,8 @@ table4, Var_loader, use_pragma_unroll, + c_array, + c_variable, ) diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py index e4d61f01c..a54609689 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py @@ -2,11 +2,7 @@ from keopscore.config import cuda from keopscore.config.chunks import dimfinalchunk from keopscore.formulas.reductions.Sum_Reduction import Sum_Reduction -from keopscore.formulas.reductions.sum_schemes import ( - block_sum, - kahan_scheme, - direct_sum, -) +from keopscore.formulas.reductions.sum_schemes import block_sum from keopscore.mapreduce.MapReduce import MapReduce from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero from keopscore.utils.code_gen_utils import ( @@ -17,6 +13,9 @@ pointer, Var_loader, use_pragma_unroll, + c_array, + c_zero_float, + c_variable, ) from keopscore.utils.messages import KeOps_Error diff --git a/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py b/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py index 8c2a0b309..16fb70ba2 100644 --- a/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py +++ b/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py @@ -25,7 +25,7 @@ def init_phase1(self): f = open(srcname, "w") f.write(self.get_pybind11_code()) f.close() - compile_command = f"{pykeopsconfig.cxx.get_cxx_compiler()} {pykeopsconfig.cxx.get_compile_options()} {pykeopsconfig.cxx.get_include_options()} {pykeopsconfig.path.get_include_options()} {pykeopsconfig.python_includes} {srcname} {pykeopsconfig.cxx.get_linking_options()} -o {dllname}" + compile_command = f"{pykeopsconfig.cxx.get_cxx_compiler()} {pykeopsconfig.cxx.get_compile_options()} {pykeopsconfig.openmp.get_compile_options()} {pykeopsconfig.openmp.get_include_options()} {pykeopsconfig.path.get_include_options()} {pykeopsconfig.python_includes} {srcname} {pykeopsconfig.cxx.get_linking_options()} {pykeopsconfig.openmp.get_linking_options()} -o {dllname}" pyKeOps_Message( "Compiling pykeops cpp " + self.params.tag + " module ... ", flush=True, From 1e6758b87d273491b6bd184b3691a3c71b1a7fd6 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Fri, 1 May 2026 13:24:26 +0200 Subject: [PATCH 19/98] remove sum_scheme eval --- .../keopscore/formulas/reductions/__init__.py | 5 +++-- .../keopscore/formulas/reductions/sum_schemes.py | 16 ++++++++++++++++ keopscore/keopscore/mapreduce/MapReduce.py | 4 ++-- .../keopscore/mapreduce/gpu/GpuReduc1D_chunks.py | 6 ++++-- .../mapreduce/gpu/GpuReduc1D_ranges_chunks.py | 6 ++++-- 5 files changed, 29 insertions(+), 8 deletions(-) diff --git a/keopscore/keopscore/formulas/reductions/__init__.py b/keopscore/keopscore/formulas/reductions/__init__.py index 4c057081d..909625ca8 100644 --- a/keopscore/keopscore/formulas/reductions/__init__.py +++ b/keopscore/keopscore/formulas/reductions/__init__.py @@ -14,7 +14,7 @@ from .Min_Reduction import Min_Reduction from .Sum_Reduction import Sum_Reduction from .Zero_Reduction import Zero_Reduction -from .sum_schemes import block_sum, kahan_scheme, direct_sum +from .sum_schemes import block_sum, kahan_scheme, direct_sum, make_sum_scheme _exports = [ Reduction, @@ -33,7 +33,8 @@ Zero_Reduction, block_sum, kahan_scheme, - direct_sum + direct_sum, + make_sum_scheme, ] diff --git a/keopscore/keopscore/formulas/reductions/sum_schemes.py b/keopscore/keopscore/formulas/reductions/sum_schemes.py index 437581a94..f1c7be7cf 100644 --- a/keopscore/keopscore/formulas/reductions/sum_schemes.py +++ b/keopscore/keopscore/formulas/reductions/sum_schemes.py @@ -5,6 +5,12 @@ c_variable, ) +""" +This module defines the different schemes for performing the reduction of a formula. + +The schemes must be added in the _SUM_SCHEME_CLASSES dictionary, and must be subclasses of Sum_Scheme, which defines the interface for the different schemes. +""" + class Sum_Scheme: def __init__(self, red_formula, dtype, dimred=None): @@ -84,3 +90,13 @@ def initialize_temporary_accumulator_first_init(self): def accumulate_result(self, acc, fout, j, hack=False): return self.red_formula.KahanScheme(acc, fout, self.tmp_acc) + +_SUM_SCHEME_CLASSES = { + "direct_sum": direct_sum, + "block_sum": block_sum, + "kahan_scheme": kahan_scheme, +} + +def make_sum_scheme(sum_scheme_string, red_formula, dtype, dimred=None): + sum_scheme_class = _SUM_SCHEME_CLASSES.get(sum_scheme_string, None) + return sum_scheme_class(red_formula, dtype, dimred=dimred) \ No newline at end of file diff --git a/keopscore/keopscore/mapreduce/MapReduce.py b/keopscore/keopscore/mapreduce/MapReduce.py index d824cb3d7..70e48df65 100644 --- a/keopscore/keopscore/mapreduce/MapReduce.py +++ b/keopscore/keopscore/mapreduce/MapReduce.py @@ -1,4 +1,4 @@ -from keopscore.formulas.reductions import block_sum, kahan_scheme, direct_sum +from keopscore.formulas.reductions import make_sum_scheme from keopscore.formulas.GetReduction import GetReduction from keopscore.utils.code_gen_utils import ( Var_loader, @@ -63,7 +63,7 @@ def get_code(self): dtype = self.dtype dtypeacc = self.dtypeacc nargs = self.nargs - self.sum_scheme = eval(self.sum_scheme_string)(red_formula, dtype) + self.sum_scheme = make_sum_scheme(self.sum_scheme_string, red_formula, dtype) self.i = i = c_variable("signed long int", "i") self.j = j = c_variable("signed long int", "j") diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py index db8295012..6b59778ea 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py @@ -1,7 +1,7 @@ from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile from keopscore.config import cuda from keopscore.config.chunks import dimchunk -from keopscore.formulas.reductions.sum_schemes import * +from keopscore.formulas.reductions import make_sum_scheme from keopscore.mapreduce.Chunk_Mode_Constants import Chunk_Mode_Constants from keopscore.mapreduce.MapReduce import MapReduce from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero @@ -155,7 +155,9 @@ def get_code(self): chk = self.chk param_loc = c_array(dtype, chk.dimp, "param_loc") acc = c_array(dtypeacc, chk.dimred, "acc") - sum_scheme = eval(self.sum_scheme_string)(red_formula, dtype, dimred=chk.dimred) + sum_scheme = make_sum_scheme( + self.sum_scheme_string, red_formula, dtype, dimred=chk.dimred + ) xi = c_array(dtype, chk.dimx, "xi") fout_chunk = c_array( dtype, self.blocksize_chunks * chk.dimout_chunk, "fout_chunk" diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py index ac4242ba4..419df228b 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py @@ -1,7 +1,7 @@ from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile from keopscore.config import cuda from keopscore.config.chunks import dimchunk -from keopscore.formulas.reductions.sum_schemes import block_sum, kahan_scheme, direct_sum +from keopscore.formulas.reductions import make_sum_scheme from keopscore.mapreduce.Chunk_Mode_Constants import Chunk_Mode_Constants from keopscore.mapreduce.MapReduce import MapReduce from keopscore.mapreduce.gpu.GpuAssignZero import GpuAssignZero @@ -247,7 +247,9 @@ def get_code(self): chk = self.chk param_loc = c_array(dtype, chk.dimp, "param_loc") acc = c_array(dtypeacc, chk.dimred, "acc") - sum_scheme = eval(self.sum_scheme_string)(red_formula, dtype, dimred=chk.dimred) + sum_scheme = make_sum_scheme( + self.sum_scheme_string, red_formula, dtype, dimred=chk.dimred + ) xi = c_array(dtype, chk.dimx, "xi") fout_chunk = c_array( dtype, self.blocksize_chunks * chk.dimout_chunk, "fout_chunk" From c6f29c16d4d97a7e9fff245b2d1d0f968d84a630 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Fri, 1 May 2026 13:27:24 +0200 Subject: [PATCH 20/98] linting --- keopscore/keopscore/__init__.py | 2 +- keopscore/keopscore/binders/LinkCompile.py | 1 + .../binders/nvrtc/Gpu_link_compile.py | 13 ++- keopscore/keopscore/config/Cuda.py | 51 ++++++------ keopscore/keopscore/config/Debug.py | 11 +-- keopscore/keopscore/config/KeOpsPath.py | 18 +++-- keopscore/keopscore/config/OpenMP.py | 27 +++++-- keopscore/keopscore/config/Platform.py | 4 +- keopscore/keopscore/config/__init__.py | 1 + keopscore/keopscore/config/chunks.py | 2 +- keopscore/keopscore/formulas/Operation.py | 6 +- .../keopscore/formulas/autodiff/__init__.py | 2 +- .../keopscore/formulas/complex/__init__.py | 2 +- .../keopscore/formulas/maths/__init__.py | 80 ++++++++++++++++--- .../keopscore/formulas/reductions/__init__.py | 2 +- .../formulas/reductions/sum_schemes.py | 4 +- .../keopscore/formulas/variables/__init__.py | 2 +- keopscore/keopscore/get_keops_dll.py | 2 +- keopscore/keopscore/mapreduce/__init__.py | 1 + keopscore/keopscore/mapreduce/cpu/__init__.py | 9 ++- .../mapreduce/gpu/GpuReduc1D_chunks.py | 4 +- .../mapreduce/gpu/GpuReduc1D_finalchunks.py | 4 +- .../mapreduce/gpu/GpuReduc1D_ranges_chunks.py | 4 +- .../gpu/GpuReduc1D_ranges_finalchunks.py | 4 +- keopscore/keopscore/mapreduce/gpu/__init__.py | 4 +- keopscore/keopscore/sandbox/laplacian.py | 1 + keopscore/keopscore/utils/Cache.py | 11 +-- keopscore/keopscore/utils/code_gen_utils.py | 1 - keopscore/keopscore/utils/gpu_utils.py | 19 +++-- keopscore/keopscore/utils/system_utils.py | 4 +- pykeops/pykeops/__init__.py | 1 + pykeops/pykeops/common/gpu_utils.py | 3 +- .../pykeops/common/keops_io/LoadKeOps_cpp.py | 2 +- pykeops/pykeops/common/operations.py | 6 +- pykeops/pykeops/config.py | 4 +- pykeops/pykeops/test/test_numpy.py | 3 +- 36 files changed, 209 insertions(+), 106 deletions(-) diff --git a/keopscore/keopscore/__init__.py b/keopscore/keopscore/__init__.py index 8ede817db..8da592cf9 100644 --- a/keopscore/keopscore/__init__.py +++ b/keopscore/keopscore/__init__.py @@ -18,4 +18,4 @@ Gpu_link_compile.compile_jit_compile_dll() # expose to the user -set_build_folder = keopscore.config.path.set_different_build_folder +set_build_folder = keopscore.config.path.set_different_build_folder diff --git a/keopscore/keopscore/binders/LinkCompile.py b/keopscore/keopscore/binders/LinkCompile.py index a94f0c01a..6f2f9b744 100644 --- a/keopscore/keopscore/binders/LinkCompile.py +++ b/keopscore/keopscore/binders/LinkCompile.py @@ -9,6 +9,7 @@ cpp_flag += keopscore.config.cuda.get_nvrtc_flags() cpp_flag += keopscore.config.cuda.get_include_options() + class LinkCompile: """ Base class for compiling the map_reduce schemes and providing the dll to KeOps bindings. diff --git a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py index 13e447ea8..81ba9587b 100644 --- a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py +++ b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py @@ -8,7 +8,6 @@ from keopscore.utils.messages import KeOps_Error, KeOps_Message from keopscore.utils.system_utils import KeOps_OS_Run - jit_source_file = os.path.join( keopscore.config.path.get_base_dir_path(), "binders", "nvrtc", "keops_nvrtc.cpp" ) @@ -27,11 +26,15 @@ def jit_compile_dll(): class Gpu_link_compile(LinkCompile): source_code_extension = "cu" - low_level_code_prefix = "cubin_" if keopscore.config.cuda.get_cuda_version() >= 11010 else "ptx_" + low_level_code_prefix = ( + "cubin_" if keopscore.config.cuda.get_cuda_version() >= 11010 else "ptx_" + ) def __init__(self): # checking that the system has a Gpu : - if not (keopscore.config.cuda.get_use_cuda() and keopscore.config.cuda.get_n_gpus()): + if not ( + keopscore.config.cuda.get_use_cuda() and keopscore.config.cuda.get_n_gpus() + ): KeOps_Error( "Trying to compile cuda code... but we detected that the system has no properly configured cuda lib." ) @@ -79,7 +82,9 @@ def generate_code(self): @staticmethod def get_compile_command( - sourcename=jit_source_file, dllname=keopscore.config.path.get_jit_binary(), extra_flags="" + sourcename=jit_source_file, + dllname=keopscore.config.path.get_jit_binary(), + extra_flags="", ): # This is about the main KeOps binary (dll) that will be used to JIT compile all formulas. # If the dll is not present, it compiles it from source, except if check_compile is False. diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index 9664f9741..88f3ab075 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -150,7 +150,7 @@ def find_install_path(self, lib_dict_info, warn=None): result["library"] = _first_matching_file( _path_candidates(candidate_roots, self.library_suffixes), - result["lib_basename_candidate"] + result["lib_basename_candidate"], ) if result["library"] is None: result["library"] = _find_library_by_names((result["name"],)) @@ -167,11 +167,7 @@ def find_install_path(self, lib_dict_info, warn=None): (result["header_basename"],), ) - if ( - result["header"] is None - and warn - and result["header_basename"] is not None - ): + if result["header"] is None and warn and result["header_basename"] is not None: KeOps_Warning(f"{result['name']} header files not found.") return result @@ -199,18 +195,20 @@ def _find_and_load_libcuda(self): False, "libcuda was detected, but driver API could not be initialized. Rebooting the system may help. Switching to CPU only.", ) - + # If we successfully loaded libcuda and initialized it, store the handle in the config for potential future use self._libcuda_info["ctype_handle"] = libcuda nGpus = ctypes.c_int() - if libcuda.cuDeviceGetCount(ctypes.byref(nGpus)) != self.CUDA_SUCCESS or nGpus.value == 0: + if ( + libcuda.cuDeviceGetCount(ctypes.byref(nGpus)) != self.CUDA_SUCCESS + or nGpus.value == 0 + ): return ( False, "libcuda was detected and driver API was initialized, but no working GPU found. Switching to CPU only.", ) - self._MaxThreadsPerBlock = [0] * nGpus.value self._SharedMemPerBlock = [0] * nGpus.value @@ -225,9 +223,9 @@ def _find_and_load_libcuda(self): + err_msg + " Switching to CPU only.", ) - + self._n_gpus = nGpus.value - + return True, "" def _find_and_load_libnvrtc(self): @@ -247,10 +245,10 @@ def _find_and_load_libnvrtc(self): False, f"Failed to load library '{os.path.basename(libnvrtc_path)}': {e}", ) - + # If we successfully loaded libnvrtc, store the handle in the config self._libnvrtc_info["ctype_handle"] = libnvrtc_handle - + return True, "" def _find_and_load_cudart(self): @@ -269,7 +267,7 @@ def _find_and_load_cudart(self): False, "libcudart not found. Make sure the CUDA toolkit is installed and accessible. Switching to CPU only.", ) - + try: libcudart_handle = ctypes.CDLL(libcudart_path) except OSError as e: @@ -277,18 +275,21 @@ def _find_and_load_cudart(self): False, f"Failed to load '{os.path.basename(libcudart_path)}': {e}", ) - + cuda_version = ctypes.c_int() - if libcudart_handle.cudaRuntimeGetVersion(ctypes.byref(cuda_version)) != self.CUDA_SUCCESS: + if ( + libcudart_handle.cudaRuntimeGetVersion(ctypes.byref(cuda_version)) + != self.CUDA_SUCCESS + ): return ( False, "libcudart was found and loaded, but failed to get CUDA runtime version. Switching to CPU only.", ) - + # If we successfully loaded libcudart, store the handle in the config self._cudart_info["ctype_handle"] = libcudart_handle self._cuda_version = int(cuda_version.value) - + return True, "" def _cuda_libraries_available(self): @@ -348,7 +349,7 @@ def set_specific_gpus(self): """Set specific GPUs from CUDA_VISIBLE_DEVICES.""" if os.getenv("CUDA_VISIBLE_DEVICES"): self._specific_gpus = os.getenv("CUDA_VISIBLE_DEVICES").replace(",", "_") - + def get_specific_gpus(self): """Get the specific GPUs.""" return self._specific_gpus @@ -361,7 +362,7 @@ def set_n_gpus(self): def get_n_gpus(self): """Get the number of GPUs detected.""" return self._n_gpus - + def print_n_gpus(self): """Print the number of GPUs detected.""" print(f"Number of GPUs Detected: {self.get_n_gpus()}") @@ -482,10 +483,14 @@ def get_preprocessing_options(self): return self._preprocessing_options def print_preprocessing_options(self): - print(f"GPU Preprocessing Options: {self.get_preprocessing_options() or not_found_str}") + print( + f"GPU Preprocessing Options: {self.get_preprocessing_options() or not_found_str}" + ) - def set_include_options(self): - self._include_options += "".join(f" -I{p}" for p in list(set(self.get_cuda_include_path()))) + def set_include_options(self): + self._include_options += "".join( + f" -I{p}" for p in list(set(self.get_cuda_include_path())) + ) def get_include_options(self): return self._include_options diff --git a/keopscore/keopscore/config/Debug.py b/keopscore/keopscore/config/Debug.py index ee2610713..12cc46ccd 100644 --- a/keopscore/keopscore/config/Debug.py +++ b/keopscore/keopscore/config/Debug.py @@ -1,5 +1,6 @@ import os + class DebugConfig: # prints information about atomic operations during code building @@ -12,21 +13,21 @@ class DebugConfig: def __init__(self): pass - + def set_debug_ops(self, debug_ops): self._debug_ops = debug_ops - + def get_debug_ops(self): return self._debug_ops def set_debug_ops_at_exec(self, debug_ops_at_exec): - self._debug_ops_at_exec = debug_ops_at_exec + self._debug_ops_at_exec = debug_ops_at_exec def get_debug_ops_at_exec(self): return self._debug_ops_at_exec def get_verbose(self): return self._verbose - + def set_verbose(self, verbose): - self._verbose = verbose \ No newline at end of file + self._verbose = verbose diff --git a/keopscore/keopscore/config/KeOpsPath.py b/keopscore/keopscore/config/KeOpsPath.py index 951e32e44..32fc1fc63 100644 --- a/keopscore/keopscore/config/KeOpsPath.py +++ b/keopscore/keopscore/config/KeOpsPath.py @@ -15,13 +15,11 @@ class KeOpsPathConfig: _default_build_folder_name = None _default_build_path = None _build_folder = None - + _jit_binary = None _include_options = "" - path_env_vars = ( - "KEOPS_CACHE_FOLDER", - ) + path_env_vars = ("KEOPS_CACHE_FOLDER",) def __init__(self, platform, cuda): @@ -107,7 +105,9 @@ def set_different_build_folder( - reset_all: If True, reset all cached formulas and recompile necessary components. """ # If path is not given, we either read the save file or use the default build path - save_file = os.path.join(self.get_keops_cache_folder(), "build_folder_location.txt") + save_file = os.path.join( + self.get_keops_cache_folder(), "build_folder_location.txt" + ) if not path: if read_save_file and os.path.isfile(save_file): with open(save_file, "r") as f: @@ -155,7 +155,9 @@ def get_build_folder(self): def set_default_build_path(self): """Set the default build path.""" self._default_build_path = ensure_directory( - os.path.join(self.get_keops_cache_folder(), self.get_default_build_folder_name()), + os.path.join( + self.get_keops_cache_folder(), self.get_default_build_folder_name() + ), add_to_syspath=True, ) # Initialize _build_path : TODO : check ?/ @@ -181,7 +183,7 @@ def get_jit_binary(self): def print_jit_binary(self): """Print the path to the JIT binary.""" print(f"JIT Binary Path: {self.get_jit_binary() or not_found_str}") - + # include options management def set_include_options(self): """Set the include options for compilation.""" @@ -228,4 +230,4 @@ def print_all(self): cuda_info.print_all() keops_info = KeOpsPathConfig(platform_info, cuda_info) - keops_info.print_all() \ No newline at end of file + keops_info.print_all() diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index c47c02825..410331c7d 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -9,7 +9,10 @@ _ordered_search_roots, _path_candidates, ) -from keopscore.utils.system_utils import _find_library_by_names, get_include_file_abspath +from keopscore.utils.system_utils import ( + _find_library_by_names, + get_include_file_abspath, +) class OpenMPConfig: @@ -42,9 +45,9 @@ class OpenMPConfig: ] openmp_basename_candidate = ( - "libomp.dylib", - "libgomp.dylib", - "libomp.so", + "libomp.dylib", + "libgomp.dylib", + "libomp.so", ) openmp_library_suffixes = ( @@ -63,7 +66,13 @@ def __init__(self, platform, cxx_compiler): self.platform = platform self.cxx_compiler = cxx_compiler - self.openmp_system_suffixes += [self.platform.get_brew_prefix(),] if self.platform.get_brew_prefix() else [] + self.openmp_system_suffixes += ( + [ + self.platform.get_brew_prefix(), + ] + if self.platform.get_brew_prefix() + else [] + ) self.set_openmplib_path() self.set_compile_options() @@ -109,7 +118,9 @@ def find_openmp_install(self): # First try to find OpenMP library using standard names via ctypes. result["library"] = _find_library_by_names(("gomp", "omp")) - result["header"] = get_include_file_abspath(self._openmp_header_basename, self.cxx_compiler.get_cxx_compiler()) + result["header"] = get_include_file_abspath( + self._openmp_header_basename, self.cxx_compiler.get_cxx_compiler() + ) # If that fails, search for OpenMP headers and libraries in common locations. if not result["library"]: @@ -231,14 +242,14 @@ def get_compile_options(self): # C++ Compiler Include Options def set_include_options(self): - self._include_options = f'-I{self.get_openmp_include_dir()}' + self._include_options = f"-I{self.get_openmp_include_dir()}" def get_include_options(self): return self._include_options # C++ linking Options def set_linking_options(self): - self._linking_options += f'-L{self.get_openmp_lib_dir()}' + self._linking_options += f"-L{self.get_openmp_lib_dir()}" def get_linking_options(self): return self._linking_options diff --git a/keopscore/keopscore/config/Platform.py b/keopscore/keopscore/config/Platform.py index a9a63bb04..2ff1bbf08 100644 --- a/keopscore/keopscore/config/Platform.py +++ b/keopscore/keopscore/config/Platform.py @@ -130,7 +130,9 @@ def set_brew_prefix(self): return out = KeOps_OS_Run(f"brew --prefix", print_warning=False) - self._brew_prefix = out.stdout.decode("utf-8").strip() if out.stderr != b"" else None + self._brew_prefix = ( + out.stdout.decode("utf-8").strip() if out.stderr != b"" else None + ) def get_brew_prefix(self): """Get Homebrew prefix path using KeOps_OS_Run""" diff --git a/keopscore/keopscore/config/__init__.py b/keopscore/keopscore/config/__init__.py index ea0efe90f..cf9ded87f 100644 --- a/keopscore/keopscore/config/__init__.py +++ b/keopscore/keopscore/config/__init__.py @@ -63,6 +63,7 @@ def clean_keops(recompile_jit_binary=True, verbose=True): Gpu_link_compile.compile_jit_compile_dll() + __all__ = [ "platform", "cxx", diff --git a/keopscore/keopscore/config/chunks.py b/keopscore/keopscore/config/chunks.py index 531e52676..c1e597969 100644 --- a/keopscore/keopscore/config/chunks.py +++ b/keopscore/keopscore/config/chunks.py @@ -9,6 +9,7 @@ enable_chunk = True + def get_enable_chunk(): global enable_chunk return enable_chunk @@ -34,7 +35,6 @@ def set_enable_finalchunk(val): enable_final_chunk = False - def get_dimfinalchunk(): global dimfinalchunk return dimfinalchunk diff --git a/keopscore/keopscore/formulas/Operation.py b/keopscore/keopscore/formulas/Operation.py index 8f71ea066..663485761 100644 --- a/keopscore/keopscore/formulas/Operation.py +++ b/keopscore/keopscore/formulas/Operation.py @@ -78,7 +78,7 @@ def __call__(self, out, table): KeOps_Print("table=", table) for v in table: KeOps_Print(f"dim of {v} : ", v.dim) - if keopscore.config.debug.get_debug_ops_at_exec(): + if keopscore.config.debug.get_debug_ops_at_exec(): string += f'printf("\\n\\nComputing {self.__repr__()} :\\n");\n' args = [] # Evaluation of the child operations @@ -104,12 +104,12 @@ def __call__(self, out, table): string += self.Op(out, table, *args) # some debugging helper : - if keopscore.config.debug.get_debug_ops_at_exec(): + if keopscore.config.debug.get_debug_ops_at_exec(): for arg in args: string += arg.c_print string += out.c_print string += f'printf("\\n\\n");\n' - if keopscore.config.debug.get_debug_ops(): + if keopscore.config.debug.get_debug_ops(): KeOps_Print(f"Finished building code block for {self.__repr__()}") string += f"\n\n// Finished code block for {self.__repr__()}.\n}}\n\n" diff --git a/keopscore/keopscore/formulas/autodiff/__init__.py b/keopscore/keopscore/formulas/autodiff/__init__.py index 6452b9a96..f5685f9a1 100644 --- a/keopscore/keopscore/formulas/autodiff/__init__.py +++ b/keopscore/keopscore/formulas/autodiff/__init__.py @@ -14,4 +14,4 @@ Divergence, ] -__all__ = [cls.__name__ for cls in _exports] \ No newline at end of file +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/formulas/complex/__init__.py b/keopscore/keopscore/formulas/complex/__init__.py index b0bf68c61..6f4add17f 100644 --- a/keopscore/keopscore/formulas/complex/__init__.py +++ b/keopscore/keopscore/formulas/complex/__init__.py @@ -36,4 +36,4 @@ Real2Complex, ] -__all__ = [cls.__name__ for cls in _exports] \ No newline at end of file +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/formulas/maths/__init__.py b/keopscore/keopscore/formulas/maths/__init__.py index 5914ca5e4..ce39f4a04 100644 --- a/keopscore/keopscore/formulas/maths/__init__.py +++ b/keopscore/keopscore/formulas/maths/__init__.py @@ -67,19 +67,73 @@ from .XLogX import XLogX _exports = [ - Abs, Acos, Add, ArgMax, ArgMin, Asin, Atan, Atan2, BSpline, - Clamp, ClampInt, Concat, Cos, DiffClampInt, Divide, - Elem, ElemT, Equal, Exp, Extract, ExtractT, Floor, - GradMatrix, IfElse, IntInv, Inv, Kron, - LessOrEqual, LessThan, Log, MatVecMult, Max, Min, Minus, Mod, Mult, - Norm2, Normalize, NotEqual, OneHot, - Pow, Powf, ReLU, Round, Rsqrt, Scalprod, Sign, - Sin, SinXDivX, SoftDTW_SqDist, - SqDist, SqNorm2, SqNormDiag, SqNormIso, - Sqrt, Square, Step, Subtract, Sum, SumT, - SymSqNorm, TensorDot, TensorProd, - VecMatMult, WeightedSqDist, WeightedSqNorm, + Abs, + Acos, + Add, + ArgMax, + ArgMin, + Asin, + Atan, + Atan2, + BSpline, + Clamp, + ClampInt, + Concat, + Cos, + DiffClampInt, + Divide, + Elem, + ElemT, + Equal, + Exp, + Extract, + ExtractT, + Floor, + GradMatrix, + IfElse, + IntInv, + Inv, + Kron, + LessOrEqual, + LessThan, + Log, + MatVecMult, + Max, + Min, + Minus, + Mod, + Mult, + Norm2, + Normalize, + NotEqual, + OneHot, + Pow, + Powf, + ReLU, + Round, + Rsqrt, + Scalprod, + Sign, + Sin, + SinXDivX, + SoftDTW_SqDist, + SqDist, + SqNorm2, + SqNormDiag, + SqNormIso, + Sqrt, + Square, + Step, + Subtract, + Sum, + SumT, + SymSqNorm, + TensorDot, + TensorProd, + VecMatMult, + WeightedSqDist, + WeightedSqNorm, XLogX, ] -__all__ = [cls.__name__ for cls in _exports] \ No newline at end of file +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/formulas/reductions/__init__.py b/keopscore/keopscore/formulas/reductions/__init__.py index 909625ca8..5ad75016f 100644 --- a/keopscore/keopscore/formulas/reductions/__init__.py +++ b/keopscore/keopscore/formulas/reductions/__init__.py @@ -38,4 +38,4 @@ ] -__all__ = [cls.__name__ for cls in _exports] \ No newline at end of file +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/formulas/reductions/sum_schemes.py b/keopscore/keopscore/formulas/reductions/sum_schemes.py index f1c7be7cf..97fba9e46 100644 --- a/keopscore/keopscore/formulas/reductions/sum_schemes.py +++ b/keopscore/keopscore/formulas/reductions/sum_schemes.py @@ -91,12 +91,14 @@ def initialize_temporary_accumulator_first_init(self): def accumulate_result(self, acc, fout, j, hack=False): return self.red_formula.KahanScheme(acc, fout, self.tmp_acc) + _SUM_SCHEME_CLASSES = { "direct_sum": direct_sum, "block_sum": block_sum, "kahan_scheme": kahan_scheme, } + def make_sum_scheme(sum_scheme_string, red_formula, dtype, dimred=None): sum_scheme_class = _SUM_SCHEME_CLASSES.get(sum_scheme_string, None) - return sum_scheme_class(red_formula, dtype, dimred=dimred) \ No newline at end of file + return sum_scheme_class(red_formula, dtype, dimred=dimred) diff --git a/keopscore/keopscore/formulas/variables/__init__.py b/keopscore/keopscore/formulas/variables/__init__.py index d39491a44..2e2bc23ba 100644 --- a/keopscore/keopscore/formulas/variables/__init__.py +++ b/keopscore/keopscore/formulas/variables/__init__.py @@ -10,4 +10,4 @@ Zero, ] -__all__ = [cls.__name__ for cls in _exports] \ No newline at end of file +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/get_keops_dll.py b/keopscore/keopscore/get_keops_dll.py index 91a4fd10c..3cd47fca5 100644 --- a/keopscore/keopscore/get_keops_dll.py +++ b/keopscore/keopscore/get_keops_dll.py @@ -112,7 +112,7 @@ def get_keops_dll_impl( rf = map_reduce_obj.red_formula - if keopscore.config.debug.get_debug_ops(): + if keopscore.config.debug.get_debug_ops(): KeOps_Print("In get_keops_dll, formula is :", rf) KeOps_Print("formula.__repr__() is : ", rf.__repr__()) rf.make_dot() diff --git a/keopscore/keopscore/mapreduce/__init__.py b/keopscore/keopscore/mapreduce/__init__.py index b490efc0d..517b05fc5 100644 --- a/keopscore/keopscore/mapreduce/__init__.py +++ b/keopscore/keopscore/mapreduce/__init__.py @@ -1,5 +1,6 @@ from keopscore.config import cuda from .cpu import * + if cuda.get_use_cuda(): from .gpu import * diff --git a/keopscore/keopscore/mapreduce/cpu/__init__.py b/keopscore/keopscore/mapreduce/cpu/__init__.py index adf1ff1ff..27dffc967 100644 --- a/keopscore/keopscore/mapreduce/cpu/__init__.py +++ b/keopscore/keopscore/mapreduce/cpu/__init__.py @@ -2,7 +2,10 @@ from .CpuReduc import CpuReduc from .CpuReduc_ranges import CpuReduc_ranges +_exports = [ + CpuAssignZero, + CpuReduc, + CpuReduc_ranges, +] -_exports = [CpuAssignZero, CpuReduc, CpuReduc_ranges,] - -__all__ = [cls.__name__ for cls in _exports] \ No newline at end of file +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py index 6b59778ea..fc644de88 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py @@ -131,7 +131,9 @@ def __init__(self, *args): self.chk = Chunk_Mode_Constants(self.red_formula) self.dimy = self.chk.dimy self.blocksize_chunks = min( - cuda.get_cuda_block_size(), 1024, 49152 // max(1, self.dimy * sizeof(self.dtype)) + cuda.get_cuda_block_size(), + 1024, + 49152 // max(1, self.dimy * sizeof(self.dtype)), ) def get_code(self): diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py index da12b3b4c..d171e9df0 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py @@ -121,7 +121,9 @@ def get_code(self): self.dimy = max(dimfinalchunk, dimy) blocksize_chunks = min( - cuda.get_cuda_block_size(), 1024, 49152 // max(1, self.dimy * sizeof(self.dtype)) + cuda.get_cuda_block_size(), + 1024, + 49152 // max(1, self.dimy * sizeof(self.dtype)), ) if not isinstance(sum_scheme, block_sum): diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py index 419df228b..b96fafedc 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py @@ -196,7 +196,9 @@ def __init__(self, *args): self.chk = Chunk_Mode_Constants(self.red_formula) self.dimy = self.chk.dimy self.blocksize_chunks = min( - cuda.get_cuda_block_size(), 1024, 49152 // max(1, self.dimy * sizeof(self.dtype)) + cuda.get_cuda_block_size(), + 1024, + 49152 // max(1, self.dimy * sizeof(self.dtype)), ) def get_code(self): diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py index a54609689..afc8aaed8 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py @@ -144,7 +144,9 @@ def get_code(self): self.dimy = max(dimfinalchunk, dimy) blocksize_chunks = min( - cuda.get_cuda_block_size(), 1024, 49152 // max(1, self.dimy * sizeof(self.dtype)) + cuda.get_cuda_block_size(), + 1024, + 49152 // max(1, self.dimy * sizeof(self.dtype)), ) if not isinstance(sum_scheme, block_sum): diff --git a/keopscore/keopscore/mapreduce/gpu/__init__.py b/keopscore/keopscore/mapreduce/gpu/__init__.py index d74c11a76..caa8a929e 100644 --- a/keopscore/keopscore/mapreduce/gpu/__init__.py +++ b/keopscore/keopscore/mapreduce/gpu/__init__.py @@ -17,7 +17,7 @@ GpuReduc1D_ranges, GpuReduc1D_ranges_chunks, GpuReduc1D_ranges_finalchunks, - GpuReduc2D + GpuReduc2D, ] -__all__ = [cls.__name__ for cls in _exports] \ No newline at end of file +__all__ = [cls.__name__ for cls in _exports] diff --git a/keopscore/keopscore/sandbox/laplacian.py b/keopscore/keopscore/sandbox/laplacian.py index 7792212e5..a3be6eaf2 100644 --- a/keopscore/keopscore/sandbox/laplacian.py +++ b/keopscore/keopscore/sandbox/laplacian.py @@ -4,6 +4,7 @@ from keopscore.formulas import * import keopscore + keopscore.config.debug.set_debug_ops(True) diff --git a/keopscore/keopscore/utils/Cache.py b/keopscore/keopscore/utils/Cache.py index e82552508..fd5a9f422 100644 --- a/keopscore/keopscore/utils/Cache.py +++ b/keopscore/keopscore/utils/Cache.py @@ -8,11 +8,12 @@ # N.B we turn this into a function because the parameters need to be read dynamically. env_param = ( lambda: keopscore.config.cxx.get_compile_options() - + keopscore.config.cxx.get_linking_options() - + " auto_factorize=" + str(keopscore.config.auto_factorize) - + keopscore.config.cuda.get_include_options() - + keopscore.config.cuda.get_preprocessing_options() - + keopscore.config.cuda.get_nvrtc_flags() + + keopscore.config.cxx.get_linking_options() + + " auto_factorize=" + + str(keopscore.config.auto_factorize) + + keopscore.config.cuda.get_include_options() + + keopscore.config.cuda.get_preprocessing_options() + + keopscore.config.cuda.get_nvrtc_flags() ) diff --git a/keopscore/keopscore/utils/code_gen_utils.py b/keopscore/keopscore/utils/code_gen_utils.py index 1102e0d4e..3dfc1d79c 100644 --- a/keopscore/keopscore/utils/code_gen_utils.py +++ b/keopscore/keopscore/utils/code_gen_utils.py @@ -822,4 +822,3 @@ def varseq_to_array(vars, vars_ptr_name): string += f""" {vars_ptr_name}[{i}] = {vars[i].id}; """ return string - diff --git a/keopscore/keopscore/utils/gpu_utils.py b/keopscore/keopscore/utils/gpu_utils.py index 6569bf983..be1e82d63 100644 --- a/keopscore/keopscore/utils/gpu_utils.py +++ b/keopscore/keopscore/utils/gpu_utils.py @@ -16,20 +16,20 @@ def orig_cuda_include_fp16_path(): # First try to find the library file using the cuda includes cuda_include_path = keopscore.config.cuda.get_cuda_include_path() cuda_fp16_h_abspath = _first_matching_file( - cuda_include_path, - "cuda_fp16.h", - ) - cuda_fp16_hpp_abspath =_first_matching_file( - cuda_include_path, - "cuda_fp16.hpp", - ) + cuda_include_path, + "cuda_fp16.h", + ) + cuda_fp16_hpp_abspath = _first_matching_file( + cuda_include_path, + "cuda_fp16.hpp", + ) if cuda_fp16_h_abspath and cuda_fp16_hpp_abspath: return os.path.dirname(cuda_fp16_h_abspath) # Second try with compiler - cuda_fp16_h_abspath = get_include_file_abspath("cuda_fp16.h") - cuda_fp16_hpp_abspath = get_include_file_abspath("cuda_fp16.hpp") + cuda_fp16_h_abspath = get_include_file_abspath("cuda_fp16.h") + cuda_fp16_hpp_abspath = get_include_file_abspath("cuda_fp16.hpp") if cuda_fp16_h_abspath and cuda_fp16_hpp_abspath: path = os.path.dirname(cuda_fp16_h_abspath) if path != os.path.dirname(cuda_fp16_hpp_abspath): @@ -57,4 +57,3 @@ def custom_cuda_include_fp16_path(): if not os.path.isfile(fp16_header_path): pack_header(fp16_header, orig_cuda_include_fp16_path(), build_folder) return build_folder - diff --git a/keopscore/keopscore/utils/system_utils.py b/keopscore/keopscore/utils/system_utils.py index 82392b1bd..ceef605c7 100644 --- a/keopscore/keopscore/utils/system_utils.py +++ b/keopscore/keopscore/utils/system_utils.py @@ -9,7 +9,9 @@ def KeOps_OS_Run(command, print_warning=True): out = subprocess.run(command, shell=True, capture_output=True) if out.stderr != b"" and print_warning: - KeOps_Warning("There were warnings or errors while executing: " + command, newline=True) + KeOps_Warning( + "There were warnings or errors while executing: " + command, newline=True + ) KeOps_Print(out.stderr.decode("utf-8")) return out diff --git a/pykeops/pykeops/__init__.py b/pykeops/pykeops/__init__.py index 68910620d..98bbf13e8 100644 --- a/pykeops/pykeops/__init__.py +++ b/pykeops/pykeops/__init__.py @@ -39,6 +39,7 @@ def set_verbose(val): compile_jit_binary() + def clean_pykeops(recompile_jit_binaries=True): r""" This function cleans the KeOps cache and recompiles the JIT binaries if necessary. diff --git a/pykeops/pykeops/common/gpu_utils.py b/pykeops/pykeops/common/gpu_utils.py index d5c45638c..49722ca6e 100644 --- a/pykeops/pykeops/common/gpu_utils.py +++ b/pykeops/pykeops/common/gpu_utils.py @@ -1,4 +1,5 @@ -import pykeops.config as pykeopsconfig +import pykeops.config as pykeopsconfig + def get_gpu_number(): return pykeopsconfig.cuda.get_n_gpus() diff --git a/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py b/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py index 16fb70ba2..7359dc1b8 100644 --- a/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py +++ b/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py @@ -183,7 +183,7 @@ def get_pybind11_code(self): LoadKeOps_cpp = Cache_partial( LoadKeOps_cpp_class, use_cache_file=True, - save_folder=pykeopsconfig.path.get_build_folder() + save_folder=pykeopsconfig.path.get_build_folder(), ) cpp_dtype = { diff --git a/pykeops/pykeops/common/operations.py b/pykeops/pykeops/common/operations.py index ace359afe..2670f4649 100644 --- a/pykeops/pykeops/common/operations.py +++ b/pykeops/pykeops/common/operations.py @@ -174,7 +174,9 @@ def ConjugateGradientSolver( else: # for loop exhausted # Return incomplete progress - KeOps_Warning("[KeOps CG]: Maximum iterations reached. Check convergence...") + KeOps_Warning( + "[KeOps CG]: Maximum iterations reached. Check convergence..." + ) it = -maxiter - 1 if verbose: @@ -189,7 +191,7 @@ def ConjugateGradientSolver( "maxiter": maxiter, "x0_provided": x0 is not None, } - + KeOps_Print(f"[KeOps CG]: {info}") return x diff --git a/pykeops/pykeops/config.py b/pykeops/pykeops/config.py index 2c9a44c1e..2892ba04a 100644 --- a/pykeops/pykeops/config.py +++ b/pykeops/pykeops/config.py @@ -26,7 +26,9 @@ def pykeops_nvrtc_name(type="src"): extension = ".cpp" if type == "src" else sysconfig.get_config_var("EXT_SUFFIX") return os.path.join( ( - os.path.join(os.path.dirname(os.path.realpath(__file__)), "common", "keops_io") + os.path.join( + os.path.dirname(os.path.realpath(__file__)), "common", "keops_io" + ) if type == "src" else get_build_folder() ), diff --git a/pykeops/pykeops/test/test_numpy.py b/pykeops/pykeops/test/test_numpy.py index dbc08ab73..f3e63b69b 100644 --- a/pykeops/pykeops/test/test_numpy.py +++ b/pykeops/pykeops/test/test_numpy.py @@ -82,7 +82,7 @@ def test_cg_solver_stops_immediately_when_x0_is_good(self): x = K_xx.solve(b, alpha=alpha, x0=self.f, eps=1e-12) self.assertTrue(np.allclose(self.f, x)) - + stream = io.StringIO() with redirect_stdout(stream): x = K_xx.solve(b, alpha=alpha, x0=self.f, eps=1e-12, verbose=True) @@ -93,7 +93,6 @@ def test_cg_solver_stops_immediately_when_x0_is_good(self): self.assertIn("'niter': 0", output) self.assertIn("'x0_provided': True", output) - ############################################################ def test_cg_solver_verbose_prints_info(self): ############################################################ From f1fd1c2bf77e307436857d7deb3fbc5c47f48fa1 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 4 May 2026 09:04:41 +0200 Subject: [PATCH 21/98] Add include path init Co-authored-by: Copilot --- keopscore/keopscore/config/Cuda.py | 1 + 1 file changed, 1 insertion(+) diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index 88f3ab075..cff772aac 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -120,6 +120,7 @@ def __init__(self): self.set_nvrtc_flags() self.set_cuda_block_size() self.set_preprocessing_options() + self.set_include_options() def find_install_path(self, lib_dict_info, warn=None): """ From 7d35354e4e3514999ce7c920592c4b639039e758 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 4 May 2026 11:08:47 +0200 Subject: [PATCH 22/98] reorder compiles flags --- keopscore/keopscore/binders/LinkCompile.py | 2 +- .../binders/nvrtc/Gpu_link_compile.py | 2 +- .../keopscore/binders/nvrtc/keops_nvrtc.cpp | 7 ++--- keopscore/keopscore/config/Cuda.py | 31 +++++++++---------- keopscore/keopscore/config/CxxCompiler.py | 2 +- keopscore/keopscore/utils/Cache.py | 4 +-- 6 files changed, 23 insertions(+), 25 deletions(-) diff --git a/keopscore/keopscore/binders/LinkCompile.py b/keopscore/keopscore/binders/LinkCompile.py index 6f2f9b744..0d010cbe0 100644 --- a/keopscore/keopscore/binders/LinkCompile.py +++ b/keopscore/keopscore/binders/LinkCompile.py @@ -6,8 +6,8 @@ cpp_flag = keopscore.config.cxx.get_compile_options() cpp_flag += keopscore.config.cxx.get_linking_options() -cpp_flag += keopscore.config.cuda.get_nvrtc_flags() cpp_flag += keopscore.config.cuda.get_include_options() +cpp_flag += keopscore.config.cuda.get_linking_options() class LinkCompile: diff --git a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py index 81ba9587b..e08738e95 100644 --- a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py +++ b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py @@ -100,7 +100,7 @@ def get_compile_command( else '\\"compute\\"' ) target_type_define = f"-DnvrtcGetTARGET={nvrtcGetTARGET} -DnvrtcGetTARGETSize={nvrtcGetTARGETSize} -DARCHTAG={arch_tag}" - return f"{keopscore.config.cxx.get_cxx_compiler()} {keopscore.config.cuda.get_preprocessing_options()} {target_type_define} {keopscore.config.cxx.get_compile_options()} {keopscore.config.cuda.get_nvrtc_flags()} {keopscore.config.cuda.get_include_options()} {keopscore.config.path.get_include_options()} {extra_flags} {sourcename} {keopscore.config.cxx.get_linking_options()} -o {dllname}" + return f"{keopscore.config.cxx.get_cxx_compiler()} {keopscore.config.cuda.get_preprocessing_options()} {target_type_define} {keopscore.config.cxx.get_compile_options()} {keopscore.config.cuda.get_include_options()} {keopscore.config.path.get_include_options()} {extra_flags} {sourcename} {keopscore.config.cxx.get_linking_options()} {keopscore.config.cuda.get_linking_options()} -o {dllname}" @staticmethod def compile_jit_compile_dll(): diff --git a/keopscore/keopscore/binders/nvrtc/keops_nvrtc.cpp b/keopscore/keopscore/binders/nvrtc/keops_nvrtc.cpp index bad02d6d7..78a62b648 100644 --- a/keopscore/keopscore/binders/nvrtc/keops_nvrtc.cpp +++ b/keopscore/keopscore/binders/nvrtc/keops_nvrtc.cpp @@ -1,16 +1,15 @@ // nvcc -shared -Xcompiler -fPIC -lnvrtc -lcuda keops_nvrtc.cu -o keops_nvrtc.so // g++ --verbose -L/opt/cuda/lib64 -L/opt/cuda/targets/x86_64-linux/lib/ // -I/opt/cuda/targets/x86_64-linux/include/ -I../../include -shared -fPIC -// -lcuda -lnvrtc -fpermissive -DMAXIDGPU=0 -DMAXTHREADSPERBLOCK0=1024 +// -lcuda -lnvrtc -DMAXIDGPU=0 -DMAXTHREADSPERBLOCK0=1024 // -DSHAREDMEMPERBLOCK0=49152 -DnvrtcGetTARGET=nvrtcGetCUBIN // -DnvrtcGetTARGETSize=nvrtcGetCUBINSize -DARCHTAG=\"sm\" keops_nvrtc.cpp -o // keops_nvrtc.so g++ -std=c++11 -shared -fPIC -O3 -fpermissive -L /usr/lib -L // /opt/cuda/lib64 -lcuda -lnvrtc -DnvrtcGetTARGET=nvrtcGetCUBIN -// -DnvrtcGetTARGETSize=nvrtcGetCUBINSize -DARCHTAG=\"sm\" -// -I/home/bcharlier/projets/keops/keops/keops/include -I/opt/cuda/include +// -DnvrtcGetTARGETSize=nvrtcGetCUBINSize -DARCHTAG=\"sm\" -I/opt/cuda/include // -I/usr/include/python3.10/ -DMAXIDGPU=0 -DMAXTHREADSPERBLOCK0=1024 // -DSHAREDMEMPERBLOCK0=49152 -// /home/bcharlier/projets/keops/keops/keops/binders/nvrtc/keops_nvrtc.cpp -o +// ./keopscore/binders/nvrtc/keops_nvrtc.cpp -o // keops_nvrtc.cpython-310-x86_64-linux-gnu.so #include diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index cff772aac..ebb1d5f8d 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -25,17 +25,16 @@ class CudaConfig: # Cuda detection variables _use_cuda = None _specific_gpus = None - - _cuda_include_path = None - _nvrtc_flags = None - _cuda_version = None - _n_gpus = 0 _MaxThreadsPerBlock = [] _SharedMemPerBlock = [] + _cuda_version = None + _cuda_include_path = None + _preprocessing_options = "" _include_options = "" + _linking_options = None _cuda_block_size = None # ------------------------ # @@ -117,7 +116,7 @@ def __init__(self): if self.get_use_cuda(): self.set_cuda_version() self.set_cuda_include_path() - self.set_nvrtc_flags() + self.set_linking_options() self.set_cuda_block_size() self.set_preprocessing_options() self.set_include_options() @@ -448,18 +447,18 @@ def print_cuda_include_path(self): ) # NVRTC Flags - def set_nvrtc_flags(self): - """Set the NVRTC flags for CUDA compilation.""" + def set_linking_options(self): + """Set the Linking option for nvrt/cuda entry point compilation.""" - self._nvrtc_flags = f" -fpermissive -L{self.get_libcuda_folder()} -L{self.get_libnvrtc_folder()} -lcuda -lnvrtc" + self._linking_options = f"-L{self.get_libcuda_folder()} -L{self.get_libnvrtc_folder()} -lcuda -lnvrtc" - def get_nvrtc_flags(self): - """Get the NVRTC flags for CUDA compilation.""" - return self._nvrtc_flags + def get_linking_options(self): + """Get the Linking option for nvrt/cuda entry point compilation.""" + return self._linking_options - def print_nvrtc_flags(self): - """Print the NVRTC flags for CUDA compilation.""" - print(f"NVRTC Flags: {self.get_nvrtc_flags()}") + def print_linking_options(self): + """Print the Linking option for nvrt/cuda entry point compilation.""" + print(f"Linking Options: {self.get_linking_options()}") # GPU compile flags def set_preprocessing_options(self): @@ -571,10 +570,10 @@ def print_all(self): self.print_libcuda_folder() self.print_libnvrtc_folder() self.print_cuda_include_path() - self.print_nvrtc_flags() self.print_preprocessing_options() self.print_include_options() + self.print_linking_options() # Print relevant environment variables. print_envs(self.cuda_env_vars) diff --git a/keopscore/keopscore/config/CxxCompiler.py b/keopscore/keopscore/config/CxxCompiler.py index 3c2a7867a..902b1c0dc 100644 --- a/keopscore/keopscore/config/CxxCompiler.py +++ b/keopscore/keopscore/config/CxxCompiler.py @@ -122,7 +122,7 @@ def set_compile_options(self): self._compile_options = self.get_cxx_env_flags() # compilation / behavior - self.add_to_compile_option("-std=c++11 -O3 -flto=auto") + self.add_to_compile_option("-std=c++11 -O3 -flto=auto -fpermissive") # code model self.add_to_compile_option("-fPIC") diff --git a/keopscore/keopscore/utils/Cache.py b/keopscore/keopscore/utils/Cache.py index fd5a9f422..ff7f46565 100644 --- a/keopscore/keopscore/utils/Cache.py +++ b/keopscore/keopscore/utils/Cache.py @@ -11,9 +11,9 @@ + keopscore.config.cxx.get_linking_options() + " auto_factorize=" + str(keopscore.config.auto_factorize) - + keopscore.config.cuda.get_include_options() + keopscore.config.cuda.get_preprocessing_options() - + keopscore.config.cuda.get_nvrtc_flags() + + keopscore.config.cuda.get_include_options() + + keopscore.config.cuda.get_linking_options() ) From f9ce2a1ef94bf1cda981a6f2afac8036c2b1995d Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 4 May 2026 11:29:36 +0200 Subject: [PATCH 23/98] explicit name in cuda options Co-authored-by: Copilot --- .../binders/nvrtc/Gpu_link_compile.py | 3 +- keopscore/keopscore/config/Cuda.py | 87 ++++++++++--------- keopscore/keopscore/config/KeOpsPath.py | 6 +- pykeops/pykeops/common/gpu_utils.py | 2 +- 4 files changed, 51 insertions(+), 47 deletions(-) diff --git a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py index e08738e95..6022c20c2 100644 --- a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py +++ b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py @@ -33,7 +33,8 @@ class Gpu_link_compile(LinkCompile): def __init__(self): # checking that the system has a Gpu : if not ( - keopscore.config.cuda.get_use_cuda() and keopscore.config.cuda.get_n_gpus() + keopscore.config.cuda.get_use_cuda() + and keopscore.config.cuda.get_n_visible_devices() ): KeOps_Error( "Trying to compile cuda code... but we detected that the system has no properly configured cuda lib." diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index ebb1d5f8d..4805b26c4 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -24,18 +24,18 @@ class CudaConfig: # Cuda detection variables _use_cuda = None - _specific_gpus = None - _n_gpus = 0 - _MaxThreadsPerBlock = [] - _SharedMemPerBlock = [] - _cuda_version = None _cuda_include_path = None + _visible_devices = None + _n_visible_devices = 0 + _MaxThreadsPerBlock = [] + _SharedMemPerBlock = [] + _cuda_block_size = None + _preprocessing_options = "" _include_options = "" _linking_options = None - _cuda_block_size = None # ------------------------ # # Search location # @@ -106,7 +106,7 @@ class CudaConfig: def __init__(self): - self.set_specific_gpus() + self.set_visible_devices() super().__init__() @@ -224,7 +224,7 @@ def _find_and_load_libcuda(self): + " Switching to CPU only.", ) - self._n_gpus = nGpus.value + self._n_visible_devices = nGpus.value return True, "" @@ -344,28 +344,28 @@ def get_cuda_block_size(self): def print_cuda_block_size(self): print(f"CUDA Block Size: {self.get_cuda_block_size()}") - # Specific GPUs - def set_specific_gpus(self): + # Visibles GPUs devices + def set_visible_devices(self): """Set specific GPUs from CUDA_VISIBLE_DEVICES.""" if os.getenv("CUDA_VISIBLE_DEVICES"): - self._specific_gpus = os.getenv("CUDA_VISIBLE_DEVICES").replace(",", "_") + self._visible_devices = os.getenv("CUDA_VISIBLE_DEVICES").replace(",", "_") - def get_specific_gpus(self): + def get_visible_devices(self): """Get the specific GPUs.""" - return self._specific_gpus + return self._visible_devices # Number of GPUs - def set_n_gpus(self): + def set_n_visible_devices(self): """Set the number of GPUs detected. This is done in _find_and_load_libcuda.""" pass - def get_n_gpus(self): + def get_n_visible_devices(self): """Get the number of GPUs detected.""" - return self._n_gpus + return self._n_visible_devices - def print_n_gpus(self): + def print_n_visible_devices(self): """Print the number of GPUs detected.""" - print(f"Number of GPUs Detected: {self.get_n_gpus()}") + print(f"Number of GPUs Detected: {self.get_n_visible_devices()}") # Libcuda folder def set_libcuda_folder(self): @@ -446,28 +446,28 @@ def print_cuda_include_path(self): f"CUDA Include Path: {":".join(self.get_cuda_include_path()) or not_found_str}" ) - # NVRTC Flags - def set_linking_options(self): - """Set the Linking option for nvrt/cuda entry point compilation.""" - - self._linking_options = f"-L{self.get_libcuda_folder()} -L{self.get_libnvrtc_folder()} -lcuda -lnvrtc" + # NVRTC include options + def set_include_options(self): + self._include_options += "".join( + f" -I{p}" for p in list(set(self.get_cuda_include_path())) + ) - def get_linking_options(self): - """Get the Linking option for nvrt/cuda entry point compilation.""" - return self._linking_options + def get_include_options(self): + return self._include_options - def print_linking_options(self): - """Print the Linking option for nvrt/cuda entry point compilation.""" - print(f"Linking Options: {self.get_linking_options()}") + def print_include_options(self): + print(f"GPU Include Options: {self.get_include_options() or not_found_str}") - # GPU compile flags + # NVRTC preprocessins options def set_preprocessing_options(self): """Set GPU preprocessing compile flags based on detected GPU properties.""" - if self.get_n_gpus() == 0: + if self.get_n_visible_devices() == 0: return - self.add_to_preprocessing_options(f"-DMAXIDGPU={self.get_n_gpus() - 1}") - for d in range(self.get_n_gpus()): + self.add_to_preprocessing_options( + f"-DMAXIDGPU={self.get_n_visible_devices() - 1}" + ) + for d in range(self.get_n_visible_devices()): self.add_to_preprocessing_options( f"-DMAXTHREADSPERBLOCK{d}={self.get_MaxThreadsPerBlock()[d]}" ) @@ -487,16 +487,19 @@ def print_preprocessing_options(self): f"GPU Preprocessing Options: {self.get_preprocessing_options() or not_found_str}" ) - def set_include_options(self): - self._include_options += "".join( - f" -I{p}" for p in list(set(self.get_cuda_include_path())) - ) + # NVRTC linking options + def set_linking_options(self): + """Set the Linking option for nvrt/cuda entry point compilation.""" - def get_include_options(self): - return self._include_options + self._linking_options = f"-L{self.get_libcuda_folder()} -L{self.get_libnvrtc_folder()} -lcuda -lnvrtc" - def print_include_options(self): - print(f"GPU Include Options: {self.get_include_options() or not_found_str}") + def get_linking_options(self): + """Get the Linking option for nvrt/cuda entry point compilation.""" + return self._linking_options + + def print_linking_options(self): + """Print the Linking option for nvrt/cuda entry point compilation.""" + print(f"Linking Options: {self.get_linking_options()}") # Helper functions for CUDA attribute detection def _get_device_attributes(self, libcuda, device_index): @@ -565,7 +568,7 @@ def print_all(self): self.print_use_cuda() if self.get_use_cuda(): - self.print_n_gpus() + self.print_n_visible_devices() self.print_cuda_version() self.print_libcuda_folder() self.print_libnvrtc_folder() diff --git a/keopscore/keopscore/config/KeOpsPath.py b/keopscore/keopscore/config/KeOpsPath.py index 32fc1fc63..ce6520e25 100644 --- a/keopscore/keopscore/config/KeOpsPath.py +++ b/keopscore/keopscore/config/KeOpsPath.py @@ -78,9 +78,9 @@ def set_default_build_folder_name(self): ] if self.cuda.get_use_cuda(): name_parts.append(f"CUDA{self.cuda.get_cuda_version()}") - specific_gpus = self.cuda.get_specific_gpus() - if specific_gpus: - name_parts.append(f"VISIBLE_DEVICES{specific_gpus}") + visible_devices = self.cuda.get_visible_devices() + if visible_devices: + name_parts.append(f"VISIBLE_DEVICES{visible_devices}") self._default_build_folder_name = "_".join(name_parts) diff --git a/pykeops/pykeops/common/gpu_utils.py b/pykeops/pykeops/common/gpu_utils.py index 49722ca6e..4b3b7df78 100644 --- a/pykeops/pykeops/common/gpu_utils.py +++ b/pykeops/pykeops/common/gpu_utils.py @@ -2,4 +2,4 @@ def get_gpu_number(): - return pykeopsconfig.cuda.get_n_gpus() + return pykeopsconfig.cuda.get_n_visible_devices() From 310c801e16ee2338c97b37161e16e3abb05d6ac0 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 4 May 2026 12:14:05 +0200 Subject: [PATCH 24/98] init options to str Co-authored-by: Copilot --- keopscore/keopscore/config/Cuda.py | 23 +++++++++++------------ keopscore/keopscore/config/CxxCompiler.py | 6 +++--- keopscore/keopscore/config/KeOpsPath.py | 12 ++++++------ keopscore/keopscore/config/OpenMP.py | 6 +++--- keopscore/keopscore/config/Platform.py | 16 ++++++++-------- keopscore/keopscore/utils/Cache.py | 3 +-- 6 files changed, 32 insertions(+), 34 deletions(-) diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index 4805b26c4..c0762d4f2 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -23,19 +23,19 @@ class CudaConfig: CUDA_BLOCK_SIZE = 192 # Cuda detection variables - _use_cuda = None - _cuda_version = None - _cuda_include_path = None + _use_cuda = False + _cuda_version = -1 + _cuda_include_path = "" - _visible_devices = None - _n_visible_devices = 0 + _visible_devices = "" + _n_visible_devices = -1 _MaxThreadsPerBlock = [] _SharedMemPerBlock = [] - _cuda_block_size = None + _cuda_block_size = -1 _preprocessing_options = "" _include_options = "" - _linking_options = None + _linking_options = "" # ------------------------ # # Search location # @@ -416,7 +416,7 @@ def get_cuda_version(self, out_type="single_value"): return self._cuda_version def print_cuda_version(self): - str = f"CUDA Version: {self.get_cuda_version(out_type="string") if self.get_cuda_version() else not_found_str}" + str = f"CUDA Version: {self.get_cuda_version(out_type='string') if self.get_cuda_version() else not_found_str}" print(str) # CUDA Include Path @@ -443,7 +443,7 @@ def get_cuda_include_path(self): def print_cuda_include_path(self): print( - f"CUDA Include Path: {":".join(self.get_cuda_include_path()) or not_found_str}" + f"CUDA Include Path: {':'.join(self.get_cuda_include_path()) or not_found_str}" ) # NVRTC include options @@ -464,9 +464,8 @@ def set_preprocessing_options(self): if self.get_n_visible_devices() == 0: return - self.add_to_preprocessing_options( - f"-DMAXIDGPU={self.get_n_visible_devices() - 1}" - ) + self._preprocessing_options = f"-DMAXIDGPU={self.get_n_visible_devices() - 1}" + for d in range(self.get_n_visible_devices()): self.add_to_preprocessing_options( f"-DMAXTHREADSPERBLOCK{d}={self.get_MaxThreadsPerBlock()[d]}" diff --git a/keopscore/keopscore/config/CxxCompiler.py b/keopscore/keopscore/config/CxxCompiler.py index 902b1c0dc..038015afe 100644 --- a/keopscore/keopscore/config/CxxCompiler.py +++ b/keopscore/keopscore/config/CxxCompiler.py @@ -11,12 +11,12 @@ class CxxCompilerConfig: Common platform and C++ compiler configuration. """ - _cxx_compiler = None - _cxx_env_flags = None + _cxx_compiler = "" + _cxx_env_flags = "" _compile_options = "" _linking_options = "" - _disable_pragma_unrolls = None + _disable_pragma_unrolls = True _use_Apple_clang = False cxx_envs = ["CXX", "CXXFLAGS"] diff --git a/keopscore/keopscore/config/KeOpsPath.py b/keopscore/keopscore/config/KeOpsPath.py index ce6520e25..721129be3 100644 --- a/keopscore/keopscore/config/KeOpsPath.py +++ b/keopscore/keopscore/config/KeOpsPath.py @@ -10,13 +10,13 @@ class KeOpsPathConfig: Class to manage the path to the KeOps library. """ - _base_dir_path = None - _keops_cache_folder = None - _default_build_folder_name = None - _default_build_path = None - _build_folder = None + _base_dir_path = "" + _keops_cache_folder = "" + _default_build_folder_name = "" + _default_build_path = "" + _build_folder = "" - _jit_binary = None + _jit_binary = "" _include_options = "" path_env_vars = ("KEOPS_CACHE_FOLDER",) diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 410331c7d..009675694 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -20,9 +20,9 @@ class OpenMPConfig: Class for OpenMP detection and configuration. """ - _use_OpenMP = None - _openmp_lib_name = None - _openmp_lib_include_dir = None + _use_OpenMP = False + _openmp_lib_name = "" + _openmp_lib_include_dir = "" _compile_options = "" _include_options = "" diff --git a/keopscore/keopscore/config/Platform.py b/keopscore/keopscore/config/Platform.py index 2ff1bbf08..05d1aaa89 100644 --- a/keopscore/keopscore/config/Platform.py +++ b/keopscore/keopscore/config/Platform.py @@ -11,14 +11,14 @@ class PlatformConfig: Class for detecting the operating system, Python version, and environment type. """ - _os = None - _platform = None - _machine = None - _uname = None - _python_version = None - _python_executable = None - _env_type = None - _brew_prefix = None + _os = "" + _platform = "" + _machine = "" + _uname = "" + _python_version = "" + _python_executable = "" + _env_type = "" + _brew_prefix = "" platform_envs = [ "PYTHONPATH", diff --git a/keopscore/keopscore/utils/Cache.py b/keopscore/keopscore/utils/Cache.py index ff7f46565..e5e7c44e8 100644 --- a/keopscore/keopscore/utils/Cache.py +++ b/keopscore/keopscore/utils/Cache.py @@ -9,8 +9,7 @@ env_param = ( lambda: keopscore.config.cxx.get_compile_options() + keopscore.config.cxx.get_linking_options() - + " auto_factorize=" - + str(keopscore.config.auto_factorize) + + f" auto_factorize={keopscore.config.auto_factorize}" + keopscore.config.cuda.get_preprocessing_options() + keopscore.config.cuda.get_include_options() + keopscore.config.cuda.get_linking_options() From 62bb4cadaf468e9e565788c1bbc185af9aca3db6 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 4 May 2026 12:15:01 +0200 Subject: [PATCH 25/98] lint --- keopscore/keopscore/config/Cuda.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index c0762d4f2..4c88eee5e 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -465,7 +465,7 @@ def set_preprocessing_options(self): return self._preprocessing_options = f"-DMAXIDGPU={self.get_n_visible_devices() - 1}" - + for d in range(self.get_n_visible_devices()): self.add_to_preprocessing_options( f"-DMAXTHREADSPERBLOCK{d}={self.get_MaxThreadsPerBlock()[d]}" From 33ddab919cf9a5fa71098683325f63e3f4f41c2f Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 4 May 2026 17:40:58 +0200 Subject: [PATCH 26/98] cleanup openMP detection --- keopscore/keopscore/config/Cuda.py | 91 ++++++++---- keopscore/keopscore/config/OpenMP.py | 212 ++++++++++++++------------- 2 files changed, 167 insertions(+), 136 deletions(-) diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index 4c88eee5e..ce61ec805 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -25,12 +25,12 @@ class CudaConfig: # Cuda detection variables _use_cuda = False _cuda_version = -1 - _cuda_include_path = "" + _cuda_include_path = [""] # str or list of str _visible_devices = "" _n_visible_devices = -1 - _MaxThreadsPerBlock = [] - _SharedMemPerBlock = [] + _MaxThreadsPerBlock = [] # list of int + _SharedMemPerBlock = [] # list of int _cuda_block_size = -1 _preprocessing_options = "" @@ -83,23 +83,23 @@ class CudaConfig: "name": "cuda", "lib_basename_candidate": ["libcuda.so.*", "libcuda.dylib", "cuda.lib"], "header_basename": "cuda.h", - "library": None, # to be filled later - "header": None, # to be filled later + "library": "", # to be filled later + "header": "", # to be filled later "ctype_handle": None, # to be filled later } _libnvrtc_info = { "name": "nvrtc", "lib_basename_candidate": ["libnvrtc.so.*", "libnvrtc.dylib", "nvrtc.lib"], "header_basename": "nvrtc.h", - "library": None, # to be filled later - "header": None, # to be filled later + "library": "", # to be filled later + "header": "", # to be filled later "ctype_handle": None, # to be filled later } _cudart_info = { "name": "cudart", "lib_basename_candidate": ["libcudart.so.*", "libcudart.dylib", "cudart.lib"], - "header_basename": None, - "library": None, # to be filled later + "header_basename": None, # not needed + "library": "", # to be filled later "header": None, # not needed "ctype_handle": None, # to be filled later } @@ -121,7 +121,7 @@ def __init__(self): self.set_preprocessing_options() self.set_include_options() - def find_install_path(self, lib_dict_info, warn=None): + def find_install_path(self, lib_dict_info): """ Locate a cuda and headers using an explicit ordered search. @@ -131,9 +131,8 @@ def find_install_path(self, lib_dict_info, warn=None): Returns: result (dict): a copy of lib_dict_info completed with the ``library`` and ``header`` keys containing the absolute paths to the library file and include directory, or None if not found. """ - result = ( - lib_dict_info.copy() - ) # Start with the provided info, which may contain names and file patterns + # Start with the provided info, which may contain names and file patterns + result = lib_dict_info.copy() candidate_roots = _ordered_search_roots( env_vars=self.cuda_env_vars, @@ -155,9 +154,6 @@ def find_install_path(self, lib_dict_info, warn=None): if result["library"] is None: result["library"] = _find_library_by_names((result["name"],)) - if result["library"] is None and warn: - KeOps_Warning(f"lib{result['name']} not found.") - # ------------------------ # # Search for header files # # ------------------------ # @@ -167,14 +163,11 @@ def find_install_path(self, lib_dict_info, warn=None): (result["header_basename"],), ) - if result["header"] is None and warn and result["header_basename"] is not None: - KeOps_Warning(f"{result['name']} header files not found.") - return result def _find_and_load_libcuda(self): """Locate, load, and initialize the CUDA driver library.""" - self._libcuda_info = self.find_install_path(self._libcuda_info, warn=True) + self._libcuda_info = self.find_install_path(self._libcuda_info) libcuda_path = self._libcuda_info["library"] if not libcuda_path: return ( @@ -230,7 +223,7 @@ def _find_and_load_libcuda(self): def _find_and_load_libnvrtc(self): """Locate and load the NVRTC runtime compilation library.""" - self._libnvrtc_info = self.find_install_path(self._libnvrtc_info, warn=True) + self._libnvrtc_info = self.find_install_path(self._libnvrtc_info) libnvrtc_path = self._libnvrtc_info["library"] if not libnvrtc_path: return ( @@ -260,7 +253,7 @@ def _find_and_load_cudart(self): """ - self._cudart_info = self.find_install_path(self._cudart_info, warn=False) + self._cudart_info = self.find_install_path(self._cudart_info) libcudart_path = self._cudart_info["library"] if not libcudart_path: return ( @@ -379,8 +372,18 @@ def get_libcuda_folder(self): self._libcuda_info["library"] ) - def print_libcuda_folder(self): - print(f"Libcuda Folder: {self.get_libcuda_folder() or not_found_str}") + # Libcuda path + def set_libcuda_path(self): + """ + Is set in _cuda_libraries_available. + """ + pass + + def get_libcuda_path(self): + return self._libcuda_info["library"] + + def print_libcuda_path(self): + print(f"Libcuda Path: {self.get_libcuda_path() or not_found_str}") # Libnvrtc folder def set_libnvrtc_folder(self): @@ -395,11 +398,36 @@ def get_libnvrtc_folder(self): self._libnvrtc_info["library"] ) - def print_libnvrtc_folder(self): - print(f"Libnvrtc Folder: {self.get_libnvrtc_folder() or not_found_str}") + # Libnvrtc path + def set_libnvrtc_path(self): + """ + Return nothing if not using cuda + self.libnvrtc_path is already set in _cuda_libraries_available. + """ + pass + + def get_libnvrtc_path(self): + return self._libnvrtc_info["library"] + + def print_libnvrtc_path(self): + print(f"Libnvrtc Path: {self.get_libnvrtc_path() or not_found_str}") + + # Libcudart path + def set_libcudart_path(self): + """ + Return nothing if not using cuda + self.libcudart_path is already set in _cuda_libraries_available. + """ + pass + + def get_libcudart_path(self): + return self._cudart_info["library"] + + def print_libcudart_path(self): + print(f"Libcudart Path: {self.get_libcudart_path() or not_found_str}") # CUDA Version - def set_cuda_version(self, warn=True): + def set_cuda_version(self): """Set the CUDA version by querying the CUDA runtime library. Done in _find_and_load_cudart.""" pass @@ -569,8 +597,9 @@ def print_all(self): if self.get_use_cuda(): self.print_n_visible_devices() self.print_cuda_version() - self.print_libcuda_folder() - self.print_libnvrtc_folder() + self.print_libcuda_path() + self.print_libnvrtc_path() + self.print_libcudart_path() self.print_cuda_include_path() self.print_preprocessing_options() @@ -592,8 +621,8 @@ def print_all(self): cxx_compiler_info = CxxCompilerConfig(platform_info) cxx_compiler_info.print_all() - openmp_info = OpenMPConfig(platform_info, cxx_compiler_info) - openmp_info.print_all() + # openmp_info = OpenMPConfig(platform_info, cxx_compiler_info) + # openmp_info.print_all() cuda_info = CudaConfig() cuda_info.print_all() diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 009675694..a728ab5c3 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -20,16 +20,15 @@ class OpenMPConfig: Class for OpenMP detection and configuration. """ - _use_OpenMP = False - _openmp_lib_name = "" - _openmp_lib_include_dir = "" + _use_omp = False + + _libomp_folder = "" + _libomp_include_path = "" _compile_options = "" _include_options = "" _linking_options = "" - _openmp_header_basename = "omp.h" - openmp_env_vars = ( "OMP_PATH", "LIBOMP_PATH", @@ -37,7 +36,7 @@ class OpenMPConfig: "OpenMP_ROOT_DIR", ) - openmp_system_suffixes = [ + _openmp_system_suffixes = [ # self.get_brew_prefix() added later on, os.path.join(os.path.sep, "usr", "local", "opt", "libomp"), os.path.join(os.path.sep, "opt", "local"), @@ -61,73 +60,60 @@ class OpenMPConfig: os.path.join("opt", "libomp", "include"), ) - def __init__(self, platform, cxx_compiler): - - self.platform = platform - self.cxx_compiler = cxx_compiler - - self.openmp_system_suffixes += ( - [ - self.platform.get_brew_prefix(), - ] - if self.platform.get_brew_prefix() - else [] - ) - self.set_openmplib_path() + _omp_info = { + "name": ["omp", "gomp"], + "lib_basename_candidate": ["libomp.*", "libgomp.*", "libm.so*"], + "header_basename": ["omp.h", "gomp.h"], + "library": "", # to be filled later + "header": "", # not needed + "ctype_handle": None, # to be filled later + } - self.set_compile_options() - self.set_include_options() - self.set_linking_options() + def __init__(self, platform, cxx): - self.set_use_OpenMP() - - # OpenMP library path - def set_openmplib_path(self): - """try to locate OpenMP libraries""" - openmp_install = self.find_openmp_install() - openmp_lib = openmp_install["library"] - if openmp_lib: - self._openmp_lib_include_dir = os.path.dirname(openmp_install["header"]) - self._openmp_lib_lib_dir = os.path.dirname(openmp_lib) - self._openmp_lib_name = openmp_lib - else: - KeOps_Warning( - "OpenMP runtime library not found. " - "Set OMP_PATH, LIBOMP_PATH, or OpenMP_ROOT if it is installed in a non-standard location." + self.platform = platform + self.cxx = cxx + + # On MacOS: Add brew prefix to OpenMP search paths if it exists + if self.platform.get_brew_prefix(): + self._openmp_system_suffixes.insert( + 0, + [ + self.platform.get_brew_prefix(), + ], ) - def print_openmplib_path(self): - if self.get_openmp_lib_name() and self.get_openmp_lib_dir(): - full_path = os.path.join( - self.get_openmp_lib_dir(), self.get_openmp_lib_name() + # Detect OpenMP and set related configuration variables + if self._omp_is_available(): + self.set_libomp_include_path() + self.set_libomp_folder() + self.set_compile_options() + self.set_include_options() + self.set_linking_options() + + # Chech if the compiler support omp + self.set_use_omp() + + def find_install_path(self, lib_dict_info): + result = lib_dict_info.copy() + + # First try to find OpenMP library using standard names via ctypes /cxx compiler. + for name, header_basename in zip( + lib_dict_info["name"], lib_dict_info["header_basename"] + ): + result["library"] = _find_library_by_names(name) + result["header"] = get_include_file_abspath( + header_basename, self.cxx.get_cxx_compiler() ) - elif self.get_openmp_lib_name(): - full_path = self.get_openmp_lib_name() - else: - full_path = None - - print(f"OpenMP Library Path: {full_path or not_found_str}") - - def find_openmp_install(self): - """ - Locate OpenMP runtime files without assuming a specific package manager. - - Returns a dict with optional ``library`` and ``header`` entries. - """ - result = {"library": None, "header": None} - - # First try to find OpenMP library using standard names via ctypes. - result["library"] = _find_library_by_names(("gomp", "omp")) - result["header"] = get_include_file_abspath( - self._openmp_header_basename, self.cxx_compiler.get_cxx_compiler() - ) + if result["library"] and result["header"]: + return result # If that fails, search for OpenMP headers and libraries in common locations. if not result["library"]: candidate_roots = _ordered_search_roots( env_vars=self.openmp_env_vars, conda_root="CONDA_PREFIX", - system_roots=self.openmp_system_suffixes, + system_roots=self._openmp_system_suffixes, ) result["library"] = _first_matching_file( @@ -143,57 +129,70 @@ def find_openmp_install(self): return result - # OpenMP library name and directory getters/setters - def set_openmp_lib_name(self): - """Set the OpenMP library name (e.g., libomp.so or libgomp.dylib).""" - # This is set in set_openmplib_path if the library is found, otherwise it remains None. - pass + def _omp_is_available(self): + self._omp_info = self.find_install_path(self._omp_info) + liomp_path = self._omp_info["library"] + if not liomp_path: + KeOps_Warning( + "libomp not found. Set OMP_PATH, LIBOMP_PATH, or OpenMP_ROOT if it is installed in a non-standard location." + ) + return False + return True + + # OpenMP support + def set_use_omp(self): + self._use_omp = self._omp_is_available() and self.check_compiler_for_openmp() + + def get_use_omp(self): + return self._use_omp + + def print_use_omp(self): + print(f"OpenMP Support: {enabled_dict[self.get_use_omp() or False]}") + + # OpenMP library path + def get_libomp_path(self): + """try to locate OpenMP libraries""" + return self._omp_info["library"] - def get_openmp_lib_name(self): - """Get the OpenMP library name (e.g., libomp.so or libgomp.dylib).""" - return self._openmp_lib_name + def print_libomp_path(self): + print(f"OpenMP Library Path: {self._omp_info['library'] or not_found_str}") - def set_openmp_lib_dir(self): + def set_libomp_folder(self): """Set the OpenMP library directory (containing .so or .dylib files).""" # This is set in set_openmplib_path if the library is found, otherwise it remains None. - pass + self._libomp_folder = self._omp_info["library"] and os.path.dirname( + self._omp_info["library"] + ) - def get_openmp_lib_dir(self): + def get_libomp_folder(self): """Get the OpenMP library directory (containing .so or .dylib files).""" - return self._openmp_lib_lib_dir + return self._libomp_folder # OpenMP header path - def set_openmp_lib_include_dir(self): + def set_libomp_include_path(self): """Set the OpenMP include directory (containing headers).""" - # This is set in set_openmplib_path if the library is found, otherwise it remains None. - pass + self._libomp_include_path = ( + os.path.dirname(self._omp_info["header"]) + if self._omp_info["header"] + else "" + ) - def get_openmp_include_dir(self): + def get_libomp_include_path(self): """Get the OpenMP include directory (containing headers).""" - return self._openmp_lib_include_dir - - def print_openmp_include_dir(self): - if self.get_openmp_include_dir(): - print(f"OpenMP Header Path: {self.get_openmp_include_dir()}") - - # OpenMP use detection - def set_use_OpenMP(self): - """Determine and set whether to use OpenMP.""" - compiler_supports_openmp = self.check_compiler_for_openmp() - self._use_OpenMP = compiler_supports_openmp and ( - self.get_openmp_lib_name() is not None - ) + return self._libomp_include_path - def get_use_OpenMP(self): - return self._use_OpenMP + def print_libomp_include_path(self): + print(f"OpenMP Include Path: {self.get_libomp_include_path() or not_found_str}") - def print_use_OpenMP(self): - print(f"OpenMP Support: {enabled_dict[self.get_use_OpenMP() or False]}") + def get_openmp_include_dir(self): + """Get the OpenMP include directory (containing headers).""" + return self._omp_info["header"] and os.path.dirname(self._omp_info["header"]) + # Helper functions def check_compiler_for_openmp(self): """Attempt to compile a simple OpenMP program to check if the compiler supports OpenMP.""" - if not self.cxx_compiler.get_cxx_compiler(): + if not self.cxx.get_cxx_compiler(): KeOps_Warning("No C++ compiler available to check for OpenMP support.") return False @@ -210,7 +209,7 @@ def check_compiler_for_openmp(self): test_file = f.name compile_command = [ - self.cxx_compiler.get_cxx_compiler(), + self.cxx.get_cxx_compiler(), test_file, self.get_include_options(), self.get_compile_options(), @@ -227,12 +226,15 @@ def check_compiler_for_openmp(self): return True except subprocess.CalledProcessError: os.remove(test_file) + KeOps_Warning( + f"{self.cxx.get_cxx_compiler()} does not support OpenMP. OpenMP support will be disabled." + ) return False # C++ Compiler Options def set_compile_options(self): # Add special fix for openMP prgama and Apple Clang. Order matters. - if self.cxx_compiler.get_use_Apple_clang() and self.platform.get_brew_prefix(): + if self.cxx.get_use_Apple_clang() and self.platform.get_brew_prefix(): self._compile_options += "-Xpreprocessor " self._compile_options += "-fopenmp" @@ -249,7 +251,7 @@ def get_include_options(self): # C++ linking Options def set_linking_options(self): - self._linking_options += f"-L{self.get_openmp_lib_dir()}" + self._linking_options += f"-L{self.get_libomp_folder()}" def get_linking_options(self): return self._linking_options @@ -264,9 +266,9 @@ def print_all(self): print(f"OpenMP Configuration") print("=" * 60) - self.print_use_OpenMP() - self.print_openmplib_path() - self.print_openmp_include_dir() + self.print_use_omp() + self.print_libomp_path() + self.print_libomp_include_path() # Print relevant environment variables. print_envs(self.openmp_env_vars) @@ -279,8 +281,8 @@ def print_all(self): platform_info = PlatformConfig() platform_info.print_all() - cxx_compiler_info = CxxCompilerConfig(platform_info) - cxx_compiler_info.print_all() + cxx_info = CxxCompilerConfig(platform_info) + cxx_info.print_all() - openmp_info = OpenMPConfig(platform_info, cxx_compiler_info) + openmp_info = OpenMPConfig(platform_info, cxx_info) openmp_info.print_all() From 3aa45f34aa53fa1de8acaf9948bbad530dfc9596 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 4 May 2026 17:57:11 +0200 Subject: [PATCH 27/98] fix openMp init --- keopscore/keopscore/config/OpenMP.py | 39 ++++++++++++++++++---------- 1 file changed, 25 insertions(+), 14 deletions(-) diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index a728ab5c3..018cd9e4b 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -20,7 +20,7 @@ class OpenMPConfig: Class for OpenMP detection and configuration. """ - _use_omp = False + _use_OpenMP = False _libomp_folder = "" _libomp_include_path = "" @@ -92,7 +92,7 @@ def __init__(self, platform, cxx): self.set_linking_options() # Chech if the compiler support omp - self.set_use_omp() + self.set_use_OpenMP() def find_install_path(self, lib_dict_info): result = lib_dict_info.copy() @@ -140,14 +140,15 @@ def _omp_is_available(self): return True # OpenMP support - def set_use_omp(self): - self._use_omp = self._omp_is_available() and self.check_compiler_for_openmp() + def set_use_OpenMP(self): + self._use_OpenMP = self._omp_is_available() and self.check_compiler_for_openmp() - def get_use_omp(self): - return self._use_omp + def get_use_OpenMP(self): + """Boolean to determine if OpenMP is available *and* can be used through cxx compiler""" + return self._use_OpenMP - def print_use_omp(self): - print(f"OpenMP Support: {enabled_dict[self.get_use_omp() or False]}") + def print_use_OpenMP(self): + print(f"OpenMP Support: {enabled_dict[self.get_use_OpenMP() or False]}") # OpenMP library path def get_libomp_path(self): @@ -171,11 +172,7 @@ def get_libomp_folder(self): # OpenMP header path def set_libomp_include_path(self): """Set the OpenMP include directory (containing headers).""" - self._libomp_include_path = ( - os.path.dirname(self._omp_info["header"]) - if self._omp_info["header"] - else "" - ) + self._libomp_include_path = self._omp_info["header"] def get_libomp_include_path(self): """Get the OpenMP include directory (containing headers).""" @@ -242,6 +239,9 @@ def set_compile_options(self): def get_compile_options(self): return self._compile_options + def print_compile_options(self): + print(f"Compile Options: {self.get_compile_options()}") + # C++ Compiler Include Options def set_include_options(self): self._include_options = f"-I{self.get_openmp_include_dir()}" @@ -249,6 +249,9 @@ def set_include_options(self): def get_include_options(self): return self._include_options + def print_include_options(self): + print(f"Include Options: {self.get_include_options()}") + # C++ linking Options def set_linking_options(self): self._linking_options += f"-L{self.get_libomp_folder()}" @@ -256,6 +259,9 @@ def set_linking_options(self): def get_linking_options(self): return self._linking_options + def print_linking_options(self): + print(f"Linking Options: {self.get_linking_options()}") + # OpenMP configuration printing def print_all(self): """ @@ -266,10 +272,15 @@ def print_all(self): print(f"OpenMP Configuration") print("=" * 60) - self.print_use_omp() + self.print_use_OpenMP() self.print_libomp_path() self.print_libomp_include_path() + if self.get_use_OpenMP(): + self.print_compile_options() + self.print_include_options() + self.print_linking_options() + # Print relevant environment variables. print_envs(self.openmp_env_vars) From 9dba8df3925b9187bd2ccd708e311ca017a082c8 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 4 May 2026 18:10:22 +0200 Subject: [PATCH 28/98] add fix to OpemMP detection on MacOS Co-authored-by: Copilot --- keopscore/keopscore/utils/system_utils.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/keopscore/keopscore/utils/system_utils.py b/keopscore/keopscore/utils/system_utils.py index ceef605c7..2cc70bc37 100644 --- a/keopscore/keopscore/utils/system_utils.py +++ b/keopscore/keopscore/utils/system_utils.py @@ -50,10 +50,16 @@ class LINKMAP(Structure): if res is None: return "" + if os.path.isabs(res): + return res + lib = CDLL(res) libdl = CDLL(find_library("dl")) - dlinfo = libdl.dlinfo + try: + dlinfo = libdl.dlinfo + except AttributeError: + return "" dlinfo.argtypes = c_void_p, c_int, c_void_p dlinfo.restype = c_int From db13b9c3621694d5066a845ef189a258441d740f Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 4 May 2026 22:56:15 +0200 Subject: [PATCH 29/98] Add forgotten init. Remove cxx options if checks fail. --- keopscore/keopscore/config/CxxCompiler.py | 1 + keopscore/keopscore/config/OpenMP.py | 4 ++++ keopscore/keopscore/config/Platform.py | 1 + 3 files changed, 6 insertions(+) diff --git a/keopscore/keopscore/config/CxxCompiler.py b/keopscore/keopscore/config/CxxCompiler.py index 038015afe..8af7f469a 100644 --- a/keopscore/keopscore/config/CxxCompiler.py +++ b/keopscore/keopscore/config/CxxCompiler.py @@ -27,6 +27,7 @@ def __init__(self, platform): self.platform = platform self.set_cxx_compiler() + self.set_use_Apple_clang() self.set_cxx_env_flags() self.set_compile_options() self.set_linking_options() diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 018cd9e4b..37704c13b 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -142,6 +142,10 @@ def _omp_is_available(self): # OpenMP support def set_use_OpenMP(self): self._use_OpenMP = self._omp_is_available() and self.check_compiler_for_openmp() + if not self._use_OpenMP: + self._compile_options = "" + self._include_options = "" + self._linking_options = "" def get_use_OpenMP(self): """Boolean to determine if OpenMP is available *and* can be used through cxx compiler""" diff --git a/keopscore/keopscore/config/Platform.py b/keopscore/keopscore/config/Platform.py index 05d1aaa89..fd91d091c 100644 --- a/keopscore/keopscore/config/Platform.py +++ b/keopscore/keopscore/config/Platform.py @@ -36,6 +36,7 @@ def __init__(self): self.set_python_version() self.set_python_executable() self.set_env_type() + self.set_brew_prefix() # OS Detection (Distribution Name and Version) def set_os(self): From f475ef1589f0aa5328174aad0cd3afa3a4f5f6cc Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Tue, 5 May 2026 12:19:20 +0200 Subject: [PATCH 30/98] relove _shared.py --- keopscore/keopscore/config/Cuda.py | 3 +- keopscore/keopscore/config/CxxCompiler.py | 3 +- keopscore/keopscore/config/KeOpsPath.py | 3 +- keopscore/keopscore/config/OpenMP.py | 3 +- keopscore/keopscore/config/Platform.py | 2 +- keopscore/keopscore/config/_shared.py | 41 ----------------------- keopscore/keopscore/utils/messages.py | 25 ++++++++++++++ keopscore/keopscore/utils/path_utils.py | 8 +++++ 8 files changed, 42 insertions(+), 46 deletions(-) delete mode 100644 keopscore/keopscore/config/_shared.py diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index ce61ec805..0d5bf168e 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -1,7 +1,8 @@ import ctypes import os -from keopscore.config._shared import print_envs, not_found_str, enabled_dict +from keopscore.utils.messages import print_envs +from keopscore.utils.messages import enabled_dict, not_found_str from keopscore.utils.messages import KeOps_Warning from keopscore.utils.path_utils import ( _first_matching_file, diff --git a/keopscore/keopscore/config/CxxCompiler.py b/keopscore/keopscore/config/CxxCompiler.py index 8af7f469a..8ca512630 100644 --- a/keopscore/keopscore/config/CxxCompiler.py +++ b/keopscore/keopscore/config/CxxCompiler.py @@ -1,7 +1,8 @@ import os import shutil -from keopscore.config._shared import print_envs, not_found_str +from keopscore.utils.messages import print_envs +from keopscore.utils.messages import not_found_str from keopscore.utils.messages import KeOps_Error from keopscore.utils.system_utils import KeOps_OS_Run diff --git a/keopscore/keopscore/config/KeOpsPath.py b/keopscore/keopscore/config/KeOpsPath.py index 721129be3..2c4c414d8 100644 --- a/keopscore/keopscore/config/KeOpsPath.py +++ b/keopscore/keopscore/config/KeOpsPath.py @@ -2,7 +2,8 @@ import sys import keopscore -from keopscore.config._shared import not_found_str, ensure_directory, print_envs +from keopscore.utils.path_utils import ensure_directory +from keopscore.utils.messages import not_found_str, print_envs class KeOpsPathConfig: diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 37704c13b..5349a79bc 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -2,7 +2,8 @@ import subprocess import tempfile -from keopscore.config._shared import print_envs, enabled_dict, not_found_str +from keopscore.utils.messages import print_envs +from keopscore.utils.messages import enabled_dict, not_found_str from keopscore.utils.messages import KeOps_Warning from keopscore.utils.path_utils import ( _first_matching_file, diff --git a/keopscore/keopscore/config/Platform.py b/keopscore/keopscore/config/Platform.py index fd91d091c..cef374386 100644 --- a/keopscore/keopscore/config/Platform.py +++ b/keopscore/keopscore/config/Platform.py @@ -2,7 +2,7 @@ import platform import sys -from ._shared import print_envs +from ..utils.messages import print_envs from keopscore.utils.system_utils import KeOps_OS_Run diff --git a/keopscore/keopscore/config/_shared.py b/keopscore/keopscore/config/_shared.py deleted file mode 100644 index 6476d52d3..000000000 --- a/keopscore/keopscore/config/_shared.py +++ /dev/null @@ -1,41 +0,0 @@ -import os -import sys - -from keopscore.utils.path_utils import ( - _env_roots, - _ordered_search_roots, - _path_candidates, - _python_package_roots, - _unique_paths, -) -from keopscore.utils.system_utils import _find_library_by_names - -CHECK_MARK = "✅" -CROSS_MARK = "❌" - -not_found_str = f"Not Found. {CROSS_MARK}" -enabled_dict = {True: f"Enabled {CHECK_MARK}", False: f"Disabled {CROSS_MARK}"} - - -def ensure_directory(path, add_to_syspath=False): - """Create a directory and optionally register it in sys.path.""" - os.makedirs(path, exist_ok=True) - if add_to_syspath and path not in sys.path: - sys.path.append(path) - return path - - -def print_envs(env_vars): - """Print the values of specified environment variables.""" - - if not env_vars: - return - - print("\nRelevant Environment Variables") - print("-" * 60) - for var in env_vars: - value = os.environ.get(var) - if value: - print(f"{var} = {value}") - else: - print(f"{var} is not set") diff --git a/keopscore/keopscore/utils/messages.py b/keopscore/keopscore/utils/messages.py index 290b0973c..f0f1cee9e 100644 --- a/keopscore/keopscore/utils/messages.py +++ b/keopscore/keopscore/utils/messages.py @@ -2,6 +2,14 @@ import sys + +CROSS_MARK = "❌" +CHECK_MARK = "✅" + +not_found_str = f"Not Found. {CROSS_MARK}" +enabled_dict = {True: f"Enabled {CHECK_MARK}", False: f"Disabled {CROSS_MARK}"} + + def _keops_verbose(): config = sys.modules.get("keopscore.config") debug = getattr(config, "debug", None) @@ -38,3 +46,20 @@ def KeOps_Error(message, show_line_number=True): frameinfo = getframeinfo(currentframe().f_back) message += f" (error at line {frameinfo.lineno} in file {frameinfo.filename})" raise ValueError(message) + + +def print_envs(env_vars): + """Print the values of specified environment variables.""" + + if not env_vars: + return + + print("\nRelevant Environment Variables") + print("-" * 60) + for var in env_vars: + value = os.environ.get(var) + if value: + print(f"{var} = {value}") + else: + print(f"{var} is not set") + diff --git a/keopscore/keopscore/utils/path_utils.py b/keopscore/keopscore/utils/path_utils.py index ed21fc170..a76a87670 100644 --- a/keopscore/keopscore/utils/path_utils.py +++ b/keopscore/keopscore/utils/path_utils.py @@ -96,3 +96,11 @@ def _first_matching_file(directories, patterns): continue return None + + +def ensure_directory(path, add_to_syspath=False): + """Create a directory and optionally register it in sys.path.""" + os.makedirs(path, exist_ok=True) + if add_to_syspath and path not in sys.path: + sys.path.append(path) + return path From 47a1054e2cf4acf871754ffadb96e8f16d052835 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Tue, 5 May 2026 16:04:12 +0200 Subject: [PATCH 31/98] remove global constant in chunks Co-authored-by: Copilot --- keopscore/keopscore/config/__init__.py | 2 + keopscore/keopscore/config/chunks.py | 147 ++++++++++-------- keopscore/keopscore/formulas/Chunkable_Op.py | 6 +- keopscore/keopscore/get_keops_dll.py | 30 ++-- .../mapreduce/Chunk_Mode_Constants.py | 8 +- .../mapreduce/gpu/GpuReduc1D_chunks.py | 11 +- .../mapreduce/gpu/GpuReduc1D_finalchunks.py | 21 ++- .../mapreduce/gpu/GpuReduc1D_ranges_chunks.py | 18 +-- .../gpu/GpuReduc1D_ranges_finalchunks.py | 24 +-- keopscore/keopscore/utils/code_gen_utils.py | 2 +- 10 files changed, 139 insertions(+), 130 deletions(-) diff --git a/keopscore/keopscore/config/__init__.py b/keopscore/keopscore/config/__init__.py index cf9ded87f..cc7c2b51d 100644 --- a/keopscore/keopscore/config/__init__.py +++ b/keopscore/keopscore/config/__init__.py @@ -8,6 +8,7 @@ from .KeOpsPath import KeOpsPathConfig from .OpenMP import OpenMPConfig from .Platform import PlatformConfig +from .chunks import ChunksConfig from keopscore.utils.messages import KeOps_Message # Instantiate the configurations once at import time to preserve the existing API. @@ -17,6 +18,7 @@ openmp = OpenMPConfig(platform, cxx) cuda = CudaConfig() path = KeOpsPathConfig(platform, cuda) +chunks = ChunksConfig() # flag for automatic factorization : apply automatic factorization for all formulas before reduction. auto_factorize = False diff --git a/keopscore/keopscore/config/chunks.py b/keopscore/keopscore/config/chunks.py index c1e597969..3880ff531 100644 --- a/keopscore/keopscore/config/chunks.py +++ b/keopscore/keopscore/config/chunks.py @@ -1,65 +1,82 @@ -# special computation scheme for dim>100 - -dimchunk = 64 -dim_treshold_chunk = 146 -specdims_use_chunk = [99, 100, 102, 120, 133, 138, 139, 140, 141, 142] - -dimfinalchunk = 64 - - -enable_chunk = True - - -def get_enable_chunk(): - global enable_chunk - return enable_chunk - - -def set_enable_chunk(val): - global enable_chunk - if val == 1: - enable_chunk = True - elif val == 0: - enable_chunk = False - - -# special mode for formula of the type sum_j k(x_i,y_j)*b_j with high dimensional b_j -enable_final_chunk = True - - -def set_enable_finalchunk(val): - global enable_final_chunk - if val == 1: - enable_final_chunk = True - elif val == 0: - enable_final_chunk = False - - -def get_dimfinalchunk(): - global dimfinalchunk - return dimfinalchunk - - -def set_dimfinalchunk(val): - global dimfinalchunk - dimfinalchunk = val - - -def use_final_chunks(red_formula): - global enable_final_chunk - global mult_var_highdim - global dim_treshold_chunk - return ( - enable_final_chunk and mult_var_highdim and red_formula.dim > dim_treshold_chunk - ) - - -mult_var_highdim = False - - -def set_mult_var_highdim(val): - global mult_var_highdim - if val == 1: - mult_var_highdim = True - elif val == 0: - mult_var_highdim = False +class ChunksConfig: + """ + Configuration and state management for chunked computation schemes. + Special computation scheme for dim>100 + """ + + # Constants for chunking strategy + _dimchunk = 64 + _dim_treshold_chunk = 146 + _specdims_use_chunk = [99, 100, 102, 120, 133, 138, 139, 140, 141, 142] + + _enable_chunks = True + _dimfinalchunk = 64 + _enable_final_chunk = True + _mult_var_highdim = False + + def __init__(self): + pass + + def get_dimchunk(self): + """Get the current chunk dimension.""" + return self._dimchunk + + def get_dim_treshold_chunk(self): + """Get the current dimension threshold for chunking.""" + return self._dim_treshold_chunk + + def get_specdims_use_chunk(self): + """Get the current list of specific dimensions for which chunking is used.""" + return self._specdims_use_chunk + + def get_enable_chunks(self): + """Get the current enable_chunk state.""" + return self._enable_chunks + + def set_enable_chunks(self, val): + """Set enable_chunk state from int (1=True, 0=False, -1=no change).""" + if val == 1: + self._enable_chunks = True + elif val == 0: + self._enable_chunks = False + # val == -1 means keep previous value + + def get_dimfinalchunk(self): + """Get the current final chunk dimension.""" + return self._dimfinalchunk + + def set_dimfinalchunk(self, val): + """Set the final chunk dimension.""" + self._dimfinalchunk = val + + def set_enable_finalchunk(self, val): + """Set enable_final_chunk state from int (1=True, 0=False, -1=no change).""" + if val == 1: + self._enable_final_chunk = True + elif val == 0: + self._enable_final_chunk = False + # val == -1 means keep previous value + + def get_enable_final_chunk(self): + """Get the current enable_final_chunk state.""" + return self._enable_final_chunk + + def set_mult_var_highdim(self, val): + """Set mult_var_highdim state from int (1=True, 0=False, -1=no change).""" + if val == 1: + self._mult_var_highdim = True + elif val == 0: + self._mult_var_highdim = False + # val == -1 means keep previous value + + def get_mult_var_highdim(self): + """Get the current mult_var_highdim state.""" + return self._mult_var_highdim + + def use_final_chunks(self, red_formula): + """Determine if final chunks mode should be used for this formula.""" + return ( + self.get_enable_final_chunk() + and self.get_mult_var_highdim() + and red_formula.dim > self.get_dim_treshold_chunk() + ) diff --git a/keopscore/keopscore/formulas/Chunkable_Op.py b/keopscore/keopscore/formulas/Chunkable_Op.py index 783edfba1..5031834e6 100644 --- a/keopscore/keopscore/formulas/Chunkable_Op.py +++ b/keopscore/keopscore/formulas/Chunkable_Op.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.formulas.variables.Var import Var -from keopscore.config.chunks import enable_chunk, dim_treshold_chunk, specdims_use_chunk +from keopscore.config import chunks class Chunkable_Op(Operation): @@ -19,9 +19,9 @@ def notchunked_vars(self, cat): @property def use_chunk(self): - test = enable_chunk & all(child.is_chunkable for child in self.children) + test = chunks.get_enable_chunks() & all(child.is_chunkable for child in self.children) child = self.children[0] - subtest = (child.dim >= dim_treshold_chunk) | (child.dim in specdims_use_chunk) + subtest = (child.dim >= chunks.get_dim_treshold_chunk()) | (child.dim in chunks.get_specdims_use_chunk()) test &= subtest return test diff --git a/keopscore/keopscore/get_keops_dll.py b/keopscore/keopscore/get_keops_dll.py index 3cd47fca5..6a1670e3a 100644 --- a/keopscore/keopscore/get_keops_dll.py +++ b/keopscore/keopscore/get_keops_dll.py @@ -52,16 +52,8 @@ import inspect import sys -import keopscore.config import keopscore.mapreduce -from keopscore.config.chunks import ( - get_enable_chunk, - set_enable_chunk, - dimchunk, - set_enable_finalchunk, - use_final_chunks, - set_mult_var_highdim, -) +from keopscore.config import path, cuda, debug, chunks from keopscore.formulas import Zero_Reduction, Sum_Reduction from keopscore.formulas.GetReduction import GetReduction from keopscore.formulas.variables.Zero import Zero @@ -84,19 +76,19 @@ def get_keops_dll_impl( # detecting the need for special chunked computation modes : use_chunk_mode = 0 if "Gpu" in map_reduce_id: - if not keopscore.config.cuda.get_use_cuda(): + if not cuda.get_use_cuda(): KeOps_Error( "You selected a Gpu reduce scheme but KeOps is in Cpu only mode." ) - set_enable_chunk(enable_chunks) - set_enable_finalchunk(enable_finalchunks) - set_mult_var_highdim(mul_var_highdim) + chunks.set_enable_chunks(enable_chunks) + chunks.set_enable_finalchunk(enable_finalchunks) + chunks.set_mult_var_highdim(mul_var_highdim) red_formula = GetReduction(red_formula_string, aliases) - if use_final_chunks(red_formula) and map_reduce_id != "GpuReduc2D": + if chunks.use_final_chunks(red_formula) and map_reduce_id != "GpuReduc2D": use_chunk_mode = 2 map_reduce_id += "_finalchunks" - elif get_enable_chunk() and map_reduce_id != "GpuReduc2D": - if len(red_formula.formula.chunked_formulas(dimchunk)) == 1: + elif chunks.get_enable_chunks() and map_reduce_id != "GpuReduc2D": + if len(red_formula.formula.chunked_formulas(chunks.get_dimchunk())) == 1: from keopscore.mapreduce.Chunk_Mode_Constants import ( Chunk_Mode_Constants, ) @@ -112,7 +104,7 @@ def get_keops_dll_impl( rf = map_reduce_obj.red_formula - if keopscore.config.debug.get_debug_ops(): + if debug.get_debug_ops(): KeOps_Print("In get_keops_dll, formula is :", rf) KeOps_Print("formula.__repr__() is : ", rf.__repr__()) rf.make_dot() @@ -141,7 +133,7 @@ def get_keops_dll_impl( tagZero, res["use_half"], res["use_fast_math"], - keopscore.config.cuda.get_cuda_block_size(), + cuda.get_cuda_block_size(), use_chunk_mode, tag1D2D, res["dimred"], @@ -159,7 +151,7 @@ def get_keops_dll_impl( get_keops_dll = Cache( get_keops_dll_impl, use_cache_file=True, - save_folder=keopscore.config.path.get_build_folder(), + save_folder=path.get_build_folder(), ) diff --git a/keopscore/keopscore/mapreduce/Chunk_Mode_Constants.py b/keopscore/keopscore/mapreduce/Chunk_Mode_Constants.py index 59e05e301..d7b3a5dc8 100644 --- a/keopscore/keopscore/mapreduce/Chunk_Mode_Constants.py +++ b/keopscore/keopscore/mapreduce/Chunk_Mode_Constants.py @@ -1,4 +1,4 @@ -from keopscore.config.chunks import dimchunk +from keopscore.config import chunks from keopscore.utils.code_gen_utils import GetDims, GetInds, Var_loader from keopscore.formulas.variables.Var import Var @@ -15,10 +15,10 @@ def __init__(self, red_formula): formula = red_formula.formula self.dimfout = formula.dim # dimension of output variable of inner function - chunked_formula = formula.chunked_formulas(dimchunk)[0] + chunked_formula = formula.chunked_formulas(chunks.get_dimchunk())[0] self.dim_org = chunked_formula["dim_org"] - self.nchunks = 1 + (self.dim_org - 1) // dimchunk - self.dimlastchunk = self.dim_org - (self.nchunks - 1) * dimchunk + self.nchunks = 1 + (self.dim_org - 1) // chunks.get_dimchunk() + self.dimlastchunk = self.dim_org - (self.nchunks - 1) * chunks.get_dimchunk() self.nminargs = varloader.nminargs self.fun_chunked = chunked_formula["formula"] self.dimout_chunk = self.fun_chunked.dim diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py index fc644de88..56c06cdc0 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py @@ -1,6 +1,5 @@ from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile -from keopscore.config import cuda -from keopscore.config.chunks import dimchunk +from keopscore.config import cuda, chunks from keopscore.formulas.reductions import make_sum_scheme from keopscore.mapreduce.Chunk_Mode_Constants import Chunk_Mode_Constants from keopscore.mapreduce.MapReduce import MapReduce @@ -55,7 +54,7 @@ def do_chunk_sub( ) load_chunks_routine_i = load_vars_chunks( indsi_chunked, - dimchunk, + chunks.get_dimchunk(), dimchunk_curr, chk.dim_org, xiloc, @@ -65,7 +64,7 @@ def do_chunk_sub( ) load_chunks_routine_j = load_vars_chunks( indsj_chunked, - dimchunk, + chunks.get_dimchunk(), dimchunk_curr, chk.dim_org, yjloc, @@ -75,7 +74,7 @@ def do_chunk_sub( ) load_chunks_routine_p = load_vars_chunks( indsp_chunked, - dimchunk, + chunks.get_dimchunk(), dimchunk_curr, chk.dim_org, param_loc, @@ -182,7 +181,7 @@ def get_code(self): dtype, red_formula, chk.fun_chunked, - dimchunk, + chunks.get_dimchunk(), chk.dimsx, chk.dimsy, chk.dimsp, diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py index d171e9df0..ea996db86 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py @@ -1,6 +1,5 @@ from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile -from keopscore.config import cuda -from keopscore.config.chunks import dimfinalchunk +from keopscore.config import cuda, chunks from keopscore.formulas.reductions.Sum_Reduction import Sum_Reduction from keopscore.formulas.reductions.sum_schemes import block_sum from keopscore.mapreduce.MapReduce import MapReduce @@ -36,10 +35,10 @@ def do_finalchunk_sub( out, ): dimout = varfinal.dim - yjloc = c_variable(pointer(dtype), f"({yj.id} + threadIdx.x * {dimfinalchunk})") + yjloc = c_variable(pointer(dtype), f"({yj.id} + threadIdx.x * {chunks.get_dimfinalchunk()})") load_chunks_routine_j = load_vars_chunks( [varfinal.ind], - dimfinalchunk, + chunks.get_dimfinalchunk(), dimfinalchunk_curr, varfinal.dim, yjloc, @@ -54,7 +53,7 @@ def do_finalchunk_sub( {load_chunks_routine_j} }} __syncthreads(); - for (signed long int jrel = 0; (jrel < blockDim.x) && (jrel < {ny.id} - {jstart.id}); jrel++, yjrel += {dimfinalchunk}) {{ + for (signed long int jrel = 0; (jrel < blockDim.x) && (jrel < {ny.id} - {jstart.id}); jrel++, yjrel += {chunks.get_dimfinalchunk()}) {{ if ({i.id} < {nx.id}) {{ // we compute only if needed {use_pragma_unroll()} for (signed long int k=0; k<{dimfinalchunk_curr}; k++) {{ @@ -66,7 +65,7 @@ def do_finalchunk_sub( if ({i.id} < {nx.id}) {{ {use_pragma_unroll()} for (signed long int k=0; k<{dimfinalchunk_curr}; k++) - {out.id}[i*{dimout}+{chunk.id}*{dimfinalchunk}+k] += {acc.id}[k]; + {out.id}[i*{dimout}+{chunk.id}*{chunks.get_dimfinalchunk()}+k] += {acc.id}[k]; }} __syncthreads(); """ @@ -101,8 +100,8 @@ def get_code(self): ) formula = fun_internal.formula varfinal = self.red_formula.formula.children[1 - ind_fun_internal] - nchunks = 1 + (varfinal.dim - 1) // dimfinalchunk - dimlastfinalchunk = varfinal.dim - (nchunks - 1) * dimfinalchunk + nchunks = 1 + (varfinal.dim - 1) // chunks.get_dimfinalchunk() + dimlastfinalchunk = varfinal.dim - (nchunks - 1) * chunks.get_dimfinalchunk() varloader = Var_loader(fun_internal) dimsx = varloader.dimsx dimsy = varloader.dimsy @@ -119,7 +118,7 @@ def get_code(self): KeOps_Error("dimfout should be 1") sum_scheme = self.sum_scheme - self.dimy = max(dimfinalchunk, dimy) + self.dimy = max(chunks.get_dimfinalchunk(), dimy) blocksize_chunks = min( cuda.get_cuda_block_size(), 1024, @@ -131,7 +130,7 @@ def get_code(self): param_loc = c_array(dtype, dimp, "param_loc") fout = c_array(dtype, dimfout * blocksize_chunks, "fout") xi = c_array(dtype, dimx, "xi") - acc = c_array(dtypeacc, dimfinalchunk, "acc") + acc = c_array(dtypeacc, chunks.get_dimfinalchunk(), "acc") yjloc = c_array(dtype, dimy, f"(yj + threadIdx.x * {dimy})") foutjrel = c_array(dtype, dimfout, f"({fout.id}+jrel*{dimfout})") yjrel = c_array(dtype, dimy, "yjrel") @@ -142,7 +141,7 @@ def get_code(self): chunk_sub_routine = do_finalchunk_sub( dtype, varfinal, - dimfinalchunk, + chunks.get_dimfinalchunk(), acc, i, j, diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py index b96fafedc..592da43d5 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py @@ -1,6 +1,6 @@ +from keopscore.config import cuda, chunks + from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile -from keopscore.config import cuda -from keopscore.config.chunks import dimchunk from keopscore.formulas.reductions import make_sum_scheme from keopscore.mapreduce.Chunk_Mode_Constants import Chunk_Mode_Constants from keopscore.mapreduce.MapReduce import MapReduce @@ -63,7 +63,7 @@ def do_chunk_sub_ranges( load_chunks_routine_i = load_vars_chunks( indsi_chunked, - dimchunk, + chunks.get_dimchunk(), dimchunk_curr, chk.dim_org, xiloc, @@ -81,7 +81,7 @@ def do_chunk_sub_ranges( load_chunks_routine_i_batches = load_vars_chunks_offsets( indsi_chunked, indsi_global, - dimchunk, + chunks.get_dimchunk(), dimchunk_curr, chk.dim_org, xiloc, @@ -93,7 +93,7 @@ def do_chunk_sub_ranges( load_chunks_routine_j = load_vars_chunks( indsj_chunked, - dimchunk, + chunks.get_dimchunk(), dimchunk_curr, chk.dim_org, yjloc, @@ -105,7 +105,7 @@ def do_chunk_sub_ranges( load_chunks_routine_j_batches = load_vars_chunks_offsets( indsj_chunked, indsj_global, - dimchunk, + chunks.get_dimchunk(), dimchunk_curr, chk.dim_org, yjloc, @@ -117,7 +117,7 @@ def do_chunk_sub_ranges( load_chunks_routine_p = load_vars_chunks( indsp_chunked, - dimchunk, + chunks.get_dimchunk(), dimchunk_curr, chk.dim_org, param_loc, @@ -128,7 +128,7 @@ def do_chunk_sub_ranges( load_chunks_routine_p_batches = load_vars_chunks_offsets( indsp_chunked, indsp_global, - dimchunk, + chunks.get_dimchunk(), dimchunk_curr, chk.dim_org, param_loc, @@ -281,7 +281,7 @@ def get_code(self): dtype, red_formula, chk.fun_chunked, - dimchunk, + chunks.get_dimchunk(), chk.dimsx, chk.dimsy, chk.dimsp, diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py index afc8aaed8..d913348c7 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py @@ -1,6 +1,6 @@ +from keopscore.config import cuda, chunks + from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile -from keopscore.config import cuda -from keopscore.config.chunks import dimfinalchunk from keopscore.formulas.reductions.Sum_Reduction import Sum_Reduction from keopscore.formulas.reductions.sum_schemes import block_sum from keopscore.mapreduce.MapReduce import MapReduce @@ -41,11 +41,11 @@ def do_finalchunk_sub_ranges( out, ): dimout = varfinal.dim - yjloc = c_variable(pointer(dtype), f"({yj.id} + threadIdx.x * {dimfinalchunk})") + yjloc = c_variable(pointer(dtype), f"({yj.id} + threadIdx.x * {chunks.get_dimfinalchunk()})") indsj_global = Var_loader(fun_global).indsj load_chunks_routine_j = load_vars_chunks( [varfinal.ind], - dimfinalchunk, + chunks.get_dimfinalchunk(), dimfinalchunk_curr, varfinal.dim, yjloc, @@ -56,7 +56,7 @@ def do_finalchunk_sub_ranges( load_chunks_routine_j_ranges = load_vars_chunks_offsets( [varfinal.ind], indsj_global, - dimfinalchunk, + chunks.get_dimfinalchunk(), dimfinalchunk_curr, varfinal.dim, yjloc, @@ -76,7 +76,7 @@ def do_finalchunk_sub_ranges( }} }} __syncthreads(); - for (signed long int jrel = 0; (jrel < blockDim.x) && (jrel < {end_y.id} - {jstart.id}); jrel++, yjrel += {dimfinalchunk}) {{ + for (signed long int jrel = 0; (jrel < blockDim.x) && (jrel < {end_y.id} - {jstart.id}); jrel++, yjrel += {chunks.get_dimfinalchunk()}) {{ if ({i.id} < {end_x.id}) {{ // we compute only if needed {use_pragma_unroll()} for (signed long int k=0; k<{dimfinalchunk_curr}; k++) {{ @@ -88,7 +88,7 @@ def do_finalchunk_sub_ranges( if ({i.id} < {end_x.id}) {{ {use_pragma_unroll()} for (signed long int k=0; k<{dimfinalchunk_curr}; k++) - {out.id}[i*{dimout}+{chunk.id}*{dimfinalchunk}+k] += {acc.id}[k]; + {out.id}[i*{dimout}+{chunk.id}*{chunks.get_dimfinalchunk()}+k] += {acc.id}[k]; }} __syncthreads(); """ @@ -124,8 +124,8 @@ def get_code(self): ) formula = fun_internal.formula varfinal = self.red_formula.formula.children[1 - ind_fun_internal] - nchunks = 1 + (varfinal.dim - 1) // dimfinalchunk - dimlastfinalchunk = varfinal.dim - (nchunks - 1) * dimfinalchunk + nchunks = 1 + (varfinal.dim - 1) // chunks.get_dimfinalchunk() + dimlastfinalchunk = varfinal.dim - (nchunks - 1) * chunks.get_dimfinalchunk() varloader = Var_loader(fun_internal) dimsx = varloader.dimsx dimsy = varloader.dimsy @@ -142,7 +142,7 @@ def get_code(self): KeOps_Error("dimfout should be 1") sum_scheme = self.sum_scheme - self.dimy = max(dimfinalchunk, dimy) + self.dimy = max(chunks.get_dimfinalchunk(), dimy) blocksize_chunks = min( cuda.get_cuda_block_size(), 1024, @@ -154,7 +154,7 @@ def get_code(self): param_loc = c_array(dtype, dimp, "param_loc") fout = c_array(dtype, dimfout * blocksize_chunks, "fout") xi = c_array(dtype, dimx, "xi") - acc = c_array(dtypeacc, dimfinalchunk, "acc") + acc = c_array(dtypeacc, chunks.get_dimfinalchunk(), "acc") yjloc = c_array(dtype, dimy, f"(yj + threadIdx.x * {dimy})") foutjrel = c_array(dtype, dimfout, f"({fout.id}+jrel*{dimfout})") yjrel = c_array(dtype, dimy, "yjrel") @@ -204,7 +204,7 @@ def get_code(self): dtype, fun_global, varfinal, - dimfinalchunk, + chunks.get_dimfinalchunk(), acc, i, j, diff --git a/keopscore/keopscore/utils/code_gen_utils.py b/keopscore/keopscore/utils/code_gen_utils.py index 3dfc1d79c..c75b09573 100644 --- a/keopscore/keopscore/utils/code_gen_utils.py +++ b/keopscore/keopscore/utils/code_gen_utils.py @@ -1,7 +1,7 @@ from hashlib import sha256 import keopscore.config -from keopscore.utils.messages import KeOps_Error, KeOps_Message +from keopscore.utils.messages import KeOps_Error def get_hash_name(*args): From 9c4a832f1ea7c0c9911639e3d769b5e3462e0bf4 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Tue, 5 May 2026 16:07:27 +0200 Subject: [PATCH 32/98] rename chunks --- keopscore/keopscore/config/{chunks.py => Chunks.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename keopscore/keopscore/config/{chunks.py => Chunks.py} (100%) diff --git a/keopscore/keopscore/config/chunks.py b/keopscore/keopscore/config/Chunks.py similarity index 100% rename from keopscore/keopscore/config/chunks.py rename to keopscore/keopscore/config/Chunks.py From 979f2e1452cfe5804931591c21463e4d18816ab8 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Tue, 5 May 2026 23:18:45 +0200 Subject: [PATCH 33/98] add adaptive verbosity level Co-authored-by: Copilot --- doc/python/installation.rst | 4 +- keopscore/keopscore/binders/LinkCompile.py | 3 +- keopscore/keopscore/config/Debug.py | 25 ++++++-- keopscore/keopscore/config/__init__.py | 2 +- keopscore/keopscore/utils/messages.py | 20 +++--- pykeops/pykeops/__init__.py | 29 +++------ pykeops/pykeops/common/utils.py | 14 ++--- pykeops/pykeops/config.py | 71 ++++++++++++++++++++-- pykeops/pykeops/numpy/test_install.py | 17 ++++-- pykeops/pykeops/torch/test_install.py | 15 ++++- 10 files changed, 143 insertions(+), 57 deletions(-) diff --git a/doc/python/installation.rst b/doc/python/installation.rst index 599a7c591..a2aaaf796 100644 --- a/doc/python/installation.rst +++ b/doc/python/installation.rst @@ -331,4 +331,6 @@ Alternatively, you can disable verbose compilation from your python script using .. code-block:: python import pykeops - pykeops.set_verbose(False) + pykeops.set_verbose(0) # no output + pykeops.set_verbose(1) # default verbosity level + pykeops.set_verbose(2) # maximum verbosity level \ No newline at end of file diff --git a/keopscore/keopscore/binders/LinkCompile.py b/keopscore/keopscore/binders/LinkCompile.py index 0d010cbe0..f7a76af74 100644 --- a/keopscore/keopscore/binders/LinkCompile.py +++ b/keopscore/keopscore/binders/LinkCompile.py @@ -102,10 +102,11 @@ def get_dll_and_params(self): "Generating code for " + self.red_formula.__str__() + " ... ", flush=True, end="", + level=2 ) self.generate_code() self.save_info() - KeOps_Message("OK", use_tag=False, flush=True) + KeOps_Message("OK", use_tag=False, flush=True, level=2) else: self.read_info() return dict( diff --git a/keopscore/keopscore/config/Debug.py b/keopscore/keopscore/config/Debug.py index 12cc46ccd..b13a11ff8 100644 --- a/keopscore/keopscore/config/Debug.py +++ b/keopscore/keopscore/config/Debug.py @@ -1,5 +1,7 @@ import os +from keopscore.utils.messages import KeOps_Warning + class DebugConfig: @@ -8,11 +10,17 @@ class DebugConfig: # adds C++ code for printing all input and output values for all atomic operations during computations _debug_ops_at_exec = False - # Verbosity level (default is False unless KEOPS_VERBOSE define and not 0) - _verbose = os.getenv("KEOPS_VERBOSE") != "0" - def __init__(self): - pass + env_val = os.getenv("KEOPS_VERBOSE") + if env_val is None: + self._verbose = 1 + else: + val = int(env_val) + if val in (0, 1, 2): + self._verbose = val + else: + KeOps_Warning(f"Invalid KEOPS_VERBOSE value: {env_val}. Verbose level must be 0, 1 or 2. Defaulting to 1.") + self._verbose = 1 def set_debug_ops(self, debug_ops): self._debug_ops = debug_ops @@ -29,5 +37,10 @@ def get_debug_ops_at_exec(self): def get_verbose(self): return self._verbose - def set_verbose(self, verbose): - self._verbose = verbose + def set_verbose(self, val): + if val in (0, 1, 2): + self._verbose = val + else: + KeOps_Warning(f"Invalid verbose value: {val}. Verbose level must be 0, 1 or 2. Keeping previous value: {self._verbose}.") + return + diff --git a/keopscore/keopscore/config/__init__.py b/keopscore/keopscore/config/__init__.py index cc7c2b51d..9543fb77b 100644 --- a/keopscore/keopscore/config/__init__.py +++ b/keopscore/keopscore/config/__init__.py @@ -8,7 +8,7 @@ from .KeOpsPath import KeOpsPathConfig from .OpenMP import OpenMPConfig from .Platform import PlatformConfig -from .chunks import ChunksConfig +from .Chunks import ChunksConfig from keopscore.utils.messages import KeOps_Message # Instantiate the configurations once at import time to preserve the existing API. diff --git a/keopscore/keopscore/utils/messages.py b/keopscore/keopscore/utils/messages.py index f0f1cee9e..fd302a29e 100644 --- a/keopscore/keopscore/utils/messages.py +++ b/keopscore/keopscore/utils/messages.py @@ -15,27 +15,29 @@ def _keops_verbose(): debug = getattr(config, "debug", None) if debug is not None: return debug.get_verbose() + else: + message = "[KeOps] Warning : Could not access to verbosity level. Defaulting to verbose level 1." + print(message) - keopscore = sys.modules.get("keopscore") - return getattr(keopscore, "verbose", os.getenv("KEOPS_VERBOSE") != "0") + return 1 # default verbose level if config or debug is not available -def KeOps_Print(*messages, force_print=False, **kwargs): - if _keops_verbose() or force_print: +def KeOps_Print(*messages, force_print=False, level=1, **kwargs): + if _keops_verbose() >= level or force_print: print(*messages, **kwargs) -def KeOps_Message(message, use_tag=True, **kwargs): - if _keops_verbose(): +def KeOps_Message(message, use_tag=True, level=1, **kwargs): + if _keops_verbose() >= level: tag = "[KeOps] " if use_tag else "" message = tag + message print(message, **kwargs) -def KeOps_Warning(message, newline=False): - if _keops_verbose(): +def KeOps_Warning(message, newline=False, level=1, **kwargs): + if _keops_verbose() >= level: message = ("\n" if newline else "") + "[KeOps] Warning : " + message - print(message) + print(message, **kwargs) def KeOps_Error(message, show_line_number=True): diff --git a/pykeops/pykeops/__init__.py b/pykeops/pykeops/__init__.py index 98bbf13e8..67d30809a 100644 --- a/pykeops/pykeops/__init__.py +++ b/pykeops/pykeops/__init__.py @@ -1,32 +1,23 @@ +import sys import os import keopscore -############################################################## -# Verbosity level (we must do this before importing keopscore) -verbose = True -if os.getenv("PYKEOPS_VERBOSE") == "0": - verbose = False - os.environ["KEOPS_VERBOSE"] = "0" - - from . import config as pykeopsconfig -def set_verbose(val): - global verbose - verbose = val - keopscore.config.debug.set_verbose(val) +########################################################### +# PykeOps version +__version__ = pykeopsconfig.get_version() -########################################################### -# Set version +############################################################## +# Verbosity level (we must do this before importing keopscore) -with open( - os.path.join(os.path.abspath(os.path.dirname(__file__)), "keops_version"), - encoding="utf-8", -) as v: - __version__ = v.read().rstrip() +verbose = pykeopsconfig.init_verbose() +def set_verbose(val): + pykeopsconfig.set_verbose(val) + sys.modules[__name__].verbose = pykeopsconfig.get_verbose() ########################################################### # Utils diff --git a/pykeops/pykeops/common/utils.py b/pykeops/pykeops/common/utils.py index 104485ee9..8d4c6230a 100644 --- a/pykeops/pykeops/common/utils.py +++ b/pykeops/pykeops/common/utils.py @@ -95,19 +95,19 @@ def check_broadcasting(dims_1, dims_2): return max_tuple(padded_dims_1, padded_dims_2) -def pyKeOps_Print(message, **kwargs): - if pykeops.verbose: +def pyKeOps_Print(message, level=1, **kwargs): + if pykeops.verbose >= level: print(message, **kwargs) -def pyKeOps_Message(message, use_tag=True, **kwargs): - if pykeops.verbose: +def pyKeOps_Message(message, use_tag=True, level=1, **kwargs): + if pykeops.verbose >= level: tag = "[pyKeOps] " if use_tag else "" message = tag + message print(message, **kwargs) -def pyKeOps_Warning(message): - if pykeops.verbose: +def pyKeOps_Warning(message, level=1, **kwargs): + if pykeops.verbose >= level: message = "[pyKeOps] Warning : " + message - print(message) + print(message, **kwargs) diff --git a/pykeops/pykeops/config.py b/pykeops/pykeops/config.py index 2892ba04a..3afefe89c 100644 --- a/pykeops/pykeops/config.py +++ b/pykeops/pykeops/config.py @@ -3,11 +3,6 @@ import sys import sysconfig -############################################################### -# Initialize some variables: the values may be redefined later - -numpy_found = importlib.util.find_spec("numpy") is not None -torch_found = importlib.util.find_spec("torch") is not None # Instantiating the keopscore.config main classes for pykeops import keopscore.config @@ -16,11 +11,77 @@ path = keopscore.config.path openmp = keopscore.config.openmp cxx = keopscore.config.cxx +debug = keopscore.config.debug get_build_folder = path.get_build_folder gpu_available = cuda.get_use_cuda() +# Initialize some variables: the values may be redefined later + +numpy_found = importlib.util.find_spec("numpy") is not None +torch_found = importlib.util.find_spec("torch") is not None + + +# Verbosity level + + +class _VerboseConfig: + def __init__(self, level=1): + self.level = level + + +_verbose_state = _VerboseConfig(1) + +def set_verbose(val): + val = int(val) + if val not in (0, 1, 2): + raise ValueError("Verbose level must be 0, 1 or 2.") + + _verbose_state.level = val + os.environ["KEOPS_VERBOSE"] = str(val) + debug.set_verbose(val) + return _verbose_state.level + +def init_verbose(): + env_val = os.getenv("PYKEOPS_VERBOSE") + if env_val in ("0", "1", "2"): + return set_verbose(int(env_val)) + + return set_verbose(1) + + +def get_verbose(): + return _verbose_state.level + + +# version +def read_version(version_file): + with open(version_file, encoding="utf-8") as v: + return v.read().rstrip() + +_version = read_version(os.path.join(os.path.abspath(os.path.dirname(__file__)), "keops_version")) + +def get_version(): + assert _version == keopscore.__version__, f"Version mismatch between pykeops and keopscore: {_version} vs {keopscore.__version__}" + return _version + + + + + + + + + + + + + + + + + def pykeops_nvrtc_name(type="src"): basename = "pykeops_nvrtc" extension = ".cpp" if type == "src" else sysconfig.get_config_var("EXT_SUFFIX") diff --git a/pykeops/pykeops/numpy/test_install.py b/pykeops/pykeops/numpy/test_install.py index bd741fb56..0682f5f30 100644 --- a/pykeops/pykeops/numpy/test_install.py +++ b/pykeops/pykeops/numpy/test_install.py @@ -9,8 +9,7 @@ def test_numpy_bindings(): """ - This function try to compile a simple keops formula using the numpy binder. - + Try to compile a simple KeOps formula using the NumPy binder. """ x = np.arange(1, 10).reshape(-1, 3).astype("float32") y = np.arange(3, 9).reshape(-1, 3).astype("float32") @@ -18,7 +17,15 @@ def test_numpy_bindings(): import pykeops.numpy as pknp my_conv = pknp.Genred(formula, var) - if np.allclose(my_conv(x, y).flatten(), expected_res): - pyKeOps_Message("pyKeOps with numpy bindings is working!", use_tag=False) + + try: + keops_res = my_conv(x, y).flatten() + except Exception as e: + raise ValueError(f"Error during computation: {e}", use_tag=False) + + if np.allclose(keops_res, expected_res): + pyKeOps_Message("pyKeOps with torch bindings is working!", use_tag=False, level=1) + return True else: - pyKeOps_Message("outputs wrong values...", use_tag=False) + pyKeOps_Message(f"outputs wrong values: expected {expected_res} but get {keops_res}", use_tag=False, level=1) + return False \ No newline at end of file diff --git a/pykeops/pykeops/torch/test_install.py b/pykeops/pykeops/torch/test_install.py index 541ec3231..119329ab8 100644 --- a/pykeops/pykeops/torch/test_install.py +++ b/pykeops/pykeops/torch/test_install.py @@ -9,7 +9,7 @@ def test_torch_bindings(): """ - This function try to compile a simple keops formula using the pytorch binder. + Try to compile a simple KeOps formula using the PyTorch binder. """ x = torch.arange(1, 10, dtype=torch.float32).view(-1, 3) y = torch.arange(3, 9, dtype=torch.float32).view(-1, 3) @@ -17,9 +17,18 @@ def test_torch_bindings(): import pykeops.torch as pktorch my_conv = pktorch.Genred(formula, var) + + try: + keops_res = my_conv(x, y).view(-1) + except Exception as e: + pyKeOps_Message(f"Error during computation: {e}", use_tag=False) + return False + if torch.allclose( - my_conv(x, y).view(-1), torch.tensor(expected_res).type(torch.float32) + keops_res, torch.tensor(expected_res, dtype=torch.float32) ): pyKeOps_Message("pyKeOps with torch bindings is working!", use_tag=False) + return True else: - pyKeOps_Message("outputs wrong values...", use_tag=False) + pyKeOps_Message(f"outputs wrong values: expected {expected_res} but get {keops_res}", use_tag=False) + return False \ No newline at end of file From dbc8c95dec3b8954f64972bbf3ecdbb9fca7d48b Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Tue, 5 May 2026 23:22:25 +0200 Subject: [PATCH 34/98] lint --- keopscore/keopscore/binders/LinkCompile.py | 2 +- keopscore/keopscore/config/Chunks.py | 28 +++++++++---------- keopscore/keopscore/config/Debug.py | 9 ++++-- keopscore/keopscore/formulas/Chunkable_Op.py | 8 ++++-- .../mapreduce/gpu/GpuReduc1D_finalchunks.py | 4 ++- .../gpu/GpuReduc1D_ranges_finalchunks.py | 4 ++- keopscore/keopscore/utils/messages.py | 3 -- pykeops/pykeops/__init__.py | 4 ++- pykeops/pykeops/config.py | 28 +++++++------------ pykeops/pykeops/numpy/test_install.py | 14 +++++++--- pykeops/pykeops/torch/test_install.py | 13 +++++---- 11 files changed, 63 insertions(+), 54 deletions(-) diff --git a/keopscore/keopscore/binders/LinkCompile.py b/keopscore/keopscore/binders/LinkCompile.py index f7a76af74..0f985526f 100644 --- a/keopscore/keopscore/binders/LinkCompile.py +++ b/keopscore/keopscore/binders/LinkCompile.py @@ -102,7 +102,7 @@ def get_dll_and_params(self): "Generating code for " + self.red_formula.__str__() + " ... ", flush=True, end="", - level=2 + level=2, ) self.generate_code() self.save_info() diff --git a/keopscore/keopscore/config/Chunks.py b/keopscore/keopscore/config/Chunks.py index 3880ff531..aadc6e800 100644 --- a/keopscore/keopscore/config/Chunks.py +++ b/keopscore/keopscore/config/Chunks.py @@ -1,22 +1,22 @@ class ChunksConfig: """ - Configuration and state management for chunked computation schemes. + Configuration and state management for chunked computation schemes. Special computation scheme for dim>100 """ - + # Constants for chunking strategy _dimchunk = 64 _dim_treshold_chunk = 146 _specdims_use_chunk = [99, 100, 102, 120, 133, 138, 139, 140, 141, 142] - + _enable_chunks = True _dimfinalchunk = 64 _enable_final_chunk = True _mult_var_highdim = False - + def __init__(self): pass - + def get_dimchunk(self): """Get the current chunk dimension.""" return self._dimchunk @@ -24,15 +24,15 @@ def get_dimchunk(self): def get_dim_treshold_chunk(self): """Get the current dimension threshold for chunking.""" return self._dim_treshold_chunk - + def get_specdims_use_chunk(self): """Get the current list of specific dimensions for which chunking is used.""" return self._specdims_use_chunk - + def get_enable_chunks(self): """Get the current enable_chunk state.""" return self._enable_chunks - + def set_enable_chunks(self, val): """Set enable_chunk state from int (1=True, 0=False, -1=no change).""" if val == 1: @@ -40,15 +40,15 @@ def set_enable_chunks(self, val): elif val == 0: self._enable_chunks = False # val == -1 means keep previous value - + def get_dimfinalchunk(self): """Get the current final chunk dimension.""" return self._dimfinalchunk - + def set_dimfinalchunk(self, val): """Set the final chunk dimension.""" self._dimfinalchunk = val - + def set_enable_finalchunk(self, val): """Set enable_final_chunk state from int (1=True, 0=False, -1=no change).""" if val == 1: @@ -60,7 +60,7 @@ def set_enable_finalchunk(self, val): def get_enable_final_chunk(self): """Get the current enable_final_chunk state.""" return self._enable_final_chunk - + def set_mult_var_highdim(self, val): """Set mult_var_highdim state from int (1=True, 0=False, -1=no change).""" if val == 1: @@ -72,11 +72,11 @@ def set_mult_var_highdim(self, val): def get_mult_var_highdim(self): """Get the current mult_var_highdim state.""" return self._mult_var_highdim - + def use_final_chunks(self, red_formula): """Determine if final chunks mode should be used for this formula.""" return ( self.get_enable_final_chunk() - and self.get_mult_var_highdim() + and self.get_mult_var_highdim() and red_formula.dim > self.get_dim_treshold_chunk() ) diff --git a/keopscore/keopscore/config/Debug.py b/keopscore/keopscore/config/Debug.py index b13a11ff8..be50213e7 100644 --- a/keopscore/keopscore/config/Debug.py +++ b/keopscore/keopscore/config/Debug.py @@ -19,7 +19,9 @@ def __init__(self): if val in (0, 1, 2): self._verbose = val else: - KeOps_Warning(f"Invalid KEOPS_VERBOSE value: {env_val}. Verbose level must be 0, 1 or 2. Defaulting to 1.") + KeOps_Warning( + f"Invalid KEOPS_VERBOSE value: {env_val}. Verbose level must be 0, 1 or 2. Defaulting to 1." + ) self._verbose = 1 def set_debug_ops(self, debug_ops): @@ -41,6 +43,7 @@ def set_verbose(self, val): if val in (0, 1, 2): self._verbose = val else: - KeOps_Warning(f"Invalid verbose value: {val}. Verbose level must be 0, 1 or 2. Keeping previous value: {self._verbose}.") + KeOps_Warning( + f"Invalid verbose value: {val}. Verbose level must be 0, 1 or 2. Keeping previous value: {self._verbose}." + ) return - diff --git a/keopscore/keopscore/formulas/Chunkable_Op.py b/keopscore/keopscore/formulas/Chunkable_Op.py index 5031834e6..b38c92751 100644 --- a/keopscore/keopscore/formulas/Chunkable_Op.py +++ b/keopscore/keopscore/formulas/Chunkable_Op.py @@ -19,9 +19,13 @@ def notchunked_vars(self, cat): @property def use_chunk(self): - test = chunks.get_enable_chunks() & all(child.is_chunkable for child in self.children) + test = chunks.get_enable_chunks() & all( + child.is_chunkable for child in self.children + ) child = self.children[0] - subtest = (child.dim >= chunks.get_dim_treshold_chunk()) | (child.dim in chunks.get_specdims_use_chunk()) + subtest = (child.dim >= chunks.get_dim_treshold_chunk()) | ( + child.dim in chunks.get_specdims_use_chunk() + ) test &= subtest return test diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py index ea996db86..6d34113cd 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py @@ -35,7 +35,9 @@ def do_finalchunk_sub( out, ): dimout = varfinal.dim - yjloc = c_variable(pointer(dtype), f"({yj.id} + threadIdx.x * {chunks.get_dimfinalchunk()})") + yjloc = c_variable( + pointer(dtype), f"({yj.id} + threadIdx.x * {chunks.get_dimfinalchunk()})" + ) load_chunks_routine_j = load_vars_chunks( [varfinal.ind], chunks.get_dimfinalchunk(), diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py index d913348c7..54c9ecd14 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py @@ -41,7 +41,9 @@ def do_finalchunk_sub_ranges( out, ): dimout = varfinal.dim - yjloc = c_variable(pointer(dtype), f"({yj.id} + threadIdx.x * {chunks.get_dimfinalchunk()})") + yjloc = c_variable( + pointer(dtype), f"({yj.id} + threadIdx.x * {chunks.get_dimfinalchunk()})" + ) indsj_global = Var_loader(fun_global).indsj load_chunks_routine_j = load_vars_chunks( [varfinal.ind], diff --git a/keopscore/keopscore/utils/messages.py b/keopscore/keopscore/utils/messages.py index fd302a29e..8635f706e 100644 --- a/keopscore/keopscore/utils/messages.py +++ b/keopscore/keopscore/utils/messages.py @@ -1,8 +1,6 @@ import os import sys - - CROSS_MARK = "❌" CHECK_MARK = "✅" @@ -64,4 +62,3 @@ def print_envs(env_vars): print(f"{var} = {value}") else: print(f"{var} is not set") - diff --git a/pykeops/pykeops/__init__.py b/pykeops/pykeops/__init__.py index 67d30809a..e6331f17e 100644 --- a/pykeops/pykeops/__init__.py +++ b/pykeops/pykeops/__init__.py @@ -5,7 +5,6 @@ from . import config as pykeopsconfig - ########################################################### # PykeOps version @@ -15,10 +14,13 @@ # Verbosity level (we must do this before importing keopscore) verbose = pykeopsconfig.init_verbose() + + def set_verbose(val): pykeopsconfig.set_verbose(val) sys.modules[__name__].verbose = pykeopsconfig.get_verbose() + ########################################################### # Utils diff --git a/pykeops/pykeops/config.py b/pykeops/pykeops/config.py index 3afefe89c..8d2085e19 100644 --- a/pykeops/pykeops/config.py +++ b/pykeops/pykeops/config.py @@ -3,7 +3,6 @@ import sys import sysconfig - # Instantiating the keopscore.config main classes for pykeops import keopscore.config @@ -33,6 +32,7 @@ def __init__(self, level=1): _verbose_state = _VerboseConfig(1) + def set_verbose(val): val = int(val) if val not in (0, 1, 2): @@ -43,6 +43,7 @@ def set_verbose(val): debug.set_verbose(val) return _verbose_state.level + def init_verbose(): env_val = os.getenv("PYKEOPS_VERBOSE") if env_val in ("0", "1", "2"): @@ -60,26 +61,17 @@ def read_version(version_file): with open(version_file, encoding="utf-8") as v: return v.read().rstrip() -_version = read_version(os.path.join(os.path.abspath(os.path.dirname(__file__)), "keops_version")) - -def get_version(): - assert _version == keopscore.__version__, f"Version mismatch between pykeops and keopscore: {_version} vs {keopscore.__version__}" - return _version - - - - - - - - - - - - +_version = read_version( + os.path.join(os.path.abspath(os.path.dirname(__file__)), "keops_version") +) +def get_version(): + assert ( + _version == keopscore.__version__ + ), f"Version mismatch between pykeops and keopscore: {_version} vs {keopscore.__version__}" + return _version def pykeops_nvrtc_name(type="src"): diff --git a/pykeops/pykeops/numpy/test_install.py b/pykeops/pykeops/numpy/test_install.py index 0682f5f30..3d78bf1a5 100644 --- a/pykeops/pykeops/numpy/test_install.py +++ b/pykeops/pykeops/numpy/test_install.py @@ -22,10 +22,16 @@ def test_numpy_bindings(): keops_res = my_conv(x, y).flatten() except Exception as e: raise ValueError(f"Error during computation: {e}", use_tag=False) - + if np.allclose(keops_res, expected_res): - pyKeOps_Message("pyKeOps with torch bindings is working!", use_tag=False, level=1) + pyKeOps_Message( + "pyKeOps with torch bindings is working!", use_tag=False, level=1 + ) return True else: - pyKeOps_Message(f"outputs wrong values: expected {expected_res} but get {keops_res}", use_tag=False, level=1) - return False \ No newline at end of file + pyKeOps_Message( + f"outputs wrong values: expected {expected_res} but get {keops_res}", + use_tag=False, + level=1, + ) + return False diff --git a/pykeops/pykeops/torch/test_install.py b/pykeops/pykeops/torch/test_install.py index 119329ab8..2bc05931a 100644 --- a/pykeops/pykeops/torch/test_install.py +++ b/pykeops/pykeops/torch/test_install.py @@ -23,12 +23,13 @@ def test_torch_bindings(): except Exception as e: pyKeOps_Message(f"Error during computation: {e}", use_tag=False) return False - - if torch.allclose( - keops_res, torch.tensor(expected_res, dtype=torch.float32) - ): + + if torch.allclose(keops_res, torch.tensor(expected_res, dtype=torch.float32)): pyKeOps_Message("pyKeOps with torch bindings is working!", use_tag=False) return True else: - pyKeOps_Message(f"outputs wrong values: expected {expected_res} but get {keops_res}", use_tag=False) - return False \ No newline at end of file + pyKeOps_Message( + f"outputs wrong values: expected {expected_res} but get {keops_res}", + use_tag=False, + ) + return False From f19ae0ac29320757d02e095acbcdc7a4fbb6fe30 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Wed, 6 May 2026 16:54:05 +0200 Subject: [PATCH 35/98] Add a size check to avoid segfault with missing (dummy) dimension in variable --- keopscore/keopscore/sandbox/Conv2D.py | 2 +- keopscore/keopscore/sandbox/chunks.py | 2 +- keopscore/keopscore/sandbox/chunks_ranges.py | 2 +- keopscore/keopscore/sandbox/finalchunks.py | 2 +- .../keopscore/sandbox/finalchunks_ranges.py | 2 +- .../sandbox/genred_auto_factorize.py | 2 +- keopscore/keopscore/sandbox/genred_float16.py | 11 +- .../keopscore/sandbox/genred_gaussian.py | 2 +- keopscore/keopscore/sandbox/genred_inplace.py | 2 +- .../keopscore/sandbox/lazytensor_gaussian.py | 2 +- .../sandbox/lazytensor_gaussian_batch.py | 2 +- .../sandbox/lazytensor_gaussian_batchdims.py | 2 +- .../sandbox/lazytensor_gaussian_inplace.py | 2 +- .../keopscore/sandbox/lazytensor_grad.py | 2 +- .../keopscore/sandbox/lazytensor_sandbox.py | 2 +- .../keopscore/sandbox/lazytensor_tensordot.py | 2 +- keopscore/keopscore/utils/messages.py | 4 +- .../pykeops/common/keops_io/LoadKeOps_cpp.py | 29 ++ .../pykeops/common/keops_io/pykeops_nvrtc.cpp | 29 ++ pykeops/pykeops/sandbox/issue_335.py | 2 +- pykeops/pykeops/sandbox/issue_335_alt.py | 2 +- pykeops/pykeops/sandbox/test_chunks.py | 2 +- pykeops/pykeops/sandbox/test_finalchunks.py | 4 +- .../sandbox/test_finalchunks_ranges.py | 2 +- .../pykeops/sandbox/test_lazytensor_clamp.py | 2 +- .../sandbox/test_lazytensor_comp_op.py | 2 +- .../sandbox/test_lazytensor_gaussian.py | 2 +- .../pykeops/sandbox/test_soft_dtw_kernel.py | 2 +- .../test_soft_dtw_kernel_dissmatrix.py | 2 +- .../sandbox/test_soft_dtw_kernel_v0.py | 2 +- pykeops/pykeops/test/__init__.py | 24 + pykeops/pykeops/test/test_chunks.py | 5 +- pykeops/pykeops/test/test_chunks_ranges.py | 5 +- pykeops/pykeops/test/test_complex.py | 3 +- pykeops/pykeops/test/test_complex_numpy.py | 3 +- pykeops/pykeops/test/test_contiguous_numpy.py | 3 +- pykeops/pykeops/test/test_contiguous_torch.py | 3 +- .../test/{conv2d.py => test_conv2d.py} | 11 +- pykeops/pykeops/test/test_finalchunks.py | 5 +- .../pykeops/test/test_finalchunks_ranges.py | 5 +- pykeops/pykeops/test/test_float16.py | 5 +- pykeops/pykeops/test/test_gpu_cpu.py | 9 +- pykeops/pykeops/test/test_kron.py | 20 +- pykeops/pykeops/test/test_lazytensor_clamp.py | 7 +- .../test/test_lazytensor_gaussian_batch.py | 5 +- .../test/test_lazytensor_gaussian_cpu.py | 7 +- .../test/test_lazytensor_gaussian_fromhost.py | 3 +- .../test/test_lazytensor_gaussian_inplace.py | 5 +- .../test_lazytensor_gaussian_numpy_inplace.py | 3 +- pykeops/pykeops/test/test_lazytensor_grad.py | 9 +- .../pykeops/test/test_lazytensor_tensordot.py | 5 +- pykeops/pykeops/test/test_numpy.py | 78 +-- pykeops/pykeops/test/test_sinc.py | 5 +- pykeops/pykeops/test/test_torch.py | 448 +++++++----------- pykeops/pykeops/test/test_torch_func.py | 15 +- .../pykeops/test/test_torch_func_logsumexp.py | 15 +- pykeops/pykeops/test/test_transpose.py | 3 +- pykeops/pykeops/test/test_vmap.py | 13 +- .../tutorials/kmeans/plot_kmeans_torch.py | 2 +- 59 files changed, 436 insertions(+), 415 deletions(-) rename pykeops/pykeops/test/{conv2d.py => test_conv2d.py} (78%) diff --git a/keopscore/keopscore/sandbox/Conv2D.py b/keopscore/keopscore/sandbox/Conv2D.py index 8b7cfd8f8..3fb2c776e 100644 --- a/keopscore/keopscore/sandbox/Conv2D.py +++ b/keopscore/keopscore/sandbox/Conv2D.py @@ -12,7 +12,7 @@ test_grad = False test_grad2 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = False x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/chunks.py b/keopscore/keopscore/sandbox/chunks.py index 8b5024054..a8cadcc72 100644 --- a/keopscore/keopscore/sandbox/chunks.py +++ b/keopscore/keopscore/sandbox/chunks.py @@ -11,7 +11,7 @@ dtype = torch.float32 sum_scheme = "block_sum" -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = False x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/chunks_ranges.py b/keopscore/keopscore/sandbox/chunks_ranges.py index 83cc28c61..419575726 100644 --- a/keopscore/keopscore/sandbox/chunks_ranges.py +++ b/keopscore/keopscore/sandbox/chunks_ranges.py @@ -16,7 +16,7 @@ pykeops.config.gpu_available = 0 -device_id = "cuda:0" if pykeops.config.gpu_available else "cpu" +device_id = "cuda" if pykeops.config.gpu_available else "cpu" do_warmup = False x = torch.rand(B1, B2, M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/finalchunks.py b/keopscore/keopscore/sandbox/finalchunks.py index 88125b7ee..0db91a099 100644 --- a/keopscore/keopscore/sandbox/finalchunks.py +++ b/keopscore/keopscore/sandbox/finalchunks.py @@ -11,7 +11,7 @@ dtype = torch.float32 sum_scheme = "block_sum" -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/finalchunks_ranges.py b/keopscore/keopscore/sandbox/finalchunks_ranges.py index a5301f2c8..eaa8f8347 100644 --- a/keopscore/keopscore/sandbox/finalchunks_ranges.py +++ b/keopscore/keopscore/sandbox/finalchunks_ranges.py @@ -11,7 +11,7 @@ dtype = torch.float32 sum_scheme = "block_sum" -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = False x = torch.rand(B1, B2, M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/genred_auto_factorize.py b/keopscore/keopscore/sandbox/genred_auto_factorize.py index c2ab468f8..48396ed51 100644 --- a/keopscore/keopscore/sandbox/genred_auto_factorize.py +++ b/keopscore/keopscore/sandbox/genred_auto_factorize.py @@ -8,7 +8,7 @@ dtype = torch.float32 -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" x = torch.rand(M, D, device=device_id, dtype=dtype) / math.sqrt(D) y = torch.rand(N, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/genred_float16.py b/keopscore/keopscore/sandbox/genred_float16.py index f43c093b6..38c117d1c 100644 --- a/keopscore/keopscore/sandbox/genred_float16.py +++ b/keopscore/keopscore/sandbox/genred_float16.py @@ -6,17 +6,12 @@ import torch from pykeops.torch import Genred -( - M, - N, -) = ( - 5, - 5, -) +M = 5 +N = 5 dtype = torch.float16 -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" x = torch.zeros(M, 1, device=device_id, dtype=dtype) b = torch.ones(N, 1, device=device_id, dtype=dtype) diff --git a/keopscore/keopscore/sandbox/genred_gaussian.py b/keopscore/keopscore/sandbox/genred_gaussian.py index 51b7cfae1..bb9d36184 100644 --- a/keopscore/keopscore/sandbox/genred_gaussian.py +++ b/keopscore/keopscore/sandbox/genred_gaussian.py @@ -10,7 +10,7 @@ dtype = torch.float64 -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = False x = torch.rand(M, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/genred_inplace.py b/keopscore/keopscore/sandbox/genred_inplace.py index 149bd1439..f2e123de0 100644 --- a/keopscore/keopscore/sandbox/genred_inplace.py +++ b/keopscore/keopscore/sandbox/genred_inplace.py @@ -10,7 +10,7 @@ dtype = torch.float64 -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = False x = torch.rand(M, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/lazytensor_gaussian.py b/keopscore/keopscore/sandbox/lazytensor_gaussian.py index 501abf387..5649fd678 100644 --- a/keopscore/keopscore/sandbox/lazytensor_gaussian.py +++ b/keopscore/keopscore/sandbox/lazytensor_gaussian.py @@ -16,7 +16,7 @@ pykeops.config.gpu_available = 0 -device_id = "cuda:0" if pykeops.config.gpu_available else "cpu" +device_id = "cuda" if pykeops.config.gpu_available else "cpu" do_warmup = True x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/lazytensor_gaussian_batch.py b/keopscore/keopscore/sandbox/lazytensor_gaussian_batch.py index e0a3aec8e..e6e8edbc1 100644 --- a/keopscore/keopscore/sandbox/lazytensor_gaussian_batch.py +++ b/keopscore/keopscore/sandbox/lazytensor_gaussian_batch.py @@ -7,7 +7,7 @@ D = 3 M, N = 2000, 3000 -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" dtype = torch.float32 x = torch.rand(B1, 1, B3, M, 1, D, dtype=dtype, device=device_id) diff --git a/keopscore/keopscore/sandbox/lazytensor_gaussian_batchdims.py b/keopscore/keopscore/sandbox/lazytensor_gaussian_batchdims.py index 5defa71c7..814980cb7 100644 --- a/keopscore/keopscore/sandbox/lazytensor_gaussian_batchdims.py +++ b/keopscore/keopscore/sandbox/lazytensor_gaussian_batchdims.py @@ -11,7 +11,7 @@ dtype = torch.float32 sum_scheme = "block_sum" -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" x = torch.rand(3, M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) y = torch.rand(1, 1, N, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/lazytensor_gaussian_inplace.py b/keopscore/keopscore/sandbox/lazytensor_gaussian_inplace.py index 1dd479273..1a1bd0a45 100644 --- a/keopscore/keopscore/sandbox/lazytensor_gaussian_inplace.py +++ b/keopscore/keopscore/sandbox/lazytensor_gaussian_inplace.py @@ -11,7 +11,7 @@ dtype = torch.float32 sum_scheme = "block_sum" -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/lazytensor_grad.py b/keopscore/keopscore/sandbox/lazytensor_grad.py index 39763fd14..b1bbc0eb5 100644 --- a/keopscore/keopscore/sandbox/lazytensor_grad.py +++ b/keopscore/keopscore/sandbox/lazytensor_grad.py @@ -19,7 +19,7 @@ test_grad = True test_grad2 = True test_grad3 = True -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" x = torch.rand(M, 1, D, requires_grad=test_grad, device=device_id, dtype=dtype) y = torch.rand(1, N, 1, device=device_id, dtype=dtype) diff --git a/keopscore/keopscore/sandbox/lazytensor_sandbox.py b/keopscore/keopscore/sandbox/lazytensor_sandbox.py index eeb431916..f4eb0f485 100644 --- a/keopscore/keopscore/sandbox/lazytensor_sandbox.py +++ b/keopscore/keopscore/sandbox/lazytensor_sandbox.py @@ -11,7 +11,7 @@ dtype = torch.float32 sum_scheme = "block_sum" -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/keopscore/keopscore/sandbox/lazytensor_tensordot.py b/keopscore/keopscore/sandbox/lazytensor_tensordot.py index d43213667..7a3ec76cd 100644 --- a/keopscore/keopscore/sandbox/lazytensor_tensordot.py +++ b/keopscore/keopscore/sandbox/lazytensor_tensordot.py @@ -17,7 +17,7 @@ # Matrix multiplication as a special case of Tensordot # ---------------------------------------------------- # -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True a = torch.randn(4 * 7, requires_grad=True, device=device_id, dtype=torch.float64) diff --git a/keopscore/keopscore/utils/messages.py b/keopscore/keopscore/utils/messages.py index 8635f706e..8dffc64c7 100644 --- a/keopscore/keopscore/utils/messages.py +++ b/keopscore/keopscore/utils/messages.py @@ -34,12 +34,12 @@ def KeOps_Message(message, use_tag=True, level=1, **kwargs): def KeOps_Warning(message, newline=False, level=1, **kwargs): if _keops_verbose() >= level: - message = ("\n" if newline else "") + "[KeOps] Warning : " + message + message = ("\n" if newline else "") + "[KeOps] Warning: " + message print(message, **kwargs) def KeOps_Error(message, show_line_number=True): - message = "[KeOps] Error : " + message + message = "[KeOps] Error: " + message if show_line_number: from inspect import currentframe, getframeinfo diff --git a/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py b/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py index 7359dc1b8..bc2a95f52 100644 --- a/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py +++ b/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py @@ -72,6 +72,7 @@ def get_pybind11_code(self): #include "{self.params.source_name}" #include +#include namespace py = pybind11; template < typename TYPE > @@ -150,6 +151,34 @@ def get_pybind11_code(self): }} + /*------------------------------------*/ + /* Shape validation */ + /*------------------------------------*/ + for (int k = 0; k < (int)indsi_v.size(); k++) {{ + int idx = indsi_v[k]; + const auto& shape = argshape_v[idx]; + if ((int)shape.size() < 2) + throw std::invalid_argument("[pyKeOps] Error: Vi argument #" + std::to_string(idx) + " requires at least 2 dimensions, got " + std::to_string(shape.size()) + "."); + if (shape.back() != dimsx_v[k]) + throw std::invalid_argument("[pyKeOps] Error: Vi argument #" + std::to_string(idx) + " has trailing dim " + std::to_string(shape.back()) + ", expected " + std::to_string(dimsx_v[k]) + "."); + }} + for (int k = 0; k < (int)indsj_v.size(); k++) {{ + int idx = indsj_v[k]; + const auto& shape = argshape_v[idx]; + if ((int)shape.size() < 2) + throw std::invalid_argument("[pyKeOps] Error: Vj argument #" + std::to_string(idx) + " requires at least 2 dimensions, got " + std::to_string(shape.size()) + "."); + if (shape.back() != dimsy_v[k]) + throw std::invalid_argument("[pyKeOps] Error: Vj argument #" + std::to_string(idx) + " has trailing dim " + std::to_string(shape.back()) + ", expected " + std::to_string(dimsy_v[k]) + "."); + }} + for (int k = 0; k < (int)indsp_v.size(); k++) {{ + int idx = indsp_v[k]; + const auto& shape = argshape_v[idx]; + if ((int)shape.size() < 1) + throw std::invalid_argument("[pyKeOps] Error: Pm argument #" + std::to_string(idx) + " requires at least 1 dimension (got a scalar)."); + if (shape.back() != dimsp_v[k]) + throw std::invalid_argument("[pyKeOps] Error: Pm argument #" + std::to_string(idx) + " has trailing dim " + std::to_string(shape.back()) + ", expected " + std::to_string(dimsp_v[k]) + "."); + }} + return launch_keops_cpu_{self.params.tag}< TYPE >(dimY, nx, ny, diff --git a/pykeops/pykeops/common/keops_io/pykeops_nvrtc.cpp b/pykeops/pykeops/common/keops_io/pykeops_nvrtc.cpp index 6805a9b3f..14feb33ab 100644 --- a/pykeops/pykeops/common/keops_io/pykeops_nvrtc.cpp +++ b/pykeops/pykeops/common/keops_io/pykeops_nvrtc.cpp @@ -7,6 +7,7 @@ #include #include +#include namespace py = pybind11; @@ -76,6 +77,34 @@ template class KeOps_module_python : public KeOps_module { argshape_v[i] = tmp_v; } + /*------------------------------------*/ + /* Shape validation */ + /*------------------------------------*/ + for (int k = 0; k < (int)indsi_v.size(); k++) { + int idx = indsi_v[k]; + const auto& shape = argshape_v[idx]; + if ((int)shape.size() < 2) + throw std::invalid_argument("[pyKeOps] Error: Vi argument #" + std::to_string(idx) + " requires at least 3 dimensions, got " + std::to_string(shape.size()) + "."); + if (shape.back() != dimsx_v[k]) + throw std::invalid_argument("[pyKeOps] Error: Vi argument #" + std::to_string(idx) + " has trailing dim " + std::to_string(shape.back()) + ", expected " + std::to_string(dimsx_v[k]) + "."); + } + for (int k = 0; k < (int)indsj_v.size(); k++) { + int idx = indsj_v[k]; + const auto& shape = argshape_v[idx]; + if ((int)shape.size() < 2) + throw std::invalid_argument("[pyKeOps] Error: Vj argument #" + std::to_string(idx) + " requires at least 2 dimensions, got " + std::to_string(shape.size()) + "."); + if (shape.back() != dimsy_v[k]) + throw std::invalid_argument("[pyKeOps] Error: Vj argument #" + std::to_string(idx) + " has trailing dim " + std::to_string(shape.back()) + ", expected " + std::to_string(dimsy_v[k]) + "."); + } + for (int k = 0; k < (int)indsp_v.size(); k++) { + int idx = indsp_v[k]; + const auto& shape = argshape_v[idx]; + if ((int)shape.size() < 1) + throw std::invalid_argument("[pyKeOps] Error: Pm argument #" + std::to_string(idx) + " requires at least 1 dimension (got a scalar)."); + if (shape.back() != dimsp_v[k]) + throw std::invalid_argument("[pyKeOps] Error: Pm argument #" + std::to_string(idx) + " has trailing dim " + std::to_string(shape.back()) + ", expected " + std::to_string(dimsp_v[k]) + "."); + } + return KeOps_module::launch_kernel( tagHostDevice, dimY, nx, ny, tagI, tagZero, use_half, tag1D2D, dimred, cuda_block_size, use_chunk_mode, indsi_v, indsj_v, indsp_v, dimout, diff --git a/pykeops/pykeops/sandbox/issue_335.py b/pykeops/pykeops/sandbox/issue_335.py index fef25a356..49b5b7d52 100644 --- a/pykeops/pykeops/sandbox/issue_335.py +++ b/pykeops/pykeops/sandbox/issue_335.py @@ -23,7 +23,7 @@ def covar_func(x1, x2=None): if __name__ == "__main__": - device = "cuda:0" + device = "cuda" train_x = torch.randn(30, 10000, 3, device=device) # covar_module = gpytorch.kernels.keops.MaternKernel(nu=2.5).to(device) diff --git a/pykeops/pykeops/sandbox/issue_335_alt.py b/pykeops/pykeops/sandbox/issue_335_alt.py index 85cd98853..fc4c2f578 100644 --- a/pykeops/pykeops/sandbox/issue_335_alt.py +++ b/pykeops/pykeops/sandbox/issue_335_alt.py @@ -25,7 +25,7 @@ def covar_func(x1, x2=None): B, M, N, D = 25, 10000, 10000, 3 if __name__ == "__main__": - device = "cuda:0" + device = "cuda" train_x = torch.randn(B, M, D, device=device) # covar_module = gpytorch.kernels.keops.MaternKernel(nu=2.5).to(device) diff --git a/pykeops/pykeops/sandbox/test_chunks.py b/pykeops/pykeops/sandbox/test_chunks.py index 3a1530246..7bb38c084 100644 --- a/pykeops/pykeops/sandbox/test_chunks.py +++ b/pykeops/pykeops/sandbox/test_chunks.py @@ -12,7 +12,7 @@ test_grad = True test_grad2 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True x = torch.rand( diff --git a/pykeops/pykeops/sandbox/test_finalchunks.py b/pykeops/pykeops/sandbox/test_finalchunks.py index c697b7cb5..3861232f8 100644 --- a/pykeops/pykeops/sandbox/test_finalchunks.py +++ b/pykeops/pykeops/sandbox/test_finalchunks.py @@ -6,13 +6,13 @@ import torch from pykeops.torch import LazyTensor -M, N, D, DV = 10000, 10000, 3, 10000 +M, N, D, DV = 1000, 1000, 3, 1000 dtype = torch.float32 test_grad = True test_grad2 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True x = torch.rand( diff --git a/pykeops/pykeops/sandbox/test_finalchunks_ranges.py b/pykeops/pykeops/sandbox/test_finalchunks_ranges.py index aac85b2ac..aca386d4a 100644 --- a/pykeops/pykeops/sandbox/test_finalchunks_ranges.py +++ b/pykeops/pykeops/sandbox/test_finalchunks_ranges.py @@ -8,7 +8,7 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" torch.manual_seed(0) x = torch.rand(B1, B2, M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/pykeops/pykeops/sandbox/test_lazytensor_clamp.py b/pykeops/pykeops/sandbox/test_lazytensor_clamp.py index 2d021f3ac..005737d1d 100644 --- a/pykeops/pykeops/sandbox/test_lazytensor_clamp.py +++ b/pykeops/pykeops/sandbox/test_lazytensor_clamp.py @@ -10,7 +10,7 @@ test_grad = True -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True diff --git a/pykeops/pykeops/sandbox/test_lazytensor_comp_op.py b/pykeops/pykeops/sandbox/test_lazytensor_comp_op.py index 6486377a1..521aa40a3 100644 --- a/pykeops/pykeops/sandbox/test_lazytensor_comp_op.py +++ b/pykeops/pykeops/sandbox/test_lazytensor_comp_op.py @@ -10,7 +10,7 @@ dtype = torch.float16 -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True diff --git a/pykeops/pykeops/sandbox/test_lazytensor_gaussian.py b/pykeops/pykeops/sandbox/test_lazytensor_gaussian.py index b7f88efcd..0a7bc16a6 100644 --- a/pykeops/pykeops/sandbox/test_lazytensor_gaussian.py +++ b/pykeops/pykeops/sandbox/test_lazytensor_gaussian.py @@ -12,7 +12,7 @@ test_grad = True test_grad2 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True x = torch.rand( diff --git a/pykeops/pykeops/sandbox/test_soft_dtw_kernel.py b/pykeops/pykeops/sandbox/test_soft_dtw_kernel.py index 294607614..8cda610bd 100644 --- a/pykeops/pykeops/sandbox/test_soft_dtw_kernel.py +++ b/pykeops/pykeops/sandbox/test_soft_dtw_kernel.py @@ -10,7 +10,7 @@ test_grad = True -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True diff --git a/pykeops/pykeops/sandbox/test_soft_dtw_kernel_dissmatrix.py b/pykeops/pykeops/sandbox/test_soft_dtw_kernel_dissmatrix.py index eeab3f0b7..b0320bd1f 100644 --- a/pykeops/pykeops/sandbox/test_soft_dtw_kernel_dissmatrix.py +++ b/pykeops/pykeops/sandbox/test_soft_dtw_kernel_dissmatrix.py @@ -10,7 +10,7 @@ test_grad = True -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True diff --git a/pykeops/pykeops/sandbox/test_soft_dtw_kernel_v0.py b/pykeops/pykeops/sandbox/test_soft_dtw_kernel_v0.py index 2aa8ec433..fe9846327 100644 --- a/pykeops/pykeops/sandbox/test_soft_dtw_kernel_v0.py +++ b/pykeops/pykeops/sandbox/test_soft_dtw_kernel_v0.py @@ -10,7 +10,7 @@ # D is size of each sample M, N, D = 1000, 1000, 50 -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" do_warmup = True diff --git a/pykeops/pykeops/test/__init__.py b/pykeops/pykeops/test/__init__.py index e69de29bb..7cd6237a6 100644 --- a/pykeops/pykeops/test/__init__.py +++ b/pykeops/pykeops/test/__init__.py @@ -0,0 +1,24 @@ +import numpy as np +import torch + + +def assert_torch_allclose(actual, expected, *, label=None, **kwargs): + # test dtype, convert long to float if needed + + if actual.dtype in [torch.int64, torch.int32]: + actual = actual.float() + + if expected.dtype in [torch.int64, torch.int32]: + expected = expected.float() + + ok = torch.allclose(actual, expected, **kwargs) + diff = torch.linalg.norm(actual - expected).item() + prefix = label if label is not None else "torch.allclose failed" + assert ok, f"{prefix}: ||ref-test||_2={diff:.6e} and ||ref||_2={torch.linalg.norm(expected).item():.6e}" + + +def assert_np_allclose(actual, expected, *, label=None, **kwargs): + ok = np.allclose(actual, expected, **kwargs) + diff = float(np.linalg.norm(np.asarray(actual) - np.asarray(expected))) + prefix = label if label is not None else "np.allclose failed" + assert ok, f"{prefix}: ||ref-test||_2={diff:.6e} and ||ref||_2={float(np.linalg.norm(expected)):.6e}" diff --git a/pykeops/pykeops/test/test_chunks.py b/pykeops/pykeops/test/test_chunks.py index b5b03d8d0..0344b2345 100644 --- a/pykeops/pykeops/test/test_chunks.py +++ b/pykeops/pykeops/test/test_chunks.py @@ -1,6 +1,7 @@ import math import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose M, N, D, DV = 200, 300, 300, 1 @@ -8,7 +9,7 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" torch.manual_seed(0) x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) @@ -38,4 +39,4 @@ def fun(x, y, b, backend): def test_chunks(): - assert torch.allclose(out[0], out[1], atol=0.0001) + assert_torch_allclose(out[0], out[1], atol=0.0001, label="chunks") diff --git a/pykeops/pykeops/test/test_chunks_ranges.py b/pykeops/pykeops/test/test_chunks_ranges.py index 19d5c0e27..d186eebc4 100644 --- a/pykeops/pykeops/test/test_chunks_ranges.py +++ b/pykeops/pykeops/test/test_chunks_ranges.py @@ -1,6 +1,7 @@ import math import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose B1, B2, M, N, D, DV = 2, 3, 200, 300, 300, 1 @@ -8,7 +9,7 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" torch.manual_seed(0) x = torch.rand(B1, B2, M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) @@ -38,4 +39,4 @@ def fun(x, y, b, backend): def test_chunk_ranges(): - assert torch.allclose(out[0], out[1], rtol=0.0001, atol=0.0001) + assert_torch_allclose(out[0], out[1], rtol=0.0001, atol=0.0001, label="chunks_ranges") diff --git a/pykeops/pykeops/test/test_complex.py b/pykeops/pykeops/test/test_complex.py index 2378163f6..9808ba70e 100644 --- a/pykeops/pykeops/test/test_complex.py +++ b/pykeops/pykeops/test/test_complex.py @@ -3,6 +3,7 @@ import math import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose dtype = torch.float32 dtype_c = torch.complex64 @@ -40,7 +41,7 @@ def fun(x, p, f, backend): def test_complex_fw(): - assert torch.allclose(out[0], out[1]) + assert_torch_allclose(out[0], out[1], label="complex_fw") # out_g = [] diff --git a/pykeops/pykeops/test/test_complex_numpy.py b/pykeops/pykeops/test/test_complex_numpy.py index fb33a507c..fe4ce25d2 100644 --- a/pykeops/pykeops/test/test_complex_numpy.py +++ b/pykeops/pykeops/test/test_complex_numpy.py @@ -1,5 +1,6 @@ import numpy as np from pykeops.numpy import LazyTensor, ComplexLazyTensor +from pykeops.test import assert_np_allclose M, N, D = 1000, 1000, 3 @@ -24,4 +25,4 @@ def fun(x, y, backend): def test_complex_numpy(): - assert np.allclose(out[0], out[1]) + assert_np_allclose(out[0], out[1], label="complex_numpy") diff --git a/pykeops/pykeops/test/test_contiguous_numpy.py b/pykeops/pykeops/test/test_contiguous_numpy.py index aa96a5125..5f75a8a38 100644 --- a/pykeops/pykeops/test/test_contiguous_numpy.py +++ b/pykeops/pykeops/test/test_contiguous_numpy.py @@ -1,5 +1,6 @@ from pykeops.numpy import LazyTensor import numpy as np +from pykeops.test import assert_np_allclose np.random.seed(0) a1 = np.random.rand(2, 1000, 5) @@ -22,4 +23,4 @@ def test_contiguous_numpy(): - assert np.allclose(d2, d1) + assert_np_allclose(d2, d1, label="contiguous_numpy") diff --git a/pykeops/pykeops/test/test_contiguous_torch.py b/pykeops/pykeops/test/test_contiguous_torch.py index c8af66361..bf789dd83 100644 --- a/pykeops/pykeops/test/test_contiguous_torch.py +++ b/pykeops/pykeops/test/test_contiguous_torch.py @@ -1,5 +1,6 @@ from pykeops.torch import LazyTensor import torch +from pykeops.test import assert_torch_allclose torch.backends.cuda.matmul.allow_tf32 = False torch.manual_seed(0) @@ -24,4 +25,4 @@ def test_contiguous_torch(): - assert torch.allclose(d2, d1) + assert_torch_allclose(d2, d1, label="contiguous_torch") diff --git a/pykeops/pykeops/test/conv2d.py b/pykeops/pykeops/test/test_conv2d.py similarity index 78% rename from pykeops/pykeops/test/conv2d.py rename to pykeops/pykeops/test/test_conv2d.py index 1713f225f..132b4a4f3 100644 --- a/pykeops/pykeops/test/conv2d.py +++ b/pykeops/pykeops/test/test_conv2d.py @@ -1,13 +1,14 @@ import math import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose -M, N, D, DV = 20000, 30000, 3, 1 +M, N, D, DV = 2000, 3000, 3, 1 dtype = torch.float32 torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) y = torch.rand(1, N, D, device=device_id, dtype=dtype) / math.sqrt(D) @@ -47,10 +48,10 @@ def fun(x, y, b, backend): class TestCase: def test_conv2d_fw(self): - assert torch.allclose(out[0], out[1]) + assert_torch_allclose(out[0], out[1], label="conv2d_fw") def test_conv2d_bw1(self): - assert torch.allclose(out_g[0], out_g[1]) + assert_torch_allclose(out_g[0], out_g[1], label="conv2d_bw1") def test_conv2d_bw2(self): - assert torch.allclose(out_g2[0], out_g2[1]) + assert_torch_allclose(out_g2[0], out_g2[1], label="conv2d_bw2") diff --git a/pykeops/pykeops/test/test_finalchunks.py b/pykeops/pykeops/test/test_finalchunks.py index 5a7f98000..4bcc05d0a 100644 --- a/pykeops/pykeops/test/test_finalchunks.py +++ b/pykeops/pykeops/test/test_finalchunks.py @@ -1,6 +1,7 @@ import math import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose M, N, D, DV = 200, 300, 3, 300 @@ -8,7 +9,7 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" torch.manual_seed(0) x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) @@ -40,4 +41,4 @@ def fun(x, y, b, backend): def test_finalchunk(): # print(out[0] - out[1]) - assert torch.allclose(out[0], out[1], atol=0.0001) + assert_torch_allclose(out[0], out[1], atol=0.0001, label="finalchunks") diff --git a/pykeops/pykeops/test/test_finalchunks_ranges.py b/pykeops/pykeops/test/test_finalchunks_ranges.py index aac85b2ac..f42c08dae 100644 --- a/pykeops/pykeops/test/test_finalchunks_ranges.py +++ b/pykeops/pykeops/test/test_finalchunks_ranges.py @@ -1,6 +1,7 @@ import math import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose B1, B2, M, N, D, DV = 2, 3, 200, 300, 3, 300 @@ -8,7 +9,7 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" torch.manual_seed(0) x = torch.rand(B1, B2, M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) @@ -38,4 +39,4 @@ def fun(x, y, b, backend): def test_finalchunks_ranges(): - assert torch.allclose(out[0], out[1], atol=0.0001) + assert_torch_allclose(out[0], out[1], atol=0.0001, label="finalchunks_ranges") diff --git a/pykeops/pykeops/test/test_float16.py b/pykeops/pykeops/test/test_float16.py index 161f1f9b7..6a7156129 100644 --- a/pykeops/pykeops/test/test_float16.py +++ b/pykeops/pykeops/test/test_float16.py @@ -2,6 +2,7 @@ import pytest import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose dtype = torch.float16 @@ -34,7 +35,7 @@ def test_float16_fw(self): for backend in ["torch", "keops"]: self.out.append(fun(x, y, backend).squeeze()) - assert torch.allclose(self.out[0], self.out[1], atol=0.001, rtol=0.001) + assert_torch_allclose(self.out[0], self.out[1], atol=0.001, rtol=0.001, label="float16_fw") @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires a GPU") def test_float16_bw(self): @@ -42,4 +43,4 @@ def test_float16_bw(self): for k, backend in enumerate(["torch", "keops"]): out_g.append(torch.autograd.grad(self.out[k][0], [x])[0]) - assert torch.allclose(out_g[0], out_g[1]) + assert_torch_allclose(out_g[0], out_g[1], label="float16_bw") diff --git a/pykeops/pykeops/test/test_gpu_cpu.py b/pykeops/pykeops/test/test_gpu_cpu.py index 505437f1a..5977e8c91 100644 --- a/pykeops/pykeops/test/test_gpu_cpu.py +++ b/pykeops/pykeops/test/test_gpu_cpu.py @@ -39,8 +39,13 @@ def fun(x, y, b, backend): class TestCase: def test_torch_keops_cpu(self): - assert torch.allclose(out[0], out[1]) + assert torch.allclose(out[0], out[1]), ( + f"torch vs keops_cpu mismatch: ||ref-test||_2={torch.norm(out[0] - out[1]).item():.6e} and ||ref||_2={torch.norm(out[0]).item():.6e}" + ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires a GPU") def test_torch_keops_gpu(self): - assert torch.allclose(out[0], fun(x, y, b, ["keops_gpu"]).squeeze()) + out_gpu = fun(x, y, b, ["keops_gpu"]).squeeze() + assert torch.allclose(out[0], out_gpu), ( + f"torch vs keops_gpu mismatch: ||ref-test||_2={torch.norm(out[0] - out_gpu).item():.6e} and ||ref||_2={torch.norm(out[0]).item():.6e}" + ) diff --git a/pykeops/pykeops/test/test_kron.py b/pykeops/pykeops/test/test_kron.py index 29fcadb17..dd6c02095 100644 --- a/pykeops/pykeops/test/test_kron.py +++ b/pykeops/pykeops/test/test_kron.py @@ -2,6 +2,7 @@ import torch from pykeops.numpy import Genred +from pykeops.test import assert_np_allclose, assert_torch_allclose M = 11 N = 150 @@ -32,7 +33,7 @@ def test_kron_lazytensor_np(): Y = LazyTensor(y[None, :, :]) gamma_keops_lazytensor_np = (X.keops_kron(Y, dimfa, dimfb)).sum(axis=axis) - assert np.allclose(gamma_keops_lazytensor_np, gamma_py, atol=1e-6) + assert_np_allclose(gamma_keops_lazytensor_np, gamma_py, atol=1e-6, label="kron_lazytensor_np") ############################################################################ @@ -43,8 +44,11 @@ def test_kron_lazytensor_torch(): Y_t = LazyTensor_torch(torch.from_numpy(y)[None, :, :]) gamma_keops_lazytensor_torch = (X_t.keops_kron(Y_t, dimfa, dimfb)).sum(axis=axis) - assert torch.allclose( - gamma_keops_lazytensor_torch, torch.from_numpy(gamma_py), atol=1e-6 + assert_torch_allclose( + gamma_keops_lazytensor_torch, + torch.from_numpy(gamma_py), + atol=1e-6, + label="kron_lazytensor_torch", ) @@ -79,8 +83,8 @@ def test_kron_genred(): myconv2 = Genred(formula2, aliases, reduction_op="Sum", axis=axis) gamma_keops_genred_TensorDot = myconv2(x0, y0) - assert np.allclose(gamma_keops_genred_Kron, gamma_py2, atol=1e-6) - assert np.allclose(gamma_keops_genred_TensorDot, gamma_py2, atol=1e-6) + assert_np_allclose(gamma_keops_genred_Kron, gamma_py2, atol=1e-6, label="kron_genred") + assert_np_allclose(gamma_keops_genred_TensorDot, gamma_py2, atol=1e-6, label="tensordot_genred") ############################################################################ @@ -123,6 +127,6 @@ def test_kron_genred2(): myconv2 = Genred(formula2, aliases, reduction_op="Sum", axis=axis) gamma_keops_genred_TensorDot2 = myconv2(x2, y2) - assert np.allclose(gamma_keops_genred_Kron2, gamma_py3, atol=1e-6) - assert np.allclose(gamma_keops_genred_TensorDot2, gamma_py3, atol=1e-6) - assert np.allclose(gamma_py4, gamma_py3, atol=1e-6) + assert_np_allclose(gamma_keops_genred_Kron2, gamma_py3, atol=1e-6, label="kron_genred2") + assert_np_allclose(gamma_keops_genred_TensorDot2, gamma_py3, atol=1e-6, label="tensordot_genred2") + assert_np_allclose(gamma_py4, gamma_py3, atol=1e-6, label="numpy_kron_einsum") diff --git a/pykeops/pykeops/test/test_lazytensor_clamp.py b/pykeops/pykeops/test/test_lazytensor_clamp.py index fa3dbf456..28b7fa631 100644 --- a/pykeops/pykeops/test/test_lazytensor_clamp.py +++ b/pykeops/pykeops/test/test_lazytensor_clamp.py @@ -1,10 +1,11 @@ import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose M, N, D = 1000, 1000, 3 torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" torch.manual_seed(0) x = torch.randn(M, 1, D, requires_grad=True, device=device_id) @@ -35,7 +36,7 @@ def fun(x, y, a, b, backend): class TestCase: def test_lazytensor_clamp_fw(self): - assert torch.allclose(out[0], out[1]) + assert_torch_allclose(out[0], out[1], label="clamp_fw") def test_lazytensor_clamp_bw(self): - assert torch.allclose(out_g[0], out_g[1], atol=0.01, rtol=0.001) + assert_torch_allclose(out_g[0], out_g[1], atol=0.01, rtol=0.001, label="clamp_bw") diff --git a/pykeops/pykeops/test/test_lazytensor_gaussian_batch.py b/pykeops/pykeops/test/test_lazytensor_gaussian_batch.py index b2c5c84cc..6f0a3779a 100644 --- a/pykeops/pykeops/test/test_lazytensor_gaussian_batch.py +++ b/pykeops/pykeops/test/test_lazytensor_gaussian_batch.py @@ -1,6 +1,7 @@ import math import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose B1, B2, M, N, D, DV = 3, 4, 20, 25, 3, 2 @@ -41,7 +42,7 @@ def fun(x, y, b, backend): class TestCase: def test_lazytensor_gaussian_batch_fw(self): # print(out[0]- out[1]) - assert torch.allclose(out[0], out[1], atol=1e-6) + assert_torch_allclose(out[0], out[1], atol=1e-6, label="gaussian_batch_fw") def test_lazytensor_gaussian_batch_bw(self): - assert torch.allclose(out_g[0], out_g[1]) + assert_torch_allclose(out_g[0], out_g[1], label="gaussian_batch_bw") diff --git a/pykeops/pykeops/test/test_lazytensor_gaussian_cpu.py b/pykeops/pykeops/test/test_lazytensor_gaussian_cpu.py index efa7712b1..f07f7be98 100644 --- a/pykeops/pykeops/test/test_lazytensor_gaussian_cpu.py +++ b/pykeops/pykeops/test/test_lazytensor_gaussian_cpu.py @@ -1,6 +1,7 @@ import math import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose M, N, D, DV = 1000, 1000, 3, 1 @@ -48,10 +49,10 @@ def fun(x, y, b, backend): class TestClass: def test_lazytensor_gaussian_cpu(self): - assert torch.allclose(out[0], out[1]) + assert_torch_allclose(out[0], out[1], label="gaussian_cpu_fw") def test_lazytensor_gaussian_cpu_bw1(self): - assert torch.allclose(out_g[0], out_g[1]) + assert_torch_allclose(out_g[0], out_g[1], label="gaussian_cpu_bw1") def test_lazytensor_gaussian_cpu_bw2(self): - assert torch.allclose(out_g2[0], out_g2[1]) + assert_torch_allclose(out_g2[0], out_g2[1], label="gaussian_cpu_bw2") diff --git a/pykeops/pykeops/test/test_lazytensor_gaussian_fromhost.py b/pykeops/pykeops/test/test_lazytensor_gaussian_fromhost.py index d78eca154..1cc3b42e4 100644 --- a/pykeops/pykeops/test/test_lazytensor_gaussian_fromhost.py +++ b/pykeops/pykeops/test/test_lazytensor_gaussian_fromhost.py @@ -1,6 +1,7 @@ import math import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose M, N, D, DV = 2500, 2000, 3, 1 @@ -38,4 +39,4 @@ def fun(x, y, b, backend): def test_lazytensor_gaussian_fromhost(): - assert torch.allclose(out[0], out[1]) + assert_torch_allclose(out[0], out[1], label="gaussian_fromhost") diff --git a/pykeops/pykeops/test/test_lazytensor_gaussian_inplace.py b/pykeops/pykeops/test/test_lazytensor_gaussian_inplace.py index 1a14d951b..01d059764 100644 --- a/pykeops/pykeops/test/test_lazytensor_gaussian_inplace.py +++ b/pykeops/pykeops/test/test_lazytensor_gaussian_inplace.py @@ -1,6 +1,7 @@ import math import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose M, N, D, DV = 20, 30, 3, 1 @@ -8,7 +9,7 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" torch.manual_seed(0) x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) @@ -39,4 +40,4 @@ def fun(x, y, b, backend, out=None): def test_lazytensor_gaussian_inplace(): - assert torch.allclose(out[0], out[1]) + assert_torch_allclose(out[0], out[1], label="gaussian_inplace") diff --git a/pykeops/pykeops/test/test_lazytensor_gaussian_numpy_inplace.py b/pykeops/pykeops/test/test_lazytensor_gaussian_numpy_inplace.py index 8ea7bcbb2..180ddc751 100644 --- a/pykeops/pykeops/test/test_lazytensor_gaussian_numpy_inplace.py +++ b/pykeops/pykeops/test/test_lazytensor_gaussian_numpy_inplace.py @@ -1,6 +1,7 @@ import math import numpy as np from pykeops.numpy import LazyTensor +from pykeops.test import assert_np_allclose M, N, D, DV = 3000, 2000, 3, 1 @@ -35,4 +36,4 @@ def fun(x, y, b, backend, out=None): def test_lazytensor_gaussian_numpy_inplace(): - assert np.allclose(out[0], out[1]) + assert_np_allclose(out[0], out[1], label="gaussian_numpy_inplace") diff --git a/pykeops/pykeops/test/test_lazytensor_grad.py b/pykeops/pykeops/test/test_lazytensor_grad.py index 7010bfd64..98aac080f 100644 --- a/pykeops/pykeops/test/test_lazytensor_grad.py +++ b/pykeops/pykeops/test/test_lazytensor_grad.py @@ -1,13 +1,14 @@ import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose M, N, D, DV = 100, 100, 3, 1 dtype = torch.float32 torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" torch.manual_seed(0) x = torch.rand(M, 1, D, requires_grad=True, device=device_id, dtype=dtype) @@ -45,10 +46,10 @@ def fun(x, y, b, backend): class TestClass: def test_lazytensor_grad(self): - assert torch.allclose(out[0], out[1], rtol=0.0001) + assert_torch_allclose(out[0], out[1], rtol=0.0001, label="lazytensor_grad_fw") def test_lazytensor_grad_bw1(self): - assert torch.allclose(out_g[0], out_g[1], rtol=0.0001) + assert_torch_allclose(out_g[0], out_g[1], rtol=0.0001, label="lazytensor_grad_bw1") def test_lazytensor_grad_bw2(self): - assert torch.allclose(out_g2[0], out_g2[1], rtol=0.001) + assert_torch_allclose(out_g2[0], out_g2[1], rtol=0.001, label="lazytensor_grad_bw2") diff --git a/pykeops/pykeops/test/test_lazytensor_tensordot.py b/pykeops/pykeops/test/test_lazytensor_tensordot.py index d1cf43348..2ca9ba8a1 100644 --- a/pykeops/pykeops/test/test_lazytensor_tensordot.py +++ b/pykeops/pykeops/test/test_lazytensor_tensordot.py @@ -1,11 +1,12 @@ import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose M, N = 2, 10 # Matrix multiplication as a special case of Tensordot torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda:0" if torch.cuda.is_available() else "cpu" +device_id = "cuda" if torch.cuda.is_available() else "cpu" torch.manual_seed(0) a = torch.randn(4 * 7, requires_grad=True, device=device_id, dtype=torch.float64) @@ -18,4 +19,4 @@ def test_tensordot(): - assert torch.allclose(c.flatten(), C.flatten()) + assert_torch_allclose(c.flatten(), C.flatten(), label="lazytensor_tensordot") diff --git a/pykeops/pykeops/test/test_numpy.py b/pykeops/pykeops/test/test_numpy.py index f3e63b69b..8a8a0bcb2 100644 --- a/pykeops/pykeops/test/test_numpy.py +++ b/pykeops/pykeops/test/test_numpy.py @@ -22,6 +22,7 @@ import pykeops import pykeops.config +from pykeops.test import assert_np_allclose from pykeops.numpy.utils import ( squared_distances, log_sum_exp, @@ -29,28 +30,29 @@ class NumpyUnitTestCase(unittest.TestCase): - A = int(4) # Batchdim 1 - B = int(6) # Batchdim 2 - M = int(10) - N = int(6) - D = int(3) - E = int(3) - nbatchdims = int(2) - - x = np.random.rand(M, D) - a = np.random.rand(M, E) - f = np.random.rand(M, 1) - y = np.random.rand(N, D) - b = np.random.rand(N, E) - g = np.random.rand(N, 1) - sigma = np.array([0.4]) - - X = np.random.rand(A, 1, M, D) - L = np.random.rand(1, B, M, 1) - Y = np.random.rand(1, B, N, D) - S = np.random.rand(A, B, 1) + 1 - - type_to_test = ["float32", "float64"] + def SetUp(self): + self.A = int(4) # Batchdim 1 + self.B = int(6) # Batchdim 2 + self.M = int(10) + self.N = int(6) + self.D = int(3) + self.E = int(3) + self.nbatchdims = int(2) + + self.x = np.random.rand(M, D) + self.a = np.random.rand(M, E) + self.f = np.random.rand(M, 1) + self.y = np.random.rand(N, D) + self.b = np.random.rand(N, E) + self.g = np.random.rand(N, 1) + self.sigma = np.array([0.4]) + + self.X = np.random.rand(A, 1, M, D) + self.L = np.random.rand(1, B, M, 1) + self.Y = np.random.rand(1, B, N, D) + self.S = np.random.rand(A, B, 1) + 1 + + self.type_to_test = ["float32", "float64"] ############################################################ def test_numpytools_function_binding(self): @@ -61,10 +63,10 @@ def test_numpytools_function_binding(self): x = self.x.astype(self.type_to_test[0]) self.assertTrue(np.array_equal(tools.copy(x), x)) - self.assertTrue(np.allclose(tools.exp(x), np.exp(x))) - self.assertTrue(np.allclose(tools.log(x + 1), np.log(x + 1))) - self.assertTrue(np.allclose(tools.norm(x), np.linalg.norm(x))) - self.assertTrue(np.allclose(tools.arraysum(x, axis=0), np.sum(x, axis=0))) + assert_np_allclose(tools.exp(x), np.exp(x)) + assert_np_allclose(tools.log(x + 1), np.log(x + 1)) + assert_np_allclose(tools.norm(x), np.linalg.norm(x)) + assert_np_allclose(tools.arraysum(x, axis=0), np.sum(x, axis=0)) ############################################################ def test_cg_solver_stops_immediately_when_x0_is_good(self): @@ -81,13 +83,13 @@ def test_cg_solver_stops_immediately_when_x0_is_good(self): b = K_xx @ self.f + alpha * self.f x = K_xx.solve(b, alpha=alpha, x0=self.f, eps=1e-12) - self.assertTrue(np.allclose(self.f, x)) + assert_np_allclose(self.f, x) stream = io.StringIO() with redirect_stdout(stream): x = K_xx.solve(b, alpha=alpha, x0=self.f, eps=1e-12, verbose=True) - self.assertTrue(np.allclose(self.f, x)) + assert_np_allclose(self.f, x) output = stream.getvalue() self.assertIn("'status': 'Converged'", output) self.assertIn("'niter': 0", output) @@ -111,7 +113,7 @@ def test_cg_solver_verbose_prints_info(self): with redirect_stdout(stream): x = K_xx.solve(b, alpha=alpha, x0=self.f, eps=1e-12, verbose=True) - self.assertTrue(np.allclose(self.f, x)) + assert_np_allclose(self.f, x) output = stream.getvalue() self.assertIn("'status': 'Converged'", output) self.assertIn("'x0_provided': True", output) @@ -150,7 +152,7 @@ def test_generic_syntax_sum(self): ).T # compare output - self.assertTrue(np.allclose(gamma_keops, gamma_py, atol=1e-6)) + assert_np_allclose(gamma_keops, gamma_py, atol=1e-6) ############################################################ def test_generic_syntax_lse(self): @@ -185,7 +187,7 @@ def test_generic_syntax_lse(self): ) # compare output - self.assertTrue(np.allclose(gamma_keops.ravel(), gamma_py, atol=1e-6)) + assert_np_allclose(gamma_keops.ravel(), gamma_py, atol=1e-6) ############################################################ def test_generic_syntax_softmax(self): @@ -231,9 +233,7 @@ def np_softmax(x, w): ) # compare output - self.assertTrue( - np.allclose(gamma_keops.ravel(), gamma_py.ravel(), atol=1e-6) - ) + assert_np_allclose(gamma_keops.ravel(), gamma_py.ravel(), atol=1e-6) ############################################################ def test_non_contiguity(self): @@ -261,7 +261,7 @@ def test_non_contiguity(self): # check output self.assertFalse(yc_tmp.flags.c_contiguous) - self.assertTrue(np.allclose(gamma_keops1, gamma_keops2)) + assert_np_allclose(gamma_keops1, gamma_keops2) ############################################################ def test_heterogeneous_var_aliases(self): @@ -290,7 +290,7 @@ def test_heterogeneous_var_aliases(self): ) # compare output - self.assertTrue(np.allclose(gamma_keops.ravel(), gamma_py, atol=1e-6)) + assert_np_allclose(gamma_keops.ravel(), gamma_py, atol=1e-6) ############################################################ def test_formula_simplification(self): @@ -316,7 +316,7 @@ def test_formula_simplification(self): gamma_py = np.zeros_like(self.x) # compare output - self.assertTrue(np.allclose(gamma_keops, gamma_py, atol=1e-6)) + assert_np_allclose(gamma_keops, gamma_py, atol=1e-6) ############################################################ def test_argkmin(self): @@ -343,7 +343,7 @@ def test_argkmin(self): np.sum((self.x[:, np.newaxis, :] - self.y[np.newaxis, :, :]) ** 2, axis=2), axis=1, )[:, :3] - self.assertTrue(np.allclose(c.ravel(), cnp.ravel())) + assert_np_allclose(c.ravel(), cnp.ravel()) ############################################################ def test_LazyTensor_sum(self): @@ -388,7 +388,7 @@ def test_LazyTensor_sum(self): for res_keops, res_numpy in zip(full_results[0], full_results[1]): self.assertTrue(res_keops.shape == res_numpy.shape) - self.assertTrue(np.allclose(res_keops, res_numpy, atol=1e-3)) + assert_np_allclose(res_keops, res_numpy, atol=1e-3) if __name__ == "__main__": diff --git a/pykeops/pykeops/test/test_sinc.py b/pykeops/pykeops/test/test_sinc.py index 15589c265..996153ff3 100644 --- a/pykeops/pykeops/test/test_sinc.py +++ b/pykeops/pykeops/test/test_sinc.py @@ -4,6 +4,7 @@ import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose device = "cuda" if torch.cuda.is_available() else "cpu" @@ -25,9 +26,9 @@ @pytest.mark.skipif(torch.__version__ < "1.8", reason="Requires torch>=1.8") def test_sinc(): - assert torch.allclose(s1, s2), torch.abs(s1 - s2) + assert_torch_allclose(s1, s2, label="sinc") @pytest.mark.skipif(torch.__version__ < "1.8", reason="Requires torch>=1.8") def test_sinc_grad(): - assert torch.allclose(x.grad, y.grad) + assert_torch_allclose(x.grad, y.grad, label="sinc_grad") diff --git a/pykeops/pykeops/test/test_torch.py b/pykeops/pykeops/test/test_torch.py index 4d272264c..47e502fc7 100644 --- a/pykeops/pykeops/test/test_torch.py +++ b/pykeops/pykeops/test/test_torch.py @@ -2,6 +2,7 @@ import sys from contextlib import redirect_stdout import io +from math import prod sys.path.append( os.path.join( @@ -17,17 +18,27 @@ ) import unittest -import numpy as np + +try: + import torch + HAS_TORCH = True + + use_cuda = torch.cuda.is_available() + if use_cuda: + device = "cuda" + torch.backends.cuda.matmul.allow_tf32 = False + else: + device = "cpu" + +except ImportError: + HAS_TORCH = False import pykeops import pykeops.config -from pykeops.numpy.utils import ( - squared_distances, - log_np_kernel, - log_sum_exp, -) +from pykeops.test import assert_torch_allclose +@unittest.skipUnless(HAS_TORCH, "torch not available") class PytorchUnitTestCase(unittest.TestCase): A = int(5) # Batchdim 1 B = int(3) # Batchdim 2 @@ -36,67 +47,41 @@ class PytorchUnitTestCase(unittest.TestCase): D = int(3) E = int(3) nbatchdims = int(2) + def SetUp(self): + self.x64 = torch.rand((M, D), dtype=torch.float64, device=device, requires_grad=True) + self.a64 = torch.rand((M, E), dtype=torch.float64, device=device, requires_grad=False) + self.e64 = torch.rand((M, E), dtype=torch.float64, device=device, requires_grad=False) + self.f64 = torch.rand((M, 1), dtype=torch.float64, device=device, requires_grad=True) + self.y64 = torch.rand((N, D), dtype=torch.float64, device=device, requires_grad=False) + self.b64 = torch.rand((N, E), dtype=torch.float64, device=device, requires_grad=False) + self.g64 = torch.rand((N, 1), dtype=torch.float64, device=device, requires_grad=True) + self.p64 = torch.rand(2, dtype=torch.float64, device=device, requires_grad=False) + + self.sigma64 = torch.tensor([0.4], dtype=torch.float64, device=device) + self.alpha64 = torch.tensor([0.1], dtype=torch.float64, device=device) + + self.X64 = torch.rand((A, B, M, D), dtype=torch.float64, device=device, requires_grad=True) + self.L64 = torch.rand((A, 1, M, 1), dtype=torch.float64, device=device, requires_grad=False) + self.Y64 = torch.rand((1, B, N, D), dtype=torch.float64, device=device, requires_grad=True) + self.S64 = 1 + torch.rand((A, B, 1), dtype=torch.float64, device=device, requires_grad=True) + + self.x32 = x64.to(torch.float32).clone().requires_grad_(True) + self.a32 = a64.to(torch.float32).clone().requires_grad_(False) + self.e32 = e64.to(torch.float32).clone().requires_grad_(False) + self.f32 = f64.to(torch.float32).clone().requires_grad_(True) + self.y32 = y64.to(torch.float32).clone().requires_grad_(False) + self.b32 = b64.to(torch.float32).clone().requires_grad_(False) + self.g32 = g64.to(torch.float32).clone().requires_grad_(True) + self.p32 = p64.to(torch.float32).clone().requires_grad_(False) + + self.sigma32 = sigma64.to(torch.float32).clone().requires_grad_(False) + self.alpha32 = alpha64.to(torch.float32).clone().requires_grad_(False) + + self.X32 = X64.to(torch.float32).clone().requires_grad_(True) + self.L32 = L64.to(torch.float32).clone().requires_grad_(False) + self.Y32 = Y64.to(torch.float32).clone().requires_grad_(True) + self.S32 = S64.to(torch.float32).clone().requires_grad_(True) - x = np.random.rand(M, D) - a = np.random.rand(M, E) - e = np.random.rand(M, E) - f = np.random.rand(M, 1) - y = np.random.rand(N, D) - b = np.random.rand(N, E) - g = np.random.rand(N, 1) - p = np.random.rand(2) - sigma = np.array([0.4]) - alpha = np.array([0.1]) - - X = np.random.rand(A, B, M, D) - L = np.random.rand(A, 1, M, 1) - Y = np.random.rand(1, B, N, D) - S = np.random.rand(A, B, 1) + 1 - - try: - import torch - - use_cuda = torch.cuda.is_available() - device = "cuda" if use_cuda else "cpu" - torch.backends.cuda.matmul.allow_tf32 = False - - dtype = torch.float32 - xc = torch.tensor(x, dtype=dtype, device=device, requires_grad=True) - ac = torch.tensor(a, dtype=dtype, device=device, requires_grad=False) - ec = torch.tensor(e, dtype=dtype, device=device, requires_grad=False) - fc = torch.tensor(f, dtype=dtype, device=device, requires_grad=True) - yc = torch.tensor(y, dtype=dtype, device=device, requires_grad=False) - bc = torch.tensor(b, dtype=dtype, device=device, requires_grad=False) - gc = torch.tensor(g, dtype=dtype, device=device, requires_grad=True) - pc = torch.tensor(p, dtype=dtype, device=device, requires_grad=False) - sigmac = torch.tensor(sigma, dtype=dtype, device=device, requires_grad=False) - alphac = torch.tensor(alpha, dtype=dtype, device=device, requires_grad=False) - - Xc = torch.tensor(X, dtype=dtype, device=device, requires_grad=True) - Lc = torch.tensor(L, dtype=dtype, device=device, requires_grad=False) - Yc = torch.tensor(Y, dtype=dtype, device=device, requires_grad=True) - Sc = torch.tensor(S, dtype=dtype, device=device, requires_grad=True) - - dtype = torch.float64 - xcd = torch.tensor(x, dtype=dtype, device=device, requires_grad=False) - acd = torch.tensor(a, dtype=dtype, device=device, requires_grad=False) - ecd = torch.tensor(e, dtype=dtype, device=device, requires_grad=False) - fcd = torch.tensor(f, dtype=dtype, device=device, requires_grad=False) - ycd = torch.tensor(y, dtype=dtype, device=device, requires_grad=False) - bcd = torch.tensor(b, dtype=dtype, device=device, requires_grad=False) - gcd = torch.tensor(g, dtype=dtype, device=device, requires_grad=False) - pcd = torch.tensor(p, dtype=dtype, device=device, requires_grad=False) - sigmacd = torch.tensor(sigma, dtype=dtype, device=device, requires_grad=False) - alphacd = torch.tensor(alpha, dtype=dtype, device=device, requires_grad=False) - Xcd = torch.tensor(X, dtype=dtype, device=device, requires_grad=True) - Lcd = torch.tensor(L, dtype=dtype, device=device, requires_grad=False) - Ycd = torch.tensor(Y, dtype=dtype, device=device, requires_grad=True) - Scd = torch.tensor(S, dtype=dtype, device=device, requires_grad=True) - - print("Running Pytorch tests.") - except: - print("Pytorch could not be loaded. Skip tests.") - pass ############################################################ def test_torchtools_function_binding(self): @@ -105,12 +90,12 @@ def test_torchtools_function_binding(self): import torch tools = torchtools() - x = self.xc.detach() + x = self.x32.detach() self.assertTrue(torch.equal(tools.copy(x), x)) - self.assertTrue(torch.allclose(tools.exp(x), torch.exp(x))) - self.assertTrue(torch.allclose(tools.log(x + 1), torch.log(x + 1))) - self.assertTrue(torch.allclose(tools.norm(x), torch.norm(x))) + assert_torch_allclose(tools.exp(x), torch.exp(x)) + assert_torch_allclose(tools.log(x + 1), torch.log(x + 1)) + assert_torch_allclose(tools.norm(x), torch.norm(x)) ############################################################ def test_generic_syntax_float(self): @@ -128,18 +113,16 @@ def test_generic_syntax_float(self): with self.subTest(b=b): # Call cuda kernel gamma_keops = Genred(formula, aliases, axis=1)( - self.sigmac, self.gc, self.xc, self.yc, backend=b + self.sigma32, self.g32, self.x32, self.y32, backend=b ) - # Numpy version - gamma_py = np.sum( - (self.sigma - self.g) ** 2 - * np.exp((self.y.T[:, :, np.newaxis] + self.x.T[:, np.newaxis, :])), - axis=1, + # Torch reference in float64 for stable comparisons. + gamma_ref = torch.sum( + (self.sigma64 - self.g64) ** 2 + * torch.exp(self.y64.T[:, :, None] + self.x64.T[:, None, :]), + dim=1, ).T # compare output - self.assertTrue( - np.allclose(gamma_keops.cpu().data.numpy(), gamma_py, atol=1e-6) - ) + assert_torch_allclose(gamma_keops.to(torch.float64), gamma_ref, atol=1e-6) ############################################################ def test_generic_syntax_double(self): @@ -157,18 +140,16 @@ def test_generic_syntax_double(self): with self.subTest(b=b): # Call cuda kernel gamma_keops = Genred(formula, aliases, axis=1)( - self.sigmacd, self.gcd, self.xcd, self.ycd, backend=b + self.sigma64, self.g64, self.x64, self.y64, backend=b ) - # Numpy version - gamma_py = np.sum( - (self.sigma - self.g) ** 2 - * np.exp((self.y.T[:, :, np.newaxis] + self.x.T[:, np.newaxis, :])), - axis=1, + # Torch reference in float64. + gamma_ref = torch.sum( + (self.sigma64 - self.g64) ** 2 + * torch.exp(self.y64.T[:, :, None] + self.x64.T[:, None, :]), + dim=1, ).T # compare output - self.assertTrue( - np.allclose(gamma_keops.cpu().data.numpy(), gamma_py, atol=1e-6) - ) + assert_torch_allclose(gamma_keops, gamma_ref, atol=1e-6) ############################################################ def test_generic_syntax_softmax(self): @@ -194,24 +175,17 @@ def test_generic_syntax_softmax(self): formula2=formula_weights, ) gamma_keops = myop( - self.sigmacd, self.gcd, self.xcd, self.ycd, backend=b + self.sigma64, self.g64, self.x64, self.y64, backend=b ) - # Numpy version - def np_softmax(x, w): - x -= np.max(x, axis=1)[:, None] # subtract the max for robustness - return np.exp(x) @ w / np.sum(np.exp(x), axis=1)[:, None] - - gamma_py = np_softmax( - (self.sigma - self.g.T) ** 2 - * np.exp(-squared_distances(self.x, self.y)), - self.y, - ) + # Torch reference + sqdist = torch.sum((self.x64[:, None, :] - self.y64[None, :, :]) ** 2, dim=2) + scores = (self.sigma64 - self.g64.T) ** 2 * torch.exp(-sqdist) + scores = scores - torch.max(scores, dim=1, keepdim=True).values + gamma_ref = torch.exp(scores) @ self.y64 / torch.sum(torch.exp(scores), dim=1, keepdim=True) # compare output - self.assertTrue( - np.allclose(gamma_keops.cpu().data.numpy(), gamma_py, atol=1e-6) - ) + assert_torch_allclose(gamma_keops, gamma_ref, atol=1e-6) ############################################################ def test_generic_syntax_simple(self): @@ -221,9 +195,9 @@ def test_generic_syntax_simple(self): aliases = [ "P = Pm(2)", # 1st argument, a parameter, dim 2. "X = Vi(" - + str(self.xc.shape[1]) + + str(self.x64.shape[1]) + ") ", # 2nd argument, indexed by i, dim D. - "Y = Vj(" + str(self.yc.shape[1]) + ") ", + "Y = Vj(" + str(self.y64.shape[1]) + ") ", ] # 3rd argument, indexed by j, dim D. formula = "Pow((X|Y),2) * ((Elem(P,0) * X) + (Elem(P,1) * Y))" @@ -236,18 +210,16 @@ def test_generic_syntax_simple(self): for b in backend_to_test: with self.subTest(b=b): my_routine = Genred(formula, aliases, reduction_op="Sum", axis=1) - gamma_keops = my_routine(self.pc, self.xc, self.yc, backend=b) + gamma_keops = my_routine(self.p64, self.x64, self.y64, backend=b) - # Numpy version - scals = (self.x @ self.y.T) ** 2 # Memory-intensive computation! - gamma_py = self.p[0] * scals.sum(1).reshape(-1, 1) * self.x + self.p[ + # Torch reference + scals = (self.x64 @ self.y64.T) ** 2 # Memory-intensive computation! + gamma_ref = self.p64[0] * scals.sum(1).reshape(-1, 1) * self.x64 + self.p64[ 1 - ] * (scals @ self.y) + ] * (scals @ self.y64) # compare output - self.assertTrue( - np.allclose(gamma_keops.cpu().data.numpy(), gamma_py, atol=1e-6) - ) + assert_torch_allclose(gamma_keops, gamma_ref, atol=1e-6) ############################################################ def test_logSumExp_kernels_feature(self): @@ -277,19 +249,24 @@ def test_logSumExp_kernels_feature(self): for k in ["gaussian", "laplacian", "cauchy", "inverse_multiquadric"]: with self.subTest(k=k): # Call cuda kernel - gamma_lazy = kernels[k](self.xc, self.yc, self.sigmac) - gamma_lazy = gamma_lazy.logsumexp(dim=1, weight=Vj(self.gc.exp())).cpu() - # gamma = kernel_product(params, self.xc, self.yc, self.gc).cpu() - - # Numpy version - log_K = log_np_kernel(self.x, self.y, self.sigma, kernel=k) - log_KP = log_K + self.g.T - gamma_py = log_sum_exp(log_KP, axis=1) + gamma_lazy = kernels[k](self.x64, self.y64, self.sigma64) + gamma_lazy = gamma_lazy.logsumexp(dim=1, weight=Vj(self.g64.exp())) + + # Torch reference + sqdist = torch.sum((self.x64[:, None, :] - self.y64[None, :, :]) ** 2, dim=2) + inv_s2 = 1 / (self.sigma64**2) + if k == "gaussian": + log_k = -(inv_s2 * sqdist) + elif k == "laplacian": + log_k = -torch.sqrt(inv_s2 * sqdist) + elif k == "cauchy": + log_k = -(1 + inv_s2 * sqdist).log() + else: + log_k = -0.5 * (1 + inv_s2 * sqdist).log() + gamma_ref = torch.logsumexp(log_k + self.g64.T, dim=1) # compare output - self.assertTrue( - np.allclose(gamma_lazy.data.numpy().ravel(), gamma_py, atol=1e-6) - ) + assert_torch_allclose(gamma_lazy.reshape(-1), gamma_ref.reshape(-1), atol=1e-6) ############################################################ def test_logSumExp_gradient_kernels_feature(self): @@ -300,38 +277,33 @@ def test_logSumExp_gradient_kernels_feature(self): aliases = [ "P = Pm(2)", # 1st argument, a parameter, dim 2. "X = Vi(" - + str(self.gc.shape[1]) + + str(self.g64.shape[1]) + ") ", # 2nd argument, indexed by i, dim D. - "Y = Vj(" + str(self.fc.shape[1]) + ") ", + "Y = Vj(" + str(self.f64.shape[1]) + ") ", ] # 3rd argument, indexed by j, dim D. formula = "(Elem(P,0) * X) + (Elem(P,1) * Y)" # Pytorch version my_routine = Genred(formula, aliases, reduction_op="LogSumExp", axis=1) - tmp = my_routine(self.pc, self.fc, self.gc, backend="auto") + tmp = my_routine(self.p64, self.f64, self.g64, backend="auto") res = torch.dot( torch.ones_like(tmp).view(-1), tmp.view(-1) ) # equivalent to tmp.sum() but avoiding contiguity pb - gamma_keops = torch.autograd.grad(res, [self.fc, self.gc], create_graph=False) - - # Numpy version - tmp = self.p[0] * self.f + self.p[1] * self.g.T - res_py = (np.exp(tmp)).sum(axis=1) - tmp2 = np.exp(tmp.T) / res_py.reshape(1, -1) - gamma_py = [np.ones(self.M) * self.p[0], self.p[1] * tmp2.T.sum(axis=0)] + gamma_keops = torch.autograd.grad(res, [self.f64, self.g64], create_graph=False) + + # Torch reference + tmp = self.p64[0] * self.f64 + self.p64[1] * self.g64.T + res_ref = torch.exp(tmp).sum(dim=1) + tmp2 = torch.exp(tmp.T) / res_ref.reshape(1, -1) + gamma_ref = [ + torch.ones(self.M, dtype=torch.float64, device=device) * self.p64[0], + self.p64[1] * tmp2.T.sum(dim=0), + ] # compare output - self.assertTrue( - np.allclose( - gamma_keops[0].cpu().data.numpy().ravel(), gamma_py[0], atol=1e-6 - ) - ) - self.assertTrue( - np.allclose( - gamma_keops[1].cpu().data.numpy().ravel(), gamma_py[1], atol=1e-6 - ) - ) + assert_torch_allclose(gamma_keops[0].reshape(-1), gamma_ref[0].reshape(-1), atol=1e-6) + assert_torch_allclose(gamma_keops[1].reshape(-1), gamma_ref[1].reshape(-1), atol=1e-6) ############################################################ def test_non_contiguity(self): @@ -341,45 +313,38 @@ def test_non_contiguity(self): aliases = [ "P = Pm(2)", # 1st argument, a parameter, dim 2. "X = Vi(" - + str(self.xc.shape[1]) + + str(self.x64.shape[1]) + ") ", # 2nd argument, indexed by i, dim D. - "Y = Vj(" + str(self.yc.shape[1]) + ") ", + "Y = Vj(" + str(self.y64.shape[1]) + ") ", ] # 3rd argument, indexed by j, dim D. formula = "Pow((X|Y),2) * ((Elem(P,0) * X) + (Elem(P,1) * Y))" my_routine = Genred(formula, aliases, reduction_op="Sum", axis=1) - yc_tmp = self.yc.t().contiguous().t() # create a non contiguous copy + yc_tmp = self.y64.t().contiguous().t() # create a non contiguous copy # check output self.assertFalse(yc_tmp.is_contiguous()) - my_routine(self.pc, self.xc, yc_tmp, backend="auto") + my_routine(self.p64, self.x64, yc_tmp, backend="auto") ############################################################ def test_heterogeneous_var_aliases(self): ############################################################ from pykeops.torch import Genred - from pykeops.numpy.utils import squared_distances aliases = ["p=Pm(0,1)", "x=Vi(1,3)", "y=Vj(2,3)"] formula = "Square(p-Var(3,1,1))*Exp(-SqNorm2(y-x))" # Call cuda kernel myconv = Genred(formula, aliases, reduction_op="Sum", axis=1) - gamma_keops = myconv(self.sigmac, self.xc, self.yc, self.gc, backend="auto") + gamma_keops = myconv(self.sigma64, self.x64, self.y64, self.g64, backend="auto") - # Numpy version - gamma_py = np.sum( - (self.sigma - self.g.T) ** 2 * np.exp(-squared_distances(self.x, self.y)), - axis=1, - ) + # Torch reference + sqdist = torch.sum((self.x64[:, None, :] - self.y64[None, :, :]) ** 2, dim=2) + gamma_ref = torch.sum((self.sigma64 - self.g64.T) ** 2 * torch.exp(-sqdist), dim=1) # compare output - self.assertTrue( - np.allclose( - gamma_keops.cpu().data.numpy().ravel(), gamma_py.ravel(), atol=1e-6 - ) - ) + assert_torch_allclose(gamma_keops.reshape(-1), gamma_ref.reshape(-1), atol=1e-6) ############################################################ def test_invkernel(self): @@ -397,33 +362,25 @@ def test_invkernel(self): Kinv = KernelSolve(formula, aliases, "b", axis=1) - c = Kinv(self.xc, self.xc, self.ac, self.sigmac, alpha=self.alphac) + c = Kinv(self.x64, self.x64, self.a64, self.sigma64, alpha=self.alpha64) if torch.__version__ >= "1.8": torchsolve = lambda A, B: torch.linalg.solve(A, B) else: torchsolve = lambda A, B: torch.solve(B, A)[0] c_ = torchsolve( - self.alphac * torch.eye(self.M, device=self.device) + self.alpha64 * torch.eye(self.M, device=device, dtype=torch.float64) + torch.exp( - -torch.sum((self.xc[:, None, :] - self.xc[None, :, :]) ** 2, dim=2) - * self.sigmac + -torch.sum((self.x64[:, None, :] - self.x64[None, :, :]) ** 2, dim=2) + * self.sigma64 ), - self.ac, + self.a64, ) - self.assertTrue( - np.allclose( - c.cpu().data.numpy().ravel(), c_.cpu().data.numpy().ravel(), atol=1e-4 - ) - ) + assert_torch_allclose(c.reshape(-1), c_.reshape(-1), atol=1e-4) - (u,) = torch.autograd.grad(c, self.xc, self.ec) - (u_,) = torch.autograd.grad(c_, self.xc, self.ec) - self.assertTrue( - np.allclose( - u.cpu().data.numpy().ravel(), u_.cpu().data.numpy().ravel(), atol=1e-4 - ) - ) + (u,) = torch.autograd.grad(c, self.x64, self.e64) + (u_,) = torch.autograd.grad(c_, self.x64, self.e64) + assert_torch_allclose(u.reshape(-1), u_.reshape(-1), atol=1e-4) ############################################################ def test_cg_solver_stops_immediately_when_x0_is_good(self): @@ -434,20 +391,20 @@ def test_cg_solver_stops_immediately_when_x0_is_good(self): alpha = 2.0 - x_i = LazyTensor(self.xc[:, None, :]) - x_j = LazyTensor(self.xc[None, :, :]) + x_i = LazyTensor(self.x64[:, None, :]) + x_j = LazyTensor(self.x64[None, :, :]) K_xx = (((x_i - x_j).abs()).sum(-1)).exp() - b = K_xx @ self.fc + alpha * self.fc + b = K_xx @ self.f64 + alpha * self.f64 - x = K_xx.solve(b, alpha=alpha, x0=self.fc, eps=1e-12) - self.assertTrue(torch.allclose(self.fc, x)) + x = K_xx.solve(b, alpha=alpha, x0=self.f64, eps=1e-12) + assert_torch_allclose(self.f64, x) stream = io.StringIO() with redirect_stdout(stream): - x = K_xx.solve(b, alpha=alpha, x0=self.fc, eps=1e-12, verbose=True) + x = K_xx.solve(b, alpha=alpha, x0=self.f64, eps=1e-12, verbose=True) - self.assertTrue(torch.allclose(self.fc, x)) + assert_torch_allclose(self.f64, x) output = stream.getvalue() self.assertIn("'status': 'Converged'", output) self.assertIn("'niter': 0", output) @@ -476,22 +433,18 @@ def test_softmax(self): formula2=formula_weights, ) - c = softmax_op(self.xc, self.yc, self.bc) + c = softmax_op(self.x64, self.y64, self.b64) # compare with direct implementation cc = 0 for k in range(self.D): - xk = self.xc[:, k][:, None] - yk = self.yc[:, k][:, None] + xk = self.x64[:, k][:, None] + yk = self.y64[:, k][:, None] cc += (xk - yk.t()) ** 2 cc -= torch.max(cc, dim=1)[0][:, None] # subtract the max for robustness - cc = torch.exp(cc) @ self.bc / torch.sum(torch.exp(cc), dim=1)[:, None] + cc = torch.exp(cc) @ self.b64 / torch.sum(torch.exp(cc), dim=1)[:, None] - self.assertTrue( - np.allclose( - c.cpu().data.numpy().ravel(), cc.cpu().data.numpy().ravel(), atol=1e-6 - ) - ) + assert_torch_allclose(c.reshape(-1), cc.reshape(-1), atol=1e-6) ############################################################ def test_pickle(self): @@ -525,19 +478,17 @@ def test_LazyTensor_sum(self): results = [] # N.B.: We could loop over float32 and float64, but this would take longer... - for x, l, y, s in [(self.Xc, self.Lc, self.Yc, self.Sc)]: # Float32 + for x, l, y, s in [(self.X32, self.L32, self.Y32, self.S32)]: # Float32 x_i = x.unsqueeze(-2) l_i = l.unsqueeze(-2) y_j = y.unsqueeze(-3) s_p = s.unsqueeze(-2).unsqueeze(-2) if use_keops: - x_i, l_i, y_j, s_p = ( - LazyTensor(x_i), - LazyTensor(l_i), - LazyTensor(y_j), - LazyTensor(s_p), - ) + x_i = LazyTensor(x_i) + l_i = LazyTensor(l_i) + y_j = LazyTensor(y_j) + s_p = LazyTensor(s_p) D_ij = (0.5 * (l_i * x_i - y_j) ** 2 / s_p).sum(-1) K_ij = (-D_ij).exp() @@ -555,18 +506,7 @@ def test_LazyTensor_sum(self): for res_keops, res_torch in zip(full_results[0], full_results[1]): self.assertTrue(res_keops.shape == res_torch.shape) - self.assertTrue( - np.allclose( - res_keops.cpu().data.numpy().ravel(), - res_torch.cpu().data.numpy().ravel(), - atol=1e-3, - ), - "KeOps:\n" - + str(res_keops) - + "\nPyTorch:\n" - + str(res_torch) - + "\nMax error: {:.2e}".format((res_keops - res_torch).abs().max()), - ) + assert_torch_allclose(res_keops.reshape(-1), res_torch.reshape(-1), atol=1e-3) ############################################################ def test_LazyTensor_logsumexp(self): @@ -579,19 +519,17 @@ def test_LazyTensor_logsumexp(self): results = [] # N.B.: We could loop over float32 and float64, but this would take longer... - for x, l, y, s in [(self.Xcd, self.Lcd, self.Ycd, self.Scd)]: # Float64 + for x, l, y, s in [(self.X64, self.L64, self.Y64, self.S64)]: # Float64 x_i = x.unsqueeze(-2) l_i = l.unsqueeze(-2) y_j = y.unsqueeze(-3) s_p = s.unsqueeze(-2).unsqueeze(-2) if use_keops: - x_i, l_i, y_j, s_p = ( - LazyTensor(x_i), - LazyTensor(l_i), - LazyTensor(y_j), - LazyTensor(s_p), - ) + x_i = LazyTensor(x_i) + l_i = LazyTensor(l_i) + y_j = LazyTensor(y_j) + s_p = LazyTensor(s_p) D_ij = ((l_i * x_i + y_j).relu() * s_p / 9).sum(-1) K_ij = -1 / (1 + D_ij) @@ -612,13 +550,7 @@ def test_LazyTensor_logsumexp(self): for res_keops, res_torch in zip(full_results[0], full_results[1]): self.assertTrue(res_keops.shape == res_torch.shape) - self.assertTrue( - np.allclose( - res_keops.cpu().data.numpy().ravel(), - res_torch.cpu().data.numpy().ravel(), - atol=1e-5, - ) - ) + assert_torch_allclose(res_keops.reshape(-1), res_torch.reshape(-1), atol=1e-5) ############################################################ # Test min reduction with chunk without batches @@ -627,25 +559,23 @@ def test_LazyTensor_min_chunked(self): from pykeops.torch import LazyTensor import torch - X = np.random.rand(self.M, 990) - Xc = torch.tensor(X, dtype=self.dtype, device=self.device) - Y = np.random.rand(self.N, 990) - Yc = torch.tensor(Y, dtype=self.dtype, device=self.device) + X64 = torch.rand((self.M, 990), dtype=torch.float64, device=device) + Y64 = torch.rand((self.N, 990), dtype=torch.float64, device=device) + + X32 = X64.to(torch.float32).clone().requires_grad_(True) + Y32 = Y64.to(torch.float32).clone().requires_grad_(True) full_results = [] for use_keops in [True, False]: results = [] - # N.B.: We could loop over float32 and float64, but this would take longer... - for x, y in [(Xc, Yc)]: # Float32 + for x, y in [(X32, Y32), (X64, Y64)]: x_i = x.unsqueeze(-2) y_j = y.unsqueeze(-3) if use_keops: - x_i, y_j = ( - LazyTensor(x_i), - LazyTensor(y_j), - ) + x_i = LazyTensor(x_i) + y_j = LazyTensor(y_j) K_ij = ((-(((x_i + y_j)) ** 2)).exp()).sum(-1, keepdim=True) @@ -660,13 +590,7 @@ def test_LazyTensor_min_chunked(self): for res_keops, res_torch in zip(full_results[0], full_results[1]): self.assertTrue(res_keops.shape == res_torch.shape) - self.assertTrue( - np.allclose( - res_keops.cpu().data.numpy().ravel(), - res_torch.cpu().data.numpy().ravel(), - atol=1e-5, - ) - ) + assert_torch_allclose(res_keops.reshape(-1), res_torch.reshape(-1), atol=1e-5) ############################################################ def test_LazyTensor_min(self): @@ -678,19 +602,17 @@ def test_LazyTensor_min(self): results = [] # N.B.: We could loop over float32 and float64, but this would take longer... - for x, l, y, s in [(self.Xc, self.Lc, self.Yc, self.Sc)]: # Float32 + for x, l, y, s in [(self.X32, self.L32, self.Y32, self.S32)]: # Float32 x_i = x.unsqueeze(-2) l_i = l.unsqueeze(-2) y_j = y.unsqueeze(-3) s_p = s.unsqueeze(-2).unsqueeze(-2) if use_keops: - x_i, l_i, y_j, s_p = ( - LazyTensor(x_i), - LazyTensor(l_i), - LazyTensor(y_j), - LazyTensor(s_p), - ) + x_i = LazyTensor(x_i) + l_i = LazyTensor(l_i) + y_j = LazyTensor(y_j) + s_p = LazyTensor(s_p) D_ij = ((1 + ((l_i * x_i + y_j).relu() * s_p) ** 2).log()).sum( -1, keepdim=True @@ -708,13 +630,7 @@ def test_LazyTensor_min(self): for res_keops, res_torch in zip(full_results[0], full_results[1]): self.assertTrue(res_keops.shape == res_torch.shape) - self.assertTrue( - np.allclose( - res_keops.cpu().data.numpy().ravel(), - res_torch.cpu().data.numpy().ravel(), - atol=1e-5, - ) - ) + assert_torch_allclose(res_keops.reshape(-1), res_torch.reshape(-1), atol=1e-5) ############################################################ def test_TensorDot_with_permute(self): @@ -726,7 +642,7 @@ def my_tensordort_perm(a, b, dims=None, perm=None): return torch.tensordot(a, b, dims=dims).sum(3).permute(perm) def invert_permutation_numpy(permutation): - return np.arange(len(permutation))[np.argsort(permutation)] + return [idx for idx, _ in sorted(enumerate(permutation), key=lambda x: x[1])] x = torch.randn(self.M, 2, 3, 2, 2, 4, requires_grad=True, dtype=torch.float64) y = torch.randn( @@ -740,17 +656,17 @@ def invert_permutation_numpy(permutation): sum_f_torch2 = my_tensordort_perm(x, y, dims=(contfa, contfb), perm=perm_torch) f_keops = LazyTensor( - x.reshape(self.M, 1, int(np.array((dimfa)).prod())) + x.reshape(self.M, 1, prod(dimfa)) ).keops_tensordot( - LazyTensor(y.reshape(1, self.N, int(np.array(dimfb).prod()))), + LazyTensor(y.reshape(1, self.N, prod(dimfb))), dimfa, dimfb, - tuple(np.array(contfa) - 1), - tuple(np.array(contfb) - 1), + tuple(i - 1 for i in contfa), + tuple(i - 1 for i in contfb), tuple(perm), ) sum_f_keops = f_keops.sum_reduction(dim=1) - self.assertTrue(torch.allclose(sum_f_keops.flatten(), sum_f_torch2.flatten())) + assert_torch_allclose(sum_f_keops.flatten(), sum_f_torch2.flatten()) e = torch.randn_like(sum_f_torch2) # checking gradients @@ -758,15 +674,11 @@ def invert_permutation_numpy(permutation): sum_f_keops, x, e.reshape(self.M, -1), retain_graph=True )[0] grad_torch = torch.autograd.grad(sum_f_torch2, x, e, retain_graph=True)[0] - self.assertTrue( - torch.allclose(grad_keops.flatten(), grad_torch.flatten(), rtol=1e-4) - ) + assert_torch_allclose(grad_keops.flatten(), grad_torch.flatten(), rtol=1e-4) grad_keops = torch.autograd.grad(sum_f_keops, y, e.reshape(self.M, -1))[0] grad_torch = torch.autograd.grad(sum_f_torch2, y, e)[0] - self.assertTrue( - torch.allclose(grad_keops.flatten(), grad_torch.flatten(), rtol=1e-4) - ) + assert_torch_allclose(grad_keops.flatten(), grad_torch.flatten(), rtol=1e-4) if __name__ == "__main__": diff --git a/pykeops/pykeops/test/test_torch_func.py b/pykeops/pykeops/test/test_torch_func.py index c518cbe1a..240a29864 100644 --- a/pykeops/pykeops/test/test_torch_func.py +++ b/pykeops/pykeops/test/test_torch_func.py @@ -1,6 +1,7 @@ import keopscore import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose torch.manual_seed(0) @@ -37,36 +38,36 @@ class TestCase: def test_torch_func_vmap(self): res1 = torch.func.vmap(fn_torch)(x_i, y_j, b_j, p) res2 = torch.func.vmap(fn_keops)(x_i, y_j, b_j, p) - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_vmap") def test_torch_func_grad(self): res1 = torch.func.grad(fn_torch, (0, 1, 2, 3))(x_i, y_j, b_j, p) res2 = torch.func.grad(fn_keops, (0, 1, 2, 3))(x_i, y_j, b_j, p) for k in range(4): - assert torch.allclose(res1[k], res2[k], atol=1e-5) + assert_torch_allclose(res1[k], res2[k], atol=1e-5, label=f"torch_func_grad[{k}]") def test_torch_func_vjp(self): res1 = torch.func.vjp(fn_torch, x_i, y_j, b_j, p)[1](torch.tensor(1.0))[0] res2 = torch.func.vjp(fn_keops, x_i, y_j, b_j, p)[1](torch.tensor(1.0))[0] - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_vjp") def test_torch_func_jvp(self): _, res1 = torch.func.jvp(fn_torch, (x_i, y_j, b_j, p), (dx_i, dy_j, db_j, dp)) _, res2 = torch.func.jvp(fn_keops, (x_i, y_j, b_j, p), (dx_i, dy_j, db_j, dp)) - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_jvp") def test_torch_func_jacrev(self): res1 = torch.func.jacrev(fn_torch, (0, 1, 2, 3))(x_i, y_j, b_j, p) res2 = torch.func.jacrev(fn_keops, (0, 1, 2, 3))(x_i, y_j, b_j, p) for k in range(4): - assert torch.allclose(res1[k], res2[k], atol=1e-5) + assert_torch_allclose(res1[k], res2[k], atol=1e-5, label=f"torch_func_jacrev[{k}]") def test_torch_func_jacfwd(self): res1 = torch.func.jacfwd(fn_torch)(x_i, y_j, b_j, p) res2 = torch.func.jacfwd(fn_keops)(x_i, y_j, b_j, p) - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_jacfwd") def test_torch_func_hessian(self): res1 = torch.func.hessian(fn_torch)(x_i, y_j, b_j, p) res2 = torch.func.hessian(fn_keops)(x_i, y_j, b_j, p) - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_hessian") diff --git a/pykeops/pykeops/test/test_torch_func_logsumexp.py b/pykeops/pykeops/test/test_torch_func_logsumexp.py index 83dde275c..d45e6df26 100644 --- a/pykeops/pykeops/test/test_torch_func_logsumexp.py +++ b/pykeops/pykeops/test/test_torch_func_logsumexp.py @@ -1,5 +1,6 @@ import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose torch.manual_seed(0) @@ -29,36 +30,36 @@ class TestCase: def test_torch_func_vmap(self): res1 = torch.func.vmap(fn_torch)(x_i, y_j, b_j) res2 = torch.func.vmap(fn_keops)(x_i, y_j, b_j) - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_lse_vmap") def test_torch_func_grad(self): res1 = torch.func.grad(fn_torch, (0, 1, 2))(x_i, y_j, b_j) res2 = torch.func.grad(fn_keops, (0, 1, 2))(x_i, y_j, b_j) for k in range(3): - assert torch.allclose(res1[k], res2[k], atol=1e-5) + assert_torch_allclose(res1[k], res2[k], atol=1e-5, label=f"torch_func_lse_grad[{k}]") def test_torch_func_vjp(self): res1 = torch.func.vjp(fn_torch, x_i, y_j, b_j)[1](torch.tensor(1.0))[0] res2 = torch.func.vjp(fn_keops, x_i, y_j, b_j)[1](torch.tensor(1.0))[0] - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_lse_vjp") def test_torch_func_jvp(self): _, res1 = torch.func.jvp(fn_torch, (x_i, y_j, b_j), (dx_i, dy_j, db_j)) _, res2 = torch.func.jvp(fn_keops, (x_i, y_j, b_j), (dx_i, dy_j, db_j)) - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_lse_jvp") def test_torch_func_jacrev(self): res1 = torch.func.jacrev(fn_torch, (0, 1, 2))(x_i, y_j, b_j) res2 = torch.func.jacrev(fn_keops, (0, 1, 2))(x_i, y_j, b_j) for k in range(3): - assert torch.allclose(res1[k], res2[k], atol=1e-5) + assert_torch_allclose(res1[k], res2[k], atol=1e-5, label=f"torch_func_lse_jacrev[{k}]") def test_torch_func_jacfwd(self): res1 = torch.func.jacfwd(fn_torch)(x_i, y_j, b_j) res2 = torch.func.jacfwd(fn_keops)(x_i, y_j, b_j) - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_lse_jacfwd") def test_torch_func_hessian(self): res1 = torch.func.hessian(fn_torch)(x_i, y_j, b_j) res2 = torch.func.hessian(fn_keops)(x_i, y_j, b_j) - assert torch.allclose(res1, res2, atol=1e-5) + assert_torch_allclose(res1, res2, atol=1e-5, label="torch_func_lse_hessian") diff --git a/pykeops/pykeops/test/test_transpose.py b/pykeops/pykeops/test/test_transpose.py index cd253b291..7c3f58f41 100644 --- a/pykeops/pykeops/test/test_transpose.py +++ b/pykeops/pykeops/test/test_transpose.py @@ -1,5 +1,6 @@ import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose x, y = torch.randn(1000, 3), torch.randn(2000, 3) x_i, y_j = LazyTensor(x[:, None, :]), LazyTensor(y[None, :, :]) @@ -16,4 +17,4 @@ def test_transpose(): - assert torch.allclose(K_keops.t() @ w, K_torch.t() @ w) + assert_torch_allclose(K_keops.t() @ w, K_torch.t() @ w, label="transpose") diff --git a/pykeops/pykeops/test/test_vmap.py b/pykeops/pykeops/test/test_vmap.py index 4f3f21e96..68444636f 100644 --- a/pykeops/pykeops/test/test_vmap.py +++ b/pykeops/pykeops/test/test_vmap.py @@ -1,5 +1,6 @@ import torch from pykeops.torch import LazyTensor +from pykeops.test import assert_torch_allclose test_grad = True torch.manual_seed(0) @@ -31,13 +32,13 @@ def test_vmap_0(self): x_i, x_j, y_j, use_keops=False ) - assert torch.allclose(res_keops, res_torch) + assert_torch_allclose(res_keops, res_torch, label="vmap_0_fw") u = torch.rand(res_torch.shape) (res_torch_grad,) = torch.autograd.grad(res_torch, x_i, u) (res_keops_grad,) = torch.autograd.grad(res_keops, x_i, u) - assert torch.allclose(res_keops_grad, res_torch_grad) + assert_torch_allclose(res_keops_grad, res_torch_grad, label="vmap_0_bw") def test_vmap_1(self): x_i = torch.randn(10, 5, 1, 2, requires_grad=test_grad, dtype=torch.float64) @@ -52,13 +53,13 @@ def test_vmap_1(self): res_torch = torch.vmap(fn, in_dims=in_dims, out_dims=out_dims)( x_i, x_j, y_j, use_keops=False ) - assert torch.allclose(res_keops, res_torch) + assert_torch_allclose(res_keops, res_torch, label="vmap_1_fw") u = torch.rand(res_torch.shape) (res_torch_grad,) = torch.autograd.grad(res_torch, x_i, u) (res_keops_grad,) = torch.autograd.grad(res_keops, x_i, u) - assert torch.allclose(res_keops_grad, res_torch_grad) + assert_torch_allclose(res_keops_grad, res_torch_grad, label="vmap_1_bw") def test_vmap_2(self): x_i = torch.randn(10, 5, 1, 2, requires_grad=test_grad, dtype=torch.float64) @@ -73,10 +74,10 @@ def test_vmap_2(self): res_torch = torch.vmap(fn, in_dims=in_dims, out_dims=out_dims)( x_i, x_j, y_j, use_keops=False ) - assert torch.allclose(res_keops, res_torch) + assert_torch_allclose(res_keops, res_torch, label="vmap_2_fw") u = torch.rand(res_torch.shape) (res_torch_grad,) = torch.autograd.grad(res_torch, x_i, u) (res_keops_grad,) = torch.autograd.grad(res_keops, x_i, u) - assert torch.allclose(res_keops_grad, res_torch_grad) + assert_torch_allclose(res_keops_grad, res_torch_grad, label="vmap_2_bw") diff --git a/pykeops/pykeops/tutorials/kmeans/plot_kmeans_torch.py b/pykeops/pykeops/tutorials/kmeans/plot_kmeans_torch.py index c09f87c7c..a2043577d 100644 --- a/pykeops/pykeops/tutorials/kmeans/plot_kmeans_torch.py +++ b/pykeops/pykeops/tutorials/kmeans/plot_kmeans_torch.py @@ -29,7 +29,7 @@ use_cuda = torch.cuda.is_available() dtype = torch.float32 if use_cuda else torch.float64 -device_id = "cuda:0" if use_cuda else "cpu" +device_id = "cuda" if use_cuda else "cpu" ######################################################################## # Simple implementation of the K-means algorithm: From 6dae0ceacbcbaa4e0a1443184ea62406d6400935 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Wed, 6 May 2026 16:55:48 +0200 Subject: [PATCH 36/98] lint --- pykeops/pykeops/test/__init__.py | 8 +- pykeops/pykeops/test/test_chunks_ranges.py | 4 +- pykeops/pykeops/test/test_float16.py | 4 +- pykeops/pykeops/test/test_gpu_cpu.py | 12 +- pykeops/pykeops/test/test_kron.py | 20 ++- pykeops/pykeops/test/test_lazytensor_clamp.py | 4 +- pykeops/pykeops/test/test_lazytensor_grad.py | 8 +- pykeops/pykeops/test/test_torch.py | 117 +++++++++++++----- pykeops/pykeops/test/test_torch_func.py | 8 +- .../pykeops/test/test_torch_func_logsumexp.py | 8 +- 10 files changed, 138 insertions(+), 55 deletions(-) diff --git a/pykeops/pykeops/test/__init__.py b/pykeops/pykeops/test/__init__.py index 7cd6237a6..6fa62a837 100644 --- a/pykeops/pykeops/test/__init__.py +++ b/pykeops/pykeops/test/__init__.py @@ -14,11 +14,15 @@ def assert_torch_allclose(actual, expected, *, label=None, **kwargs): ok = torch.allclose(actual, expected, **kwargs) diff = torch.linalg.norm(actual - expected).item() prefix = label if label is not None else "torch.allclose failed" - assert ok, f"{prefix}: ||ref-test||_2={diff:.6e} and ||ref||_2={torch.linalg.norm(expected).item():.6e}" + assert ( + ok + ), f"{prefix}: ||ref-test||_2={diff:.6e} and ||ref||_2={torch.linalg.norm(expected).item():.6e}" def assert_np_allclose(actual, expected, *, label=None, **kwargs): ok = np.allclose(actual, expected, **kwargs) diff = float(np.linalg.norm(np.asarray(actual) - np.asarray(expected))) prefix = label if label is not None else "np.allclose failed" - assert ok, f"{prefix}: ||ref-test||_2={diff:.6e} and ||ref||_2={float(np.linalg.norm(expected)):.6e}" + assert ( + ok + ), f"{prefix}: ||ref-test||_2={diff:.6e} and ||ref||_2={float(np.linalg.norm(expected)):.6e}" diff --git a/pykeops/pykeops/test/test_chunks_ranges.py b/pykeops/pykeops/test/test_chunks_ranges.py index d186eebc4..c48a03725 100644 --- a/pykeops/pykeops/test/test_chunks_ranges.py +++ b/pykeops/pykeops/test/test_chunks_ranges.py @@ -39,4 +39,6 @@ def fun(x, y, b, backend): def test_chunk_ranges(): - assert_torch_allclose(out[0], out[1], rtol=0.0001, atol=0.0001, label="chunks_ranges") + assert_torch_allclose( + out[0], out[1], rtol=0.0001, atol=0.0001, label="chunks_ranges" + ) diff --git a/pykeops/pykeops/test/test_float16.py b/pykeops/pykeops/test/test_float16.py index 6a7156129..1d2484a32 100644 --- a/pykeops/pykeops/test/test_float16.py +++ b/pykeops/pykeops/test/test_float16.py @@ -35,7 +35,9 @@ def test_float16_fw(self): for backend in ["torch", "keops"]: self.out.append(fun(x, y, backend).squeeze()) - assert_torch_allclose(self.out[0], self.out[1], atol=0.001, rtol=0.001, label="float16_fw") + assert_torch_allclose( + self.out[0], self.out[1], atol=0.001, rtol=0.001, label="float16_fw" + ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires a GPU") def test_float16_bw(self): diff --git a/pykeops/pykeops/test/test_gpu_cpu.py b/pykeops/pykeops/test/test_gpu_cpu.py index 5977e8c91..fe49cb7d0 100644 --- a/pykeops/pykeops/test/test_gpu_cpu.py +++ b/pykeops/pykeops/test/test_gpu_cpu.py @@ -39,13 +39,13 @@ def fun(x, y, b, backend): class TestCase: def test_torch_keops_cpu(self): - assert torch.allclose(out[0], out[1]), ( - f"torch vs keops_cpu mismatch: ||ref-test||_2={torch.norm(out[0] - out[1]).item():.6e} and ||ref||_2={torch.norm(out[0]).item():.6e}" - ) + assert torch.allclose( + out[0], out[1] + ), f"torch vs keops_cpu mismatch: ||ref-test||_2={torch.norm(out[0] - out[1]).item():.6e} and ||ref||_2={torch.norm(out[0]).item():.6e}" @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires a GPU") def test_torch_keops_gpu(self): out_gpu = fun(x, y, b, ["keops_gpu"]).squeeze() - assert torch.allclose(out[0], out_gpu), ( - f"torch vs keops_gpu mismatch: ||ref-test||_2={torch.norm(out[0] - out_gpu).item():.6e} and ||ref||_2={torch.norm(out[0]).item():.6e}" - ) + assert torch.allclose( + out[0], out_gpu + ), f"torch vs keops_gpu mismatch: ||ref-test||_2={torch.norm(out[0] - out_gpu).item():.6e} and ||ref||_2={torch.norm(out[0]).item():.6e}" diff --git a/pykeops/pykeops/test/test_kron.py b/pykeops/pykeops/test/test_kron.py index dd6c02095..22104db09 100644 --- a/pykeops/pykeops/test/test_kron.py +++ b/pykeops/pykeops/test/test_kron.py @@ -33,7 +33,9 @@ def test_kron_lazytensor_np(): Y = LazyTensor(y[None, :, :]) gamma_keops_lazytensor_np = (X.keops_kron(Y, dimfa, dimfb)).sum(axis=axis) - assert_np_allclose(gamma_keops_lazytensor_np, gamma_py, atol=1e-6, label="kron_lazytensor_np") + assert_np_allclose( + gamma_keops_lazytensor_np, gamma_py, atol=1e-6, label="kron_lazytensor_np" + ) ############################################################################ @@ -83,8 +85,12 @@ def test_kron_genred(): myconv2 = Genred(formula2, aliases, reduction_op="Sum", axis=axis) gamma_keops_genred_TensorDot = myconv2(x0, y0) - assert_np_allclose(gamma_keops_genred_Kron, gamma_py2, atol=1e-6, label="kron_genred") - assert_np_allclose(gamma_keops_genred_TensorDot, gamma_py2, atol=1e-6, label="tensordot_genred") + assert_np_allclose( + gamma_keops_genred_Kron, gamma_py2, atol=1e-6, label="kron_genred" + ) + assert_np_allclose( + gamma_keops_genred_TensorDot, gamma_py2, atol=1e-6, label="tensordot_genred" + ) ############################################################################ @@ -127,6 +133,10 @@ def test_kron_genred2(): myconv2 = Genred(formula2, aliases, reduction_op="Sum", axis=axis) gamma_keops_genred_TensorDot2 = myconv2(x2, y2) - assert_np_allclose(gamma_keops_genred_Kron2, gamma_py3, atol=1e-6, label="kron_genred2") - assert_np_allclose(gamma_keops_genred_TensorDot2, gamma_py3, atol=1e-6, label="tensordot_genred2") + assert_np_allclose( + gamma_keops_genred_Kron2, gamma_py3, atol=1e-6, label="kron_genred2" + ) + assert_np_allclose( + gamma_keops_genred_TensorDot2, gamma_py3, atol=1e-6, label="tensordot_genred2" + ) assert_np_allclose(gamma_py4, gamma_py3, atol=1e-6, label="numpy_kron_einsum") diff --git a/pykeops/pykeops/test/test_lazytensor_clamp.py b/pykeops/pykeops/test/test_lazytensor_clamp.py index 28b7fa631..c9d962b03 100644 --- a/pykeops/pykeops/test/test_lazytensor_clamp.py +++ b/pykeops/pykeops/test/test_lazytensor_clamp.py @@ -39,4 +39,6 @@ def test_lazytensor_clamp_fw(self): assert_torch_allclose(out[0], out[1], label="clamp_fw") def test_lazytensor_clamp_bw(self): - assert_torch_allclose(out_g[0], out_g[1], atol=0.01, rtol=0.001, label="clamp_bw") + assert_torch_allclose( + out_g[0], out_g[1], atol=0.01, rtol=0.001, label="clamp_bw" + ) diff --git a/pykeops/pykeops/test/test_lazytensor_grad.py b/pykeops/pykeops/test/test_lazytensor_grad.py index 98aac080f..350742ed6 100644 --- a/pykeops/pykeops/test/test_lazytensor_grad.py +++ b/pykeops/pykeops/test/test_lazytensor_grad.py @@ -49,7 +49,11 @@ def test_lazytensor_grad(self): assert_torch_allclose(out[0], out[1], rtol=0.0001, label="lazytensor_grad_fw") def test_lazytensor_grad_bw1(self): - assert_torch_allclose(out_g[0], out_g[1], rtol=0.0001, label="lazytensor_grad_bw1") + assert_torch_allclose( + out_g[0], out_g[1], rtol=0.0001, label="lazytensor_grad_bw1" + ) def test_lazytensor_grad_bw2(self): - assert_torch_allclose(out_g2[0], out_g2[1], rtol=0.001, label="lazytensor_grad_bw2") + assert_torch_allclose( + out_g2[0], out_g2[1], rtol=0.001, label="lazytensor_grad_bw2" + ) diff --git a/pykeops/pykeops/test/test_torch.py b/pykeops/pykeops/test/test_torch.py index 47e502fc7..0ed581de0 100644 --- a/pykeops/pykeops/test/test_torch.py +++ b/pykeops/pykeops/test/test_torch.py @@ -21,6 +21,7 @@ try: import torch + HAS_TORCH = True use_cuda = torch.cuda.is_available() @@ -47,23 +48,48 @@ class PytorchUnitTestCase(unittest.TestCase): D = int(3) E = int(3) nbatchdims = int(2) + def SetUp(self): - self.x64 = torch.rand((M, D), dtype=torch.float64, device=device, requires_grad=True) - self.a64 = torch.rand((M, E), dtype=torch.float64, device=device, requires_grad=False) - self.e64 = torch.rand((M, E), dtype=torch.float64, device=device, requires_grad=False) - self.f64 = torch.rand((M, 1), dtype=torch.float64, device=device, requires_grad=True) - self.y64 = torch.rand((N, D), dtype=torch.float64, device=device, requires_grad=False) - self.b64 = torch.rand((N, E), dtype=torch.float64, device=device, requires_grad=False) - self.g64 = torch.rand((N, 1), dtype=torch.float64, device=device, requires_grad=True) - self.p64 = torch.rand(2, dtype=torch.float64, device=device, requires_grad=False) + self.x64 = torch.rand( + (M, D), dtype=torch.float64, device=device, requires_grad=True + ) + self.a64 = torch.rand( + (M, E), dtype=torch.float64, device=device, requires_grad=False + ) + self.e64 = torch.rand( + (M, E), dtype=torch.float64, device=device, requires_grad=False + ) + self.f64 = torch.rand( + (M, 1), dtype=torch.float64, device=device, requires_grad=True + ) + self.y64 = torch.rand( + (N, D), dtype=torch.float64, device=device, requires_grad=False + ) + self.b64 = torch.rand( + (N, E), dtype=torch.float64, device=device, requires_grad=False + ) + self.g64 = torch.rand( + (N, 1), dtype=torch.float64, device=device, requires_grad=True + ) + self.p64 = torch.rand( + 2, dtype=torch.float64, device=device, requires_grad=False + ) self.sigma64 = torch.tensor([0.4], dtype=torch.float64, device=device) self.alpha64 = torch.tensor([0.1], dtype=torch.float64, device=device) - self.X64 = torch.rand((A, B, M, D), dtype=torch.float64, device=device, requires_grad=True) - self.L64 = torch.rand((A, 1, M, 1), dtype=torch.float64, device=device, requires_grad=False) - self.Y64 = torch.rand((1, B, N, D), dtype=torch.float64, device=device, requires_grad=True) - self.S64 = 1 + torch.rand((A, B, 1), dtype=torch.float64, device=device, requires_grad=True) + self.X64 = torch.rand( + (A, B, M, D), dtype=torch.float64, device=device, requires_grad=True + ) + self.L64 = torch.rand( + (A, 1, M, 1), dtype=torch.float64, device=device, requires_grad=False + ) + self.Y64 = torch.rand( + (1, B, N, D), dtype=torch.float64, device=device, requires_grad=True + ) + self.S64 = 1 + torch.rand( + (A, B, 1), dtype=torch.float64, device=device, requires_grad=True + ) self.x32 = x64.to(torch.float32).clone().requires_grad_(True) self.a32 = a64.to(torch.float32).clone().requires_grad_(False) @@ -82,7 +108,6 @@ def SetUp(self): self.Y32 = Y64.to(torch.float32).clone().requires_grad_(True) self.S32 = S64.to(torch.float32).clone().requires_grad_(True) - ############################################################ def test_torchtools_function_binding(self): ############################################################ @@ -122,7 +147,9 @@ def test_generic_syntax_float(self): dim=1, ).T # compare output - assert_torch_allclose(gamma_keops.to(torch.float64), gamma_ref, atol=1e-6) + assert_torch_allclose( + gamma_keops.to(torch.float64), gamma_ref, atol=1e-6 + ) ############################################################ def test_generic_syntax_double(self): @@ -179,10 +206,16 @@ def test_generic_syntax_softmax(self): ) # Torch reference - sqdist = torch.sum((self.x64[:, None, :] - self.y64[None, :, :]) ** 2, dim=2) + sqdist = torch.sum( + (self.x64[:, None, :] - self.y64[None, :, :]) ** 2, dim=2 + ) scores = (self.sigma64 - self.g64.T) ** 2 * torch.exp(-sqdist) scores = scores - torch.max(scores, dim=1, keepdim=True).values - gamma_ref = torch.exp(scores) @ self.y64 / torch.sum(torch.exp(scores), dim=1, keepdim=True) + gamma_ref = ( + torch.exp(scores) + @ self.y64 + / torch.sum(torch.exp(scores), dim=1, keepdim=True) + ) # compare output assert_torch_allclose(gamma_keops, gamma_ref, atol=1e-6) @@ -214,9 +247,9 @@ def test_generic_syntax_simple(self): # Torch reference scals = (self.x64 @ self.y64.T) ** 2 # Memory-intensive computation! - gamma_ref = self.p64[0] * scals.sum(1).reshape(-1, 1) * self.x64 + self.p64[ - 1 - ] * (scals @ self.y64) + gamma_ref = self.p64[0] * scals.sum(1).reshape( + -1, 1 + ) * self.x64 + self.p64[1] * (scals @ self.y64) # compare output assert_torch_allclose(gamma_keops, gamma_ref, atol=1e-6) @@ -253,7 +286,9 @@ def test_logSumExp_kernels_feature(self): gamma_lazy = gamma_lazy.logsumexp(dim=1, weight=Vj(self.g64.exp())) # Torch reference - sqdist = torch.sum((self.x64[:, None, :] - self.y64[None, :, :]) ** 2, dim=2) + sqdist = torch.sum( + (self.x64[:, None, :] - self.y64[None, :, :]) ** 2, dim=2 + ) inv_s2 = 1 / (self.sigma64**2) if k == "gaussian": log_k = -(inv_s2 * sqdist) @@ -266,7 +301,9 @@ def test_logSumExp_kernels_feature(self): gamma_ref = torch.logsumexp(log_k + self.g64.T, dim=1) # compare output - assert_torch_allclose(gamma_lazy.reshape(-1), gamma_ref.reshape(-1), atol=1e-6) + assert_torch_allclose( + gamma_lazy.reshape(-1), gamma_ref.reshape(-1), atol=1e-6 + ) ############################################################ def test_logSumExp_gradient_kernels_feature(self): @@ -302,8 +339,12 @@ def test_logSumExp_gradient_kernels_feature(self): ] # compare output - assert_torch_allclose(gamma_keops[0].reshape(-1), gamma_ref[0].reshape(-1), atol=1e-6) - assert_torch_allclose(gamma_keops[1].reshape(-1), gamma_ref[1].reshape(-1), atol=1e-6) + assert_torch_allclose( + gamma_keops[0].reshape(-1), gamma_ref[0].reshape(-1), atol=1e-6 + ) + assert_torch_allclose( + gamma_keops[1].reshape(-1), gamma_ref[1].reshape(-1), atol=1e-6 + ) ############################################################ def test_non_contiguity(self): @@ -341,7 +382,9 @@ def test_heterogeneous_var_aliases(self): # Torch reference sqdist = torch.sum((self.x64[:, None, :] - self.y64[None, :, :]) ** 2, dim=2) - gamma_ref = torch.sum((self.sigma64 - self.g64.T) ** 2 * torch.exp(-sqdist), dim=1) + gamma_ref = torch.sum( + (self.sigma64 - self.g64.T) ** 2 * torch.exp(-sqdist), dim=1 + ) # compare output assert_torch_allclose(gamma_keops.reshape(-1), gamma_ref.reshape(-1), atol=1e-6) @@ -506,7 +549,9 @@ def test_LazyTensor_sum(self): for res_keops, res_torch in zip(full_results[0], full_results[1]): self.assertTrue(res_keops.shape == res_torch.shape) - assert_torch_allclose(res_keops.reshape(-1), res_torch.reshape(-1), atol=1e-3) + assert_torch_allclose( + res_keops.reshape(-1), res_torch.reshape(-1), atol=1e-3 + ) ############################################################ def test_LazyTensor_logsumexp(self): @@ -550,7 +595,9 @@ def test_LazyTensor_logsumexp(self): for res_keops, res_torch in zip(full_results[0], full_results[1]): self.assertTrue(res_keops.shape == res_torch.shape) - assert_torch_allclose(res_keops.reshape(-1), res_torch.reshape(-1), atol=1e-5) + assert_torch_allclose( + res_keops.reshape(-1), res_torch.reshape(-1), atol=1e-5 + ) ############################################################ # Test min reduction with chunk without batches @@ -569,7 +616,7 @@ def test_LazyTensor_min_chunked(self): for use_keops in [True, False]: results = [] - for x, y in [(X32, Y32), (X64, Y64)]: + for x, y in [(X32, Y32), (X64, Y64)]: x_i = x.unsqueeze(-2) y_j = y.unsqueeze(-3) @@ -590,7 +637,9 @@ def test_LazyTensor_min_chunked(self): for res_keops, res_torch in zip(full_results[0], full_results[1]): self.assertTrue(res_keops.shape == res_torch.shape) - assert_torch_allclose(res_keops.reshape(-1), res_torch.reshape(-1), atol=1e-5) + assert_torch_allclose( + res_keops.reshape(-1), res_torch.reshape(-1), atol=1e-5 + ) ############################################################ def test_LazyTensor_min(self): @@ -630,7 +679,9 @@ def test_LazyTensor_min(self): for res_keops, res_torch in zip(full_results[0], full_results[1]): self.assertTrue(res_keops.shape == res_torch.shape) - assert_torch_allclose(res_keops.reshape(-1), res_torch.reshape(-1), atol=1e-5) + assert_torch_allclose( + res_keops.reshape(-1), res_torch.reshape(-1), atol=1e-5 + ) ############################################################ def test_TensorDot_with_permute(self): @@ -642,7 +693,9 @@ def my_tensordort_perm(a, b, dims=None, perm=None): return torch.tensordot(a, b, dims=dims).sum(3).permute(perm) def invert_permutation_numpy(permutation): - return [idx for idx, _ in sorted(enumerate(permutation), key=lambda x: x[1])] + return [ + idx for idx, _ in sorted(enumerate(permutation), key=lambda x: x[1]) + ] x = torch.randn(self.M, 2, 3, 2, 2, 4, requires_grad=True, dtype=torch.float64) y = torch.randn( @@ -655,9 +708,7 @@ def invert_permutation_numpy(permutation): perm_torch = (0,) + tuple([(i + 1) for i in invert_permutation_numpy(perm)]) sum_f_torch2 = my_tensordort_perm(x, y, dims=(contfa, contfb), perm=perm_torch) - f_keops = LazyTensor( - x.reshape(self.M, 1, prod(dimfa)) - ).keops_tensordot( + f_keops = LazyTensor(x.reshape(self.M, 1, prod(dimfa))).keops_tensordot( LazyTensor(y.reshape(1, self.N, prod(dimfb))), dimfa, dimfb, diff --git a/pykeops/pykeops/test/test_torch_func.py b/pykeops/pykeops/test/test_torch_func.py index 240a29864..332b1968a 100644 --- a/pykeops/pykeops/test/test_torch_func.py +++ b/pykeops/pykeops/test/test_torch_func.py @@ -44,7 +44,9 @@ def test_torch_func_grad(self): res1 = torch.func.grad(fn_torch, (0, 1, 2, 3))(x_i, y_j, b_j, p) res2 = torch.func.grad(fn_keops, (0, 1, 2, 3))(x_i, y_j, b_j, p) for k in range(4): - assert_torch_allclose(res1[k], res2[k], atol=1e-5, label=f"torch_func_grad[{k}]") + assert_torch_allclose( + res1[k], res2[k], atol=1e-5, label=f"torch_func_grad[{k}]" + ) def test_torch_func_vjp(self): res1 = torch.func.vjp(fn_torch, x_i, y_j, b_j, p)[1](torch.tensor(1.0))[0] @@ -60,7 +62,9 @@ def test_torch_func_jacrev(self): res1 = torch.func.jacrev(fn_torch, (0, 1, 2, 3))(x_i, y_j, b_j, p) res2 = torch.func.jacrev(fn_keops, (0, 1, 2, 3))(x_i, y_j, b_j, p) for k in range(4): - assert_torch_allclose(res1[k], res2[k], atol=1e-5, label=f"torch_func_jacrev[{k}]") + assert_torch_allclose( + res1[k], res2[k], atol=1e-5, label=f"torch_func_jacrev[{k}]" + ) def test_torch_func_jacfwd(self): res1 = torch.func.jacfwd(fn_torch)(x_i, y_j, b_j, p) diff --git a/pykeops/pykeops/test/test_torch_func_logsumexp.py b/pykeops/pykeops/test/test_torch_func_logsumexp.py index d45e6df26..6795938c1 100644 --- a/pykeops/pykeops/test/test_torch_func_logsumexp.py +++ b/pykeops/pykeops/test/test_torch_func_logsumexp.py @@ -36,7 +36,9 @@ def test_torch_func_grad(self): res1 = torch.func.grad(fn_torch, (0, 1, 2))(x_i, y_j, b_j) res2 = torch.func.grad(fn_keops, (0, 1, 2))(x_i, y_j, b_j) for k in range(3): - assert_torch_allclose(res1[k], res2[k], atol=1e-5, label=f"torch_func_lse_grad[{k}]") + assert_torch_allclose( + res1[k], res2[k], atol=1e-5, label=f"torch_func_lse_grad[{k}]" + ) def test_torch_func_vjp(self): res1 = torch.func.vjp(fn_torch, x_i, y_j, b_j)[1](torch.tensor(1.0))[0] @@ -52,7 +54,9 @@ def test_torch_func_jacrev(self): res1 = torch.func.jacrev(fn_torch, (0, 1, 2))(x_i, y_j, b_j) res2 = torch.func.jacrev(fn_keops, (0, 1, 2))(x_i, y_j, b_j) for k in range(3): - assert_torch_allclose(res1[k], res2[k], atol=1e-5, label=f"torch_func_lse_jacrev[{k}]") + assert_torch_allclose( + res1[k], res2[k], atol=1e-5, label=f"torch_func_lse_jacrev[{k}]" + ) def test_torch_func_jacfwd(self): res1 = torch.func.jacfwd(fn_torch)(x_i, y_j, b_j) From 7deb341bddbc2abfd3a3f522f28b76d15fb46046 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Wed, 6 May 2026 17:13:50 +0200 Subject: [PATCH 37/98] Fix test. Add manual random seed. Shape tests. --- pykeops/pykeops/test/test_numpy.py | 25 ++-- pykeops/pykeops/test/test_shape_regression.py | 129 ++++++++++++++++++ pykeops/pykeops/test/test_torch.py | 84 +++++++----- 3 files changed, 191 insertions(+), 47 deletions(-) create mode 100644 pykeops/pykeops/test/test_shape_regression.py diff --git a/pykeops/pykeops/test/test_numpy.py b/pykeops/pykeops/test/test_numpy.py index 8a8a0bcb2..f1005a25d 100644 --- a/pykeops/pykeops/test/test_numpy.py +++ b/pykeops/pykeops/test/test_numpy.py @@ -28,9 +28,12 @@ log_sum_exp, ) +np.random.seed(42) + class NumpyUnitTestCase(unittest.TestCase): - def SetUp(self): + + def setUp(self): self.A = int(4) # Batchdim 1 self.B = int(6) # Batchdim 2 self.M = int(10) @@ -39,18 +42,18 @@ def SetUp(self): self.E = int(3) self.nbatchdims = int(2) - self.x = np.random.rand(M, D) - self.a = np.random.rand(M, E) - self.f = np.random.rand(M, 1) - self.y = np.random.rand(N, D) - self.b = np.random.rand(N, E) - self.g = np.random.rand(N, 1) + self.x = np.random.rand(self.M, self.D) + self.a = np.random.rand(self.M, self.E) + self.f = np.random.rand(self.M, 1) + self.y = np.random.rand(self.N, self.D) + self.b = np.random.rand(self.N, self.E) + self.g = np.random.rand(self.N, 1) self.sigma = np.array([0.4]) - self.X = np.random.rand(A, 1, M, D) - self.L = np.random.rand(1, B, M, 1) - self.Y = np.random.rand(1, B, N, D) - self.S = np.random.rand(A, B, 1) + 1 + self.X = np.random.rand(self.A, 1, self.M, self.D) + self.L = np.random.rand(1, self.B, self.M, 1) + self.Y = np.random.rand(1, self.B, self.N, self.D) + self.S = np.random.rand(self.A, self.B, 1) + 1 self.type_to_test = ["float32", "float64"] diff --git a/pykeops/pykeops/test/test_shape_regression.py b/pykeops/pykeops/test/test_shape_regression.py new file mode 100644 index 000000000..a6e27e572 --- /dev/null +++ b/pykeops/pykeops/test/test_shape_regression.py @@ -0,0 +1,129 @@ +import re +import unittest +import numpy as np + +from pykeops.test import assert_np_allclose, assert_torch_allclose + + +class ShapeNumpyTestCase(unittest.TestCase): + + def setUp(self): + self.param = np.array([0.4], dtype=np.float32) + self.param_scalar = np.array(0.4, dtype=np.float32) + + self.x = np.array([[0.1], [0.2], [0.3]], dtype=np.float32) # shape (3,1) + self.x_vect = np.array([0.1, 0.2, 0.3], dtype=np.float32) # shape (3,) + + def test_numpy_pm1_shape_regression(self): + from pykeops.numpy import Genred + + out = Genred("param", ["param=Pm(1)"], axis=1)(self.param, backend="auto") + ref = np.full_like(out, self.param[0]) + assert_np_allclose(out, ref, atol=1e-7) + + def test_numpy_pm1_scalar_rejected(self): + from pykeops.numpy import Genred + + with self.assertRaisesRegex( + ValueError, + re.escape( + "[pyKeOps] Error: Pm argument #0 requires at least 1 dimension (got a scalar)." + ), + ): + Genred("param", ["param=Pm(1)"], axis=1)(self.param_scalar, backend="auto") + + def test_numpy_vj_single_dim_formula_x(self): + from pykeops.numpy import Genred + + out = Genred("x", ["x=Vj(1)"], axis=1)(self.x, backend="auto") + ref = np.sum(self.x, axis=0, keepdims=True) + assert_np_allclose(out, ref, atol=1e-7) + + def test_numpy_vj_single_dim_formula_x_1d_rejected_python(self): + from pykeops.numpy import Genred + + # The error is raised in the Python code (parse_type.py), before even calling the C++ code + with self.assertRaisesRegex(IndexError, "tuple index out of range"): + Genred("y", ["y=Vj(1)"], axis=1)(self.x_vect, backend="auto") + + def test_numpy_vj_single_dim_formula_x_1d_rejected_cpp(self): + from pykeops.numpy import Genred + + # The error is raised in the C++ code (LoadKeOps_cpp.py or pykeops_nvrtc.py), after the Python code has accepted the input shape but before launching the kernel + with self.assertRaisesRegex( + ValueError, + re.escape( + "[pyKeOps] Error: Vj argument #2 requires at least 2 dimensions, got 1." + ), + ): + Genred("y + z + l", ["y=Vj(1)", "z=Vi(1)", "l=Vj(1)"], axis=1)( + self.x, self.x, self.x_vect, backend="auto" + ) + + +try: + import torch + + HAS_TORCH = True +except ImportError: + HAS_TORCH = False + + +@unittest.skipUnless(HAS_TORCH, "torch not available") +class ShapeTorchTestCase(unittest.TestCase): + + def setUp(self): + device = "cuda" if torch.cuda.is_available() else "cpu" + self.param = torch.tensor([0.4], dtype=torch.float32, device=device) + self.param_scalar = torch.tensor(0.4, dtype=torch.float32, device=device) + self.x = torch.tensor([[0.1], [0.2], [0.3]], dtype=torch.float32, device=device) + self.x_vect = torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32, device=device) + + def test_torch_pm1_shape_regression(self): + from pykeops.torch import Genred + + out = Genred("param", ["param=Pm(1)"], axis=1)(self.param, backend="auto") + ref = torch.full_like(out, self.param[0]) + assert_torch_allclose(out, ref, atol=1e-7) + + def test_torch_pm1_scalar_rejected(self): + from pykeops.torch import Genred + + with self.assertRaisesRegex( + ValueError, + re.escape( + "[pyKeOps] Error: Pm argument #0 requires at least 1 dimension (got a scalar)." + ), + ): + Genred("param", ["param=Pm(1)"], axis=1)(self.param_scalar, backend="auto") + + def test_torch_vj_single_dim_formula_x(self): + from pykeops.torch import Genred + + out = Genred("x", ["x=Vj(1)"], axis=1)(self.x, backend="auto") + ref = torch.sum(self.x, dim=0, keepdim=True) + assert_torch_allclose(out, ref, atol=1e-7) + + def test_torch_vj_single_dim_formula_x_1d_rejected_python(self): + from pykeops.torch import Genred + + with self.assertRaisesRegex(IndexError, "tuple index out of range"): + Genred("y", ["y=Vj(1)"], axis=1)(self.x_vect, backend="auto") + + def test_torch_vj_single_dim_formula_x_1d_rejected_cpp(self): + from pykeops.torch import Genred + + # The error is raised in the C++ code (LoadKeOps_cpp.py or pykeops_nvrtc.py), after the Python code has accepted the input shape but before launching the kernel + with self.assertRaisesRegex( + ValueError, + re.escape( + "[pyKeOps] Error: Vj argument #2 requires at least 2 dimensions, got 1." + ), + ): + Genred("y + z + l", ["y=Vj(1)", "z=Vi(1)", "l=Vj(1)"], axis=1)( + self.x, self.x, self.x_vect, backend="auto" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/pykeops/pykeops/test/test_torch.py b/pykeops/pykeops/test/test_torch.py index 0ed581de0..7fd86ec25 100644 --- a/pykeops/pykeops/test/test_torch.py +++ b/pykeops/pykeops/test/test_torch.py @@ -24,6 +24,8 @@ HAS_TORCH = True + torch.manual_seed(42) + use_cuda = torch.cuda.is_available() if use_cuda: device = "cuda" @@ -41,35 +43,36 @@ @unittest.skipUnless(HAS_TORCH, "torch not available") class PytorchUnitTestCase(unittest.TestCase): - A = int(5) # Batchdim 1 - B = int(3) # Batchdim 2 - M = int(10) - N = int(6) - D = int(3) - E = int(3) - nbatchdims = int(2) - - def SetUp(self): + + def setUp(self): + self.A = int(5) # Batchdim 1 + self.B = int(3) # Batchdim 2 + self.M = int(10) + self.N = int(6) + self.D = int(3) + self.E = int(3) + self.nbatchdims = int(2) + self.x64 = torch.rand( - (M, D), dtype=torch.float64, device=device, requires_grad=True + (self.M, self.D), dtype=torch.float64, device=device, requires_grad=True ) self.a64 = torch.rand( - (M, E), dtype=torch.float64, device=device, requires_grad=False + (self.M, self.E), dtype=torch.float64, device=device, requires_grad=False ) self.e64 = torch.rand( - (M, E), dtype=torch.float64, device=device, requires_grad=False + (self.M, self.E), dtype=torch.float64, device=device, requires_grad=False ) self.f64 = torch.rand( - (M, 1), dtype=torch.float64, device=device, requires_grad=True + (self.M, 1), dtype=torch.float64, device=device, requires_grad=True ) self.y64 = torch.rand( - (N, D), dtype=torch.float64, device=device, requires_grad=False + (self.N, self.D), dtype=torch.float64, device=device, requires_grad=False ) self.b64 = torch.rand( - (N, E), dtype=torch.float64, device=device, requires_grad=False + (self.N, self.E), dtype=torch.float64, device=device, requires_grad=False ) self.g64 = torch.rand( - (N, 1), dtype=torch.float64, device=device, requires_grad=True + (self.N, 1), dtype=torch.float64, device=device, requires_grad=True ) self.p64 = torch.rand( 2, dtype=torch.float64, device=device, requires_grad=False @@ -79,34 +82,43 @@ def SetUp(self): self.alpha64 = torch.tensor([0.1], dtype=torch.float64, device=device) self.X64 = torch.rand( - (A, B, M, D), dtype=torch.float64, device=device, requires_grad=True + (self.A, self.B, self.M, self.D), + dtype=torch.float64, + device=device, + requires_grad=True, ) self.L64 = torch.rand( - (A, 1, M, 1), dtype=torch.float64, device=device, requires_grad=False + (self.A, 1, self.M, 1), + dtype=torch.float64, + device=device, + requires_grad=False, ) self.Y64 = torch.rand( - (1, B, N, D), dtype=torch.float64, device=device, requires_grad=True + (1, self.B, self.N, self.D), + dtype=torch.float64, + device=device, + requires_grad=True, ) self.S64 = 1 + torch.rand( - (A, B, 1), dtype=torch.float64, device=device, requires_grad=True + (self.A, self.B, 1), dtype=torch.float64, device=device, requires_grad=True ) - self.x32 = x64.to(torch.float32).clone().requires_grad_(True) - self.a32 = a64.to(torch.float32).clone().requires_grad_(False) - self.e32 = e64.to(torch.float32).clone().requires_grad_(False) - self.f32 = f64.to(torch.float32).clone().requires_grad_(True) - self.y32 = y64.to(torch.float32).clone().requires_grad_(False) - self.b32 = b64.to(torch.float32).clone().requires_grad_(False) - self.g32 = g64.to(torch.float32).clone().requires_grad_(True) - self.p32 = p64.to(torch.float32).clone().requires_grad_(False) - - self.sigma32 = sigma64.to(torch.float32).clone().requires_grad_(False) - self.alpha32 = alpha64.to(torch.float32).clone().requires_grad_(False) - - self.X32 = X64.to(torch.float32).clone().requires_grad_(True) - self.L32 = L64.to(torch.float32).clone().requires_grad_(False) - self.Y32 = Y64.to(torch.float32).clone().requires_grad_(True) - self.S32 = S64.to(torch.float32).clone().requires_grad_(True) + self.x32 = self.x64.to(torch.float32).clone().requires_grad_(True) + self.a32 = self.a64.to(torch.float32).clone().requires_grad_(False) + self.e32 = self.e64.to(torch.float32).clone().requires_grad_(False) + self.f32 = self.f64.to(torch.float32).clone().requires_grad_(True) + self.y32 = self.y64.to(torch.float32).clone().requires_grad_(False) + self.b32 = self.b64.to(torch.float32).clone().requires_grad_(False) + self.g32 = self.g64.to(torch.float32).clone().requires_grad_(True) + self.p32 = self.p64.to(torch.float32).clone().requires_grad_(False) + + self.sigma32 = self.sigma64.to(torch.float32).clone().requires_grad_(False) + self.alpha32 = self.alpha64.to(torch.float32).clone().requires_grad_(False) + + self.X32 = self.X64.to(torch.float32).clone().requires_grad_(True) + self.L32 = self.L64.to(torch.float32).clone().requires_grad_(False) + self.Y32 = self.Y64.to(torch.float32).clone().requires_grad_(True) + self.S32 = self.S64.to(torch.float32).clone().requires_grad_(True) ############################################################ def test_torchtools_function_binding(self): From 9c093a2d5649e7ba9753b4ef8a0d110110e1686e Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Wed, 6 May 2026 17:53:21 +0200 Subject: [PATCH 38/98] Skip conv2d test if not cuda --- pykeops/pykeops/test/test_conv2d.py | 41 +++++++++++++++++------------ 1 file changed, 24 insertions(+), 17 deletions(-) diff --git a/pykeops/pykeops/test/test_conv2d.py b/pykeops/pykeops/test/test_conv2d.py index 132b4a4f3..63bcbf62b 100644 --- a/pykeops/pykeops/test/test_conv2d.py +++ b/pykeops/pykeops/test/test_conv2d.py @@ -1,4 +1,5 @@ import math +import unittest import torch from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -7,13 +8,11 @@ dtype = torch.float32 +torch.manual_seed(42) + torch.backends.cuda.matmul.allow_tf32 = False device_id = "cuda" if torch.cuda.is_available() else "cpu" -x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) -y = torch.rand(1, N, D, device=device_id, dtype=dtype) / math.sqrt(D) -b = torch.randn(N, DV, requires_grad=True, device=device_id, dtype=dtype) - def fun(x, y, b, backend): if "keops" in backend: @@ -33,25 +32,33 @@ def fun(x, y, b, backend): backends = ["keops2D", "torch"] -out = [] -for backend in backends: - out.append(fun(x, y, b, backend).squeeze()) -out_g = [] -for k, backend in enumerate(backends): - out_g.append(torch.autograd.grad((out[k] ** 2).sum(), [b], create_graph=True)[0]) +@unittest.skipUnless(torch.cuda.is_available(), "CUDA not available") +class TestCase(unittest.TestCase): + @classmethod + def setUpClass(cls): + x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) + y = torch.rand(1, N, D, device=device_id, dtype=dtype) / math.sqrt(D) + b = torch.randn(N, DV, requires_grad=True, device=device_id, dtype=dtype) + cls.out = [] + for backend in backends: + cls.out.append(fun(x, y, b, backend).squeeze()) -out_g2 = [] -for k, backend in enumerate(backends): - out_g2.append(torch.autograd.grad((out_g[k] ** 2).sum(), [b])[0]) + cls.out_g = [] + for k in range(len(backends)): + cls.out_g.append( + torch.autograd.grad((cls.out[k] ** 2).sum(), [b], create_graph=True)[0] + ) + cls.out_g2 = [] + for k in range(len(backends)): + cls.out_g2.append(torch.autograd.grad((cls.out_g[k] ** 2).sum(), [b])[0]) -class TestCase: def test_conv2d_fw(self): - assert_torch_allclose(out[0], out[1], label="conv2d_fw") + assert_torch_allclose(self.out[0], self.out[1], label="conv2d_fw") def test_conv2d_bw1(self): - assert_torch_allclose(out_g[0], out_g[1], label="conv2d_bw1") + assert_torch_allclose(self.out_g[0], self.out_g[1], label="conv2d_bw1") def test_conv2d_bw2(self): - assert_torch_allclose(out_g2[0], out_g2[1], label="conv2d_bw2") + assert_torch_allclose(self.out_g2[0], self.out_g2[1], label="conv2d_bw2") From 8cb79383c6828e59e85f3b659851d63f765dffde Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Thu, 7 May 2026 14:46:54 +0200 Subject: [PATCH 39/98] move autofactorize to chunkConfig. Rename ChunkConfig to ReductionTuningConfig. Co-authored-by: Copilot --- .../config/{Chunks.py => ReductionTuning.py} | 13 ++++++++++- keopscore/keopscore/config/__init__.py | 9 +++----- keopscore/keopscore/formulas/Chunkable_Op.py | 8 +++---- keopscore/keopscore/formulas/GetReduction.py | 6 +++-- keopscore/keopscore/get_keops_dll.py | 14 ++++++------ .../mapreduce/Chunk_Mode_Constants.py | 8 +++---- .../mapreduce/gpu/GpuReduc1D_chunks.py | 10 ++++----- .../mapreduce/gpu/GpuReduc1D_finalchunks.py | 20 ++++++++--------- .../mapreduce/gpu/GpuReduc1D_ranges_chunks.py | 16 +++++++------- .../gpu/GpuReduc1D_ranges_finalchunks.py | 22 +++++++++---------- keopscore/keopscore/sandbox/formula.py | 2 +- .../keopscore/sandbox/lazytensor_grad.py | 2 +- keopscore/keopscore/utils/Cache.py | 2 +- .../sandbox/test_torch_func_hessian.py | 2 +- pykeops/pykeops/test/test_torch_func.py | 2 +- 15 files changed, 73 insertions(+), 63 deletions(-) rename keopscore/keopscore/config/{Chunks.py => ReductionTuning.py} (88%) diff --git a/keopscore/keopscore/config/Chunks.py b/keopscore/keopscore/config/ReductionTuning.py similarity index 88% rename from keopscore/keopscore/config/Chunks.py rename to keopscore/keopscore/config/ReductionTuning.py index aadc6e800..895bde41b 100644 --- a/keopscore/keopscore/config/Chunks.py +++ b/keopscore/keopscore/config/ReductionTuning.py @@ -1,4 +1,4 @@ -class ChunksConfig: +class ReductionTuningConfig: """ Configuration and state management for chunked computation schemes. Special computation scheme for dim>100 @@ -14,6 +14,9 @@ class ChunksConfig: _enable_final_chunk = True _mult_var_highdim = False + # Automatic factorization formulas + _auto_factorize = False + def __init__(self): pass @@ -73,6 +76,14 @@ def get_mult_var_highdim(self): """Get the current mult_var_highdim state.""" return self._mult_var_highdim + def set_auto_factorize(self, val): + """Set automatic factorization state.""" + self._auto_factorize = bool(val) + + def get_auto_factorize(self): + """Get automatic factorization state.""" + return self._auto_factorize + def use_final_chunks(self, red_formula): """Determine if final chunks mode should be used for this formula.""" return ( diff --git a/keopscore/keopscore/config/__init__.py b/keopscore/keopscore/config/__init__.py index 9543fb77b..e3ae6c7b2 100644 --- a/keopscore/keopscore/config/__init__.py +++ b/keopscore/keopscore/config/__init__.py @@ -8,7 +8,7 @@ from .KeOpsPath import KeOpsPathConfig from .OpenMP import OpenMPConfig from .Platform import PlatformConfig -from .Chunks import ChunksConfig +from .ReductionTuning import ReductionTuningConfig from keopscore.utils.messages import KeOps_Message # Instantiate the configurations once at import time to preserve the existing API. @@ -18,10 +18,7 @@ openmp = OpenMPConfig(platform, cxx) cuda = CudaConfig() path = KeOpsPathConfig(platform, cuda) -chunks = ChunksConfig() - -# flag for automatic factorization : apply automatic factorization for all formulas before reduction. -auto_factorize = False +reduction = ReductionTuningConfig() def check_health(infos="all"): @@ -73,7 +70,7 @@ def clean_keops(recompile_jit_binary=True, verbose=True): "cuda", "path", "debug", - "auto_factorize", + "reduction", "check_health", "clean_keops", ] diff --git a/keopscore/keopscore/formulas/Chunkable_Op.py b/keopscore/keopscore/formulas/Chunkable_Op.py index b38c92751..23cfed7ef 100644 --- a/keopscore/keopscore/formulas/Chunkable_Op.py +++ b/keopscore/keopscore/formulas/Chunkable_Op.py @@ -1,6 +1,6 @@ from keopscore.formulas.Operation import Operation from keopscore.formulas.variables.Var import Var -from keopscore.config import chunks +from keopscore.config import reduction class Chunkable_Op(Operation): @@ -19,12 +19,12 @@ def notchunked_vars(self, cat): @property def use_chunk(self): - test = chunks.get_enable_chunks() & all( + test = reduction.get_enable_chunks() & all( child.is_chunkable for child in self.children ) child = self.children[0] - subtest = (child.dim >= chunks.get_dim_treshold_chunk()) | ( - child.dim in chunks.get_specdims_use_chunk() + subtest = (child.dim >= reduction.get_dim_treshold_chunk()) | ( + child.dim in reduction.get_specdims_use_chunk() ) test &= subtest return test diff --git a/keopscore/keopscore/formulas/GetReduction.py b/keopscore/keopscore/formulas/GetReduction.py index b1a3ad88d..cd3658326 100644 --- a/keopscore/keopscore/formulas/GetReduction.py +++ b/keopscore/keopscore/formulas/GetReduction.py @@ -11,7 +11,9 @@ class GetReduction: def __new__(self, red_formula_string, aliases=[]): string_id_hash = get_hash_name( - red_formula_string, aliases, keopscore.config.auto_factorize + red_formula_string, + aliases, + keopscore.config.reduction.get_auto_factorize(), ) if string_id_hash in GetReduction.library: return GetReduction.library[string_id_hash] @@ -24,7 +26,7 @@ def __new__(self, red_formula_string, aliases=[]): varname, var = alias.split("=") aliases_dict[varname] = eval(var) reduction = eval(red_formula_string, globals(), aliases_dict) - if keopscore.config.auto_factorize: + if keopscore.config.reduction.get_auto_factorize(): formula = reduction.children[0] new_formula = AutoFactorize(formula) reduction.children[0] = new_formula diff --git a/keopscore/keopscore/get_keops_dll.py b/keopscore/keopscore/get_keops_dll.py index 6a1670e3a..2f303292d 100644 --- a/keopscore/keopscore/get_keops_dll.py +++ b/keopscore/keopscore/get_keops_dll.py @@ -53,7 +53,7 @@ import sys import keopscore.mapreduce -from keopscore.config import path, cuda, debug, chunks +from keopscore.config import path, cuda, debug, reduction from keopscore.formulas import Zero_Reduction, Sum_Reduction from keopscore.formulas.GetReduction import GetReduction from keopscore.formulas.variables.Zero import Zero @@ -80,15 +80,15 @@ def get_keops_dll_impl( KeOps_Error( "You selected a Gpu reduce scheme but KeOps is in Cpu only mode." ) - chunks.set_enable_chunks(enable_chunks) - chunks.set_enable_finalchunk(enable_finalchunks) - chunks.set_mult_var_highdim(mul_var_highdim) + reduction.set_enable_chunks(enable_chunks) + reduction.set_enable_finalchunk(enable_finalchunks) + reduction.set_mult_var_highdim(mul_var_highdim) red_formula = GetReduction(red_formula_string, aliases) - if chunks.use_final_chunks(red_formula) and map_reduce_id != "GpuReduc2D": + if reduction.use_final_chunks(red_formula) and map_reduce_id != "GpuReduc2D": use_chunk_mode = 2 map_reduce_id += "_finalchunks" - elif chunks.get_enable_chunks() and map_reduce_id != "GpuReduc2D": - if len(red_formula.formula.chunked_formulas(chunks.get_dimchunk())) == 1: + elif reduction.get_enable_chunks() and map_reduce_id != "GpuReduc2D": + if len(red_formula.formula.chunked_formulas(reduction.get_dimchunk())) == 1: from keopscore.mapreduce.Chunk_Mode_Constants import ( Chunk_Mode_Constants, ) diff --git a/keopscore/keopscore/mapreduce/Chunk_Mode_Constants.py b/keopscore/keopscore/mapreduce/Chunk_Mode_Constants.py index d7b3a5dc8..1f2f8d6ec 100644 --- a/keopscore/keopscore/mapreduce/Chunk_Mode_Constants.py +++ b/keopscore/keopscore/mapreduce/Chunk_Mode_Constants.py @@ -1,4 +1,4 @@ -from keopscore.config import chunks +from keopscore.config import reduction from keopscore.utils.code_gen_utils import GetDims, GetInds, Var_loader from keopscore.formulas.variables.Var import Var @@ -15,10 +15,10 @@ def __init__(self, red_formula): formula = red_formula.formula self.dimfout = formula.dim # dimension of output variable of inner function - chunked_formula = formula.chunked_formulas(chunks.get_dimchunk())[0] + chunked_formula = formula.chunked_formulas(reduction.get_dimchunk())[0] self.dim_org = chunked_formula["dim_org"] - self.nchunks = 1 + (self.dim_org - 1) // chunks.get_dimchunk() - self.dimlastchunk = self.dim_org - (self.nchunks - 1) * chunks.get_dimchunk() + self.nchunks = 1 + (self.dim_org - 1) // reduction.get_dimchunk() + self.dimlastchunk = self.dim_org - (self.nchunks - 1) * reduction.get_dimchunk() self.nminargs = varloader.nminargs self.fun_chunked = chunked_formula["formula"] self.dimout_chunk = self.fun_chunked.dim diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py index 56c06cdc0..a394c06b6 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py @@ -1,5 +1,5 @@ from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile -from keopscore.config import cuda, chunks +from keopscore.config import cuda, reduction from keopscore.formulas.reductions import make_sum_scheme from keopscore.mapreduce.Chunk_Mode_Constants import Chunk_Mode_Constants from keopscore.mapreduce.MapReduce import MapReduce @@ -54,7 +54,7 @@ def do_chunk_sub( ) load_chunks_routine_i = load_vars_chunks( indsi_chunked, - chunks.get_dimchunk(), + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, xiloc, @@ -64,7 +64,7 @@ def do_chunk_sub( ) load_chunks_routine_j = load_vars_chunks( indsj_chunked, - chunks.get_dimchunk(), + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, yjloc, @@ -74,7 +74,7 @@ def do_chunk_sub( ) load_chunks_routine_p = load_vars_chunks( indsp_chunked, - chunks.get_dimchunk(), + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, param_loc, @@ -181,7 +181,7 @@ def get_code(self): dtype, red_formula, chk.fun_chunked, - chunks.get_dimchunk(), + reduction.get_dimchunk(), chk.dimsx, chk.dimsy, chk.dimsp, diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py index 6d34113cd..0d77b2db6 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_finalchunks.py @@ -1,5 +1,5 @@ from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile -from keopscore.config import cuda, chunks +from keopscore.config import cuda, reduction from keopscore.formulas.reductions.Sum_Reduction import Sum_Reduction from keopscore.formulas.reductions.sum_schemes import block_sum from keopscore.mapreduce.MapReduce import MapReduce @@ -36,11 +36,11 @@ def do_finalchunk_sub( ): dimout = varfinal.dim yjloc = c_variable( - pointer(dtype), f"({yj.id} + threadIdx.x * {chunks.get_dimfinalchunk()})" + pointer(dtype), f"({yj.id} + threadIdx.x * {reduction.get_dimfinalchunk()})" ) load_chunks_routine_j = load_vars_chunks( [varfinal.ind], - chunks.get_dimfinalchunk(), + reduction.get_dimfinalchunk(), dimfinalchunk_curr, varfinal.dim, yjloc, @@ -55,7 +55,7 @@ def do_finalchunk_sub( {load_chunks_routine_j} }} __syncthreads(); - for (signed long int jrel = 0; (jrel < blockDim.x) && (jrel < {ny.id} - {jstart.id}); jrel++, yjrel += {chunks.get_dimfinalchunk()}) {{ + for (signed long int jrel = 0; (jrel < blockDim.x) && (jrel < {ny.id} - {jstart.id}); jrel++, yjrel += {reduction.get_dimfinalchunk()}) {{ if ({i.id} < {nx.id}) {{ // we compute only if needed {use_pragma_unroll()} for (signed long int k=0; k<{dimfinalchunk_curr}; k++) {{ @@ -67,7 +67,7 @@ def do_finalchunk_sub( if ({i.id} < {nx.id}) {{ {use_pragma_unroll()} for (signed long int k=0; k<{dimfinalchunk_curr}; k++) - {out.id}[i*{dimout}+{chunk.id}*{chunks.get_dimfinalchunk()}+k] += {acc.id}[k]; + {out.id}[i*{dimout}+{chunk.id}*{reduction.get_dimfinalchunk()}+k] += {acc.id}[k]; }} __syncthreads(); """ @@ -102,8 +102,8 @@ def get_code(self): ) formula = fun_internal.formula varfinal = self.red_formula.formula.children[1 - ind_fun_internal] - nchunks = 1 + (varfinal.dim - 1) // chunks.get_dimfinalchunk() - dimlastfinalchunk = varfinal.dim - (nchunks - 1) * chunks.get_dimfinalchunk() + nchunks = 1 + (varfinal.dim - 1) // reduction.get_dimfinalchunk() + dimlastfinalchunk = varfinal.dim - (nchunks - 1) * reduction.get_dimfinalchunk() varloader = Var_loader(fun_internal) dimsx = varloader.dimsx dimsy = varloader.dimsy @@ -120,7 +120,7 @@ def get_code(self): KeOps_Error("dimfout should be 1") sum_scheme = self.sum_scheme - self.dimy = max(chunks.get_dimfinalchunk(), dimy) + self.dimy = max(reduction.get_dimfinalchunk(), dimy) blocksize_chunks = min( cuda.get_cuda_block_size(), 1024, @@ -132,7 +132,7 @@ def get_code(self): param_loc = c_array(dtype, dimp, "param_loc") fout = c_array(dtype, dimfout * blocksize_chunks, "fout") xi = c_array(dtype, dimx, "xi") - acc = c_array(dtypeacc, chunks.get_dimfinalchunk(), "acc") + acc = c_array(dtypeacc, reduction.get_dimfinalchunk(), "acc") yjloc = c_array(dtype, dimy, f"(yj + threadIdx.x * {dimy})") foutjrel = c_array(dtype, dimfout, f"({fout.id}+jrel*{dimfout})") yjrel = c_array(dtype, dimy, "yjrel") @@ -143,7 +143,7 @@ def get_code(self): chunk_sub_routine = do_finalchunk_sub( dtype, varfinal, - chunks.get_dimfinalchunk(), + reduction.get_dimfinalchunk(), acc, i, j, diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py index 592da43d5..41bcdffd0 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_chunks.py @@ -1,4 +1,4 @@ -from keopscore.config import cuda, chunks +from keopscore.config import cuda, reduction from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile from keopscore.formulas.reductions import make_sum_scheme @@ -63,7 +63,7 @@ def do_chunk_sub_ranges( load_chunks_routine_i = load_vars_chunks( indsi_chunked, - chunks.get_dimchunk(), + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, xiloc, @@ -81,7 +81,7 @@ def do_chunk_sub_ranges( load_chunks_routine_i_batches = load_vars_chunks_offsets( indsi_chunked, indsi_global, - chunks.get_dimchunk(), + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, xiloc, @@ -93,7 +93,7 @@ def do_chunk_sub_ranges( load_chunks_routine_j = load_vars_chunks( indsj_chunked, - chunks.get_dimchunk(), + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, yjloc, @@ -105,7 +105,7 @@ def do_chunk_sub_ranges( load_chunks_routine_j_batches = load_vars_chunks_offsets( indsj_chunked, indsj_global, - chunks.get_dimchunk(), + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, yjloc, @@ -117,7 +117,7 @@ def do_chunk_sub_ranges( load_chunks_routine_p = load_vars_chunks( indsp_chunked, - chunks.get_dimchunk(), + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, param_loc, @@ -128,7 +128,7 @@ def do_chunk_sub_ranges( load_chunks_routine_p_batches = load_vars_chunks_offsets( indsp_chunked, indsp_global, - chunks.get_dimchunk(), + reduction.get_dimchunk(), dimchunk_curr, chk.dim_org, param_loc, @@ -281,7 +281,7 @@ def get_code(self): dtype, red_formula, chk.fun_chunked, - chunks.get_dimchunk(), + reduction.get_dimchunk(), chk.dimsx, chk.dimsy, chk.dimsp, diff --git a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py index 54c9ecd14..e727ba3d6 100644 --- a/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py +++ b/keopscore/keopscore/mapreduce/gpu/GpuReduc1D_ranges_finalchunks.py @@ -1,4 +1,4 @@ -from keopscore.config import cuda, chunks +from keopscore.config import cuda, reduction from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile from keopscore.formulas.reductions.Sum_Reduction import Sum_Reduction @@ -42,12 +42,12 @@ def do_finalchunk_sub_ranges( ): dimout = varfinal.dim yjloc = c_variable( - pointer(dtype), f"({yj.id} + threadIdx.x * {chunks.get_dimfinalchunk()})" + pointer(dtype), f"({yj.id} + threadIdx.x * {reduction.get_dimfinalchunk()})" ) indsj_global = Var_loader(fun_global).indsj load_chunks_routine_j = load_vars_chunks( [varfinal.ind], - chunks.get_dimfinalchunk(), + reduction.get_dimfinalchunk(), dimfinalchunk_curr, varfinal.dim, yjloc, @@ -58,7 +58,7 @@ def do_finalchunk_sub_ranges( load_chunks_routine_j_ranges = load_vars_chunks_offsets( [varfinal.ind], indsj_global, - chunks.get_dimfinalchunk(), + reduction.get_dimfinalchunk(), dimfinalchunk_curr, varfinal.dim, yjloc, @@ -78,7 +78,7 @@ def do_finalchunk_sub_ranges( }} }} __syncthreads(); - for (signed long int jrel = 0; (jrel < blockDim.x) && (jrel < {end_y.id} - {jstart.id}); jrel++, yjrel += {chunks.get_dimfinalchunk()}) {{ + for (signed long int jrel = 0; (jrel < blockDim.x) && (jrel < {end_y.id} - {jstart.id}); jrel++, yjrel += {reduction.get_dimfinalchunk()}) {{ if ({i.id} < {end_x.id}) {{ // we compute only if needed {use_pragma_unroll()} for (signed long int k=0; k<{dimfinalchunk_curr}; k++) {{ @@ -90,7 +90,7 @@ def do_finalchunk_sub_ranges( if ({i.id} < {end_x.id}) {{ {use_pragma_unroll()} for (signed long int k=0; k<{dimfinalchunk_curr}; k++) - {out.id}[i*{dimout}+{chunk.id}*{chunks.get_dimfinalchunk()}+k] += {acc.id}[k]; + {out.id}[i*{dimout}+{chunk.id}*{reduction.get_dimfinalchunk()}+k] += {acc.id}[k]; }} __syncthreads(); """ @@ -126,8 +126,8 @@ def get_code(self): ) formula = fun_internal.formula varfinal = self.red_formula.formula.children[1 - ind_fun_internal] - nchunks = 1 + (varfinal.dim - 1) // chunks.get_dimfinalchunk() - dimlastfinalchunk = varfinal.dim - (nchunks - 1) * chunks.get_dimfinalchunk() + nchunks = 1 + (varfinal.dim - 1) // reduction.get_dimfinalchunk() + dimlastfinalchunk = varfinal.dim - (nchunks - 1) * reduction.get_dimfinalchunk() varloader = Var_loader(fun_internal) dimsx = varloader.dimsx dimsy = varloader.dimsy @@ -144,7 +144,7 @@ def get_code(self): KeOps_Error("dimfout should be 1") sum_scheme = self.sum_scheme - self.dimy = max(chunks.get_dimfinalchunk(), dimy) + self.dimy = max(reduction.get_dimfinalchunk(), dimy) blocksize_chunks = min( cuda.get_cuda_block_size(), 1024, @@ -156,7 +156,7 @@ def get_code(self): param_loc = c_array(dtype, dimp, "param_loc") fout = c_array(dtype, dimfout * blocksize_chunks, "fout") xi = c_array(dtype, dimx, "xi") - acc = c_array(dtypeacc, chunks.get_dimfinalchunk(), "acc") + acc = c_array(dtypeacc, reduction.get_dimfinalchunk(), "acc") yjloc = c_array(dtype, dimy, f"(yj + threadIdx.x * {dimy})") foutjrel = c_array(dtype, dimfout, f"({fout.id}+jrel*{dimfout})") yjrel = c_array(dtype, dimy, "yjrel") @@ -206,7 +206,7 @@ def get_code(self): dtype, fun_global, varfinal, - chunks.get_dimfinalchunk(), + reduction.get_dimfinalchunk(), acc, i, j, diff --git a/keopscore/keopscore/sandbox/formula.py b/keopscore/keopscore/sandbox/formula.py index 8ec57b364..2b26ea517 100644 --- a/keopscore/keopscore/sandbox/formula.py +++ b/keopscore/keopscore/sandbox/formula.py @@ -3,7 +3,7 @@ import keopscore.config from keopscore.formulas import * -keopscore.config.auto_factorize = True +keopscore.config.reduction.set_auto_factorize(True) print("********************************") print("test 1") diff --git a/keopscore/keopscore/sandbox/lazytensor_grad.py b/keopscore/keopscore/sandbox/lazytensor_grad.py index b1bbc0eb5..591972b22 100644 --- a/keopscore/keopscore/sandbox/lazytensor_grad.py +++ b/keopscore/keopscore/sandbox/lazytensor_grad.py @@ -8,7 +8,7 @@ import keopscore -keopscore.auto_factorize = False +keopscore.config.chunks.set_auto_factorize(False) M, N, D, DV = ( (100000, 100000, 3, 1) if torch.cuda.is_available() else (10000, 10000, 3, 1) diff --git a/keopscore/keopscore/utils/Cache.py b/keopscore/keopscore/utils/Cache.py index e5e7c44e8..e5a5c10be 100644 --- a/keopscore/keopscore/utils/Cache.py +++ b/keopscore/keopscore/utils/Cache.py @@ -9,7 +9,7 @@ env_param = ( lambda: keopscore.config.cxx.get_compile_options() + keopscore.config.cxx.get_linking_options() - + f" auto_factorize={keopscore.config.auto_factorize}" + + f" auto_factorize={keopscore.config.reduction.get_auto_factorize()}" + keopscore.config.cuda.get_preprocessing_options() + keopscore.config.cuda.get_include_options() + keopscore.config.cuda.get_linking_options() diff --git a/pykeops/pykeops/sandbox/test_torch_func_hessian.py b/pykeops/pykeops/sandbox/test_torch_func_hessian.py index 7e574e2dd..d2dd2209c 100644 --- a/pykeops/pykeops/sandbox/test_torch_func_hessian.py +++ b/pykeops/pykeops/sandbox/test_torch_func_hessian.py @@ -2,7 +2,7 @@ import torch from pykeops.torch import LazyTensor -keopscore.config.auto_factorize = False +keopscore.config.chunks.set_auto_factorize(False) def fn_torch(x_i): diff --git a/pykeops/pykeops/test/test_torch_func.py b/pykeops/pykeops/test/test_torch_func.py index 332b1968a..f50ff880a 100644 --- a/pykeops/pykeops/test/test_torch_func.py +++ b/pykeops/pykeops/test/test_torch_func.py @@ -5,7 +5,7 @@ torch.manual_seed(0) -keopscore.config.auto_factorize = False +keopscore.config.chunks.set_auto_factorize(False) B1, B2, M, N, D = 5, 4, 10, 20, 2 From d0c978dae6c2ac7679b711ee67d8bccf0b57a471 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 10 May 2026 18:00:07 +0200 Subject: [PATCH 40/98] Split cpp compilation : create entry point with pybind11. remove evals. Co-authored-by: Copilot --- keopscore/keopscore/binders/LinkCompile.py | 25 +- .../keopscore/binders/cpp/Cpu_link_compile.py | 86 +++++- .../binders/cpp/keops_cpu_runtime.cpp | 66 +++++ .../keopscore/binders/cpp/keops_cpu_runtime.h | 65 +++++ .../binders/nvrtc/Gpu_link_compile.py | 56 ++-- keopscore/keopscore/config/Cuda.py | 45 ++- keopscore/keopscore/config/CxxCompiler.py | 9 + keopscore/keopscore/config/KeOpsPath.py | 21 +- keopscore/keopscore/get_keops_dll.py | 4 - .../keopscore/sandbox/lazytensor_grad.py | 2 +- keopscore/keopscore/utils/Cache.py | 12 +- keopscore/keopscore/utils/math_functions.py | 2 +- pykeops/pykeops/__init__.py | 18 +- pykeops/pykeops/common/keops_io/LoadKeOps.py | 5 +- .../pykeops/common/keops_io/LoadKeOps_cpp.py | 223 -------------- pykeops/pykeops/common/keops_io/__init__.py | 5 +- .../common/keops_io/cpp/LoadKeOps_cpp.py | 106 +++++++ .../pykeops/common/keops_io/cpp/__init__.py | 0 .../common/keops_io/cpp/pykeops_cpp.cpp | 273 ++++++++++++++++++ .../keops_io/{ => nvrtc}/LoadKeOps_nvrtc.py | 19 +- .../pykeops/common/keops_io/nvrtc/__init__.py | 0 .../keops_io/{ => nvrtc}/pykeops_nvrtc.cpp | 0 pykeops/pykeops/config.py | 45 ++- .../sandbox/test_torch_func_hessian.py | 2 +- pykeops/pykeops/test/test_torch_func.py | 2 +- 25 files changed, 775 insertions(+), 316 deletions(-) create mode 100644 keopscore/keopscore/binders/cpp/keops_cpu_runtime.cpp create mode 100644 keopscore/keopscore/binders/cpp/keops_cpu_runtime.h delete mode 100644 pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py create mode 100644 pykeops/pykeops/common/keops_io/cpp/LoadKeOps_cpp.py create mode 100644 pykeops/pykeops/common/keops_io/cpp/__init__.py create mode 100644 pykeops/pykeops/common/keops_io/cpp/pykeops_cpp.cpp rename pykeops/pykeops/common/keops_io/{ => nvrtc}/LoadKeOps_nvrtc.py (84%) create mode 100644 pykeops/pykeops/common/keops_io/nvrtc/__init__.py rename pykeops/pykeops/common/keops_io/{ => nvrtc}/pykeops_nvrtc.cpp (100%) diff --git a/keopscore/keopscore/binders/LinkCompile.py b/keopscore/keopscore/binders/LinkCompile.py index 0f985526f..9781f7d8c 100644 --- a/keopscore/keopscore/binders/LinkCompile.py +++ b/keopscore/keopscore/binders/LinkCompile.py @@ -54,23 +54,21 @@ def save_info(self): # dimy (sum of dimensions of j-indexed vectors) f = open(self.info_file, "w") f.write( - f"red_formula={self.red_formula_string}\ndim={self.dim}\ntagI={self.tagI}\ndimy={self.dimy}" + f"red_formula={self.red_formula_string}\ndim={self.dim}\ntagI={self.tagI}\ndimy={self.dimy}\n\n" + + f"dtype={self.dtype}\nsum_scheme={self.sum_scheme_string} in {self.dtypeacc}\n" + + f"tagHostDevice={self.tagHostDevice}\ntagCpuGpu={self.tagCpuGpu}\ntag1D2D={self.tag1D2D}\n\ndevice_id={self.device_id}\n" ) f.close() def read_info(self): # read info_file to retreive dim, tagI, dimy f = open(self.info_file, "r") - string = f.read() + f.readline() # skip line 0 + tmp_dim = f.readline().rstrip("\n").split("=") + tmp_tag = f.readline().rstrip("\n").split("=") + tmp_dimy = f.readline().rstrip("\n").split("=") f.close() - tmp = string.split("\n") - if len(tmp) != 4: - KeOps_Error("Incorrect info file") - tmp_dim, tmp_tag, tmp_dimy = ( - tmp[1].split("="), - tmp[2].split("="), - tmp[3].split("="), - ) + if ( len(tmp_dim) != 2 or tmp_dim[0] != "dim" @@ -80,9 +78,9 @@ def read_info(self): or tmp_dimy[0] != "dimy" ): KeOps_Error("Incorrect info file") - self.dim = eval(tmp_dim[1]) - self.tagI = eval(tmp_tag[1]) - self.dimy = eval(tmp_dimy[1]) + self.dim = int(tmp_dim[1]) + self.tagI = int(tmp_tag[1]) + self.dimy = int(tmp_dimy[1]) def write_code(self): # write the generated code in the source file ; this is used as a subfunction of compile_code @@ -111,7 +109,6 @@ def get_dll_and_params(self): self.read_info() return dict( tag=self.gencode_filename, - source_file=self.true_dllname, low_level_code_file=self.low_level_code_file, tagI=self.tagI, use_half=self.use_half, diff --git a/keopscore/keopscore/binders/cpp/Cpu_link_compile.py b/keopscore/keopscore/binders/cpp/Cpu_link_compile.py index 1328f94e8..3d89a6223 100644 --- a/keopscore/keopscore/binders/cpp/Cpu_link_compile.py +++ b/keopscore/keopscore/binders/cpp/Cpu_link_compile.py @@ -1,4 +1,22 @@ +import os +import sysconfig + +import keopscore.config from keopscore.binders.LinkCompile import LinkCompile +from keopscore.utils.messages import KeOps_Error, KeOps_Message +from keopscore.utils.system_utils import KeOps_OS_Run + +cpp_dtype = { + "float": "float", + "double": "double", +} + +cpu_runtime_src = os.path.join( + keopscore.config.path.get_base_dir_path(), + "binders", + "cpp", + "keops_cpu_runtime.cpp", +) class Cpu_link_compile(LinkCompile): @@ -6,20 +24,78 @@ class Cpu_link_compile(LinkCompile): def __init__(self): LinkCompile.__init__(self) - # these are used for command line compiling mode - self.low_level_code_file = "".encode("utf-8") + self.dllname = os.path.join( + keopscore.config.path.get_build_folder(), + self.gencode_filename + "_cpp" + sysconfig.get_config_var("SHLIB_SUFFIX"), + ) + self.low_level_code_file = self.dllname - # actual dll to be called - self.true_dllname = self.gencode_file # file to check for existence to detect compilation is needed - self.file_to_check = self.gencode_file + self.file_to_check = self.dllname def generate_code(self): # method to generate the code and compile it # generate the code and save it in self.code, by calling get_code method from CpuReduc class : self.get_code() + self.code += self.get_cpu_dll_entrypoint_code() # write the code in the source file self.write_code() + self.compile_dll() # retreive some parameters that will be saved into info_file. self.tagI = self.red_formula.tagI self.dim = self.red_formula.dim + + def compile_dll(self): + compile_command = ( + f"{keopscore.config.cxx.get_cxx_compiler()} " + f"{keopscore.config.cxx.get_compile_options()} " + f"{keopscore.config.openmp.get_compile_options()} " + f"{keopscore.config.openmp.get_include_options()} " + f"{keopscore.config.path.get_include_options()} " + f"{self.gencode_file} " + f"{cpu_runtime_src} " + f"{keopscore.config.cxx.get_linking_options()} " + f"{keopscore.config.openmp.get_linking_options()} " + f"-o {self.dllname}" + ) + KeOps_Message( + "Compiling cpp formula " + self.gencode_filename + " module", + flush=True, + end="", + level=1, + ) + KeOps_Message( + " in cache folder " + keopscore.config.path.get_build_folder(), + flush=True, + end="", + level=2, + ) + KeOps_Message(" ... ", flush=True, end="", level=1) + + out = KeOps_OS_Run(compile_command) + if out.returncode != 0: + KeOps_Error( + f"Error compiling cpp formula {self.gencode_filename}. See compiler output above." + ) + else: + KeOps_Message("OK", use_tag=False, flush=True, level=1) + + def get_cpu_dll_entrypoint_code(self): + if self.dtype not in cpp_dtype: + KeOps_Error( + "The cpp backend only supports float and double formulas, got " + + self.dtype + + "." + ) + + scalar_type = cpp_dtype[self.dtype] + + return f""" + +#include "binders/cpp/keops_cpu_runtime.h" + +extern "C" KEOPS_CPU_EXPORT int keops_cpu_launch(const KeOpsCpuArgs *runtime_args) {{ + return keops_cpu_launch_impl< {scalar_type} >( + runtime_args, &launch_keops_cpu_{self.gencode_filename}< {scalar_type} >); +}} +""" diff --git a/keopscore/keopscore/binders/cpp/keops_cpu_runtime.cpp b/keopscore/keopscore/binders/cpp/keops_cpu_runtime.cpp new file mode 100644 index 000000000..ae7385813 --- /dev/null +++ b/keopscore/keopscore/binders/cpp/keops_cpu_runtime.cpp @@ -0,0 +1,66 @@ +#include "binders/cpp/keops_cpu_runtime.h" + +namespace { + +template std::vector keops_cpu_make_vector(const T *data, int size) { + if (data == nullptr || size <= 0) { + return std::vector(); + } + return std::vector(data, data + size); +} + +} // namespace + +template +int keops_cpu_launch_impl(const KeOpsCpuArgs *runtime_args, + KeOpsCpuFormulaLaunchFn formula_launcher) { + if (runtime_args == nullptr) { + return -1; + } + + std::vector indsi_v = + keops_cpu_make_vector(runtime_args->indsi, runtime_args->n_indsi); + std::vector indsj_v = + keops_cpu_make_vector(runtime_args->indsj, runtime_args->n_indsj); + std::vector indsp_v = + keops_cpu_make_vector(runtime_args->indsp, runtime_args->n_indsp); + + std::vector dimsx_v = + keops_cpu_make_vector(runtime_args->dimsx, runtime_args->n_dimsx); + std::vector dimsy_v = + keops_cpu_make_vector(runtime_args->dimsy, runtime_args->n_dimsy); + std::vector dimsp_v = + keops_cpu_make_vector(runtime_args->dimsp, runtime_args->n_dimsp); + std::vector shapeout_v = + keops_cpu_make_vector(runtime_args->shapeout, runtime_args->n_shapeout); + + std::vector arg_v(runtime_args->nargs); + for (int i = 0; i < runtime_args->nargs; ++i) { + arg_v[i] = reinterpret_cast(runtime_args->arg[i]); + } + + std::vector> argshape_v(runtime_args->nargs); + for (int i = 0; i < runtime_args->nargs; ++i) { + const signed long int *shape = runtime_args->argshape[i]; + int shape_size = runtime_args->argshape_sizes[i]; + if (shape != nullptr && shape_size > 0) { + argshape_v[i].assign(shape, shape + shape_size); + } + } + + return formula_launcher( + runtime_args->dimY, runtime_args->nx, runtime_args->ny, runtime_args->tagI, + runtime_args->tagZero, runtime_args->use_half, runtime_args->dimred, + runtime_args->use_chunk_mode, indsi_v, indsj_v, indsp_v, + runtime_args->dimout, dimsx_v, dimsy_v, dimsp_v, runtime_args->ranges, + shapeout_v, reinterpret_cast(runtime_args->out), arg_v.data(), + argshape_v); +} + +template int keops_cpu_launch_impl( + const KeOpsCpuArgs *runtime_args, + KeOpsCpuFormulaLaunchFn formula_launcher); + +template int keops_cpu_launch_impl( + const KeOpsCpuArgs *runtime_args, + KeOpsCpuFormulaLaunchFn formula_launcher); diff --git a/keopscore/keopscore/binders/cpp/keops_cpu_runtime.h b/keopscore/keopscore/binders/cpp/keops_cpu_runtime.h new file mode 100644 index 000000000..aec14e7a0 --- /dev/null +++ b/keopscore/keopscore/binders/cpp/keops_cpu_runtime.h @@ -0,0 +1,65 @@ +#pragma once + +#if defined(_WIN32) +#define KEOPS_CPU_EXPORT __declspec(dllexport) +#define KEOPS_CPU_INTERNAL +#else +#define KEOPS_CPU_EXPORT __attribute__((visibility("default"))) +#define KEOPS_CPU_INTERNAL __attribute__((visibility("hidden"))) +#endif + +#include + +struct KeOpsCpuArgs { + signed long int dimY; + signed long int nx; + signed long int ny; + int tagI; + int tagZero; + int use_half; + signed long int dimred; + int use_chunk_mode; + + int n_indsi; + const int *indsi; + int n_indsj; + const int *indsj; + int n_indsp; + const int *indsp; + + signed long int dimout; + int n_dimsx; + const signed long int *dimsx; + int n_dimsy; + const signed long int *dimsy; + int n_dimsp; + const signed long int *dimsp; + + signed long int **ranges; + + int n_shapeout; + const signed long int *shapeout; + void *out; + + int nargs; + void **arg; + const signed long int **argshape; + const int *argshape_sizes; +}; + +typedef int (*KeOpsCpuLaunchFn)(const KeOpsCpuArgs *); + +template +using KeOpsCpuFormulaLaunchFn = int (*)( + signed long int dimY, signed long int nx, signed long int ny, int tagI, + int tagZero, int use_half, signed long int dimred, int use_chunk_mode, + std::vector indsi, std::vector indsj, std::vector indsp, + signed long int dimout, std::vector dimsx, + std::vector dimsy, std::vector dimsp, + signed long int **ranges, std::vector shapeout, TYPE *out, + TYPE **arg, std::vector> argshape); + +template +KEOPS_CPU_INTERNAL int +keops_cpu_launch_impl(const KeOpsCpuArgs *runtime_args, + KeOpsCpuFormulaLaunchFn formula_launcher); diff --git a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py index 6022c20c2..e349e71d2 100644 --- a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py +++ b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py @@ -26,9 +26,6 @@ def jit_compile_dll(): class Gpu_link_compile(LinkCompile): source_code_extension = "cu" - low_level_code_prefix = ( - "cubin_" if keopscore.config.cuda.get_cuda_version() >= 11010 else "ptx_" - ) def __init__(self): # checking that the system has a Gpu : @@ -43,15 +40,14 @@ def __init__(self): LinkCompile.__init__(self) # these are used for JIT compiling mode # low_level_code_file is filename of low level code (PTX for Cuda) or binary (CUBIN for Cuda) - # generated by the JIT compiler, e.g. ptx_7b9a611f7e + # generated by the JIT compiler, e.g. 7b9a611f7e_nvrtc.{ptx,cubin} self.low_level_code_file = os.path.join( keopscore.config.path.get_build_folder(), - self.low_level_code_prefix + self.gencode_filename, + self.gencode_filename + "_nvrtc." + keopscore.config.cuda.get_ir_type(), ).encode("utf-8") self.my_c_dll = ctypes.CDLL(jit_compile_dll(), mode=os.RTLD_LAZY) - # actual dll to be called is the jit binary, TODO: check if this is relevent - self.true_dllname = keopscore.config.path.get_jit_binary() + # file to check for existence to detect compilation is needed self.file_to_check = self.low_level_code_file @@ -59,7 +55,7 @@ def generate_code(self): # method to generate the code and compile it # generate the code and save it in self.code, by calling get_code method from GpuReduc class : self.get_code() - # write the code in the source file + # write the code in the source .cu file self.write_code() # we execute the main dll, passing the code as argument, and the name of the low level code file to save the assembly instructions @@ -73,9 +69,9 @@ def generate_code(self): (custom_cuda_include_fp16_path() + os.path.sep).encode("utf-8") ), ) - if res != 0: + if res != keopscore.config.cuda.CUDA_SUCCESS: KeOps_Error( - f"Error when compiling formula (error in nvrtcCompileProgram, nvrtcResult={res})" + f"Error when compiling {self.low_level_code_file} formula (error in nvrtcCompileProgram, nvrtcResult={res})" ) # retreive some parameters that will be saved into info_file. self.tagI = self.red_formula.tagI @@ -90,24 +86,36 @@ def get_compile_command( # This is about the main KeOps binary (dll) that will be used to JIT compile all formulas. # If the dll is not present, it compiles it from source, except if check_compile is False. - target_tag = ( - "CUBIN" if Gpu_link_compile.low_level_code_prefix == "cubin_" else "PTX" - ) - nvrtcGetTARGET = "nvrtcGet" + target_tag - nvrtcGetTARGETSize = nvrtcGetTARGET + "Size" - arch_tag = ( - '\\"sm\\"' - if Gpu_link_compile.low_level_code_prefix == "cubin_" - else '\\"compute\\"' + return ( + f"{keopscore.config.cxx.get_cxx_compiler()} " + f"{keopscore.config.cuda.get_preprocessing_options()} " + f"{keopscore.config.cxx.get_compile_options()} " + f"{keopscore.config.cuda.get_include_options()} " + f"{keopscore.config.path.get_include_options()} " + f" {extra_flags} " + f"{sourcename} " + f"{keopscore.config.cxx.get_linking_options()} " + f"{keopscore.config.cuda.get_linking_options()} " + f"-o {dllname}" ) - target_type_define = f"-DnvrtcGetTARGET={nvrtcGetTARGET} -DnvrtcGetTARGETSize={nvrtcGetTARGETSize} -DARCHTAG={arch_tag}" - return f"{keopscore.config.cxx.get_cxx_compiler()} {keopscore.config.cuda.get_preprocessing_options()} {target_type_define} {keopscore.config.cxx.get_compile_options()} {keopscore.config.cuda.get_include_options()} {keopscore.config.path.get_include_options()} {extra_flags} {sourcename} {keopscore.config.cxx.get_linking_options()} {keopscore.config.cuda.get_linking_options()} -o {dllname}" @staticmethod def compile_jit_compile_dll(): - KeOps_Message("Compiling cuda jit compiler engine ... ", flush=True, end="") + KeOps_Message("Compiling cuda jit compiler engine", flush=True, end="", level=1) + KeOps_Message( + " in cache folder " + keopscore.config.path.get_build_folder(), + flush=True, + end="", + level=2, + ) + KeOps_Message(" ... ", flush=True, end="", level=1) command = Gpu_link_compile.get_compile_command( sourcename=jit_compile_src, dllname=jit_compile_dll() ) - KeOps_OS_Run(command) - KeOps_Message("OK", use_tag=False, flush=True) + out = KeOps_OS_Run(command) + if out.returncode != 0: + KeOps_Error( + f"Error compiling cuda/nvrtc binder {jit_compile_dll()}. See compiler output above." + ) + else: + KeOps_Message("OK", use_tag=False, flush=True, level=1) diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index 0d5bf168e..7bb92d603 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -26,6 +26,7 @@ class CudaConfig: # Cuda detection variables _use_cuda = False _cuda_version = -1 + _ir_type = "" # "ptx" or "cubin" _cuda_include_path = [""] # str or list of str _visible_devices = "" @@ -116,6 +117,7 @@ def __init__(self): # If cuda is enabled, then we finalize the rest of the config if self.get_use_cuda(): self.set_cuda_version() + self.set_ir_type() self.set_cuda_include_path() self.set_linking_options() self.set_cuda_block_size() @@ -187,7 +189,7 @@ def _find_and_load_libcuda(self): if libcuda.cuInit(0) != self.CUDA_SUCCESS: return ( False, - "libcuda was detected, but driver API could not be initialized. Rebooting the system may help. Switching to CPU only.", + "libcuda was detected, but driver API could not be initialized (flushing KeOps caches and/or rebooting the system may help). Switching to CPU only.", ) # If we successfully loaded libcuda and initialized it, store the handle in the config for potential future use @@ -296,6 +298,13 @@ def _cuda_libraries_available(self): This is also where we handle one single warning if needed. """ + # If CUDA_VISIBLE_DEVICES is explicitly set to empty, no GPU is requested. + if os.getenv("CUDA_VISIBLE_DEVICES") == "": + KeOps_Warning( + "CUDA_VISIBLE_DEVICES is set to empty, no GPU will be used. Switching to CPU only." + ) + return False + # Libcuda (driver) loaded globally so it is available to KeOps shared objects. success_cuda, err_cuda = self._find_and_load_libcuda() if not success_cuda: @@ -341,8 +350,16 @@ def print_cuda_block_size(self): # Visibles GPUs devices def set_visible_devices(self): """Set specific GPUs from CUDA_VISIBLE_DEVICES.""" - if os.getenv("CUDA_VISIBLE_DEVICES"): - self._visible_devices = os.getenv("CUDA_VISIBLE_DEVICES").replace(",", "_") + cuda_visible = os.getenv("CUDA_VISIBLE_DEVICES") + print( + f"CUDA_VISIBLE_DEVICES: {cuda_visible if cuda_visible is not None else not_found_str}" + ) + + if cuda_visible is not None: + self._visible_devices = ( + cuda_visible.replace(",", "_") if cuda_visible else "empty" + ) + print(f"Parsed visible devices: {self._visible_devices}") def get_visible_devices(self): """Get the specific GPUs.""" @@ -475,6 +492,18 @@ def print_cuda_include_path(self): f"CUDA Include Path: {':'.join(self.get_cuda_include_path()) or not_found_str}" ) + # IR type + + def set_ir_type(self): + """Set the IR type to be used for nvrtc compilation based on the CUDA version.""" + if self.get_cuda_version() >= 11010: + self._ir_type = "cubin" + else: + self._ir_type = "ptx" + + def get_ir_type(self): + return self._ir_type + # NVRTC include options def set_include_options(self): self._include_options += "".join( @@ -503,6 +532,16 @@ def set_preprocessing_options(self): f"-DSHAREDMEMPERBLOCK{d}={self.get_SharedMemPerBlock()[d]}" ) + target_tag = "CUBIN" if self.get_ir_type == "cubin" else "PTX" + nvrtcGetTARGET = "nvrtcGet" + target_tag + self.add_to_preprocessing_options(f"-DnvrtcGetTARGET={nvrtcGetTARGET}") + + nvrtcGetTARGETSize = nvrtcGetTARGET + "Size" + self.add_to_preprocessing_options(f"-DnvrtcGetTARGETSize={nvrtcGetTARGETSize}") + + arch_tag = '\\"sm\\"' if self.get_ir_type == "cubin" else '\\"compute\\"' + self.add_to_preprocessing_options(f"-DARCHTAG={arch_tag}") + def add_to_preprocessing_options(self, flags): self._preprocessing_options += " " + flags diff --git a/keopscore/keopscore/config/CxxCompiler.py b/keopscore/keopscore/config/CxxCompiler.py index 8ca512630..a157a0352 100644 --- a/keopscore/keopscore/config/CxxCompiler.py +++ b/keopscore/keopscore/config/CxxCompiler.py @@ -173,6 +173,15 @@ def get_linking_options(self): def print_linking_options(self): print(f"Linking Options: {self.get_linking_options()}") + def get_dynamic_loader_linking_options(self): + """Return extra linker options needed for dlopen/dlsym support.""" + if self.platform.get_platform() == "Linux": + return "-ldl" + elif self.platform.get_platform() == "Darwin": + # macOS: dlopen/dlsym are in standard C library, no extra flag needed + return "" + return "" + # C++ Environment Flags. Unused yet... def set_cxx_env_flags(self): """Recover the C++ environment flags.""" diff --git a/keopscore/keopscore/config/KeOpsPath.py b/keopscore/keopscore/config/KeOpsPath.py index 2c4c414d8..7c459cd81 100644 --- a/keopscore/keopscore/config/KeOpsPath.py +++ b/keopscore/keopscore/config/KeOpsPath.py @@ -1,5 +1,6 @@ import os import sys +import sysconfig import keopscore from keopscore.utils.path_utils import ensure_directory @@ -79,9 +80,10 @@ def set_default_build_folder_name(self): ] if self.cuda.get_use_cuda(): name_parts.append(f"CUDA{self.cuda.get_cuda_version()}") - visible_devices = self.cuda.get_visible_devices() - if visible_devices: - name_parts.append(f"VISIBLE_DEVICES{visible_devices}") + + visible_devices = self.cuda.get_visible_devices() + if visible_devices: + name_parts.append(f"VISIBLE_DEVICES{visible_devices}") self._default_build_folder_name = "_".join(name_parts) @@ -152,6 +154,19 @@ def set_different_build_folder( def get_build_folder(self): return self._build_folder + def get_build_folder_file(self, filename): + return os.path.join(self.get_build_folder(), filename) + + def get_python_extension_path(self, basename): + return self.get_build_folder_file( + basename + sysconfig.get_config_var("EXT_SUFFIX") + ) + + def target_needs_update(self, target, source): + return not os.path.exists(target) or os.path.getmtime( + source + ) > os.path.getmtime(target) + # Default Build Path def set_default_build_path(self): """Set the default build path.""" diff --git a/keopscore/keopscore/get_keops_dll.py b/keopscore/keopscore/get_keops_dll.py index 2f303292d..595f147bf 100644 --- a/keopscore/keopscore/get_keops_dll.py +++ b/keopscore/keopscore/get_keops_dll.py @@ -19,9 +19,6 @@ It returns - tag : string, hash code used as id for the input formula and parameters - - source_file : string, either : - - in CPU mode : name of the source file to be compiled - - in GPU mode : name of the main dll to be called - low_level_code_file : string, either : - in CPU mode : the empty string "" - in GPU mode : name of the low level code (ptx) or binary file (cubin) to be passed to the main dll. @@ -127,7 +124,6 @@ def get_keops_dll_impl( return ( res["tag"], - res["source_file"], res["low_level_code_file"], res["tagI"], tagZero, diff --git a/keopscore/keopscore/sandbox/lazytensor_grad.py b/keopscore/keopscore/sandbox/lazytensor_grad.py index 591972b22..798fac664 100644 --- a/keopscore/keopscore/sandbox/lazytensor_grad.py +++ b/keopscore/keopscore/sandbox/lazytensor_grad.py @@ -8,7 +8,7 @@ import keopscore -keopscore.config.chunks.set_auto_factorize(False) +keopscore.config.reduction.set_auto_factorize(False) M, N, D, DV = ( (100000, 100000, 3, 1) if torch.cuda.is_available() else (10000, 10000, 3, 1) diff --git a/keopscore/keopscore/utils/Cache.py b/keopscore/keopscore/utils/Cache.py index e5a5c10be..151235fdd 100644 --- a/keopscore/keopscore/utils/Cache.py +++ b/keopscore/keopscore/utils/Cache.py @@ -21,6 +21,7 @@ def __init__(self, fun, use_cache_file=False, save_folder="."): self.fun = fun self.library = {} self.use_cache_file = use_cache_file + self.save_folder = save_folder if use_cache_file: self.cache_file = os.path.join(save_folder, fun.__name__ + "_cache.pkl") if os.path.isfile(self.cache_file) and os.path.getsize(self.cache_file) > 0: @@ -41,6 +42,10 @@ def reset(self, new_save_folder=None): self.library = {} if new_save_folder: self.save_folder = new_save_folder + if self.use_cache_file: + self.cache_file = os.path.join( + self.save_folder, self.fun.__name__ + "_cache.pkl" + ) def save_cache(self): f = open(self.cache_file, "wb") @@ -73,6 +78,7 @@ def __init__(self, cls, use_cache_file=False, save_folder="."): self.cls = cls self.library = {} self.use_cache_file = use_cache_file + self.save_folder = save_folder if self.use_cache_file: self.cache_file = os.path.join(save_folder, cls.__name__ + "_cache.pkl") if os.path.isfile(self.cache_file): @@ -86,7 +92,7 @@ def __init__(self, cls, use_cache_file=False, save_folder="."): atexit.register(self.save_cache) def __call__(self, *args): - str_id = "".join(list(str(arg) for arg in args)) + str(env_param) + str_id = "".join(list(str(arg) for arg in args)) + str(env_param()) if not str_id in self.library: if self.use_cache_file: if str_id in self.library_params: @@ -106,6 +112,10 @@ def reset(self, new_save_folder=None): self.library_params = {} if new_save_folder: self.save_folder = new_save_folder + if self.use_cache_file: + self.cache_file = os.path.join( + self.save_folder, self.cls.__name__ + "_cache.pkl" + ) def save_cache(self): f = open(self.cache_file, "wb") diff --git a/keopscore/keopscore/utils/math_functions.py b/keopscore/keopscore/utils/math_functions.py index b93aa2b1a..48de80ba6 100644 --- a/keopscore/keopscore/utils/math_functions.py +++ b/keopscore/keopscore/utils/math_functions.py @@ -160,7 +160,7 @@ def call(*args): ) keops_round = math_function( cpu_code=lambda x, d: ( - f"round({x})" if eval(d) == 0 else f"(round({x}*{10**eval(d)})/{10**eval(d)})" + f"round({x})" if int(d) == 0 else f"(round({x}*{10**int(d)})/{10**int(d)})" ), gpu_half2_code="NA", ) diff --git a/pykeops/pykeops/__init__.py b/pykeops/pykeops/__init__.py index e6331f17e..796106f95 100644 --- a/pykeops/pykeops/__init__.py +++ b/pykeops/pykeops/__init__.py @@ -26,9 +26,17 @@ def set_verbose(val): default_device_id = 0 # default Gpu device number +from .common.keops_io.cpp.LoadKeOps_cpp import ( + compile_jit_binary as compile_cpp_jit_binary, + should_compile_binder as should_compile_cpp_binder, +) + +if should_compile_cpp_binder(): + compile_cpp_jit_binary() + if pykeopsconfig.cuda.get_use_cuda(): if not os.path.exists(pykeopsconfig.pykeops_nvrtc_name(type="target")): - from .common.keops_io.LoadKeOps_nvrtc import compile_jit_binary + from .common.keops_io.nvrtc.LoadKeOps_nvrtc import compile_jit_binary compile_jit_binary() @@ -46,8 +54,10 @@ def clean_pykeops(recompile_jit_binaries=True): keops_binder = pykeops.common.keops_io.keops_binder for key in keops_binder: keops_binder[key].reset() - if recompile_jit_binaries and pykeopsconfig.cuda.get_use_cuda(): - pykeops.common.keops_io.LoadKeOps_nvrtc.compile_jit_binary() + if recompile_jit_binaries: + pykeops.common.keops_io.LoadKeOps_cpp.compile_jit_binary() + if pykeopsconfig.cuda.get_use_cuda(): + pykeops.common.keops_io.LoadKeOps_nvrtc.compile_jit_binary() def check_health(infos="all"): @@ -74,6 +84,8 @@ def set_build_folder(path=None): keops_binder = pykeops.common.keops_io.keops_binder for key in keops_binder: keops_binder[key].reset(new_save_folder=get_build_folder()) + if should_compile_cpp_binder(): + pykeops.common.keops_io.LoadKeOps_cpp.compile_jit_binary() if pykeopsconfig.cuda.get_use_cuda() and not os.path.exists( pykeopsconfig.pykeops_nvrtc_name(type="target") ): diff --git a/pykeops/pykeops/common/keops_io/LoadKeOps.py b/pykeops/pykeops/common/keops_io/LoadKeOps.py index 57e00f6b5..26786b94a 100644 --- a/pykeops/pykeops/common/keops_io/LoadKeOps.py +++ b/pykeops/pykeops/common/keops_io/LoadKeOps.py @@ -57,9 +57,9 @@ def init( cat = 2 alias_args = var[3:-1].split(",") if len(alias_args) == 1: - ind, dim = k, eval(alias_args[0]) + ind, dim = k, int(alias_args[0]) elif len(alias_args) == 2: - ind, dim = eval(alias_args[0]), eval(alias_args[1]) + ind, dim = int(alias_args[0]), int(alias_args[1]) alias = f"{varname}=Var({ind},{dim},{cat})" aliases_new.append(alias) @@ -106,7 +106,6 @@ def init( ( self.params.tag, - self.params.source_name, self.params.low_level_code_file, self.params.tagI, self.params.tagZero, diff --git a/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py b/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py deleted file mode 100644 index bc2a95f52..000000000 --- a/pykeops/pykeops/common/keops_io/LoadKeOps_cpp.py +++ /dev/null @@ -1,223 +0,0 @@ -import os -import sysconfig - -import pykeops.config as pykeopsconfig - -from keopscore.utils.Cache import Cache_partial -from keopscore.utils.system_utils import KeOps_OS_Run - -from pykeops.common.keops_io.LoadKeOps import LoadKeOps -from pykeops.common.utils import pyKeOps_Message - - -class LoadKeOps_cpp_class(LoadKeOps): - def __init__(self, *args, fast_init=False): - super().__init__(*args, fast_init=fast_init) - - def init_phase1(self): - srcname = pykeopsconfig.pykeops_cpp_name(tag=self.params.tag, extension=".cpp") - - dllname = pykeopsconfig.pykeops_cpp_name( - tag=self.params.tag, extension=sysconfig.get_config_var("EXT_SUFFIX") - ) - - if not os.path.exists(dllname): - f = open(srcname, "w") - f.write(self.get_pybind11_code()) - f.close() - compile_command = f"{pykeopsconfig.cxx.get_cxx_compiler()} {pykeopsconfig.cxx.get_compile_options()} {pykeopsconfig.openmp.get_compile_options()} {pykeopsconfig.openmp.get_include_options()} {pykeopsconfig.path.get_include_options()} {pykeopsconfig.python_includes} {srcname} {pykeopsconfig.cxx.get_linking_options()} {pykeopsconfig.openmp.get_linking_options()} -o {dllname}" - pyKeOps_Message( - "Compiling pykeops cpp " + self.params.tag + " module ... ", - flush=True, - end="", - ) - KeOps_OS_Run(compile_command) - pyKeOps_Message("OK", use_tag=False, flush=True) - - def init_phase2(self): - import importlib - - mylib = importlib.import_module( - os.path.basename(pykeopsconfig.pykeops_cpp_name(tag=self.params.tag)) - ) - - self.launch_keops_cpu = mylib.launch_pykeops_cpu - - def call_keops(self, nx, ny): - self.launch_keops_cpu( - self.params.dimy, - nx, - ny, - self.params.tagI, - self.params.tagZero, - self.params.use_half, - self.params.dimred, - self.params.use_chunk_mode, - self.params.indsi, - self.params.indsj, - self.params.indsp, - self.params.dim, - self.params.dimsx, - self.params.dimsy, - self.params.dimsp, - self.ranges_ptr_new, - self.outshape, - self.out_ptr, - self.args_ptr_new, - self.argshapes_new, - ) - - def get_pybind11_code(self): - return f""" -#include "{self.params.source_name}" - -#include -#include -namespace py = pybind11; - -template < typename TYPE > -int launch_pykeops_{self.params.tag}_cpu(signed long int dimY, signed long int nx, signed long int ny, - int tagI, int tagZero, int use_half, - signed long int dimred, - int use_chunk_mode, - py::tuple py_indsi, py::tuple py_indsj, py::tuple py_indsp, - signed long int dimout, - py::tuple py_dimsx, py::tuple py_dimsy, py::tuple py_dimsp, - py::tuple py_ranges, - py::tuple py_shapeout, - long out_void, - py::tuple py_arg, - py::tuple py_argshape){{ - - /*------------------------------------*/ - /* Cast input args */ - /*------------------------------------*/ - - std::vector< int > indsi_v(py_indsi.size()); - for (auto i = 0; i < py_indsi.size(); i++) - indsi_v[i] = py::cast< int >(py_indsi[i]); - - - std::vector< int > indsj_v(py_indsj.size()); - for (auto i = 0; i < py_indsj.size(); i++) - indsj_v[i] = py::cast< int >(py_indsj[i]); - - - std::vector< int > indsp_v(py_indsp.size()); - for (auto i = 0; i < py_indsp.size(); i++) - indsp_v[i] = py::cast< int >(py_indsp[i]); - - - std::vector< signed long int > dimsx_v(py_dimsx.size()); - for (auto i = 0; i < py_dimsx.size(); i++) - dimsx_v[i] = py::cast< signed long int >(py_dimsx[i]); - - - std::vector< signed long int > dimsy_v(py_dimsy.size()); - for (auto i = 0; i < py_dimsy.size(); i++) - dimsy_v[i] = py::cast< signed long int >(py_dimsy[i]); - - - std::vector< signed long int > dimsp_v(py_dimsp.size()); - for (auto i = 0; i < py_dimsp.size(); i++) - dimsp_v[i] = py::cast< signed long int >(py_dimsp[i]); - - - // Cast the ranges arrays - std::vector< signed long int* > ranges_v(py_ranges.size()); - for (signed long int i = 0; i < py_ranges.size(); i++) - ranges_v[i] = (signed long int*) py::cast< signed long int >(py_ranges[i]); - signed long int **ranges = (signed long int**) ranges_v.data(); - - std::vector< signed long int > shapeout_v(py_shapeout.size()); - for (auto i = 0; i < py_shapeout.size(); i++) - shapeout_v[i] = py::cast< signed long int >(py_shapeout[i]); - - TYPE *out = (TYPE*) out_void; - // std::cout << "out_ptr : " << (long) out << std::endl; - - std::vector< TYPE* > arg_v(py_arg.size()); - for (int i = 0; i < py_arg.size(); i++) - arg_v[i] = (TYPE*) py::cast< long >(py_arg[i]); - TYPE **arg = (TYPE**) arg_v.data(); - - std::vector< std::vector< signed long int > > argshape_v(py_argshape.size()); - for (auto i = 0; i < py_argshape.size(); i++){{ - py::tuple tmp = py_argshape[i]; - std::vector< signed long int > tmp_v(tmp.size()); - for (auto j =0; j < tmp.size(); j++) - tmp_v[j] = py::cast< signed long int >(tmp[j]); - argshape_v[i] = tmp_v; - }} - - - /*------------------------------------*/ - /* Shape validation */ - /*------------------------------------*/ - for (int k = 0; k < (int)indsi_v.size(); k++) {{ - int idx = indsi_v[k]; - const auto& shape = argshape_v[idx]; - if ((int)shape.size() < 2) - throw std::invalid_argument("[pyKeOps] Error: Vi argument #" + std::to_string(idx) + " requires at least 2 dimensions, got " + std::to_string(shape.size()) + "."); - if (shape.back() != dimsx_v[k]) - throw std::invalid_argument("[pyKeOps] Error: Vi argument #" + std::to_string(idx) + " has trailing dim " + std::to_string(shape.back()) + ", expected " + std::to_string(dimsx_v[k]) + "."); - }} - for (int k = 0; k < (int)indsj_v.size(); k++) {{ - int idx = indsj_v[k]; - const auto& shape = argshape_v[idx]; - if ((int)shape.size() < 2) - throw std::invalid_argument("[pyKeOps] Error: Vj argument #" + std::to_string(idx) + " requires at least 2 dimensions, got " + std::to_string(shape.size()) + "."); - if (shape.back() != dimsy_v[k]) - throw std::invalid_argument("[pyKeOps] Error: Vj argument #" + std::to_string(idx) + " has trailing dim " + std::to_string(shape.back()) + ", expected " + std::to_string(dimsy_v[k]) + "."); - }} - for (int k = 0; k < (int)indsp_v.size(); k++) {{ - int idx = indsp_v[k]; - const auto& shape = argshape_v[idx]; - if ((int)shape.size() < 1) - throw std::invalid_argument("[pyKeOps] Error: Pm argument #" + std::to_string(idx) + " requires at least 1 dimension (got a scalar)."); - if (shape.back() != dimsp_v[k]) - throw std::invalid_argument("[pyKeOps] Error: Pm argument #" + std::to_string(idx) + " has trailing dim " + std::to_string(shape.back()) + ", expected " + std::to_string(dimsp_v[k]) + "."); - }} - - return launch_keops_cpu_{self.params.tag}< TYPE >(dimY, - nx, - ny, - tagI, - tagZero, - use_half, - dimred, - use_chunk_mode, - indsi_v, - indsj_v, - indsp_v, - dimout, - dimsx_v, - dimsy_v, - dimsp_v, - ranges, - shapeout_v, - out, - arg, - argshape_v); - -}} - -PYBIND11_MODULE(pykeops_cpp_{self.params.tag}, m) {{ - m.doc() = "pyKeOps: KeOps for pytorch through pybind11 (pytorch flavour)."; - m.def("launch_pykeops_cpu", &launch_pykeops_{self.params.tag}_cpu < {cpp_dtype[self.params.dtype]} >, "Entry point to keops."); -}} - """ - - -LoadKeOps_cpp = Cache_partial( - LoadKeOps_cpp_class, - use_cache_file=True, - save_folder=pykeopsconfig.path.get_build_folder(), -) - -cpp_dtype = { - "float": "float", - "float32": "float", - "double": "double", - "float64": "double", -} diff --git a/pykeops/pykeops/common/keops_io/__init__.py b/pykeops/pykeops/common/keops_io/__init__.py index 15271bf6e..658775c37 100644 --- a/pykeops/pykeops/common/keops_io/__init__.py +++ b/pykeops/pykeops/common/keops_io/__init__.py @@ -1,13 +1,14 @@ import pykeops.config as pykeopsconfig +from .nvrtc import LoadKeOps_nvrtc if pykeopsconfig.cuda.get_use_cuda(): - from . import LoadKeOps_nvrtc, LoadKeOps_cpp + from .cpp import LoadKeOps_cpp keops_binder = { "nvrtc": LoadKeOps_nvrtc.LoadKeOps_nvrtc, "cpp": LoadKeOps_cpp.LoadKeOps_cpp, } else: - from . import LoadKeOps_cpp + from .cpp import LoadKeOps_cpp keops_binder = {"cpp": LoadKeOps_cpp.LoadKeOps_cpp} diff --git a/pykeops/pykeops/common/keops_io/cpp/LoadKeOps_cpp.py b/pykeops/pykeops/common/keops_io/cpp/LoadKeOps_cpp.py new file mode 100644 index 000000000..fbbd84da6 --- /dev/null +++ b/pykeops/pykeops/common/keops_io/cpp/LoadKeOps_cpp.py @@ -0,0 +1,106 @@ +import pykeops.config as pykeopsconfig + +from keopscore.utils.Cache import Cache_partial +from keopscore.utils.system_utils import KeOps_OS_Run +from keopscore.utils.messages import KeOps_Error + +from pykeops.common.keops_io.LoadKeOps import LoadKeOps +from pykeops.common.utils import pyKeOps_Message + + +class LoadKeOps_cpp_class(LoadKeOps): + def __init__(self, *args, fast_init=False): + super().__init__(*args, fast_init=fast_init) + + def init_phase2(self): + import importlib + + if should_compile_binder(): + compile_jit_binary() + + pykeops_cpp = importlib.import_module("pykeops_cpp") + formula_library = self.params.low_level_code_file + if not formula_library: + raise RuntimeError( + "[pyKeOps] Error: empty CPU formula library path. " + "Please clean the KeOps cache and recompile formulas." + ) + + if self.params.c_dtype == "float": + self.launch_keops_cpu = pykeops_cpp.KeOps_module_float(formula_library) + elif self.params.c_dtype == "double": + self.launch_keops_cpu = pykeops_cpp.KeOps_module_double(formula_library) + else: + raise ValueError( + "The cpp backend only supports float32 and float64 inputs." + ) + + def call_keops(self, nx, ny): + self.launch_keops_cpu( + self.params.dimy, + nx, + ny, + self.params.tagI, + self.params.tagZero, + self.params.use_half, + self.params.dimred, + self.params.use_chunk_mode, + self.params.indsi, + self.params.indsj, + self.params.indsp, + self.params.dim, + self.params.dimsx, + self.params.dimsy, + self.params.dimsp, + self.ranges_ptr_new, + self.outshape, + self.out_ptr, + self.args_ptr_new, + self.argshapes_new, + ) + + +def should_compile_binder(): + return pykeopsconfig.path.target_needs_update( + pykeopsconfig.pykeops_cpp_binder_name(type="target"), + pykeopsconfig.pykeops_cpp_binder_name(type="src"), + ) + + +def compile_jit_binary(): + """ + Compile the reusable pybind11 entry point for the cpp backend. + """ + compile_command = ( + f"{pykeopsconfig.cxx.get_cxx_compiler()} " + f"{pykeopsconfig.cxx.get_compile_options()} " + f"{pykeopsconfig.path.get_include_options()} " + f"{pykeopsconfig.python_includes} " + f"{pykeopsconfig.pykeops_cpp_binder_name(type='src')} " + f"{pykeopsconfig.cxx.get_linking_options()} " + f"{pykeopsconfig.cxx.get_dynamic_loader_linking_options()} " + f"-o {pykeopsconfig.pykeops_cpp_binder_name(type='target')}" + ) + pyKeOps_Message("Compiling cpp binder for python", flush=True, end="", level=1) + pyKeOps_Message( + " in cache folder " + pykeopsconfig.path.get_build_folder(), + flush=True, + end="", + level=2, + ) + pyKeOps_Message(" ... ", flush=True, end="", level=1) + + out = KeOps_OS_Run(compile_command) + if out.returncode != 0: + KeOps_Error( + f"Error compiling cpp binder {pykeopsconfig.pykeops_cpp_binder_name(type='target')} for python. See compiler output above." + ) + else: + pyKeOps_Message("OK", use_tag=False, flush=True, level=1) + + +LoadKeOps_cpp = Cache_partial( + LoadKeOps_cpp_class, + use_cache_file=True, + save_folder=pykeopsconfig.path.get_build_folder(), +) diff --git a/pykeops/pykeops/common/keops_io/cpp/__init__.py b/pykeops/pykeops/common/keops_io/cpp/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/pykeops/pykeops/common/keops_io/cpp/pykeops_cpp.cpp b/pykeops/pykeops/common/keops_io/cpp/pykeops_cpp.cpp new file mode 100644 index 000000000..e23d704bd --- /dev/null +++ b/pykeops/pykeops/common/keops_io/cpp/pykeops_cpp.cpp @@ -0,0 +1,273 @@ +#include "binders/cpp/keops_cpu_runtime.h" + +#include + +#include +#include +#include + +#if defined(_WIN32) +#include +#else +#include +#endif + +namespace py = pybind11; + +namespace pykeops_cpp_backend { + +class DynamicLibrary { +public: + explicit DynamicLibrary(const char *path) : path_(path), handle_(nullptr) { +#if defined(_WIN32) + handle_ = LoadLibraryA(path_.c_str()); + if (handle_ == nullptr) { + throw std::runtime_error("[pyKeOps] Error: could not load CPU formula library " + + path_ + ": " + last_error()); + } +#else + dlerror(); + handle_ = dlopen(path_.c_str(), RTLD_NOW | RTLD_LOCAL); + if (handle_ == nullptr) { + const char *error = dlerror(); + throw std::runtime_error("[pyKeOps] Error: could not load CPU formula library " + + path_ + ": " + (error ? error : "unknown error")); + } +#endif + } + + ~DynamicLibrary() { + if (handle_ != nullptr) { +#if defined(_WIN32) + FreeLibrary(static_cast(handle_)); +#else + dlclose(handle_); +#endif + } + } + + DynamicLibrary(const DynamicLibrary &) = delete; + DynamicLibrary &operator=(const DynamicLibrary &) = delete; + + void *symbol(const char *name) { +#if defined(_WIN32) + void *sym = reinterpret_cast( + GetProcAddress(static_cast(handle_), name)); + if (sym == nullptr) { + throw std::runtime_error("[pyKeOps] Error: could not find symbol " + + std::string(name) + " in " + path_ + ": " + + last_error()); + } + return sym; +#else + dlerror(); + void *sym = dlsym(handle_, name); + const char *error = dlerror(); + if (error != nullptr) { + throw std::runtime_error("[pyKeOps] Error: could not find symbol " + + std::string(name) + " in " + path_ + ": " + + error); + } + return sym; +#endif + } + +private: +#if defined(_WIN32) + static std::string last_error() { + DWORD error = GetLastError(); + if (error == 0) { + return "unknown error"; + } + + LPSTR buffer = nullptr; + size_t size = FormatMessageA( + FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | + FORMAT_MESSAGE_IGNORE_INSERTS, + nullptr, error, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), + reinterpret_cast(&buffer), 0, nullptr); + + std::string message(buffer, size); + LocalFree(buffer); + return message; + } +#endif + + std::string path_; + void *handle_; +}; + +template class KeOps_module_python { +public: + explicit KeOps_module_python(const char *formula_library_name) + : library_(formula_library_name), launch_(nullptr) { + launch_ = reinterpret_cast( + library_.symbol("keops_cpu_launch")); + } + + int operator()(signed long int dimY, signed long int nx, signed long int ny, + int tagI, int tagZero, int use_half, + signed long int dimred, int use_chunk_mode, + py::tuple py_indsi, py::tuple py_indsj, py::tuple py_indsp, + signed long int dimout, py::tuple py_dimsx, + py::tuple py_dimsy, py::tuple py_dimsp, py::tuple py_ranges, + py::tuple py_shapeout, long out_void, py::tuple py_arg, + py::tuple py_argshape) { + + std::vector indsi_v(py_indsi.size()); + for (py::ssize_t i = 0; i < py_indsi.size(); i++) + indsi_v[i] = py::cast(py_indsi[i]); + + std::vector indsj_v(py_indsj.size()); + for (py::ssize_t i = 0; i < py_indsj.size(); i++) + indsj_v[i] = py::cast(py_indsj[i]); + + std::vector indsp_v(py_indsp.size()); + for (py::ssize_t i = 0; i < py_indsp.size(); i++) + indsp_v[i] = py::cast(py_indsp[i]); + + std::vector dimsx_v(py_dimsx.size()); + for (py::ssize_t i = 0; i < py_dimsx.size(); i++) + dimsx_v[i] = py::cast(py_dimsx[i]); + + std::vector dimsy_v(py_dimsy.size()); + for (py::ssize_t i = 0; i < py_dimsy.size(); i++) + dimsy_v[i] = py::cast(py_dimsy[i]); + + std::vector dimsp_v(py_dimsp.size()); + for (py::ssize_t i = 0; i < py_dimsp.size(); i++) + dimsp_v[i] = py::cast(py_dimsp[i]); + + std::vector ranges_v(py_ranges.size()); + for (py::ssize_t i = 0; i < py_ranges.size(); i++) + ranges_v[i] = + reinterpret_cast(py::cast(py_ranges[i])); + + std::vector shapeout_v(py_shapeout.size()); + for (py::ssize_t i = 0; i < py_shapeout.size(); i++) + shapeout_v[i] = py::cast(py_shapeout[i]); + + TYPE *out = reinterpret_cast(out_void); + + std::vector typed_arg_v(py_arg.size()); + std::vector arg_v(py_arg.size()); + for (py::ssize_t i = 0; i < py_arg.size(); i++) { + typed_arg_v[i] = reinterpret_cast(py::cast(py_arg[i])); + arg_v[i] = typed_arg_v[i]; + } + + std::vector> argshape_v(py_argshape.size()); + for (py::ssize_t i = 0; i < py_argshape.size(); i++) { + py::tuple tmp = py_argshape[i]; + std::vector tmp_v(tmp.size()); + for (py::ssize_t j = 0; j < tmp.size(); j++) + tmp_v[j] = py::cast(tmp[j]); + argshape_v[i] = tmp_v; + } + + for (int k = 0; k < (int)indsi_v.size(); k++) { + int idx = indsi_v[k]; + const auto &shape = argshape_v[idx]; + if ((int)shape.size() < 2) + throw std::invalid_argument( + "[pyKeOps] Error: Vi argument #" + std::to_string(idx) + + " requires at least 2 dimensions, got " + + std::to_string(shape.size()) + "."); + if (shape.back() != dimsx_v[k]) + throw std::invalid_argument( + "[pyKeOps] Error: Vi argument #" + std::to_string(idx) + + " has trailing dim " + std::to_string(shape.back()) + + ", expected " + std::to_string(dimsx_v[k]) + "."); + } + for (int k = 0; k < (int)indsj_v.size(); k++) { + int idx = indsj_v[k]; + const auto &shape = argshape_v[idx]; + if ((int)shape.size() < 2) + throw std::invalid_argument( + "[pyKeOps] Error: Vj argument #" + std::to_string(idx) + + " requires at least 2 dimensions, got " + + std::to_string(shape.size()) + "."); + if (shape.back() != dimsy_v[k]) + throw std::invalid_argument( + "[pyKeOps] Error: Vj argument #" + std::to_string(idx) + + " has trailing dim " + std::to_string(shape.back()) + + ", expected " + std::to_string(dimsy_v[k]) + "."); + } + for (int k = 0; k < (int)indsp_v.size(); k++) { + int idx = indsp_v[k]; + const auto &shape = argshape_v[idx]; + if ((int)shape.size() < 1) + throw std::invalid_argument( + "[pyKeOps] Error: Pm argument #" + std::to_string(idx) + + " requires at least 1 dimension (got a scalar)."); + if (shape.back() != dimsp_v[k]) + throw std::invalid_argument( + "[pyKeOps] Error: Pm argument #" + std::to_string(idx) + + " has trailing dim " + std::to_string(shape.back()) + + ", expected " + std::to_string(dimsp_v[k]) + "."); + } + + std::vector argshape_ptr_v(argshape_v.size()); + std::vector argshape_size_v(argshape_v.size()); + for (std::size_t i = 0; i < argshape_v.size(); ++i) { + argshape_ptr_v[i] = argshape_v[i].data(); + argshape_size_v[i] = static_cast(argshape_v[i].size()); + } + + KeOpsCpuArgs runtime_args; + runtime_args.dimY = dimY; + runtime_args.nx = nx; + runtime_args.ny = ny; + runtime_args.tagI = tagI; + runtime_args.tagZero = tagZero; + runtime_args.use_half = use_half; + runtime_args.dimred = dimred; + runtime_args.use_chunk_mode = use_chunk_mode; + runtime_args.n_indsi = static_cast(indsi_v.size()); + runtime_args.indsi = indsi_v.data(); + runtime_args.n_indsj = static_cast(indsj_v.size()); + runtime_args.indsj = indsj_v.data(); + runtime_args.n_indsp = static_cast(indsp_v.size()); + runtime_args.indsp = indsp_v.data(); + runtime_args.dimout = dimout; + runtime_args.n_dimsx = static_cast(dimsx_v.size()); + runtime_args.dimsx = dimsx_v.data(); + runtime_args.n_dimsy = static_cast(dimsy_v.size()); + runtime_args.dimsy = dimsy_v.data(); + runtime_args.n_dimsp = static_cast(dimsp_v.size()); + runtime_args.dimsp = dimsp_v.data(); + runtime_args.ranges = ranges_v.data(); + runtime_args.n_shapeout = static_cast(shapeout_v.size()); + runtime_args.shapeout = shapeout_v.data(); + runtime_args.out = out; + runtime_args.nargs = static_cast(arg_v.size()); + runtime_args.arg = arg_v.data(); + runtime_args.argshape = argshape_ptr_v.data(); + runtime_args.argshape_sizes = argshape_size_v.data(); + + py::gil_scoped_release release; + return launch_(&runtime_args); + } + +private: + DynamicLibrary library_; + KeOpsCpuLaunchFn launch_; +}; + +} // namespace pykeops_cpp_backend + +PYBIND11_MODULE(pykeops_cpp, m) { + m.doc() = "pyKeOps: reusable CPU frontend through pybind11."; + + py::class_>( + m, "KeOps_module_float", py::module_local()) + .def(py::init()) + .def("__call__", + &pykeops_cpp_backend::KeOps_module_python::operator()); + + py::class_>( + m, "KeOps_module_double", py::module_local()) + .def(py::init()) + .def("__call__", + &pykeops_cpp_backend::KeOps_module_python::operator()); +} diff --git a/pykeops/pykeops/common/keops_io/LoadKeOps_nvrtc.py b/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py similarity index 84% rename from pykeops/pykeops/common/keops_io/LoadKeOps_nvrtc.py rename to pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py index d51187637..2806fce99 100644 --- a/pykeops/pykeops/common/keops_io/LoadKeOps_nvrtc.py +++ b/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py @@ -1,4 +1,3 @@ -import os import sys import pykeops.config as pykeopsconfig @@ -6,6 +5,7 @@ from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile from keopscore.utils.Cache import Cache_partial from keopscore.utils.system_utils import KeOps_OS_Run +from keopscore.utils.messages import KeOps_Error from pykeops.common.keops_io.LoadKeOps import LoadKeOps from pykeops.common.utils import pyKeOps_Message @@ -85,7 +85,22 @@ def compile_jit_binary(): dllname=pykeopsconfig.pykeops_nvrtc_name(type="target"), ) pyKeOps_Message("Compiling nvrtc binder for python ... ", flush=True, end="") - KeOps_OS_Run(compile_command) + pyKeOps_Message( + " in cache folder " + pykeopsconfig.path.get_build_folder(), + flush=True, + end="", + level=2, + ) + pyKeOps_Message(" ... ", flush=True, end="", level=1) + + out = KeOps_OS_Run(compile_command) + if out.returncode != 0: + KeOps_Error( + f"Error compiling nvrtc binder {pykeopsconfig.pykeops_nvrtc_name(type='target')} for python. See compiler output above." + ) + else: + pyKeOps_Message("OK", use_tag=False, flush=True, level=1) + pyKeOps_Message("OK", use_tag=False, flush=True) diff --git a/pykeops/pykeops/common/keops_io/nvrtc/__init__.py b/pykeops/pykeops/common/keops_io/nvrtc/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/pykeops/pykeops/common/keops_io/pykeops_nvrtc.cpp b/pykeops/pykeops/common/keops_io/nvrtc/pykeops_nvrtc.cpp similarity index 100% rename from pykeops/pykeops/common/keops_io/pykeops_nvrtc.cpp rename to pykeops/pykeops/common/keops_io/nvrtc/pykeops_nvrtc.cpp diff --git a/pykeops/pykeops/config.py b/pykeops/pykeops/config.py index 8d2085e19..e9a2e4fe1 100644 --- a/pykeops/pykeops/config.py +++ b/pykeops/pykeops/config.py @@ -1,7 +1,6 @@ import importlib.util import os import sys -import sysconfig # Instantiating the keopscore.config main classes for pykeops import keopscore.config @@ -23,8 +22,6 @@ # Verbosity level - - class _VerboseConfig: def __init__(self, level=1): self.level = level @@ -57,14 +54,16 @@ def get_verbose(): # version + +base_dir_path = os.path.abspath(os.path.dirname(__file__)) + + def read_version(version_file): with open(version_file, encoding="utf-8") as v: return v.read().rstrip() -_version = read_version( - os.path.join(os.path.abspath(os.path.dirname(__file__)), "keops_version") -) +_version = read_version(os.path.join(base_dir_path, "keops_version")) def get_version(): @@ -74,27 +73,23 @@ def get_version(): return _version +# path +def get_pykeops_io_folder(): + return os.path.join(base_dir_path, "common", "keops_io") + + def pykeops_nvrtc_name(type="src"): basename = "pykeops_nvrtc" - extension = ".cpp" if type == "src" else sysconfig.get_config_var("EXT_SUFFIX") - return os.path.join( - ( - os.path.join( - os.path.dirname(os.path.realpath(__file__)), "common", "keops_io" - ) - if type == "src" - else get_build_folder() - ), - basename + extension, - ) - - -def pykeops_cpp_name(tag="", extension=""): - basename = "pykeops_cpp_" - return os.path.join( - get_build_folder(), - basename + tag + extension, - ) + if type == "src": + return os.path.join(get_pykeops_io_folder(), "nvrtc", basename + ".cpp") + return path.get_python_extension_path(basename) + + +def pykeops_cpp_binder_name(type="src"): + basename = "pykeops_cpp" + if type == "src": + return os.path.join(get_pykeops_io_folder(), "cpp", basename + ".cpp") + return path.get_python_extension_path(basename) python_includes = "$({python3} -m pybind11 --includes)".format(python3=sys.executable) diff --git a/pykeops/pykeops/sandbox/test_torch_func_hessian.py b/pykeops/pykeops/sandbox/test_torch_func_hessian.py index d2dd2209c..fcb6e503c 100644 --- a/pykeops/pykeops/sandbox/test_torch_func_hessian.py +++ b/pykeops/pykeops/sandbox/test_torch_func_hessian.py @@ -2,7 +2,7 @@ import torch from pykeops.torch import LazyTensor -keopscore.config.chunks.set_auto_factorize(False) +keopscore.config.reduction.set_auto_factorize(False) def fn_torch(x_i): diff --git a/pykeops/pykeops/test/test_torch_func.py b/pykeops/pykeops/test/test_torch_func.py index f50ff880a..476c7a4dc 100644 --- a/pykeops/pykeops/test/test_torch_func.py +++ b/pykeops/pykeops/test/test_torch_func.py @@ -5,7 +5,7 @@ torch.manual_seed(0) -keopscore.config.chunks.set_auto_factorize(False) +keopscore.config.reduction.set_auto_factorize(False) B1, B2, M, N, D = 5, 4, 10, 20, 2 From f25e1e803134a1f4d7807825e5295d4a56c86f6c Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 10 May 2026 18:22:02 +0200 Subject: [PATCH 41/98] fix messages --- keopscore/keopscore/binders/cpp/Cpu_link_compile.py | 2 +- keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py | 2 +- keopscore/keopscore/config/Cuda.py | 3 --- pykeops/pykeops/common/keops_io/cpp/LoadKeOps_cpp.py | 2 +- pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py | 2 +- 5 files changed, 4 insertions(+), 7 deletions(-) diff --git a/keopscore/keopscore/binders/cpp/Cpu_link_compile.py b/keopscore/keopscore/binders/cpp/Cpu_link_compile.py index 3d89a6223..17305196b 100644 --- a/keopscore/keopscore/binders/cpp/Cpu_link_compile.py +++ b/keopscore/keopscore/binders/cpp/Cpu_link_compile.py @@ -70,7 +70,7 @@ def compile_dll(self): end="", level=2, ) - KeOps_Message(" ... ", flush=True, end="", level=1) + KeOps_Message(" ... ", flush=True, end="", use_tag=False, level=1) out = KeOps_OS_Run(compile_command) if out.returncode != 0: diff --git a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py index e349e71d2..51f767680 100644 --- a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py +++ b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py @@ -108,7 +108,7 @@ def compile_jit_compile_dll(): end="", level=2, ) - KeOps_Message(" ... ", flush=True, end="", level=1) + KeOps_Message(" ... ", flush=True, end="", use_tag=False, level=1) command = Gpu_link_compile.get_compile_command( sourcename=jit_compile_src, dllname=jit_compile_dll() ) diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index 7bb92d603..655c8c8c7 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -351,9 +351,6 @@ def print_cuda_block_size(self): def set_visible_devices(self): """Set specific GPUs from CUDA_VISIBLE_DEVICES.""" cuda_visible = os.getenv("CUDA_VISIBLE_DEVICES") - print( - f"CUDA_VISIBLE_DEVICES: {cuda_visible if cuda_visible is not None else not_found_str}" - ) if cuda_visible is not None: self._visible_devices = ( diff --git a/pykeops/pykeops/common/keops_io/cpp/LoadKeOps_cpp.py b/pykeops/pykeops/common/keops_io/cpp/LoadKeOps_cpp.py index fbbd84da6..1eef91c66 100644 --- a/pykeops/pykeops/common/keops_io/cpp/LoadKeOps_cpp.py +++ b/pykeops/pykeops/common/keops_io/cpp/LoadKeOps_cpp.py @@ -88,7 +88,7 @@ def compile_jit_binary(): end="", level=2, ) - pyKeOps_Message(" ... ", flush=True, end="", level=1) + pyKeOps_Message(" ... ", flush=True, end="", use_tag=False, level=1) out = KeOps_OS_Run(compile_command) if out.returncode != 0: diff --git a/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py b/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py index 2806fce99..30051dbf7 100644 --- a/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py +++ b/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py @@ -91,7 +91,7 @@ def compile_jit_binary(): end="", level=2, ) - pyKeOps_Message(" ... ", flush=True, end="", level=1) + pyKeOps_Message(" ... ", flush=True, end="", use_tag=False, level=1) out = KeOps_OS_Run(compile_command) if out.returncode != 0: From c8d8dbe7d9cad2fd2f2c9a8be0bba4668486fcd4 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 10 May 2026 18:48:31 +0200 Subject: [PATCH 42/98] fix omp path in brew --- keopscore/keopscore/config/OpenMP.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 5349a79bc..9ce77f1ad 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -39,8 +39,8 @@ class OpenMPConfig: _openmp_system_suffixes = [ # self.get_brew_prefix() added later on, - os.path.join(os.path.sep, "usr", "local", "opt", "libomp"), - os.path.join(os.path.sep, "opt", "local"), + os.path.join(os.path.sep, "opt", "homebrew", "opt", "libomp"), + os.path.join(os.path.sep, "usr", "local"), os.path.join(os.path.sep, "usr"), ] @@ -53,12 +53,12 @@ class OpenMPConfig: openmp_library_suffixes = ( "lib", "lib64", - os.path.join("opt", "libomp", "lib"), + os.path.join("libomp", "lib"), ) _openmp_include_sufixes = ( "include", - os.path.join("opt", "libomp", "include"), + os.path.join("libomp", "include"), ) _omp_info = { From 7e0b3ad2d1c9870166f182a43454e45a122eaca6 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 10 May 2026 19:11:40 +0200 Subject: [PATCH 43/98] test openMP on MacOS Co-authored-by: Copilot --- keopscore/keopscore/config/OpenMP.py | 32 +++++++++++++++------------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 9ce77f1ad..9adfad951 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -1,5 +1,4 @@ import os -import subprocess import tempfile from keopscore.utils.messages import print_envs @@ -13,6 +12,7 @@ from keopscore.utils.system_utils import ( _find_library_by_names, get_include_file_abspath, + KeOps_OS_Run, ) @@ -210,27 +210,29 @@ def check_compiler_for_openmp(self): f.write(test_program) test_file = f.name - compile_command = [ - self.cxx.get_cxx_compiler(), - test_file, - self.get_include_options(), - self.get_compile_options(), - ] - if self.get_linking_options(): - compile_command.append(self.get_linking_options()) - compile_command.extend(["-o", f"{test_file}.out"]) - - try: - # Warning : subprocess is used below to compile the test program (using subprocess.check_output to capture stderr) - subprocess.check_output(compile_command, stderr=subprocess.STDOUT) + # Try to compile with compiler's OpenMP options directly + compile_command_str = ( + f"{self.cxx.get_cxx_compiler()} " + f"{self.get_compile_options()} " + f"{test_file} " + f"{self.get_include_options()} " + f"{self.get_linking_options()} " + f"-o {test_file}.out" + ) + + out = KeOps_OS_Run(compile_command_str, print_warning=True) + if out.returncode == 0: os.remove(test_file) os.remove(test_file + ".out") return True - except subprocess.CalledProcessError: + else: os.remove(test_file) + if os.path.exists(test_file + ".out"): + os.remove(test_file + ".out") KeOps_Warning( f"{self.cxx.get_cxx_compiler()} does not support OpenMP. OpenMP support will be disabled." ) + self.print_all() return False # C++ Compiler Options From a35a669a4fd1aa69ea9c8d8c4c55259aeebd35f3 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 10 May 2026 19:47:36 +0200 Subject: [PATCH 44/98] try debub openmp detection Co-authored-by: Copilot --- keopscore/keopscore/config/OpenMP.py | 97 ++++++++++++++++++++-------- 1 file changed, 69 insertions(+), 28 deletions(-) diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 9adfad951..4d4dab29c 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -37,19 +37,13 @@ class OpenMPConfig: "OpenMP_ROOT_DIR", ) - _openmp_system_suffixes = [ + _openmp_system_roots = [ # self.get_brew_prefix() added later on, os.path.join(os.path.sep, "opt", "homebrew", "opt", "libomp"), os.path.join(os.path.sep, "usr", "local"), os.path.join(os.path.sep, "usr"), ] - openmp_basename_candidate = ( - "libomp.dylib", - "libgomp.dylib", - "libomp.so", - ) - openmp_library_suffixes = ( "lib", "lib64", @@ -63,7 +57,7 @@ class OpenMPConfig: _omp_info = { "name": ["omp", "gomp"], - "lib_basename_candidate": ["libomp.*", "libgomp.*", "libm.so*"], + "lib_basename_candidate": ["libomp.*", "libgomp.*", "libm.*"], "header_basename": ["omp.h", "gomp.h"], "library": "", # to be filled later "header": "", # not needed @@ -77,7 +71,7 @@ def __init__(self, platform, cxx): # On MacOS: Add brew prefix to OpenMP search paths if it exists if self.platform.get_brew_prefix(): - self._openmp_system_suffixes.insert( + self._openmp_system_roots.insert( 0, [ self.platform.get_brew_prefix(), @@ -106,27 +100,68 @@ def find_install_path(self, lib_dict_info): result["header"] = get_include_file_abspath( header_basename, self.cxx.get_cxx_compiler() ) + + #### + print("OpenMP library search using standard names:") + print(f" Trying library name: {name}") + print(f" Found library path: {result['library'] or not_found_str}") + print(f" Trying header name: {header_basename}") + print(f" Found header path: {result['header'] or not_found_str}") + print() + print() + #### + if result["library"] and result["header"]: + #### + print() + print() + print( + "Successfully found OpenMP library and header using standard names." + ) + print(f" Library: {result['library']}") + print(f" Header: {result['header']}") + #### return result # If that fails, search for OpenMP headers and libraries in common locations. - if not result["library"]: - candidate_roots = _ordered_search_roots( - env_vars=self.openmp_env_vars, - conda_root="CONDA_PREFIX", - system_roots=self._openmp_system_suffixes, - ) + candidate_roots = _ordered_search_roots( + env_vars=self.openmp_env_vars, + conda_root="CONDA_PREFIX", + system_roots=self._openmp_system_roots, + ) - result["library"] = _first_matching_file( - _path_candidates(candidate_roots, self.openmp_library_suffixes), - self.openmp_basename_candidate, - ) + result["library"] = _first_matching_file( + _path_candidates(candidate_roots, self.openmp_library_suffixes), + lib_dict_info["lib_basename_candidate"], + ) - # Finally, search for OpenMP headers in common locations. - result["header"] = _first_matching_file( - _path_candidates(candidate_roots, self._openmp_include_sufixes), - (self._openmp_header_basename,), - ) + #### + print() + print() + print("OpenMP library search in common locations:") + print(f" Candidate roots: {candidate_roots}") + print(f" Library search suffixes: {self.openmp_library_suffixes}") + print(f" Found library path: {result['library'] or not_found_str}") + print() + print() + #### + + # Finally, search for OpenMP headers in common locations. + result["header"] = _first_matching_file( + _path_candidates(candidate_roots, self._openmp_include_sufixes), + lib_dict_info["header_basename"], + ) + + #### + print() + print() + print("OpenMP header search in common locations:") + print(f" Candidate roots: {candidate_roots}") + print(f" Header search suffixes: {self._openmp_include_sufixes}") + print(f" Found header path: {result['header'] or not_found_str}") + print() + print() + #### return result @@ -237,8 +272,8 @@ def check_compiler_for_openmp(self): # C++ Compiler Options def set_compile_options(self): - # Add special fix for openMP prgama and Apple Clang. Order matters. - if self.cxx.get_use_Apple_clang() and self.platform.get_brew_prefix(): + # Apple clang does not support -fopenmp directly; -Xpreprocessor is required. + if self.cxx.get_use_Apple_clang(): self._compile_options += "-Xpreprocessor " self._compile_options += "-fopenmp" @@ -251,7 +286,11 @@ def print_compile_options(self): # C++ Compiler Include Options def set_include_options(self): - self._include_options = f"-I{self.get_openmp_include_dir()}" + self._include_options = ( + f"-I{self.get_openmp_include_dir()}" + if self.get_openmp_include_dir() + else "" + ) def get_include_options(self): return self._include_options @@ -261,7 +300,9 @@ def print_include_options(self): # C++ linking Options def set_linking_options(self): - self._linking_options += f"-L{self.get_libomp_folder()}" + self._linking_options += ( + f"-L{self.get_libomp_folder()}" if self.get_libomp_folder() else "" + ) def get_linking_options(self): return self._linking_options From d0f99653696776ddf5bbc20a98085f8f79b74e3a Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 10 May 2026 20:07:47 +0200 Subject: [PATCH 45/98] debugging in progress... --- keopscore/keopscore/config/OpenMP.py | 11 ++++++++++- keopscore/keopscore/config/Platform.py | 10 ++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 4d4dab29c..81b4df31a 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -57,7 +57,14 @@ class OpenMPConfig: _omp_info = { "name": ["omp", "gomp"], - "lib_basename_candidate": ["libomp.*", "libgomp.*", "libm.*"], + "lib_basename_candidate": [ + "libomp.dylib", + "libomp.so*", + "libgomp.dylib", + "libgomp.so*", + "libm.dylib", + "libm.so*" + ], "header_basename": ["omp.h", "gomp.h"], "library": "", # to be filled later "header": "", # not needed @@ -267,6 +274,8 @@ def check_compiler_for_openmp(self): KeOps_Warning( f"{self.cxx.get_cxx_compiler()} does not support OpenMP. OpenMP support will be disabled." ) + self.platform.print_all() + self.cxx.print_all() self.print_all() return False diff --git a/keopscore/keopscore/config/Platform.py b/keopscore/keopscore/config/Platform.py index cef374386..f16bd368f 100644 --- a/keopscore/keopscore/config/Platform.py +++ b/keopscore/keopscore/config/Platform.py @@ -138,6 +138,14 @@ def set_brew_prefix(self): def get_brew_prefix(self): """Get Homebrew prefix path using KeOps_OS_Run""" return self._brew_prefix + + def print_brew_prefix(self): + if self.get_platform() == "Darwin": + brew_prefix = self.get_brew_prefix() + if brew_prefix: + print(f"Homebrew Prefix: {brew_prefix}") + else: + print("Homebrew not found or not installed.") @staticmethod def detect_env_type(): @@ -165,6 +173,8 @@ def print_all(self): self.print_env_type() self.print_python_executable() + self.print_brew_prefix() + # Print relevant environment variables. print_envs(self.platform_envs) From b9c2100e3021138d70af23e48327f2603750edab Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 10 May 2026 20:09:29 +0200 Subject: [PATCH 46/98] debugging in progress... --- keopscore/keopscore/config/OpenMP.py | 4 ++-- keopscore/keopscore/config/Platform.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 81b4df31a..73c62f3ea 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -63,8 +63,8 @@ class OpenMPConfig: "libgomp.dylib", "libgomp.so*", "libm.dylib", - "libm.so*" - ], + "libm.so*", + ], "header_basename": ["omp.h", "gomp.h"], "library": "", # to be filled later "header": "", # not needed diff --git a/keopscore/keopscore/config/Platform.py b/keopscore/keopscore/config/Platform.py index f16bd368f..5dc0e90f8 100644 --- a/keopscore/keopscore/config/Platform.py +++ b/keopscore/keopscore/config/Platform.py @@ -138,7 +138,7 @@ def set_brew_prefix(self): def get_brew_prefix(self): """Get Homebrew prefix path using KeOps_OS_Run""" return self._brew_prefix - + def print_brew_prefix(self): if self.get_platform() == "Darwin": brew_prefix = self.get_brew_prefix() From 96147b5b61b48d3ed7c1018ca40dfe614dc47d68 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 10 May 2026 20:30:38 +0200 Subject: [PATCH 47/98] add arch flag to openmp test --- keopscore/keopscore/config/OpenMP.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 73c62f3ea..ed43fc7f6 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -174,7 +174,9 @@ def find_install_path(self, lib_dict_info): def _omp_is_available(self): self._omp_info = self.find_install_path(self._omp_info) - liomp_path = self._omp_info["library"] + liomp_path = os.path.exists(self._omp_info["library"]) and os.path.exists( + self._omp_info["header"] + ) if not liomp_path: KeOps_Warning( "libomp not found. Set OMP_PATH, LIBOMP_PATH, or OpenMP_ROOT if it is installed in a non-standard location." @@ -255,6 +257,7 @@ def check_compiler_for_openmp(self): # Try to compile with compiler's OpenMP options directly compile_command_str = ( f"{self.cxx.get_cxx_compiler()} " + f"{self.cxx.get_compile_options()} " f"{self.get_compile_options()} " f"{test_file} " f"{self.get_include_options()} " From 92e34b7474aca9f219694f043af8bf28c1f48982 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 10 May 2026 21:02:28 +0200 Subject: [PATCH 48/98] polish messages. CI pass on all systems --- keopscore/keopscore/config/Cuda.py | 15 +++-- keopscore/keopscore/config/OpenMP.py | 63 ++++++++----------- keopscore/keopscore/utils/system_utils.py | 2 +- .../common/keops_io/nvrtc/LoadKeOps_nvrtc.py | 2 - 4 files changed, 35 insertions(+), 47 deletions(-) diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index 655c8c8c7..3cd8a5e21 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -175,7 +175,7 @@ def _find_and_load_libcuda(self): if not libcuda_path: return ( False, - "libcuda not found. Make sure the CUDA driver is installed and accessible. Switching to CPU only.", + "libcuda not found. Make sure the CUDA driver is installed and accessible.", ) try: @@ -189,7 +189,7 @@ def _find_and_load_libcuda(self): if libcuda.cuInit(0) != self.CUDA_SUCCESS: return ( False, - "libcuda was detected, but driver API could not be initialized (flushing KeOps caches and/or rebooting the system may help). Switching to CPU only.", + "libcuda was detected, but driver API could not be initialized (flushing KeOps caches and/or rebooting the system may help).", ) # If we successfully loaded libcuda and initialized it, store the handle in the config for potential future use @@ -202,7 +202,7 @@ def _find_and_load_libcuda(self): ): return ( False, - "libcuda was detected and driver API was initialized, but no working GPU found. Switching to CPU only.", + "libcuda was detected and driver API was initialized, but no working GPU found.", ) self._MaxThreadsPerBlock = [0] * nGpus.value @@ -216,8 +216,7 @@ def _find_and_load_libcuda(self): return ( False, f"libcuda was detected and driver API was initialized, but " - + err_msg - + " Switching to CPU only.", + + err_msg, ) self._n_visible_devices = nGpus.value @@ -231,7 +230,7 @@ def _find_and_load_libnvrtc(self): if not libnvrtc_path: return ( False, - "libnvrtc not found. Make sure the CUDA toolkit is installed and accessible. Switching to CPU only.", + "libnvrtc not found. Make sure the CUDA toolkit is installed and accessible.", ) try: @@ -261,7 +260,7 @@ def _find_and_load_cudart(self): if not libcudart_path: return ( False, - "libcudart not found. Make sure the CUDA toolkit is installed and accessible. Switching to CPU only.", + "libcudart not found. Make sure the CUDA toolkit is installed and accessible.", ) try: @@ -279,7 +278,7 @@ def _find_and_load_cudart(self): ): return ( False, - "libcudart was found and loaded, but failed to get CUDA runtime version. Switching to CPU only.", + "libcudart was found and loaded, but failed to get CUDA runtime version.", ) # If we successfully loaded libcudart, store the handle in the config diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index ed43fc7f6..32139b131 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -1,7 +1,7 @@ import os import tempfile -from keopscore.utils.messages import print_envs +from keopscore.utils.messages import KeOps_Message, print_envs from keopscore.utils.messages import enabled_dict, not_found_str from keopscore.utils.messages import KeOps_Warning from keopscore.utils.path_utils import ( @@ -109,25 +109,19 @@ def find_install_path(self, lib_dict_info): ) #### - print("OpenMP library search using standard names:") - print(f" Trying library name: {name}") - print(f" Found library path: {result['library'] or not_found_str}") - print(f" Trying header name: {header_basename}") - print(f" Found header path: {result['header'] or not_found_str}") - print() - print() + KeOps_Message("OpenMP library search using standard names:", level=2) + KeOps_Message(f" Trying library name: {name}", level=2) + KeOps_Message( + f" Found library path: {result['library'] or not_found_str}", level=2 + ) + KeOps_Message(f" Trying header name: {header_basename}", level=2) + KeOps_Message( + f" Found header path: {result['header'] or not_found_str}", level=2 + ) + KeOps_Message() #### if result["library"] and result["header"]: - #### - print() - print() - print( - "Successfully found OpenMP library and header using standard names." - ) - print(f" Library: {result['library']}") - print(f" Header: {result['header']}") - #### return result # If that fails, search for OpenMP headers and libraries in common locations. @@ -143,14 +137,14 @@ def find_install_path(self, lib_dict_info): ) #### - print() - print() - print("OpenMP library search in common locations:") - print(f" Candidate roots: {candidate_roots}") - print(f" Library search suffixes: {self.openmp_library_suffixes}") - print(f" Found library path: {result['library'] or not_found_str}") - print() - print() + KeOps_Message("OpenMP library search in common locations:", level=2) + KeOps_Message(f" Candidate roots: {candidate_roots}", level=2) + KeOps_Message( + f" Library search suffixes: {self.openmp_library_suffixes}", level=2 + ) + KeOps_Message( + f" Found library path: {result['library'] or not_found_str}", level=2 + ) #### # Finally, search for OpenMP headers in common locations. @@ -160,14 +154,14 @@ def find_install_path(self, lib_dict_info): ) #### - print() - print() - print("OpenMP header search in common locations:") - print(f" Candidate roots: {candidate_roots}") - print(f" Header search suffixes: {self._openmp_include_sufixes}") - print(f" Found header path: {result['header'] or not_found_str}") - print() - print() + KeOps_Message("OpenMP header search in common locations:", level=2) + KeOps_Message(f" Candidate roots: {candidate_roots}", level=2) + KeOps_Message( + f" Header search suffixes: {self._openmp_include_sufixes}", level=2 + ) + KeOps_Message( + f" Found header path: {result['header'] or not_found_str}", level=2 + ) #### return result @@ -277,9 +271,6 @@ def check_compiler_for_openmp(self): KeOps_Warning( f"{self.cxx.get_cxx_compiler()} does not support OpenMP. OpenMP support will be disabled." ) - self.platform.print_all() - self.cxx.print_all() - self.print_all() return False # C++ Compiler Options diff --git a/keopscore/keopscore/utils/system_utils.py b/keopscore/keopscore/utils/system_utils.py index 2cc70bc37..ca0ba4516 100644 --- a/keopscore/keopscore/utils/system_utils.py +++ b/keopscore/keopscore/utils/system_utils.py @@ -20,7 +20,7 @@ def get_include_file_abspath(filename, compiler): """Return the full path of the header filename using compiler.""" cmd = f'echo "#include <{filename}>" | {compiler} -M -E -x c++ -' - out = KeOps_OS_Run(cmd) + out = KeOps_OS_Run(cmd, print_warning=False) text = out.stdout.decode("utf8") text = text.replace("\\\n", " ") diff --git a/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py b/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py index 30051dbf7..7d0e0927f 100644 --- a/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py +++ b/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py @@ -101,8 +101,6 @@ def compile_jit_binary(): else: pyKeOps_Message("OK", use_tag=False, flush=True, level=1) - pyKeOps_Message("OK", use_tag=False, flush=True) - LoadKeOps_nvrtc = Cache_partial( LoadKeOps_nvrtc_class, From 53feaf57e611b50a8b69259f5093a9dc99a1d754 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 10 May 2026 22:19:26 +0200 Subject: [PATCH 49/98] try to add crt header locally... Co-authored-by: Copilot --- .../binders/nvrtc/Gpu_link_compile.py | 2 +- keopscore/keopscore/config/Cuda.py | 17 ++++---- keopscore/keopscore/config/KeOpsPath.py | 5 +++ keopscore/keopscore/config/OpenMP.py | 1 - keopscore/keopscore/utils/gpu_utils.py | 41 +++++++++++++++++++ .../common/keops_io/nvrtc/LoadKeOps_nvrtc.py | 2 +- 6 files changed, 57 insertions(+), 11 deletions(-) diff --git a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py index 51f767680..646162526 100644 --- a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py +++ b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py @@ -90,8 +90,8 @@ def get_compile_command( f"{keopscore.config.cxx.get_cxx_compiler()} " f"{keopscore.config.cuda.get_preprocessing_options()} " f"{keopscore.config.cxx.get_compile_options()} " - f"{keopscore.config.cuda.get_include_options()} " f"{keopscore.config.path.get_include_options()} " + f"{keopscore.config.cuda.get_include_options()} " f" {extra_flags} " f"{sourcename} " f"{keopscore.config.cxx.get_linking_options()} " diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index 3cd8a5e21..74f6604d7 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -466,15 +466,16 @@ def set_cuda_include_path(self): """Set the CUDA include path by searching for cuda.h and nvrtc.h.""" # This is done in find_cuda_install since it relies on the cuda installation info which is only available after checking library availability. if not self.get_use_cuda(): - self._cuda_include_path = None - self._cuda_include_path = list( - set( - [ - os.path.dirname(self._libnvrtc_info["header"]), - os.path.dirname(self._libcuda_info["header"]), - ] + self._cuda_include_path = "" + else: + self._cuda_include_path = list( + set( + [ + os.path.dirname(self._libnvrtc_info["header"]), + os.path.dirname(self._libcuda_info["header"]), + ] + ) ) - ) def get_cuda_include_path(self): """ diff --git a/keopscore/keopscore/config/KeOpsPath.py b/keopscore/keopscore/config/KeOpsPath.py index 7c459cd81..510febfdb 100644 --- a/keopscore/keopscore/config/KeOpsPath.py +++ b/keopscore/keopscore/config/KeOpsPath.py @@ -3,6 +3,7 @@ import sysconfig import keopscore +from keopscore.utils.gpu_utils import add_crt_symlink_to_cuda_include_path from keopscore.utils.path_utils import ensure_directory from keopscore.utils.messages import not_found_str, print_envs @@ -178,6 +179,10 @@ def set_default_build_path(self): ) # Initialize _build_path : TODO : check ?/ self._build_folder = self._default_build_path + #### Add a symlink to the crt folder in keops include path if needed, to handle the case of cuda_fp16.h including crt/host_config.h and crt/device_runtime_api.h without proper reference to their location in the cuda include path (issue in cudatoolkit pip package) + add_crt_symlink_to_cuda_include_path( + self.cuda.get_cuda_include_path(), self._build_folder + ) def get_default_build_path(self): """Get the default build path.""" diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 32139b131..6260d4360 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -118,7 +118,6 @@ def find_install_path(self, lib_dict_info): KeOps_Message( f" Found header path: {result['header'] or not_found_str}", level=2 ) - KeOps_Message() #### if result["library"] and result["header"]: diff --git a/keopscore/keopscore/utils/gpu_utils.py b/keopscore/keopscore/utils/gpu_utils.py index be1e82d63..b1cbde01d 100644 --- a/keopscore/keopscore/utils/gpu_utils.py +++ b/keopscore/keopscore/utils/gpu_utils.py @@ -57,3 +57,44 @@ def custom_cuda_include_fp16_path(): if not os.path.isfile(fp16_header_path): pack_header(fp16_header, orig_cuda_include_fp16_path(), build_folder) return build_folder + + +def add_crt_symlink_to_cuda_include_path(cuda_include_path, build_folder): + """ + The cuda_fp16.h header includes crt/host_config.h and crt/device_runtime_api.h, but these + headers are not properly referenced in the cuda include path ship in the pip version of + cudatoolkit. This function adds a symlink to the crt folder in keops include path + (included before the cuda toolkit) that mimick the correct behavior... + """ + + files_to_check = ["host_config.h"] + cuda_include_path = [p for p in list(cuda_include_path) if os.path.isdir(p)] + + check = [ + any([os.path.isfile(os.path.join(p, "crt", file)) for p in cuda_include_path]) + for file in files_to_check + ] + if all(check): + return + + crt_folder = os.path.join(build_folder, "crt") + os.makedirs(crt_folder, exist_ok=True) + + for file in files_to_check: + crt_symlink = os.path.join(crt_folder, file) + if os.path.exists(crt_symlink): + continue + # Find the actual file in one of the cuda include paths + target = next( + ( + os.path.join(p, file) + for p in cuda_include_path + if os.path.isfile(os.path.join(p, file)) + ), + None, + ) + if target is not None: + os.symlink(target, crt_symlink) + print( + f"Created symlink for {file} in crt folder: {crt_symlink} -> {target}" + ) diff --git a/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py b/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py index 7d0e0927f..681d725d0 100644 --- a/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py +++ b/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py @@ -84,7 +84,7 @@ def compile_jit_binary(): sourcename=pykeopsconfig.pykeops_nvrtc_name(type="src"), dllname=pykeopsconfig.pykeops_nvrtc_name(type="target"), ) - pyKeOps_Message("Compiling nvrtc binder for python ... ", flush=True, end="") + pyKeOps_Message("Compiling nvrtc binder for python", flush=True, end="", level=1) pyKeOps_Message( " in cache folder " + pykeopsconfig.path.get_build_folder(), flush=True, From c09aa8fe932d51d269df40cd8a794485e4ccbc81 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 11 May 2026 01:39:47 +0200 Subject: [PATCH 50/98] Fix set_build_folder Co-authored-by: Copilot --- keopscore/keopscore/__init__.py | 12 +- .../binders/nvrtc/Gpu_link_compile.py | 43 +++-- keopscore/keopscore/config/KeOpsPath.py | 149 +++++++++--------- keopscore/keopscore/config/__init__.py | 35 ++-- keopscore/keopscore/utils/gpu_utils.py | 3 - .../common/keops_io/nvrtc/LoadKeOps_nvrtc.py | 4 +- 6 files changed, 121 insertions(+), 125 deletions(-) diff --git a/keopscore/keopscore/__init__.py b/keopscore/keopscore/__init__.py index 8da592cf9..453d93962 100644 --- a/keopscore/keopscore/__init__.py +++ b/keopscore/keopscore/__init__.py @@ -8,14 +8,8 @@ # Config import keopscore.config -# Initialize CUDA libraries if CUDA is used -if keopscore.config.cuda.get_use_cuda(): - # Initialize CUDA libraries if necessary - from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile - from keopscore.binders.nvrtc.Gpu_link_compile import jit_compile_dll - - if not os.path.exists(jit_compile_dll()): - Gpu_link_compile.compile_jit_compile_dll() +# Create build folder if it does not exist and set it as the default build folder +keopscore.config.path.set_build_folder(reset_all=False) # expose to the user -set_build_folder = keopscore.config.path.set_different_build_folder +set_build_folder = keopscore.config.path.set_build_folder diff --git a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py index 646162526..dabfa79e1 100644 --- a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py +++ b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py @@ -1,6 +1,5 @@ import ctypes import os -import sysconfig import keopscore.config from keopscore.binders.LinkCompile import LinkCompile @@ -8,23 +7,12 @@ from keopscore.utils.messages import KeOps_Error, KeOps_Message from keopscore.utils.system_utils import KeOps_OS_Run -jit_source_file = os.path.join( - keopscore.config.path.get_base_dir_path(), "binders", "nvrtc", "keops_nvrtc.cpp" -) - -jit_compile_src = os.path.join( - keopscore.config.path.get_base_dir_path(), "binders", "nvrtc", "nvrtc_jit.cpp" -) - - -def jit_compile_dll(): - return os.path.join( - keopscore.config.path.get_build_folder(), - "nvrtc_jit" + sysconfig.get_config_var("SHLIB_SUFFIX"), - ) - class Gpu_link_compile(LinkCompile): + """ + This class is used to compile the main dll that will be used to JIT compile the cuda code for each formula. + """ + source_code_extension = "cu" def __init__(self): @@ -46,7 +34,10 @@ def __init__(self): self.gencode_filename + "_nvrtc." + keopscore.config.cuda.get_ir_type(), ).encode("utf-8") - self.my_c_dll = ctypes.CDLL(jit_compile_dll(), mode=os.RTLD_LAZY) + # Load nvrtc_jit dll, that will be used to compile the cuda code in jit mode + self.my_c_dll = ctypes.CDLL( + keopscore.config.keops_jit_compile_name(type="target"), mode=os.RTLD_LAZY + ) # file to check for existence to detect compilation is needed self.file_to_check = self.low_level_code_file @@ -79,8 +70,8 @@ def generate_code(self): @staticmethod def get_compile_command( - sourcename=jit_source_file, - dllname=keopscore.config.path.get_jit_binary(), + sourcename, + dllname, extra_flags="", ): # This is about the main KeOps binary (dll) that will be used to JIT compile all formulas. @@ -100,7 +91,14 @@ def get_compile_command( ) @staticmethod - def compile_jit_compile_dll(): + def compile_jit_compile_dll(force_recompile=False): + + if ( + os.path.exists(keopscore.config.keops_jit_compile_name(type="target")) + and not force_recompile + ): + return + KeOps_Message("Compiling cuda jit compiler engine", flush=True, end="", level=1) KeOps_Message( " in cache folder " + keopscore.config.path.get_build_folder(), @@ -110,12 +108,13 @@ def compile_jit_compile_dll(): ) KeOps_Message(" ... ", flush=True, end="", use_tag=False, level=1) command = Gpu_link_compile.get_compile_command( - sourcename=jit_compile_src, dllname=jit_compile_dll() + keopscore.config.keops_jit_compile_name(type="src"), + keopscore.config.keops_jit_compile_name(type="target"), ) out = KeOps_OS_Run(command) if out.returncode != 0: KeOps_Error( - f"Error compiling cuda/nvrtc binder {jit_compile_dll()}. See compiler output above." + f"Error compiling cuda/nvrtc binder {keopscore.config.keops_jit_compile_name(type="target")}. See compiler output above." ) else: KeOps_Message("OK", use_tag=False, flush=True, level=1) diff --git a/keopscore/keopscore/config/KeOpsPath.py b/keopscore/keopscore/config/KeOpsPath.py index 510febfdb..4538a9273 100644 --- a/keopscore/keopscore/config/KeOpsPath.py +++ b/keopscore/keopscore/config/KeOpsPath.py @@ -1,11 +1,12 @@ import os import sys import sysconfig +import warnings import keopscore from keopscore.utils.gpu_utils import add_crt_symlink_to_cuda_include_path from keopscore.utils.path_utils import ensure_directory -from keopscore.utils.messages import not_found_str, print_envs +from keopscore.utils.messages import KeOps_Warning, not_found_str, print_envs class KeOpsPathConfig: @@ -14,10 +15,12 @@ class KeOpsPathConfig: """ _base_dir_path = "" + _keops_cache_folder = "" + _build_folder = "" + _default_build_folder_name = "" _default_build_path = "" - _build_folder = "" _jit_binary = "" _include_options = "" @@ -35,7 +38,6 @@ def __init__(self, platform, cuda): self.set_default_build_folder_name() self.set_default_build_path() # should be done at the end.. - self.set_jit_binary() self.set_include_options() # Base Directory Path @@ -56,11 +58,14 @@ def print_base_dir_path(self): # KeOps Cache Folder def set_keops_cache_folder(self): """Set the KeOps cache folder.""" - cache_folder = os.getenv("KEOPS_CACHE_FOLDER") - if cache_folder is None: + + if os.getenv("KEOPS_CACHE_FOLDER"): + cache_folder = os.getenv("KEOPS_CACHE_FOLDER") + else: # fallback to default cache folder in user home directory cache_folder = os.path.join( os.path.expanduser("~"), ".cache", f"keops{keopscore.__version__}" ) + os.makedirs(cache_folder, exist_ok=True) self._keops_cache_folder = cache_folder @@ -73,30 +78,7 @@ def print_keops_cache_folder(self): print(f"KeOps Cache Folder: {self.get_keops_cache_folder() or not_found_str}") # Build Folder Management - def set_default_build_folder_name(self): - """Set the default build folder name.""" - name_parts = [ - "_".join(self.platform.get_uname()[:3]), - f"python{self.platform.get_python_version()}", - ] - if self.cuda.get_use_cuda(): - name_parts.append(f"CUDA{self.cuda.get_cuda_version()}") - - visible_devices = self.cuda.get_visible_devices() - if visible_devices: - name_parts.append(f"VISIBLE_DEVICES{visible_devices}") - - self._default_build_folder_name = "_".join(name_parts) - - def get_default_build_folder_name(self): - """Return the platform-specific default build folder suffix.""" - return self._default_build_folder_name - - def print_default_build_folder_name(self): - """Print the default build folder name.""" - print(f"Default Build Folder Name: {self.get_default_build_folder_name()}") - - def set_different_build_folder( + def set_build_folder( self, path=None, read_save_file=False, write_save_file=True, reset_all=True ): """ @@ -106,8 +88,8 @@ def set_different_build_folder( - path: The new build folder path. If None, it will be determined based on saved settings or defaults. - read_save_file: If True, read the build folder path from a save file if path is not provided. - write_save_file: If True, write the new build folder path to the save file. - - reset_all: If True, reset all cached formulas and recompile necessary components. """ + # If path is not given, we either read the save file or use the default build path save_file = os.path.join( self.get_keops_cache_folder(), "build_folder_location.txt" @@ -119,17 +101,15 @@ def set_different_build_folder( else: path = self.get_default_build_path() - # Create the folder if not yet done - os.makedirs(path, exist_ok=True) + # Remove the old build path from sys.path if it's there. + old_build_folder = self.get_build_folder() + if old_build_folder and old_build_folder in sys.path: + sys.path.remove(old_build_folder) - # Remove the old build path from sys.path if it's there - if self._build_folder and self._build_folder in sys.path: - sys.path.remove(self._build_folder) + # Create the folder if not yet done and add the new build path to sys.path + ensure_directory(path, add_to_syspath=True) # Update _build_folder to the new path self._build_folder = path - # Add the new build path to sys.path - if self._build_folder not in sys.path: - sys.path.append(self._build_folder) # Saving the location of the build path in a file if write_save_file: @@ -142,33 +122,45 @@ def set_different_build_folder( keopscore.get_keops_dll.get_keops_dll.reset( new_save_folder=self._build_folder ) - # Handle CUDA-specific recompilation if CUDA is used - if hasattr(self, "_use_cuda") and self._use_cuda: - from keopscore.binders.nvrtc.Gpu_link_compile import ( - Gpu_link_compile, - jit_compile_dll, - ) - if not os.path.exists(jit_compile_dll()): - Gpu_link_compile.compile_jit_compile_dll() + # Handle CUDA-specific recompilation if CUDA is used + if self.cuda.get_use_cuda(): + from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile + + Gpu_link_compile.compile_jit_compile_dll(force_recompile=reset_all) + #### Add a symlink to the crt folder in keops include path if needed, to handle + # cudatoolkit pip package + add_crt_symlink_to_cuda_include_path( + self.cuda.get_cuda_include_path(), self._build_folder + ) def get_build_folder(self): return self._build_folder - def get_build_folder_file(self, filename): - return os.path.join(self.get_build_folder(), filename) + # Default Build Path + def set_default_build_folder_name(self): + """Set the default build folder name.""" + name_parts = [ + "_".join(self.platform.get_uname()[:3]), + f"python{self.platform.get_python_version()}", + ] + if self.cuda.get_use_cuda(): + name_parts.append(f"CUDA{self.cuda.get_cuda_version()}") - def get_python_extension_path(self, basename): - return self.get_build_folder_file( - basename + sysconfig.get_config_var("EXT_SUFFIX") - ) + visible_devices = self.cuda.get_visible_devices() + if visible_devices: + name_parts.append(f"VISIBLE_DEVICES{visible_devices}") - def target_needs_update(self, target, source): - return not os.path.exists(target) or os.path.getmtime( - source - ) > os.path.getmtime(target) + self._default_build_folder_name = "_".join(name_parts) + + def get_default_build_folder_name(self): + """Return the platform-specific default build folder suffix.""" + return self._default_build_folder_name + + def print_default_build_folder_name(self): + """Print the default build folder name.""" + print(f"Default Build Folder Name: {self.get_default_build_folder_name()}") - # Default Build Path def set_default_build_path(self): """Set the default build path.""" self._default_build_path = ensure_directory( @@ -177,12 +169,6 @@ def set_default_build_path(self): ), add_to_syspath=True, ) - # Initialize _build_path : TODO : check ?/ - self._build_folder = self._default_build_path - #### Add a symlink to the crt folder in keops include path if needed, to handle the case of cuda_fp16.h including crt/host_config.h and crt/device_runtime_api.h without proper reference to their location in the cuda include path (issue in cudatoolkit pip package) - add_crt_symlink_to_cuda_include_path( - self.cuda.get_cuda_include_path(), self._build_folder - ) def get_default_build_path(self): """Get the default build path.""" @@ -192,28 +178,34 @@ def print_default_build_path(self): """Print the default build path.""" print(f"Default Build Path: {self.get_default_build_path() or not_found_str}") - # JIT Binary Path - def set_jit_binary(self): - """Set the path to the JIT binary.""" - self._jit_binary = self.get_default_build_path() - - def get_jit_binary(self): - """Get the path to the JIT binary.""" - return self._jit_binary - - def print_jit_binary(self): - """Print the path to the JIT binary.""" - print(f"JIT Binary Path: {self.get_jit_binary() or not_found_str}") - # include options management def set_include_options(self): """Set the include options for compilation.""" - self._include_options = f" -I{self.get_base_dir_path()}" + self.add_to_include_option(f" -I{self.get_base_dir_path()}") + + def add_to_include_option(self, option): + """Add an include option for compilation.""" + self._include_options += f" {option}" def get_include_options(self): """Get the include options for compilation.""" return self._include_options + def print_include_options(self): + """Print the include options for compilation.""" + print(f"Include Options: {self.get_include_options() or not_found_str}") + + # helpers + def get_python_extension_path(self, basename, suffix="EXT_SUFFIX"): + return os.path.join( + self.get_build_folder(), basename + sysconfig.get_config_var(suffix) + ) + + def target_needs_update(self, target, source): + return not os.path.exists(target) or os.path.getmtime( + source + ) > os.path.getmtime(target) + # Comprehensive Path Information def print_all(self): """Print all path-related information.""" @@ -226,7 +218,8 @@ def print_all(self): self.print_base_dir_path() self.print_keops_cache_folder() self.print_default_build_folder_name() - self.print_jit_binary() + self.print_default_build_path() + self.print_include_options() # Print relevant environment variables. print_envs(self.path_env_vars) diff --git a/keopscore/keopscore/config/__init__.py b/keopscore/keopscore/config/__init__.py index e3ae6c7b2..458501a46 100644 --- a/keopscore/keopscore/config/__init__.py +++ b/keopscore/keopscore/config/__init__.py @@ -9,7 +9,7 @@ from .OpenMP import OpenMPConfig from .Platform import PlatformConfig from .ReductionTuning import ReductionTuningConfig -from keopscore.utils.messages import KeOps_Message +from keopscore.utils.messages import KeOps_Error, KeOps_Message # Instantiate the configurations once at import time to preserve the existing API. debug = DebugConfig() @@ -21,6 +21,16 @@ reduction = ReductionTuningConfig() +def keops_jit_compile_name(type="src"): + basename = "nvrtc_jit" + if type == "src": + return os.path.join( + path.get_base_dir_path(), "binders", "nvrtc", basename + ".cpp" + ) + + return path.get_python_extension_path(basename, suffix="SHLIB_SUFFIX") + + def check_health(infos="all"): """ Check the health of the specified configuration. @@ -40,12 +50,20 @@ def check_health(infos="all"): def clean_keops(recompile_jit_binary=True, verbose=True): build_path = path.get_build_folder() - use_cuda = cuda.get_use_cuda() - jit_binary = path.get_jit_binary() if use_cuda else None + default_build_path = path.get_default_build_path() + + if os.path.abspath(build_path) != os.path.abspath(default_build_path): + KeOps_Error( + f"Your build folder is set to {build_path}, which is not the default build folder. For safety reasons, the clean_keops function will not delete files in this folder. If you want to clean this folder, please do it manually." + ) if build_path and os.path.isdir(build_path): + jit_binary = keops_jit_compile_name(type="target") + # TODO: Add a safety check to prevent accidental deletion of important directories. for entry in os.scandir(build_path): - if recompile_jit_binary or entry.path != jit_binary: + if recompile_jit_binary or os.path.abspath(entry.path) != os.path.abspath( + jit_binary + ): if entry.is_dir(follow_symlinks=False): shutil.rmtree(entry.path) else: @@ -54,13 +72,8 @@ def clean_keops(recompile_jit_binary=True, verbose=True): if verbose: KeOps_Message(f"{build_path} has been cleaned.") - from keopscore.get_keops_dll import get_keops_dll - - get_keops_dll.reset() - if use_cuda and recompile_jit_binary: - from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile - - Gpu_link_compile.compile_jit_compile_dll() + # Re-initialize the build folder after cleaning. + path.set_build_folder(reset_all=True) __all__ = [ diff --git a/keopscore/keopscore/utils/gpu_utils.py b/keopscore/keopscore/utils/gpu_utils.py index b1cbde01d..e9682a1e4 100644 --- a/keopscore/keopscore/utils/gpu_utils.py +++ b/keopscore/keopscore/utils/gpu_utils.py @@ -95,6 +95,3 @@ def add_crt_symlink_to_cuda_include_path(cuda_include_path, build_folder): ) if target is not None: os.symlink(target, crt_symlink) - print( - f"Created symlink for {file} in crt folder: {crt_symlink} -> {target}" - ) diff --git a/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py b/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py index 681d725d0..5dd210046 100644 --- a/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py +++ b/pykeops/pykeops/common/keops_io/nvrtc/LoadKeOps_nvrtc.py @@ -80,9 +80,9 @@ def compile_jit_binary(): This function compile the main .so entry point to keops_nvrt binder... """ compile_command = Gpu_link_compile.get_compile_command( + pykeopsconfig.pykeops_nvrtc_name(type="src"), + pykeopsconfig.pykeops_nvrtc_name(type="target"), extra_flags=pykeopsconfig.python_includes, - sourcename=pykeopsconfig.pykeops_nvrtc_name(type="src"), - dllname=pykeopsconfig.pykeops_nvrtc_name(type="target"), ) pyKeOps_Message("Compiling nvrtc binder for python", flush=True, end="", level=1) pyKeOps_Message( From 28f94d943dfe812ebd22ba19044c8e08daa9b77f Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 11 May 2026 01:58:53 +0200 Subject: [PATCH 51/98] Fix set_build_folder in pykeops --- pykeops/pykeops/__init__.py | 53 ++++++++++++++++++++----------------- 1 file changed, 28 insertions(+), 25 deletions(-) diff --git a/pykeops/pykeops/__init__.py b/pykeops/pykeops/__init__.py index 796106f95..72a0554a9 100644 --- a/pykeops/pykeops/__init__.py +++ b/pykeops/pykeops/__init__.py @@ -26,20 +26,6 @@ def set_verbose(val): default_device_id = 0 # default Gpu device number -from .common.keops_io.cpp.LoadKeOps_cpp import ( - compile_jit_binary as compile_cpp_jit_binary, - should_compile_binder as should_compile_cpp_binder, -) - -if should_compile_cpp_binder(): - compile_cpp_jit_binary() - -if pykeopsconfig.cuda.get_use_cuda(): - if not os.path.exists(pykeopsconfig.pykeops_nvrtc_name(type="target")): - from .common.keops_io.nvrtc.LoadKeOps_nvrtc import compile_jit_binary - - compile_jit_binary() - def clean_pykeops(recompile_jit_binaries=True): r""" @@ -77,19 +63,33 @@ def check_health(infos="all"): keopscore.config.check_health(infos=infos) -def set_build_folder(path=None): - import pykeops +def set_build_folder(path=None, reset_all=True): + from .common.keops_io import keops_binder + from .common.keops_io.cpp.LoadKeOps_cpp import ( + should_compile_binder as should_compile_cpp_binder, + ) - keopscore.set_build_folder(path) - keops_binder = pykeops.common.keops_io.keops_binder + # Set the build folder, reset keopscore cache and recompile JIT binaries if needed + keopscore.set_build_folder(path, reset_all=reset_all) + + # Reset the cache of all pykeops binders to ensure they will be reloaded from the new build folder for key in keops_binder: keops_binder[key].reset(new_save_folder=get_build_folder()) - if should_compile_cpp_binder(): - pykeops.common.keops_io.LoadKeOps_cpp.compile_jit_binary() - if pykeopsconfig.cuda.get_use_cuda() and not os.path.exists( - pykeopsconfig.pykeops_nvrtc_name(type="target") + + # Recompile pykeops binder binaries if needed + + if reset_all or should_compile_cpp_binder(): + from .common.keops_io.cpp import LoadKeOps_cpp + + LoadKeOps_cpp.compile_jit_binary() + + if reset_all or ( + pykeopsconfig.cuda.get_use_cuda() + and not os.path.exists(pykeopsconfig.pykeops_nvrtc_name(type="target")) ): - pykeops.common.keops_io.LoadKeOps_nvrtc.compile_jit_binary() + from .common.keops_io.nvrtc import LoadKeOps_nvrtc + + LoadKeOps_nvrtc.compile_jit_binary() def get_build_folder(): @@ -99,8 +99,11 @@ def get_build_folder(): if pykeopsconfig.numpy_found: from .numpy.test_install import test_numpy_bindings + if pykeopsconfig.torch_found: from .torch.test_install import test_torch_bindings -# next line is to ensure that cache file for formulas is loaded at import -from .common import keops_io + +# set the build folder and ensure it is in the python path + +set_build_folder(reset_all=False) From ee4f11496f009495ea7682dc14aca57cb62ec781 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 11 May 2026 02:08:22 +0200 Subject: [PATCH 52/98] add build path to include when symlink are created --- keopscore/keopscore/config/KeOpsPath.py | 5 +++-- keopscore/keopscore/utils/gpu_utils.py | 6 +++++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/keopscore/keopscore/config/KeOpsPath.py b/keopscore/keopscore/config/KeOpsPath.py index 4538a9273..bb87d36e6 100644 --- a/keopscore/keopscore/config/KeOpsPath.py +++ b/keopscore/keopscore/config/KeOpsPath.py @@ -130,9 +130,10 @@ def set_build_folder( Gpu_link_compile.compile_jit_compile_dll(force_recompile=reset_all) #### Add a symlink to the crt folder in keops include path if needed, to handle # cudatoolkit pip package - add_crt_symlink_to_cuda_include_path( + if add_crt_symlink_to_cuda_include_path( self.cuda.get_cuda_include_path(), self._build_folder - ) + ): + self.add_to_include_option(f" -I{self._build_folder}") def get_build_folder(self): return self._build_folder diff --git a/keopscore/keopscore/utils/gpu_utils.py b/keopscore/keopscore/utils/gpu_utils.py index e9682a1e4..0d89ce821 100644 --- a/keopscore/keopscore/utils/gpu_utils.py +++ b/keopscore/keopscore/utils/gpu_utils.py @@ -75,11 +75,12 @@ def add_crt_symlink_to_cuda_include_path(cuda_include_path, build_folder): for file in files_to_check ] if all(check): - return + return False # no need to create the symlink crt_folder = os.path.join(build_folder, "crt") os.makedirs(crt_folder, exist_ok=True) + created = False for file in files_to_check: crt_symlink = os.path.join(crt_folder, file) if os.path.exists(crt_symlink): @@ -95,3 +96,6 @@ def add_crt_symlink_to_cuda_include_path(cuda_include_path, build_folder): ) if target is not None: os.symlink(target, crt_symlink) + created = True + + return created From c88ca57cdf79b67fe3e7cc78c6e953ff831b6623 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 11 May 2026 02:20:38 +0200 Subject: [PATCH 53/98] typo in fstring --- keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py index dabfa79e1..9e58683b4 100644 --- a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py +++ b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py @@ -114,7 +114,7 @@ def compile_jit_compile_dll(force_recompile=False): out = KeOps_OS_Run(command) if out.returncode != 0: KeOps_Error( - f"Error compiling cuda/nvrtc binder {keopscore.config.keops_jit_compile_name(type="target")}. See compiler output above." + f"Error compiling cuda/nvrtc binder {keopscore.config.keops_jit_compile_name(type='target')}. See compiler output above." ) else: KeOps_Message("OK", use_tag=False, flush=True, level=1) From 1c9586a75b3ec7afa86c4c255b5aab1e07558641 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 11 May 2026 02:30:36 +0200 Subject: [PATCH 54/98] protect first compile with a try --- pykeops/pykeops/__init__.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/pykeops/pykeops/__init__.py b/pykeops/pykeops/__init__.py index 72a0554a9..ee0c44025 100644 --- a/pykeops/pykeops/__init__.py +++ b/pykeops/pykeops/__init__.py @@ -105,5 +105,11 @@ def get_build_folder(): # set the build folder and ensure it is in the python path - -set_build_folder(reset_all=False) +try: + set_build_folder(reset_all=False) +except Exception as e: + from .common.utils import pyKeOps_Warning + pyKeOps_Warning( + f"An error occurred while setting up KeOps: {e}. Use pykeops.check_health() to get details on the current configuration.", + level=1, + ) From 293dcadf174d17b67667a83bb1d73536faf41eef Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 11 May 2026 12:26:42 +0200 Subject: [PATCH 55/98] update version number --- keops_version | 2 +- pykeops/pykeops/__init__.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/keops_version b/keops_version index bb576dbde..7451fe42d 100644 --- a/keops_version +++ b/keops_version @@ -1 +1 @@ -2.3 +2.4b diff --git a/pykeops/pykeops/__init__.py b/pykeops/pykeops/__init__.py index ee0c44025..441eec74c 100644 --- a/pykeops/pykeops/__init__.py +++ b/pykeops/pykeops/__init__.py @@ -109,6 +109,7 @@ def get_build_folder(): set_build_folder(reset_all=False) except Exception as e: from .common.utils import pyKeOps_Warning + pyKeOps_Warning( f"An error occurred while setting up KeOps: {e}. Use pykeops.check_health() to get details on the current configuration.", level=1, From f9c042adc87a9569cbde3c221da324f2808b8719 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 11 May 2026 12:37:42 +0200 Subject: [PATCH 56/98] protect keopcore compilation with a try --- keopscore/keopscore/__init__.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/keopscore/keopscore/__init__.py b/keopscore/keopscore/__init__.py index 453d93962..741b7158d 100644 --- a/keopscore/keopscore/__init__.py +++ b/keopscore/keopscore/__init__.py @@ -9,7 +9,15 @@ import keopscore.config # Create build folder if it does not exist and set it as the default build folder -keopscore.config.path.set_build_folder(reset_all=False) +try: + keopscore.config.path.set_build_folder(reset_all=False) +except Exception as e: + from .utils.messages import KeOps_Warning + + KeOps_Warning( + f"An error occurred while setting up keopcore: {e}. Use keopscore.config.check_health() to get details on the current configuration.", + level=1, + ) # expose to the user set_build_folder = keopscore.config.path.set_build_folder From 31e1ba9be76c50379f4747ffba503d818241717a Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 11 May 2026 14:11:42 +0200 Subject: [PATCH 57/98] add new path in setup.py --- keopscore/setup.py | 2 ++ pykeops/setup.py | 5 ++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/keopscore/setup.py b/keopscore/setup.py index 5a4f04b16..5c8fccdc6 100644 --- a/keopscore/setup.py +++ b/keopscore/setup.py @@ -67,6 +67,8 @@ "readme.md", "licence.txt", "keops_version", + "binders/cpp/keops_cpu_runtime.cpp", + "binders/cpp/keops_cpu_runtime.h", "binders/nvrtc/keops_nvrtc.cpp", "binders/nvrtc/nvrtc_jit.cpp", "include/CudaSizes.h", diff --git a/pykeops/setup.py b/pykeops/setup.py index 822c6d599..166551a70 100644 --- a/pykeops/setup.py +++ b/pykeops/setup.py @@ -49,6 +49,8 @@ "pykeops", "pykeops.common", "pykeops.common.keops_io", + "pykeops.common.keops_io.nvrtc", + "pykeops.common.keops_io.cpp", "pykeops.numpy", "pykeops.numpy.cluster", "pykeops.numpy.generic", @@ -64,7 +66,8 @@ "readme.md", "licence.txt", "keops_version", - "common/keops_io/pykeops_nvrtc.cpp", + "common/keops_io/cpp/pykeops_cpp.cpp", + "common/keops_io/nvrtc/pykeops_nvrtc.cpp", ], }, install_requires=["numpy", "pybind11", "keopscore"], From 054a82d465e1ebab1d6ec8fa124f272c90a27a82 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 11 May 2026 14:30:29 +0200 Subject: [PATCH 58/98] remove unused method. Add verbosity to debuf colab... --- keopscore/keopscore/config/KeOpsPath.py | 1 - keopscore/keopscore/utils/gpu_utils.py | 4 +++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/keopscore/keopscore/config/KeOpsPath.py b/keopscore/keopscore/config/KeOpsPath.py index bb87d36e6..b88f0595f 100644 --- a/keopscore/keopscore/config/KeOpsPath.py +++ b/keopscore/keopscore/config/KeOpsPath.py @@ -22,7 +22,6 @@ class KeOpsPathConfig: _default_build_folder_name = "" _default_build_path = "" - _jit_binary = "" _include_options = "" path_env_vars = ("KEOPS_CACHE_FOLDER",) diff --git a/keopscore/keopscore/utils/gpu_utils.py b/keopscore/keopscore/utils/gpu_utils.py index 0d89ce821..fc5822bdc 100644 --- a/keopscore/keopscore/utils/gpu_utils.py +++ b/keopscore/keopscore/utils/gpu_utils.py @@ -67,7 +67,7 @@ def add_crt_symlink_to_cuda_include_path(cuda_include_path, build_folder): (included before the cuda toolkit) that mimick the correct behavior... """ - files_to_check = ["host_config.h"] + files_to_check = ["host_defines.h", "host_config.h", "device_functions.h"] cuda_include_path = [p for p in list(cuda_include_path) if os.path.isdir(p)] check = [ @@ -75,6 +75,7 @@ def add_crt_symlink_to_cuda_include_path(cuda_include_path, build_folder): for file in files_to_check ] if all(check): + print("crt headers already in cuda include path, no need to create symlink.") return False # no need to create the symlink crt_folder = os.path.join(build_folder, "crt") @@ -97,5 +98,6 @@ def add_crt_symlink_to_cuda_include_path(cuda_include_path, build_folder): if target is not None: os.symlink(target, crt_symlink) created = True + print(f"Created symlink for {file} in keops include path to handle cudatoolkit pip package: {crt_symlink} -> {target}") return created From 626ceaca344a48efa651dda6a4fe213ce58a9793 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Mon, 11 May 2026 14:34:51 +0200 Subject: [PATCH 59/98] change order in compilation / symlink --- keopscore/keopscore/config/KeOpsPath.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/keopscore/keopscore/config/KeOpsPath.py b/keopscore/keopscore/config/KeOpsPath.py index b88f0595f..0d58136ec 100644 --- a/keopscore/keopscore/config/KeOpsPath.py +++ b/keopscore/keopscore/config/KeOpsPath.py @@ -124,15 +124,18 @@ def set_build_folder( # Handle CUDA-specific recompilation if CUDA is used if self.cuda.get_use_cuda(): - from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile - - Gpu_link_compile.compile_jit_compile_dll(force_recompile=reset_all) #### Add a symlink to the crt folder in keops include path if needed, to handle # cudatoolkit pip package if add_crt_symlink_to_cuda_include_path( self.cuda.get_cuda_include_path(), self._build_folder ): self.add_to_include_option(f" -I{self._build_folder}") + + # Recompile the nvrtc binder if needed + from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile + + Gpu_link_compile.compile_jit_compile_dll(force_recompile=reset_all) + def get_build_folder(self): return self._build_folder From ad9b989daeaf2f3bb44c081945b49828fd55bf0f Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Fri, 22 May 2026 19:25:01 +0200 Subject: [PATCH 60/98] add checks for cuda header --- .../keopscore/binders/nvrtc/keops_nvrtc.cpp | 10 ++-- .../keopscore/binders/nvrtc/nvrtc_jit.cpp | 7 ++- keopscore/keopscore/config/Cuda.py | 39 ++++++++++---- keopscore/keopscore/config/OpenMP.py | 2 +- keopscore/keopscore/utils/path_utils.py | 54 +++++++++++++++---- 5 files changed, 83 insertions(+), 29 deletions(-) diff --git a/keopscore/keopscore/binders/nvrtc/keops_nvrtc.cpp b/keopscore/keopscore/binders/nvrtc/keops_nvrtc.cpp index 78a62b648..bf87be432 100644 --- a/keopscore/keopscore/binders/nvrtc/keops_nvrtc.cpp +++ b/keopscore/keopscore/binders/nvrtc/keops_nvrtc.cpp @@ -12,16 +12,19 @@ // ./keopscore/binders/nvrtc/keops_nvrtc.cpp -o // keops_nvrtc.cpython-310-x86_64-linux-gnu.so -#include + #include #include #include -#include #include #include #include #include -// #include + + +#include +#include +#include #define C_CONTIGUOUS 1 #define USE_HALF 0 @@ -32,7 +35,6 @@ #include "include/utils_pe.h" #include "include/CudaSizes.h" -#include // ------------------------------------------------------------------------ // DevicePointer wrapper diff --git a/keopscore/keopscore/binders/nvrtc/nvrtc_jit.cpp b/keopscore/keopscore/binders/nvrtc/nvrtc_jit.cpp index 439481110..e0215366d 100644 --- a/keopscore/keopscore/binders/nvrtc/nvrtc_jit.cpp +++ b/keopscore/keopscore/binders/nvrtc/nvrtc_jit.cpp @@ -13,16 +13,16 @@ // /home/bcharlier/projets/keops/keops/keops/binders/nvrtc/keops_nvrtc.cpp -o // keops_nvrtc.cpython-310-x86_64-linux-gnu.so -#include #include #include -#include #include #include #include #include #include -// #include + +#include +#include #define C_CONTIGUOUS 1 #define USE_HALF 0 @@ -33,7 +33,6 @@ #include "include/utils_pe.h" #include "include/CudaSizes.h" -#include extern "C" int Compile(const char *target_file_name, const char *cu_code, int use_half, int use_fast_math, int device_id, diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index 74f6604d7..1388328a7 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -52,8 +52,9 @@ class CudaConfig: ] pip_suffixes = ( - "nvidia/cuda_runtime", - "nvidia/cuda_nvrtc", + os.path.join("nvidia", "cu13"), # for CUDA 13.X + # os.path.join("nvidia", "cuda_runtime"), # for CUDA 12.X but broken due to missing crt subfolder + # os.path.join("nvidia", "cuda_nvrtc") # for CUDA 12.X but broken due to missing crt subfolder ) system_suffixes = ( @@ -177,6 +178,11 @@ def _find_and_load_libcuda(self): False, "libcuda not found. Make sure the CUDA driver is installed and accessible.", ) + if not self._libcuda_info.get("header"): + return ( + False, + f"{self._libcuda_info['header_basename']} not found. Make sure the CUDA headers are installed and accessible.", + ) try: libcuda = ctypes.CDLL(libcuda_path, mode=ctypes.RTLD_GLOBAL) @@ -232,6 +238,11 @@ def _find_and_load_libnvrtc(self): False, "libnvrtc not found. Make sure the CUDA toolkit is installed and accessible.", ) + if not self._libnvrtc_info.get("header"): + return ( + False, + f"{self._libnvrtc_info['header_basename']} not found. Make sure the CUDA headers are installed and accessible.", + ) try: libnvrtc_handle = ctypes.CDLL(libnvrtc_path, mode=ctypes.RTLD_GLOBAL) @@ -468,14 +479,15 @@ def set_cuda_include_path(self): if not self.get_use_cuda(): self._cuda_include_path = "" else: - self._cuda_include_path = list( - set( - [ - os.path.dirname(self._libnvrtc_info["header"]), - os.path.dirname(self._libcuda_info["header"]), - ] + include_dirs = [ + os.path.dirname(header) + for header in ( + self._libnvrtc_info.get("header"), + self._libcuda_info.get("header"), ) - ) + if header + ] + self._cuda_include_path = list(set(include_dirs)) if include_dirs else "" def get_cuda_include_path(self): """ @@ -555,7 +567,14 @@ def print_preprocessing_options(self): def set_linking_options(self): """Set the Linking option for nvrt/cuda entry point compilation.""" - self._linking_options = f"-L{self.get_libcuda_folder()} -L{self.get_libnvrtc_folder()} -lcuda -lnvrtc" + link_options = [] + for library_path, fallback_name in ( + (self.get_libcuda_path(), "cuda"), + (self.get_libnvrtc_path(), "nvrtc"), + ): + link_options.append(library_path if library_path else f"-l{fallback_name}") + + self._linking_options = " ".join(link_options) def get_linking_options(self): """Get the Linking option for nvrt/cuda entry point compilation.""" diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 6260d4360..72c021c5a 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -276,7 +276,7 @@ def check_compiler_for_openmp(self): def set_compile_options(self): # Apple clang does not support -fopenmp directly; -Xpreprocessor is required. if self.cxx.get_use_Apple_clang(): - self._compile_options += "-Xpreprocessor " + self._compile_options += "-Xpreprocessor" self._compile_options += "-fopenmp" diff --git a/keopscore/keopscore/utils/path_utils.py b/keopscore/keopscore/utils/path_utils.py index a76a87670..1da2babd8 100644 --- a/keopscore/keopscore/utils/path_utils.py +++ b/keopscore/keopscore/utils/path_utils.py @@ -54,21 +54,55 @@ def _python_package_roots(): def _ordered_search_roots( - env_vars=(), pip_suffixes=(), conda_root=None, system_roots=() + env_vars=(), pip_suffixes=(), conda_root=None, system_roots=(), order=None ): - """Return roots ordered by explicit env vars, pip, conda, then system paths.""" + """Return normalized search roots in a caller-defined category order. + + Args: + env_vars (tuple[str] | list[str]): Environment variable names whose values + should be considered as root directories. + pip_suffixes (tuple[str] | list[str]): Relative suffixes appended to Python + package roots discovered from the current interpreter. + conda_root (str | None): Name of an environment variable that points to a + Conda root directory. + system_roots (tuple[str] | list[str]): Fallback root directories to append. + order (tuple[str] | list[str] | str | None): Ordered categories to apply. + Allowed items are ``env_vars``, ``pip_suffixes``, ``conda_root``, and + ``system_roots``. A comma-separated string is also accepted. + + Returns: + list[pathlib.Path]: Existing candidates, deduplicated in first-seen order. + """ + + if order is None: + order = ("env_vars", "conda_root", "system_roots", "pip_suffixes") + else: + isinstance(order, str) and (order := tuple(order.split(","))) + for item in order: + if item not in ( + "env_vars", + "pip_suffixes", + "conda_root", + "system_roots", + ): + raise ValueError(f"Invalid order item: {item}") + roots = [] - if env_vars: - roots.extend(_env_roots(env_vars)) + for item in order: + if item == "env_vars" and env_vars: + roots.extend(_env_roots(env_vars)) + + if item == "pip_suffixes" and pip_suffixes: + # Pip wheels such as nvidia-cuda-runtime expose libraries under site-packages. + roots.extend(_path_candidates(_python_package_roots(), pip_suffixes)) + + if item == "conda_root" and conda_root: + roots.extend(_env_roots((conda_root,))) - if pip_suffixes: - # Pip wheels such as nvidia-cuda-runtime expose libraries under site-packages. - roots.extend(_path_candidates(_python_package_roots(), pip_suffixes)) + if item == "system_roots" and system_roots: + roots.extend(system_roots) - if conda_root: - roots.extend(_env_roots((conda_root,))) - roots.extend(system_roots) return _unique_paths(roots) From 0f5f582fcacbb9b4fec382e4372186f06e04a639 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Fri, 22 May 2026 19:27:34 +0200 Subject: [PATCH 61/98] remove crt folder: pypi cudatoolkit 13 includes it --- keopscore/keopscore/config/KeOpsPath.py | 8 ----- keopscore/keopscore/utils/gpu_utils.py | 44 ------------------------- 2 files changed, 52 deletions(-) diff --git a/keopscore/keopscore/config/KeOpsPath.py b/keopscore/keopscore/config/KeOpsPath.py index 0d58136ec..4f73afd23 100644 --- a/keopscore/keopscore/config/KeOpsPath.py +++ b/keopscore/keopscore/config/KeOpsPath.py @@ -124,18 +124,10 @@ def set_build_folder( # Handle CUDA-specific recompilation if CUDA is used if self.cuda.get_use_cuda(): - #### Add a symlink to the crt folder in keops include path if needed, to handle - # cudatoolkit pip package - if add_crt_symlink_to_cuda_include_path( - self.cuda.get_cuda_include_path(), self._build_folder - ): - self.add_to_include_option(f" -I{self._build_folder}") - # Recompile the nvrtc binder if needed from keopscore.binders.nvrtc.Gpu_link_compile import Gpu_link_compile Gpu_link_compile.compile_jit_compile_dll(force_recompile=reset_all) - def get_build_folder(self): return self._build_folder diff --git a/keopscore/keopscore/utils/gpu_utils.py b/keopscore/keopscore/utils/gpu_utils.py index fc5822bdc..be1e82d63 100644 --- a/keopscore/keopscore/utils/gpu_utils.py +++ b/keopscore/keopscore/utils/gpu_utils.py @@ -57,47 +57,3 @@ def custom_cuda_include_fp16_path(): if not os.path.isfile(fp16_header_path): pack_header(fp16_header, orig_cuda_include_fp16_path(), build_folder) return build_folder - - -def add_crt_symlink_to_cuda_include_path(cuda_include_path, build_folder): - """ - The cuda_fp16.h header includes crt/host_config.h and crt/device_runtime_api.h, but these - headers are not properly referenced in the cuda include path ship in the pip version of - cudatoolkit. This function adds a symlink to the crt folder in keops include path - (included before the cuda toolkit) that mimick the correct behavior... - """ - - files_to_check = ["host_defines.h", "host_config.h", "device_functions.h"] - cuda_include_path = [p for p in list(cuda_include_path) if os.path.isdir(p)] - - check = [ - any([os.path.isfile(os.path.join(p, "crt", file)) for p in cuda_include_path]) - for file in files_to_check - ] - if all(check): - print("crt headers already in cuda include path, no need to create symlink.") - return False # no need to create the symlink - - crt_folder = os.path.join(build_folder, "crt") - os.makedirs(crt_folder, exist_ok=True) - - created = False - for file in files_to_check: - crt_symlink = os.path.join(crt_folder, file) - if os.path.exists(crt_symlink): - continue - # Find the actual file in one of the cuda include paths - target = next( - ( - os.path.join(p, file) - for p in cuda_include_path - if os.path.isfile(os.path.join(p, file)) - ), - None, - ) - if target is not None: - os.symlink(target, crt_symlink) - created = True - print(f"Created symlink for {file} in keops include path to handle cudatoolkit pip package: {crt_symlink} -> {target}") - - return created From ef72d813b2db92f9bfb658c8770cfa8f08ab5412 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Fri, 22 May 2026 19:54:04 +0200 Subject: [PATCH 62/98] fix imports --- keopscore/keopscore/config/Cuda.py | 22 +++++++++------------- keopscore/keopscore/config/KeOpsPath.py | 4 +--- 2 files changed, 10 insertions(+), 16 deletions(-) diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index 1388328a7..019134a7a 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -475,19 +475,15 @@ def print_cuda_version(self): # CUDA Include Path def set_cuda_include_path(self): """Set the CUDA include path by searching for cuda.h and nvrtc.h.""" - # This is done in find_cuda_install since it relies on the cuda installation info which is only available after checking library availability. - if not self.get_use_cuda(): - self._cuda_include_path = "" - else: - include_dirs = [ - os.path.dirname(header) - for header in ( - self._libnvrtc_info.get("header"), - self._libcuda_info.get("header"), - ) - if header - ] - self._cuda_include_path = list(set(include_dirs)) if include_dirs else "" + include_dirs = [ + os.path.dirname(header) + for header in ( + self._libnvrtc_info.get("header"), + self._libcuda_info.get("header"), + ) + if header + ] + self._cuda_include_path = list(set(include_dirs)) if include_dirs else "" def get_cuda_include_path(self): """ diff --git a/keopscore/keopscore/config/KeOpsPath.py b/keopscore/keopscore/config/KeOpsPath.py index 4f73afd23..abb2e8463 100644 --- a/keopscore/keopscore/config/KeOpsPath.py +++ b/keopscore/keopscore/config/KeOpsPath.py @@ -1,12 +1,10 @@ import os import sys import sysconfig -import warnings import keopscore -from keopscore.utils.gpu_utils import add_crt_symlink_to_cuda_include_path from keopscore.utils.path_utils import ensure_directory -from keopscore.utils.messages import KeOps_Warning, not_found_str, print_envs +from keopscore.utils.messages import not_found_str, print_envs class KeOpsPathConfig: From 73e81ede541a0b3170dab9f744e546da70d15849 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sat, 23 May 2026 16:44:13 +0200 Subject: [PATCH 63/98] test mac python3.12 --- .cirrus.yml | 26 +++++++++++--------------- pytest.sh | 15 ++++++--------- 2 files changed, 17 insertions(+), 24 deletions(-) diff --git a/.cirrus.yml b/.cirrus.yml index fc68ca1b0..d9ad2a98e 100644 --- a/.cirrus.yml +++ b/.cirrus.yml @@ -1,25 +1,21 @@ -task: - name: Tests (linux) - container: - image: ubuntu:latest - install_script: - - apt-get update - - apt-get -y install python3 python3-venv python3-dev g++ - test_script: - - bash ./pytest.sh - +#task: +# name: Tests (linux) +# container: +# image: ubuntu:latest +# install_script: +# - apt-get update +# - apt-get -y install python3 python3-venv python3-dev g++ +# test_script: +# - bash ./pytest.sh task: name: Tests (macOS) macos_instance: - image: ghcr.io/cirruslabs/macos-ventura-xcode:latest + image: ghcr.io/cirruslabs/macos-runner:tahoe script: - brew install libomp - - brew install python@3.11 - - $(brew --prefix python@3.11)/bin/python3.11 -m venv .test_venv - - source .test_venv/bin/activate - - pip install --upgrade pip + - python --version - sh ./pytest.sh diff --git a/pytest.sh b/pytest.sh index a4a733fe3..635e68313 100755 --- a/pytest.sh +++ b/pytest.sh @@ -65,17 +65,13 @@ install_editable_package() { } clean_pykeops_cache() { - local tmp_dir - log_verbose "-- Cleaning pykeops..." - tmp_dir="$(mktemp -d)" - - ( - cd "${tmp_dir}" - "${PYTHON_BIN}" -c 'import pykeops; pykeops.clean_pykeops()' - ) + "${PYTHON_BIN}" -c 'import pykeops; pykeops.clean_pykeops()' +} - rmdir "${tmp_dir}" +run_pykeops_health_check() { + echo "-- Running pykeops.check_health()..." + "${PYTHON_BIN}" -c 'import pykeops; pykeops.check_health()' } run_test_suite() { @@ -91,6 +87,7 @@ main() { prepare_python_environment install_editable_package "keopscore" "${SCRIPT_DIR}/keopscore" install_editable_package "pykeops" "${SCRIPT_DIR}/pykeops[test]" + run_pykeops_health_check clean_pykeops_cache run_test_suite "keopscore" "keopscore/keopscore/test/" run_test_suite "pykeops" "pykeops/pykeops/test/" From ce4e5b3f41304333de8f5266fda30c77217815ed Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sat, 23 May 2026 16:44:13 +0200 Subject: [PATCH 64/98] test mac python3.12 --- .cirrus.yml | 26 +++++++++++--------------- pytest.sh | 15 ++++++--------- 2 files changed, 17 insertions(+), 24 deletions(-) diff --git a/.cirrus.yml b/.cirrus.yml index fc68ca1b0..680af5406 100644 --- a/.cirrus.yml +++ b/.cirrus.yml @@ -1,25 +1,21 @@ -task: - name: Tests (linux) - container: - image: ubuntu:latest - install_script: - - apt-get update - - apt-get -y install python3 python3-venv python3-dev g++ - test_script: - - bash ./pytest.sh - +#task: +# name: Tests (linux) +# container: +# image: ubuntu:latest +# install_script: +# - apt-get update +# - apt-get -y install python3 python3-venv python3-dev g++ +# test_script: +# - bash ./pytest.sh task: name: Tests (macOS) macos_instance: - image: ghcr.io/cirruslabs/macos-ventura-xcode:latest + image: ghcr.io/cirruslabs/macos-runner:tahoe script: - brew install libomp - - brew install python@3.11 - - $(brew --prefix python@3.11)/bin/python3.11 -m venv .test_venv - - source .test_venv/bin/activate - - pip install --upgrade pip + - python3 --version - sh ./pytest.sh diff --git a/pytest.sh b/pytest.sh index a4a733fe3..635e68313 100755 --- a/pytest.sh +++ b/pytest.sh @@ -65,17 +65,13 @@ install_editable_package() { } clean_pykeops_cache() { - local tmp_dir - log_verbose "-- Cleaning pykeops..." - tmp_dir="$(mktemp -d)" - - ( - cd "${tmp_dir}" - "${PYTHON_BIN}" -c 'import pykeops; pykeops.clean_pykeops()' - ) + "${PYTHON_BIN}" -c 'import pykeops; pykeops.clean_pykeops()' +} - rmdir "${tmp_dir}" +run_pykeops_health_check() { + echo "-- Running pykeops.check_health()..." + "${PYTHON_BIN}" -c 'import pykeops; pykeops.check_health()' } run_test_suite() { @@ -91,6 +87,7 @@ main() { prepare_python_environment install_editable_package "keopscore" "${SCRIPT_DIR}/keopscore" install_editable_package "pykeops" "${SCRIPT_DIR}/pykeops[test]" + run_pykeops_health_check clean_pykeops_cache run_test_suite "keopscore" "keopscore/keopscore/test/" run_test_suite "pykeops" "pykeops/pykeops/test/" From 4d7408b21312b8749d53bbe5698a7b469c1a716a Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sat, 23 May 2026 16:59:57 +0200 Subject: [PATCH 65/98] fix pytest --- pytest.sh | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/pytest.sh b/pytest.sh index 635e68313..d7dbe1bc3 100755 --- a/pytest.sh +++ b/pytest.sh @@ -64,14 +64,22 @@ install_editable_package() { pip install -e "${package_path}" } +run_python_outside_repo() { + local python_code="$1" + ( + cd /tmp + "${PYTHON_BIN}" -c "${python_code}" + ) +} + clean_pykeops_cache() { log_verbose "-- Cleaning pykeops..." - "${PYTHON_BIN}" -c 'import pykeops; pykeops.clean_pykeops()' + run_python_outside_repo 'import pykeops; pykeops.clean_pykeops()' } run_pykeops_health_check() { echo "-- Running pykeops.check_health()..." - "${PYTHON_BIN}" -c 'import pykeops; pykeops.check_health()' + run_python_outside_repo 'import pykeops; pykeops.check_health()' } run_test_suite() { From 2833979db0a752c292fcad2f64a974122414a167 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sat, 23 May 2026 17:27:36 +0200 Subject: [PATCH 66/98] fix malformed compile flag. --- keopscore/keopscore/config/OpenMP.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 72c021c5a..f687f8012 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -276,9 +276,9 @@ def check_compiler_for_openmp(self): def set_compile_options(self): # Apple clang does not support -fopenmp directly; -Xpreprocessor is required. if self.cxx.get_use_Apple_clang(): - self._compile_options += "-Xpreprocessor" + self._compile_options += " -Xpreprocessor" - self._compile_options += "-fopenmp" + self._compile_options += " -fopenmp" def get_compile_options(self): return self._compile_options From c3d0b45d9171f5ac808e857f32e39e8f67057baf Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sat, 23 May 2026 19:00:56 +0200 Subject: [PATCH 67/98] add filter to CXXFLAGS. Fix brew detection message. --- keopscore/keopscore/config/CxxCompiler.py | 26 +++++++++++++++++------ keopscore/keopscore/config/KeOpsPath.py | 3 ++- keopscore/keopscore/config/Platform.py | 2 +- 3 files changed, 23 insertions(+), 8 deletions(-) diff --git a/keopscore/keopscore/config/CxxCompiler.py b/keopscore/keopscore/config/CxxCompiler.py index a157a0352..25cc75d86 100644 --- a/keopscore/keopscore/config/CxxCompiler.py +++ b/keopscore/keopscore/config/CxxCompiler.py @@ -1,4 +1,5 @@ import os +import re import shutil from keopscore.utils.messages import print_envs @@ -177,15 +178,28 @@ def get_dynamic_loader_linking_options(self): """Return extra linker options needed for dlopen/dlsym support.""" if self.platform.get_platform() == "Linux": return "-ldl" - elif self.platform.get_platform() == "Darwin": - # macOS: dlopen/dlsym are in standard C library, no extra flag needed - return "" + # macOS: dlopen/dlsym are in standard C library, no extra flag needed return "" - # C++ Environment Flags. Unused yet... + # C++ Environment Flags. def set_cxx_env_flags(self): - """Recover the C++ environment flags.""" - self._cxx_env_flags = os.getenv("CXXFLAGS") if "CXXFLAGS" in os.environ else "" + """Recover CXXFLAGS and sanitize incompatible flags when needed.""" + cxxflags = os.getenv("CXXFLAGS", "") + + # Apple clang does not support some GCC tuning flags that may leak in + # from external environments (e.g. conda, torch build stacks). + if ( + cxxflags + and self.platform.get_platform() == "Darwin" + and self.get_use_Apple_clang() + ): + cxxflags = re.sub( + r"(?:^|\s)-(?:march|mtune|mfpmath)(?:=\S+|\s+\S+)?", + "", + cxxflags, + ).strip() + + self._cxx_env_flags = cxxflags def get_cxx_env_flags(self): """Get the C++ environment flags.""" diff --git a/keopscore/keopscore/config/KeOpsPath.py b/keopscore/keopscore/config/KeOpsPath.py index abb2e8463..6ad097f7a 100644 --- a/keopscore/keopscore/config/KeOpsPath.py +++ b/keopscore/keopscore/config/KeOpsPath.py @@ -59,8 +59,9 @@ def set_keops_cache_folder(self): if os.getenv("KEOPS_CACHE_FOLDER"): cache_folder = os.getenv("KEOPS_CACHE_FOLDER") else: # fallback to default cache folder in user home directory + version = getattr(keopscore, "__version__", "") cache_folder = os.path.join( - os.path.expanduser("~"), ".cache", f"keops{keopscore.__version__}" + os.path.expanduser("~"), ".cache", f"keops{version}" ) os.makedirs(cache_folder, exist_ok=True) diff --git a/keopscore/keopscore/config/Platform.py b/keopscore/keopscore/config/Platform.py index 5dc0e90f8..ab17145b7 100644 --- a/keopscore/keopscore/config/Platform.py +++ b/keopscore/keopscore/config/Platform.py @@ -132,7 +132,7 @@ def set_brew_prefix(self): out = KeOps_OS_Run(f"brew --prefix", print_warning=False) self._brew_prefix = ( - out.stdout.decode("utf-8").strip() if out.stderr != b"" else None + out.stdout.decode("utf-8").strip() if out.returncode == 0 else None ) def get_brew_prefix(self): From 23deddf75724d33df0468a55a81d9387fce310ff Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sat, 23 May 2026 19:21:32 +0200 Subject: [PATCH 68/98] move ci to github action. --- .cirrus.yml | 21 --------------------- .github/workflows/black.yml | 19 +++++++++++++------ .github/workflows/cuda_test.yml | 15 --------------- .github/workflows/linux_cpu_test.yml | 21 +++++++++++++++++++++ .github/workflows/linux_cuda_test.yml | 16 ++++++++++++++++ .github/workflows/macOs_test.yml | 24 ++++++++++++++++++++++++ 6 files changed, 74 insertions(+), 42 deletions(-) delete mode 100644 .cirrus.yml delete mode 100644 .github/workflows/cuda_test.yml create mode 100644 .github/workflows/linux_cpu_test.yml create mode 100644 .github/workflows/linux_cuda_test.yml create mode 100644 .github/workflows/macOs_test.yml diff --git a/.cirrus.yml b/.cirrus.yml deleted file mode 100644 index 680af5406..000000000 --- a/.cirrus.yml +++ /dev/null @@ -1,21 +0,0 @@ -#task: -# name: Tests (linux) -# container: -# image: ubuntu:latest -# install_script: -# - apt-get update -# - apt-get -y install python3 python3-venv python3-dev g++ -# test_script: -# - bash ./pytest.sh - -task: - name: Tests (macOS) - macos_instance: - image: ghcr.io/cirruslabs/macos-runner:tahoe - - script: - - brew install libomp - - python3 --version - - sh ./pytest.sh - - diff --git a/.github/workflows/black.yml b/.github/workflows/black.yml index 69934f854..27bb56b8b 100644 --- a/.github/workflows/black.yml +++ b/.github/workflows/black.yml @@ -1,12 +1,19 @@ -name: Lint +name: Format Check -on: [push, pull_request] +on: + push: + pull_request: jobs: - lint: + format: + name: Format (Black) runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 - - uses: actions/setup-python@v2 - - uses: psf/black@stable + - name: Checkout + uses: actions/checkout@v4 + + - name: Black format check + uses: psf/black@stable + with: + options: "--check ." diff --git a/.github/workflows/cuda_test.yml b/.github/workflows/cuda_test.yml deleted file mode 100644 index 1d727b2f1..000000000 --- a/.github/workflows/cuda_test.yml +++ /dev/null @@ -1,15 +0,0 @@ -# Launch test on a cuda machine -name: Cuda test - -on: - push - -jobs: - test: - runs-on: self-hosted - steps: - - name: Clone Repository - uses: actions/checkout@v2 - - name: Test with pytest - run: | - sh ./pytest.sh diff --git a/.github/workflows/linux_cpu_test.yml b/.github/workflows/linux_cpu_test.yml new file mode 100644 index 000000000..51821b78b --- /dev/null +++ b/.github/workflows/linux_cpu_test.yml @@ -0,0 +1,21 @@ +name: Linux CPU Test + +on: + push: + pull_request: + +jobs: + tests-linux: + name: Tests (Linux CPU) + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.x" + + - name: Run tests + run: bash ./pytest.sh \ No newline at end of file diff --git a/.github/workflows/linux_cuda_test.yml b/.github/workflows/linux_cuda_test.yml new file mode 100644 index 000000000..9cef38911 --- /dev/null +++ b/.github/workflows/linux_cuda_test.yml @@ -0,0 +1,16 @@ +name: Linux CUDA Test + +on: + push: + pull_request: + +jobs: + tests-linux-cuda: + name: Tests (Linux CUDA) + runs-on: self-hosted + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Run tests + run: bash ./pytest.sh diff --git a/.github/workflows/macOs_test.yml b/.github/workflows/macOs_test.yml new file mode 100644 index 000000000..aa36f9f44 --- /dev/null +++ b/.github/workflows/macOs_test.yml @@ -0,0 +1,24 @@ +name: macOS Test + +on: + push: + pull_request: + +jobs: + tests-macos: + name: Tests (macOS CPU) + runs-on: macos-15 + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.x" + + - name: Install dependencies + run: brew install libomp + + - name: Run tests + run: bash ./pytest.sh From 3d6cda1a0a3432e04cba0b3bb5fa05a677c1ed9a Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sat, 23 May 2026 19:26:59 +0200 Subject: [PATCH 69/98] fix path insertion to include list --- keopscore/keopscore/config/OpenMP.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index f687f8012..aa8382357 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -80,9 +80,7 @@ def __init__(self, platform, cxx): if self.platform.get_brew_prefix(): self._openmp_system_roots.insert( 0, - [ - self.platform.get_brew_prefix(), - ], + self.platform.get_brew_prefix(), ) # Detect OpenMP and set related configuration variables From 4bc2a20b65988bd8fcb6a151b8433f7570dc261d Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sat, 23 May 2026 19:50:05 +0200 Subject: [PATCH 70/98] check macOS --- .cirrus.yml | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 .cirrus.yml diff --git a/.cirrus.yml b/.cirrus.yml new file mode 100644 index 000000000..909977052 --- /dev/null +++ b/.cirrus.yml @@ -0,0 +1,18 @@ +linux_task: + name: Tests (linux) + container: + image: ubuntu:latest + install_script: + - apt-get update + - apt-get -y install python3 python3-venv python3-dev g++ + test_script: + - bash ./pytest.sh + +macos_task: + name: Tests (macOS) + macos_instance: + image: ghcr.io/cirruslabs/macos-runner:tahoe + script: + - brew install libomp + - python3 --version + - bash ./pytest.sh From 98b959be11188651bc4e8bdda19b1a426c2e5e24 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sat, 23 May 2026 20:09:39 +0200 Subject: [PATCH 71/98] add both intel and arm for macOs test --- .github/workflows/macOs_test.yml | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/.github/workflows/macOs_test.yml b/.github/workflows/macOs_test.yml index aa36f9f44..ec66f4ed6 100644 --- a/.github/workflows/macOs_test.yml +++ b/.github/workflows/macOs_test.yml @@ -6,8 +6,15 @@ on: jobs: tests-macos: - name: Tests (macOS CPU) - runs-on: macos-15 + name: Tests (macOS ${{ matrix.arch }}) + strategy: + matrix: + include: + - arch: ARM + os: macos-latest + - arch: Intel + os: macos-26-intel + runs-on: ${{ matrix.os }} steps: - name: Checkout uses: actions/checkout@v4 From dd48a25f2114e40a532364599dc90bcf3b30f3bd Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sat, 23 May 2026 20:30:24 +0200 Subject: [PATCH 72/98] pin python versoin on CI macOs intel --- .github/workflows/macOs_test.yml | 31 +++++++++++++++++++++---------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/.github/workflows/macOs_test.yml b/.github/workflows/macOs_test.yml index ec66f4ed6..2a9ad6810 100644 --- a/.github/workflows/macOs_test.yml +++ b/.github/workflows/macOs_test.yml @@ -5,16 +5,9 @@ on: pull_request: jobs: - tests-macos: - name: Tests (macOS ${{ matrix.arch }}) - strategy: - matrix: - include: - - arch: ARM - os: macos-latest - - arch: Intel - os: macos-26-intel - runs-on: ${{ matrix.os }} + tests-macos-latest: + name: Tests (macOS latest) + runs-on: macos-latest steps: - name: Checkout uses: actions/checkout@v4 @@ -29,3 +22,21 @@ jobs: - name: Run tests run: bash ./pytest.sh + + tests-macos-intel: + name: Tests (macOS Intel) + runs-on: macos-26-intel + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install dependencies + run: brew install libomp + + - name: Run tests + run: bash ./pytest.sh From 5c88413dc6cc14bdeb21868c29d3606d5407e9af Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 24 May 2026 00:38:32 +0200 Subject: [PATCH 73/98] adapt openmp discovery strategy with conda --- .github/workflows/macOs_test.yml | 2 +- keopscore/keopscore/config/Cuda.py | 11 +++--- keopscore/keopscore/config/OpenMP.py | 57 ++++++++++++++++------------ 3 files changed, 38 insertions(+), 32 deletions(-) diff --git a/.github/workflows/macOs_test.yml b/.github/workflows/macOs_test.yml index 2a9ad6810..29c8e9207 100644 --- a/.github/workflows/macOs_test.yml +++ b/.github/workflows/macOs_test.yml @@ -33,7 +33,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v5 with: - python-version: "3.12" + python-version: "3.10" - name: Install dependencies run: brew install libomp diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index 019134a7a..a9f205633 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -564,11 +564,10 @@ def set_linking_options(self): """Set the Linking option for nvrt/cuda entry point compilation.""" link_options = [] - for library_path, fallback_name in ( - (self.get_libcuda_path(), "cuda"), - (self.get_libnvrtc_path(), "nvrtc"), - ): - link_options.append(library_path if library_path else f"-l{fallback_name}") + for lib_info in [self._libcuda_info, self._libnvrtc_info]: + link_options.append( + lib_info["library"] if lib_info["library"] else f"-l{lib_info['name']}" + ) self._linking_options = " ".join(link_options) @@ -665,7 +664,7 @@ def print_all(self): if __name__ == "__main__": from keopscore.config.Platform import PlatformConfig from keopscore.config.CxxCompiler import CxxCompilerConfig - from keopscore.config.OpenMP import OpenMPConfig + # from keopscore.config.OpenMP import OpenMPConfig platform_info = PlatformConfig() platform_info.print_all() diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index aa8382357..023298c11 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -98,28 +98,29 @@ def find_install_path(self, lib_dict_info): result = lib_dict_info.copy() # First try to find OpenMP library using standard names via ctypes /cxx compiler. - for name, header_basename in zip( - lib_dict_info["name"], lib_dict_info["header_basename"] - ): - result["library"] = _find_library_by_names(name) - result["header"] = get_include_file_abspath( - header_basename, self.cxx.get_cxx_compiler() - ) - - #### - KeOps_Message("OpenMP library search using standard names:", level=2) - KeOps_Message(f" Trying library name: {name}", level=2) - KeOps_Message( - f" Found library path: {result['library'] or not_found_str}", level=2 - ) - KeOps_Message(f" Trying header name: {header_basename}", level=2) - KeOps_Message( - f" Found header path: {result['header'] or not_found_str}", level=2 - ) - #### - - if result["library"] and result["header"]: - return result + if not self.platform.get_env_type().startswith("conda"): + for name, header_basename in zip( + lib_dict_info["name"], lib_dict_info["header_basename"] + ): + result["library"] = _find_library_by_names((name,)) + result["header"] = get_include_file_abspath( + header_basename, self.cxx.get_cxx_compiler() + ) + + #### + KeOps_Message("OpenMP library search using standard names:", level=2) + KeOps_Message(f" Trying library name: {name}", level=2) + KeOps_Message( + f" Found library path: {result['library'] or not_found_str}", level=2 + ) + KeOps_Message(f" Trying header name: {header_basename}", level=2) + KeOps_Message( + f" Found header path: {result['header'] or not_found_str}", level=2 + ) + #### + + if result["library"] and result["header"]: + return result # If that fails, search for OpenMP headers and libraries in common locations. candidate_roots = _ordered_search_roots( @@ -128,6 +129,7 @@ def find_install_path(self, lib_dict_info): system_roots=self._openmp_system_roots, ) + # libraries result["library"] = _first_matching_file( _path_candidates(candidate_roots, self.openmp_library_suffixes), lib_dict_info["lib_basename_candidate"], @@ -144,7 +146,7 @@ def find_install_path(self, lib_dict_info): ) #### - # Finally, search for OpenMP headers in common locations. + # headers result["header"] = _first_matching_file( _path_candidates(candidate_roots, self._openmp_include_sufixes), lib_dict_info["header_basename"], @@ -165,8 +167,13 @@ def find_install_path(self, lib_dict_info): def _omp_is_available(self): self._omp_info = self.find_install_path(self._omp_info) - liomp_path = os.path.exists(self._omp_info["library"]) and os.path.exists( - self._omp_info["header"] + library_path = self._omp_info.get("library") + header_path = self._omp_info.get("header") + liomp_path = bool( + library_path + and header_path + and os.path.exists(library_path) + and os.path.exists(header_path) ) if not liomp_path: KeOps_Warning( From d207e95cf2acc5bb2b70fb4d84599b0eaa13cd33 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 24 May 2026 09:59:13 +0200 Subject: [PATCH 74/98] Add brew prefix for intel macOS. Pin np version --- .github/workflows/macOs_test.yml | 6 ++++++ keopscore/keopscore/config/OpenMP.py | 1 + 2 files changed, 7 insertions(+) diff --git a/.github/workflows/macOs_test.yml b/.github/workflows/macOs_test.yml index 29c8e9207..0d36cecf4 100644 --- a/.github/workflows/macOs_test.yml +++ b/.github/workflows/macOs_test.yml @@ -26,6 +26,8 @@ jobs: tests-macos-intel: name: Tests (macOS Intel) runs-on: macos-26-intel + env: + PIP_CONSTRAINT: /tmp/pip-constraints.txt steps: - name: Checkout uses: actions/checkout@v4 @@ -38,5 +40,9 @@ jobs: - name: Install dependencies run: brew install libomp + - name: Pin NumPy for Intel CI + run: | + echo "numpy<2" > /tmp/pip-constraints.txt + - name: Run tests run: bash ./pytest.sh diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 023298c11..d6752a481 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -40,6 +40,7 @@ class OpenMPConfig: _openmp_system_roots = [ # self.get_brew_prefix() added later on, os.path.join(os.path.sep, "opt", "homebrew", "opt", "libomp"), + os.path.join(os.path.sep, "usr", "local", "opt", "libomp"), os.path.join(os.path.sep, "usr", "local"), os.path.join(os.path.sep, "usr"), ] From d3e86a67a910147755367ae565b44f21b298f1ff Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 24 May 2026 10:10:27 +0200 Subject: [PATCH 75/98] fix CI macOS --- .github/workflows/macOs_test.yml | 8 ++++---- keopscore/keopscore/config/Cuda.py | 1 + keopscore/keopscore/config/OpenMP.py | 10 +++++----- 3 files changed, 10 insertions(+), 9 deletions(-) diff --git a/.github/workflows/macOs_test.yml b/.github/workflows/macOs_test.yml index 0d36cecf4..6aef0e1b5 100644 --- a/.github/workflows/macOs_test.yml +++ b/.github/workflows/macOs_test.yml @@ -32,6 +32,10 @@ jobs: - name: Checkout uses: actions/checkout@v4 + - name: Pin NumPy for Intel CI + run: | + echo "numpy<2" > /tmp/pip-constraints.txt + - name: Set up Python uses: actions/setup-python@v5 with: @@ -40,9 +44,5 @@ jobs: - name: Install dependencies run: brew install libomp - - name: Pin NumPy for Intel CI - run: | - echo "numpy<2" > /tmp/pip-constraints.txt - - name: Run tests run: bash ./pytest.sh diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index a9f205633..aaa4bf200 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -664,6 +664,7 @@ def print_all(self): if __name__ == "__main__": from keopscore.config.Platform import PlatformConfig from keopscore.config.CxxCompiler import CxxCompilerConfig + # from keopscore.config.OpenMP import OpenMPConfig platform_info = PlatformConfig() diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index d6752a481..f75e62c9d 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -39,8 +39,7 @@ class OpenMPConfig: _openmp_system_roots = [ # self.get_brew_prefix() added later on, - os.path.join(os.path.sep, "opt", "homebrew", "opt", "libomp"), - os.path.join(os.path.sep, "usr", "local", "opt", "libomp"), + os.path.join(os.path.sep, "opt", "homebrew"), os.path.join(os.path.sep, "usr", "local"), os.path.join(os.path.sep, "usr"), ] @@ -48,12 +47,12 @@ class OpenMPConfig: openmp_library_suffixes = ( "lib", "lib64", - os.path.join("libomp", "lib"), + os.path.join("opt", "libomp", "lib"), ) _openmp_include_sufixes = ( "include", - os.path.join("libomp", "include"), + os.path.join("opt", "libomp", "include"), ) _omp_info = { @@ -112,7 +111,8 @@ def find_install_path(self, lib_dict_info): KeOps_Message("OpenMP library search using standard names:", level=2) KeOps_Message(f" Trying library name: {name}", level=2) KeOps_Message( - f" Found library path: {result['library'] or not_found_str}", level=2 + f" Found library path: {result['library'] or not_found_str}", + level=2, ) KeOps_Message(f" Trying header name: {header_basename}", level=2) KeOps_Message( From dd0d39d078b318df070a17ffc0e6da75fc981c70 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 24 May 2026 10:24:11 +0200 Subject: [PATCH 76/98] update pip constraint CI env --- .github/workflows/macOs_test.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/macOs_test.yml b/.github/workflows/macOs_test.yml index 6aef0e1b5..3dec77f0e 100644 --- a/.github/workflows/macOs_test.yml +++ b/.github/workflows/macOs_test.yml @@ -27,7 +27,7 @@ jobs: name: Tests (macOS Intel) runs-on: macos-26-intel env: - PIP_CONSTRAINT: /tmp/pip-constraints.txt + PIP_BUILD_CONSTRAINT: /tmp/pip-constraints.txt steps: - name: Checkout uses: actions/checkout@v4 From 001b86235d45af78a814f4a3e958e006faa718b1 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 24 May 2026 13:38:16 +0200 Subject: [PATCH 77/98] clean pip constraints --- .github/workflows/macOs_test.yml | 4 +- pybuild.sh | 216 ++++++++++++++----------------- pydoc.sh | 8 +- pytest.sh | 54 ++++++-- 4 files changed, 141 insertions(+), 141 deletions(-) diff --git a/.github/workflows/macOs_test.yml b/.github/workflows/macOs_test.yml index 3dec77f0e..ae2021e66 100644 --- a/.github/workflows/macOs_test.yml +++ b/.github/workflows/macOs_test.yml @@ -26,8 +26,6 @@ jobs: tests-macos-intel: name: Tests (macOS Intel) runs-on: macos-26-intel - env: - PIP_BUILD_CONSTRAINT: /tmp/pip-constraints.txt steps: - name: Checkout uses: actions/checkout@v4 @@ -45,4 +43,4 @@ jobs: run: brew install libomp - name: Run tests - run: bash ./pytest.sh + run: bash ./pytest.sh --pip-constraint=/tmp/pip-constraints.txt diff --git a/pybuild.sh b/pybuild.sh index f49b3f406..4674cfb94 100755 --- a/pybuild.sh +++ b/pybuild.sh @@ -1,4 +1,4 @@ -#!/bin/bash +#!/usr/bin/env bash ################################################################################ # Instructions for Creating a New Release @@ -6,7 +6,7 @@ # 0) Generate a Twine API token and configure your `.pypirc` file # Ensure both TestPyPI and PyPI are properly set up. # -# 1) Update the version number in the file: ./keops_version and Changes log: ./CHANGELOG.md +# 1) Update the version number in the file: ./keops_version and Changes log: ./CHANGELOG.md # # 2) Build the packages using the build script: # sh ./pybuild.sh @@ -20,157 +20,129 @@ # twine upload ./build/dist/pykeops-XXXXX.tar.gz --repository testpypi # pip install -i https://test.pypi.org/simple/ pykeops # -# ⚠️ Note: TestPyPI may have dependency resolution issues. -# If problems occur, install pykeops from PyPI, uninstall it, -# then reinstall pykeops from TestPyPI. +# Note: TestPyPI may have dependency resolution issues. +# If problems occur, install pykeops from PyPI, uninstall it, +# then reinstall pykeops from TestPyPI. # -# ⚠️ Note: Do not forget to remove the install from TestPyPI... +# Note: Do not forget to remove the install from TestPyPI. # # 5) Once validated, upload to the official PyPI: # twine upload ./build/dist/keopscore-XXXXX.tar.gz # twine upload ./build/dist/pykeops-XXXXX.tar.gz ################################################################################ +set -euo pipefail +readonly PYTHON_BIN="python3" +readonly SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +readonly BUILD_VENV="${SCRIPT_DIR}/.build_venv" +readonly BUILD_REQUIREMENTS=(pip build pyclean) +readonly VERSION="$(<"${SCRIPT_DIR}/keops_version")" -# exit in case of any errors -set -e - -################################################################################ -# help # -################################################################################ -function print_help() { - # Display Help - echo "Build script for keopscore/pykeops packages." - echo - echo "Usage: $0 [option...]" - echo - echo " -h Print the help" - echo " -l Build in local mode (without hard-coded keopscore version requirement in pykeops)" - echo " -v Verbose mode" - echo - echo "Note: by default, the keopscore version requirement is hard-coded in pykeops." - echo - exit 1 -} - -################################################################################ -# utils # -################################################################################ - -# log with verbosity management -function logging() { - if [[ ${PYBUILD_VERBOSE} == 1 ]]; then - echo -e $1 - fi -} - -################################################################################ -# process script options # -################################################################################ - -# default options LOCAL_PYBUILD=0 PYBUILD_VERBOSE=0 -# Get the options -while getopts 'hlv' option; do - case $option in - h) # display Help - print_help - ;; - l) # local build (no hard-coded keopscore version requirements) - LOCAL_PYBUILD=1 - logging "## local build (keopscore version requirements is NOT hard-coded in pykeops)" - ;; - v) # enable verbosity - PYBUILD_VERBOSE=1 - logging "## verbose mode" - ;; - \?) # Invalid option - echo "Error: Invalid option" - exit 1 - ;; - esac -done - -################################################################################ -# script setup # -################################################################################ - -# project root directory -PROJDIR=$(git rev-parse --show-toplevel) - -# python exec -PYTHON="python3" - -# python environment for build -BUILD_VENV=${PROJDIR}/.build_venv +print_help() { + cat < + Constrain build dependencies using the given constraints file. EOF - exit 1 } log_verbose() { @@ -28,21 +30,49 @@ log_verbose() { } parse_options() { - while getopts ":hv" option; do - case "${option}" in - h) + while [[ $# -gt 0 ]]; do + case "$1" in + -h) print_help + exit 0 ;; - v) + -v) PYTEST_VERBOSE=1 log_verbose "## verbose mode" ;; - \?) - echo "Error: Invalid option" + --pip-constraint) + shift + if [[ $# -eq 0 ]]; then + echo "Error: --pip-constraint requires a file path" + exit 1 + fi + PIP_CONSTRAINT_FILE="$1" + ;; + --pip-constraint=*) + PIP_CONSTRAINT_FILE="${1#*=}" + ;; + *) + echo "Error: Invalid option: $1" exit 1 ;; esac + shift done + + if [[ -n "${PIP_CONSTRAINT_FILE}" && ! -f "${PIP_CONSTRAINT_FILE}" ]]; then + echo "Error: Constraint file not found: ${PIP_CONSTRAINT_FILE}" + exit 1 + fi +} + +pip_install() { + local args=() + + if [[ -n "${PIP_CONSTRAINT_FILE}" ]]; then + args+=(--constraint "${PIP_CONSTRAINT_FILE}") + fi + + "${PYTHON_BIN}" -m pip install "${args[@]}" "$@" } prepare_python_environment() { @@ -52,8 +82,8 @@ prepare_python_environment() { # shellcheck disable=SC1091 source "${TEST_VENV}/bin/activate" - log_verbose "---- Python version = $(python -V)" - pip install -U "${TEST_REQUIREMENTS[@]}" + log_verbose "---- Python version = $(${PYTHON_BIN} -V)" + pip_install -U "${TEST_REQUIREMENTS[@]}" } install_editable_package() { @@ -61,7 +91,7 @@ install_editable_package() { local package_path="$2" log_verbose "-- Installing ${package_name}..." - pip install -e "${package_path}" + pip_install -e "${package_path}" } run_python_outside_repo() { From 0996bda9460f41742bc4ba54032d25739910f88c Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 24 May 2026 14:06:19 +0200 Subject: [PATCH 78/98] add verbose option to pytest --- .github/workflows/macOs_test.yml | 2 +- pytest.sh | 37 ++++++++++++++++++++++++-------- 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/.github/workflows/macOs_test.yml b/.github/workflows/macOs_test.yml index ae2021e66..75d159032 100644 --- a/.github/workflows/macOs_test.yml +++ b/.github/workflows/macOs_test.yml @@ -43,4 +43,4 @@ jobs: run: brew install libomp - name: Run tests - run: bash ./pytest.sh --pip-constraint=/tmp/pip-constraints.txt + run: bash ./pytest.sh --pip-constraint=/tmp/pip-constraints.txt -v 2 diff --git a/pytest.sh b/pytest.sh index 5068a07b1..c1c8be310 100755 --- a/pytest.sh +++ b/pytest.sh @@ -7,7 +7,7 @@ readonly SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 readonly TEST_VENV="${SCRIPT_DIR}/.test_venv_pytest" readonly TEST_REQUIREMENTS=(pip) -PYTEST_VERBOSE=0 +KEOPS_VERBOSE_LEVEL=-1 PIP_CONSTRAINT_FILE="" print_help() { @@ -17,14 +17,15 @@ Test script for keopscore/pykeops packages. Usage: $0 [option...] -h Print the help - -v Verbose mode + -v <0|1|2> + Verbosity level forwarded to KEOPS_VERBOSE and PYKEOPS_VERBOSE --pip-constraint - Constrain build dependencies using the given constraints file. + Constrain pip installs using the given constraints file. EOF } log_verbose() { - if [[ "${PYTEST_VERBOSE}" -eq 1 ]]; then + if [[ "${KEOPS_VERBOSE_LEVEL}" -ge 1 ]]; then printf '%b\n' "$1" fi } @@ -37,8 +38,16 @@ parse_options() { exit 0 ;; -v) - PYTEST_VERBOSE=1 - log_verbose "## verbose mode" + shift + if [[ $# -eq 0 ]]; then + echo "Error: -v requires a level (0, 1 or 2)" + exit 1 + fi + if ! [[ "$1" =~ ^[0-2]$ ]]; then + echo "Error: Invalid -v value: $1 (expected 0, 1 or 2)" + exit 1 + fi + KEOPS_VERBOSE_LEVEL="$1" ;; --pip-constraint) shift @@ -63,6 +72,16 @@ parse_options() { echo "Error: Constraint file not found: ${PIP_CONSTRAINT_FILE}" exit 1 fi + + log_verbose "## verbose mode (level=${KEOPS_VERBOSE_LEVEL})" +} + +run_with_keops_verbose() { + if [[ "${KEOPS_VERBOSE_LEVEL}" -ne -1 ]]; then + KEOPS_VERBOSE="${KEOPS_VERBOSE_LEVEL}" PYKEOPS_VERBOSE="${KEOPS_VERBOSE_LEVEL}" "$@" + else + "$@" + fi } pip_install() { @@ -72,7 +91,7 @@ pip_install() { args+=(--constraint "${PIP_CONSTRAINT_FILE}") fi - "${PYTHON_BIN}" -m pip install "${args[@]}" "$@" + run_with_keops_verbose "${PYTHON_BIN}" -m pip install "${args[@]}" "$@" } prepare_python_environment() { @@ -98,7 +117,7 @@ run_python_outside_repo() { local python_code="$1" ( cd /tmp - "${PYTHON_BIN}" -c "${python_code}" + run_with_keops_verbose "${PYTHON_BIN}" -c "${python_code}" ) } @@ -117,7 +136,7 @@ run_test_suite() { local suite_path="$2" log_verbose "-- Running ${suite_name} tests..." - pytest -v "${suite_path}" + run_with_keops_verbose pytest -v "${suite_path}" } main() { From 7c590efdcd0918e2964384f517918334b337a28d Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 24 May 2026 14:15:13 +0200 Subject: [PATCH 79/98] add verbose option to pytest --- pytest.sh | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/pytest.sh b/pytest.sh index c1c8be310..2fa155bf2 100755 --- a/pytest.sh +++ b/pytest.sh @@ -85,13 +85,12 @@ run_with_keops_verbose() { } pip_install() { - local args=() - if [[ -n "${PIP_CONSTRAINT_FILE}" ]]; then - args+=(--constraint "${PIP_CONSTRAINT_FILE}") + run_with_keops_verbose "${PYTHON_BIN}" -m pip install --constraint "${PIP_CONSTRAINT_FILE}" "$@" + return fi - run_with_keops_verbose "${PYTHON_BIN}" -m pip install "${args[@]}" "$@" + run_with_keops_verbose "${PYTHON_BIN}" -m pip install "$@" } prepare_python_environment() { From bc970d9e04ef0afcc8270236bab4a9f7a24677b0 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 24 May 2026 18:14:50 +0200 Subject: [PATCH 80/98] fix openMp discovery. Add a unit test. --- keopscore/keopscore/config/Cuda.py | 8 ++-- keopscore/keopscore/config/OpenMP.py | 63 ++++++++++++++++--------- keopscore/keopscore/test/test_openmp.py | 26 ++++++++++ 3 files changed, 72 insertions(+), 25 deletions(-) create mode 100644 keopscore/keopscore/test/test_openmp.py diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index aaa4bf200..e3a5dc4e4 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -408,7 +408,7 @@ def get_libcuda_path(self): return self._libcuda_info["library"] def print_libcuda_path(self): - print(f"Libcuda Path: {self.get_libcuda_path() or not_found_str}") + print(f"Libcuda Path: {self.get_libcuda_path() or not_found_str}") # Libnvrtc folder def set_libnvrtc_folder(self): @@ -435,7 +435,7 @@ def get_libnvrtc_path(self): return self._libnvrtc_info["library"] def print_libnvrtc_path(self): - print(f"Libnvrtc Path: {self.get_libnvrtc_path() or not_found_str}") + print(f"Libnvrtc Path: {self.get_libnvrtc_path() or not_found_str}") # Libcudart path def set_libcudart_path(self): @@ -516,10 +516,10 @@ def set_include_options(self): ) def get_include_options(self): - return self._include_options + return self._include_options.strip() def print_include_options(self): - print(f"GPU Include Options: {self.get_include_options() or not_found_str}") + print(f"CUDA Include Options: {self.get_include_options() or not_found_str}") # NVRTC preprocessins options def set_preprocessing_options(self): diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index f75e62c9d..1e866b798 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -65,7 +65,7 @@ class OpenMPConfig: "libm.dylib", "libm.so*", ], - "header_basename": ["omp.h", "gomp.h"], + "header_basename": "omp.h", "library": "", # to be filled later "header": "", # not needed "ctype_handle": None, # to be filled later @@ -83,44 +83,62 @@ def __init__(self, platform, cxx): self.platform.get_brew_prefix(), ) - # Detect OpenMP and set related configuration variables - if self._omp_is_available(): + # Detect OpenMP once and set related configuration variables. + omp_available = self._omp_is_available() + if omp_available: self.set_libomp_include_path() self.set_libomp_folder() self.set_compile_options() self.set_include_options() self.set_linking_options() - # Chech if the compiler support omp - self.set_use_OpenMP() + # Check if the compiler supports OpenMP. + self.set_use_OpenMP(omp_available) def find_install_path(self, lib_dict_info): result = lib_dict_info.copy() # First try to find OpenMP library using standard names via ctypes /cxx compiler. if not self.platform.get_env_type().startswith("conda"): - for name, header_basename in zip( - lib_dict_info["name"], lib_dict_info["header_basename"] - ): - result["library"] = _find_library_by_names((name,)) - result["header"] = get_include_file_abspath( - header_basename, self.cxx.get_cxx_compiler() + header_path = get_include_file_abspath( + lib_dict_info["header_basename"], self.cxx.get_cxx_compiler() + ) + + #### + KeOps_Message("OpenMP header search using standard names:", level=2) + KeOps_Message( + f" Trying header name: {lib_dict_info['header_basename']}", level=2 + ) + KeOps_Message( + f" Found header path: {header_path or not_found_str}", level=2 + ) + #### + + if not header_path: + KeOps_Message( + " Standard-name library probing skipped: omp.h was not found.", + level=2, ) + else: + result["header"] = header_path + + for name in lib_dict_info["name"]: + library_path = _find_library_by_names((name,)) #### KeOps_Message("OpenMP library search using standard names:", level=2) KeOps_Message(f" Trying library name: {name}", level=2) KeOps_Message( - f" Found library path: {result['library'] or not_found_str}", + f" Found library path: {library_path or not_found_str}", level=2, ) - KeOps_Message(f" Trying header name: {header_basename}", level=2) - KeOps_Message( - f" Found header path: {result['header'] or not_found_str}", level=2 - ) #### - if result["library"] and result["header"]: + if not library_path: + continue + + if header_path: + result["library"] = library_path return result # If that fails, search for OpenMP headers and libraries in common locations. @@ -184,8 +202,10 @@ def _omp_is_available(self): return True # OpenMP support - def set_use_OpenMP(self): - self._use_OpenMP = self._omp_is_available() and self.check_compiler_for_openmp() + def set_use_OpenMP(self, omp_available=None): + if omp_available is None: + omp_available = self._omp_is_available() + self._use_OpenMP = omp_available and self.check_compiler_for_openmp() if not self._use_OpenMP: self._compile_options = "" self._include_options = "" @@ -227,7 +247,7 @@ def get_libomp_include_path(self): return self._libomp_include_path def print_libomp_include_path(self): - print(f"OpenMP Include Path: {self.get_libomp_include_path() or not_found_str}") + print(f"OpenMP Header Path: {self.get_libomp_include_path() or not_found_str}") def get_openmp_include_dir(self): """Get the OpenMP include directory (containing headers).""" @@ -285,9 +305,10 @@ def set_compile_options(self): self._compile_options += " -Xpreprocessor" self._compile_options += " -fopenmp" + self._compile_options def get_compile_options(self): - return self._compile_options + return self._compile_options.strip() def print_compile_options(self): print(f"Compile Options: {self.get_compile_options()}") diff --git a/keopscore/keopscore/test/test_openmp.py b/keopscore/keopscore/test/test_openmp.py new file mode 100644 index 000000000..63c50c70e --- /dev/null +++ b/keopscore/keopscore/test/test_openmp.py @@ -0,0 +1,26 @@ +import os +import warnings + +import keopscore.config +import pytest + + +def test_openmp_detection_non_critical(): + """OpenMP is optional: warn when unavailable, validate when available.""" + openmp = keopscore.config.openmp + assert isinstance(openmp.get_use_OpenMP(), bool) + + if not openmp.get_use_OpenMP(): + with pytest.warns(UserWarning, match="OpenMP is not available"): + warnings.warn( + "OpenMP is not available on this platform/configuration.", + UserWarning, + ) + return + + lib_path = openmp.get_libomp_path() + header_path = openmp.get_libomp_include_path() + + assert lib_path and os.path.exists(lib_path) + assert header_path and os.path.exists(header_path) + assert "-fopenmp" in openmp.get_compile_options() From 1982ee10cfa79d6ef564ad5b847fac56b5400b16 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sun, 24 May 2026 23:51:36 +0200 Subject: [PATCH 81/98] add install documentation --- doc/python/api/common/Config.rst | 19 ++ doc/python/api/common/index.rst | 1 + doc/python/installation.rst | 418 +++++++++++++++++++++-------- keopscore/keopscore/config/Cuda.py | 10 +- 4 files changed, 342 insertions(+), 106 deletions(-) create mode 100644 doc/python/api/common/Config.rst diff --git a/doc/python/api/common/Config.rst b/doc/python/api/common/Config.rst new file mode 100644 index 000000000..51aa10c2f --- /dev/null +++ b/doc/python/api/common/Config.rst @@ -0,0 +1,19 @@ +Configuration +------------- + +This section documents lower-level KeOps configuration helpers that are useful +when troubleshooting PyKeOps installations. + +.. rubric:: Summary + +.. currentmodule:: keopscore.config.Cuda + +.. autosummary:: + :nosignatures: + + CudaConfig + +.. rubric:: Syntax + +.. autoclass:: CudaConfig + :members: diff --git a/doc/python/api/common/index.rst b/doc/python/api/common/index.rst index 263f096e2..ed97c421f 100644 --- a/doc/python/api/common/index.rst +++ b/doc/python/api/common/index.rst @@ -6,5 +6,6 @@ Common Python API .. toctree:: GenericLazyTensor + Config Utils diff --git a/doc/python/installation.rst b/doc/python/installation.rst index a2aaaf796..986c1e2f8 100644 --- a/doc/python/installation.rst +++ b/doc/python/installation.rst @@ -1,69 +1,159 @@ -Python install -############## +Python installation +################### -PyKeOps is a **Python 3 wrapper** around the low-level KeOpsCore library which is written in **C++/CUDA**. -It provides functions that can be used in any **NumPy** or **PyTorch** script. +The ``pykeops`` Python module provides the NumPy and PyTorch bindings for +KeOps. It relies on the ``keopscore`` Python module, the KeOps +metaprogramming engine, to generate and compile the C++/CUDA routines that +evaluate symbolic kernel formulas on the fly. + + +.. _`part.PyKeOpsRequirements`: Requirements ============ -- **Python** (>= 3.8) with the **numpy** package. -- A C++ compiler compatible with ``std=c++11``: **g++** version >=7 or **clang++** version >=8. -- The **Cuda** toolkit: version >=10.0 is recommended. -- **PyTorch** (optional): version >= 1.5. +Required dependencies: + +- **Python** (>= 3.8) with the **NumPy** package. +- A C++ compiler such as ``gcc`` or ``clang``, compatible with ``std=c++11``. + +Optional, but highly recommended: + +- An NVIDIA GPU with the **NVIDIA drivers** and the **CUDA** toolkit. +- **PyTorch** (version >= 2 is recommended). +- The **OpenMP** libraries and headers for CPU-only systems. Using pip (recommended) ======================= -1. Just in case: in a terminal, check the **consistency** of the outputs of the commands ``which python``, ``python --version``, ``which pip`` and ``pip --version``. +1. In a terminal, check that ``python`` and ``pip`` point to the same + environment with: -2. In a terminal, type: + .. prompt:: bash $ - .. prompt:: bash $ + which python + python --version + which pip + pip --version - pip install pykeops + You can also create a fresh Python virtual environment: - Note that compiled shared objects (``.so`` files on Unix, ``.dylib`` on macOS) will be stored in the folder ``~/.cache/keops/``, where ``~`` is the path to your home folder. If you want to change this default location, define the environment variable ``KEOPS_CACHE_FOLDER`` to another folder prior to importing pykeops. + .. prompt:: bash $ + + python -m venv keops_venv + source keops_venv/bin/activate + +2. **Install PyKeOps:** + + .. prompt:: bash $ + + pip install pykeops + + Compiled shared objects (``.so`` files on Unix, ``.dylib`` files on + macOS) are stored in ``~/.cache/keops``, where ``~`` is your + home folder and ```` is the installed ``pykeops`` version. To + change this location, define the ``KEOPS_CACHE_FOLDER`` environment + variable before importing ``pykeops``. + +3. Test your installation: + + .. prompt:: bash $ + + python -c "import pykeops; pykeops.test_numpy_bindings()" + +More details are in the :ref:`dedicated section `. -3. Test your installation, as described in the :ref:`next section `. On Google Colab =============== -Google provides free virtual machines where KeOps runs -out-of-the-box. -In a new `Colab notebook `_, typing: +Google provides free virtual machines, running on Ubuntu Linux, where KeOps +runs out of the box. In a new +`Colab notebook `_, typing: -.. prompt:: bash $ +.. prompt:: python >>> !pip install pykeops > install.log + import pykeops + pykeops.test_numpy_bindings() should allow you to get a working version of KeOps in less than twenty seconds. +.. _`part.CondaConfig`: + +Using a Conda/Miniconda/Mamba environment +========================================= + +Conda environments can be a convenient way to get a working configuration +without the root permissions that may be needed to install system dependencies +such as the CUDA toolkit. They are not fully isolated from every other +installation mechanism, though. If things do not behave as expected, check +whether older packages are already installed somewhere else on your system. + +The dependencies for PyKeOps are listed :ref:`above `, +but the exact packages needed depend on your system. For instance, as of May +2026, on Ubuntu 24.04 LTS with an NVIDIA GPU and driver 580 installed, but +without ``nvidia-cuda-dev`` or ``nvidia-cuda-toolkit`` installed, the following +commands create a working environment: + +.. prompt:: bash $ + + conda create --name keops_env python=3.14 + conda activate keops_env + + conda install conda-forge::openmp + conda install nvidia::cuda-toolkit==12.9.2 + pip install pykeops + + python -c "import pykeops; pykeops.test_numpy_bindings()" + +On Ubuntu 24.04 LTS, pinning ``cuda-toolkit`` to version 12 may be necessary. +This pin may no longer be needed on later Ubuntu releases. + + +On macOS +======== + +We recommend installing the OpenMP libraries with Homebrew: + +.. prompt:: bash $ + + brew install libomp + pip install pykeops + +You should now be able to test your installation: + +.. prompt:: bash $ + + python -c "import pykeops; pykeops.test_numpy_bindings()" + +More help can be found in the :ref:`dedicated section `. + Using Docker or Singularity ============================ -We provide a reference -`Dockerfile `_ -and publish full containers on our -`DockerHub channel `_ -using the -`docker-images.sh `_ script. -These environments contain a full installation of CUDA, NumPy, PyTorch, R, KeOps and GeomLoss. -Their PYTHONPATH are configured to ensure that git installations of KeOps or GeomLoss -mounted in ``/opt/keops`` or ``/opt/geomloss`` take precedence over the -pre-installed Pip versions. - -As an example, here are the steps that we follow to render this website on the +We provide a reference +`Dockerfile `_ and +publish full containers on our +`DockerHub channel `_ +using the `docker-images.sh `_ +script. These environments contain full installations of CUDA, NumPy, PyTorch, +R, KeOps (for Python and R) and GeomLoss. + +The container's ``PYTHONPATH`` environment variable is configured so that Git +installations of KeOps or GeomLoss mounted in ``/opt/keops`` or +``/opt/geomloss`` take precedence over the pre-installed pip versions. + +As an example, here are the steps that we follow to render this website on the `Jean Zay `_ scientific cluster: -.. code-block:: bash +.. prompt:: bash $ # First, clone the latest release of the KeOps repository in ~/code/keops: - mkdir ~/code - cd ~/code + mkdir ~/code + cd ~/code git clone git@github.com:getkeops/keops.git # Load singularity in our environment: @@ -76,8 +166,9 @@ As an example, here are the steps that we follow to render this website on the # Download the Docker image and store it as an immutable Singularity Image File: # N.B.: Our image is pretty heavy (~7 Gb), so it is safer to create # cache folders on the hard drive instead of relying on the RAM-only tmpfs: - # N.B.: This step may take 15mn to 60mn, so you may prefer to execute it on your - # local computer and then copy the resulting file `keops-full.sif` to the cluster. + # N.B.: This step may take 15mn to 60mn, so you may prefer to execute it on + # your local computer and then copy the resulting file `keops-full.sif` + # to the cluster. # Alternatively, on the Jean Zay cluster, you may use the `prepost` partition # to have access to both a large RAM and an internet connection. mkdir cache @@ -88,9 +179,9 @@ As an example, here are the steps that we follow to render this website on the # At this point, on the Jean Zay cluster, you should use a command like: # idrcontmgr cp keops-full.sif - # to add our new environment to the cluster's container registry as explained here: + # to add our new environment to the cluster's container registry as explained here: # http://www.idris.fr/jean-zay/cpu/jean-zay-utilisation-singularity.html - + # Then, create a separate home folder for this image. This is to ensure # that we won't see any conflict between different versions of the KeOps binaries, # stored in the ~/.cache folder of the virtual machine: @@ -106,16 +197,16 @@ Where ``keops-doc.batch`` is an executable file that contains: #!/bin/bash - #SBATCH -A dvd@a100 # Use a A100 GPU - dvd@v100 is also available - #SBATCH -C a100 + #SBATCH -A dvd@a100 # Use an A100 GPU - dvd@v100 is also available + #SBATCH -C a100 #SBATCH --partition=gpu_p5 #SBATCH --job-name=keops_doc # create a short name for your job #SBATCH --mail-type=ALL # Mail events (NONE, BEGIN, END, FAIL, ALL) - #SBATCH --mail-user=your.name@inria.fr # Where to send mail + #SBATCH --mail-user=your.name@inria.fr # Where to send mail #SBATCH --nodes=1 # node count #SBATCH --ntasks=1 # total number of tasks across all nodes #SBATCH --cpus-per-task=8 # cpu-cores per task (>1 if multi-threaded tasks) - #SBATCH --gres=gpu:1 # GPU nodes are only available in gpu partition + #SBATCH --gres=gpu:1 # GPU nodes are only available in gpu partition #SBATCH --time=03:00:00 # total run time limit (HH:MM:SS) #SBATCH --output=logs/keops_doc.out # output file name #SBATCH --error=logs/keops_doc.err # error file name @@ -179,53 +270,51 @@ And ``keops-doc.sh`` is an executable file that contains: -From source using git +From source using Git ===================== - -The simplest way of installing a specific version -of KeOps is to use `some advanced pip syntax `_: - +The simplest way to install a specific version of KeOps is to use pip's +`Git URL syntax `_: .. prompt:: bash $ - pip install git+https://github.com/getkeops/keops.git@main#subdirectory=keopscore - pip install git+https://github.com/getkeops/keops.git@main#subdirectory=pykeops + pip install git+https://github.com/getkeops/keops.git@main#subdirectory=keopscore + pip install git+https://github.com/getkeops/keops.git@main#subdirectory=pykeops Alternatively, you may: -1. Clone the KeOps repo at a location of your choice (denoted here as ``/path/to``): +1. Clone the KeOps repository at a location of your choice: - .. prompt:: bash $ + .. prompt:: bash $ - git clone --recursive https://github.com/getkeops/keops.git /path/to/libkeops + git clone https://github.com/getkeops/keops.git /path/to/keops_cloned_repo - Note that compiled **.so** routines will be stored in the folder ``/path/to/libkeops/pykeops/build``: this directory must have **write permission**. +2. Install the Python packages in editable mode: + .. prompt:: bash $ -2. Install via pip in editable mode as follows : - - .. prompt:: bash $ + pip install -e /path/to/keops_cloned_repo/keopscore -e /path/to/keops_cloned_repo/pykeops - pip install -e /path/to/libkeops/keopscore -e /path/to/libkeops/pykeops + If you prefer not to install the packages, you can add + ``/path/to/keops_cloned_repo/keopscore`` and + ``/path/to/keops_cloned_repo/pykeops`` to your Python path. To do this once + and for all, add the paths to your ``~/.bashrc``: - + Otherwise you may add the directories ``/path/to/libkeops/keopscore`` and ``/path/to/libkeops/pykeops`` to your python path. This can be done once and for all, by adding the path to to your ``~/.bashrc``. In a terminal, type: - - .. prompt:: bash $ + .. prompt:: bash $ - echo "export PYTHONPATH=$PYTHONPATH:/path/to/libkeops/keopscore:/path/to/libkeops/pykeops" >> ~/.bashrc + echo "export PYTHONPATH=$PYTHONPATH:/path/to/keops_cloned_repo/keopscore:/path/to/keops_cloned_repo/pykeops" >> ~/.bashrc - + Alternatively, you may add the following line to the beginning of your python scripts: - - .. code-block:: python + Alternatively, add these lines at the beginning of your Python scripts: - import os.path - import sys - sys.path.append('/path/to/libkeops/keopscore') - sys.path.append('/path/to/libkeops/pykeops') + .. code-block:: python -3. Test your installation, as described in the :ref:`next section. ` + import sys + + sys.path.append("/path/to/keops_cloned_repo/keopscore") + sys.path.append("/path/to/keops_cloned_repo/pykeops") + +3. Test your installation, as described in the :ref:`next section `. .. _`part.checkPython`: @@ -233,37 +322,40 @@ Alternatively, you may: Testing your installation ========================= -You can use the following test functions to compile and run simple KeOps formulas. If the compilation fails, it returns the full log. +You can use the following test functions to compile and run simple KeOps +formulas. If compilation fails, they return the full log. -1. In a python terminal, run :func:`pykeops.test_numpy_bindings `. +1. In a Python terminal, run + :func:`pykeops.test_numpy_bindings `. - .. code-block:: python + .. prompt:: python >>> - import pykeops - pykeops.test_numpy_bindings() # perform the compilation - - should return: + import pykeops + assert pykeops.test_numpy_bindings() # perform the compilation - .. code-block:: text + It should print: - pyKeOps with numpy bindings is working! + .. code-block:: text -2. If you use PyTorch, run :func:`pykeops.test_torch_bindings `. + pyKeOps with numpy bindings is working! - .. code-block:: python +2. If you use PyTorch, run + :func:`pykeops.test_torch_bindings `. - import pykeops - pykeops.test_torch_bindings() # perform the compilation - - should return: + .. prompt:: python >>> + + import pykeops + assert pykeops.test_torch_bindings() # perform the compilation - .. code-block:: text + It should print: - pyKeOps with torch bindings is working! + .. code-block:: text + pyKeOps with torch bindings is working! -Please note that running ``pytest -v`` in a copy of our git repository will also -let you perform an in-depth test of the entire KeOps codebase. + +Running ``pytest -v`` in a copy of our Git repository will also let you perform +an in-depth test of the entire KeOps codebase. Troubleshooting @@ -272,24 +364,135 @@ Troubleshooting KeOps health check ------------------ -To get an overview of your KeOps installation (along with any related issues), including relevant paths, environments, compilation flags, and more, it’s recommended to run the :func:`pykeops.check_health ` function. Simply type the following in a Python shell: +To get an overview of your KeOps installation, including relevant paths, +environments, compilation flags and possible issues, we recommend running the +:func:`pykeops.check_health ` function: -.. code-block:: python +.. prompt:: python >>> import pykeops + pykeops.clean_pykeops() pykeops.check_health() +You can inspect the paths found by KeOps to check whether the external +libraries are properly detected. Compilation issues ------------------ -First of all, make sure that you are using a C++ compiler which is compatible with the **C++11 revision**. Otherwise, compilation of formulas may fail in unexpected ways. Depending on your system, you can: +KeOps compiles small code fragments to compute kernel operations on a device +(CPU or GPU). Most installation issues occur during this compilation step. They +are usually caused by misconfigured environments, overlapping package +installations or non-standard installation paths. The common failure modes have +evolved over the years, and the +`KeOps issue tracker `_ is a good +up-to-date starting point. + +We detail some common issues below, along with generic recommendations. + + +CUDA toolkit detection +....................... + +Detecting the CUDA toolkit is not always straightforward because several CUDA +versions may be present on the same system. For instance, PyTorch installations +may bring NVIDIA runtime packages into the active Python environment. + +KeOps does not ship its own CUDA toolkit. It tries to detect a working toolkit +with :class:`keopscore.config.CudaConfig ` in +the following order: + +1. *Environment variables:* ``CUDA_PATH``, ``CUDA_HOME``, ``CUDA_ROOT`` and + ``CUDA_TOOLKIT_ROOT_DIR`` (in this order). + +2. *Conda installation:* if you use Conda, as described + :ref:`above `. + +3. *System installation from your distribution* **(recommended)**: this is the + best way to ensure that the CUDA toolkit and driver versions match and that + all paths are set consistently. + + On common Linux distributions, system CUDA packages can be installed for instance with: + + .. prompt:: bash $ + + # Debian/Ubuntu, with distribution packages: + sudo apt install nvidia-cuda-dev nvidia-cuda-toolkit + + # Debian/Ubuntu, with NVIDIA CUDA repositories enabled: + sudo apt install cuda-dev nvidia-cuda-toolkit + + # Arch Linux and derivatives: + yay -S cuda + +4. *`NVIDIA PyPI packages `_*: installed + with the ``cuda-toolkit[all]`` module from PyPI. Some of these packages are also pulled + in by PyTorch, but the solution is still fragile. Indeed, CUDA 12 packages + have historically missed the ``crt`` header directory needed by PyKeOps, so + they do not work. CUDA 13 packages fix this layout, but may require a recent + distribution and can still lead to cryptic runtime errors on Ubuntu 24.04 LTS. + + +Using NVIDIA PyPI packages works on Arch Linux. The following command: + +.. prompt:: bash $ + + PYKEOPS_VERBOSE=0 python -c "import pykeops; pykeops.config.cuda.print_all()" | head -n 9 + +yields the detection of a system-wide CUDA installation under ``/opt/cuda``: + +.. code-block:: text + + ============================================================ + CUDA Support + ============================================================ + CUDA Support: Enabled ✅ + Number of GPUs Detected: 4 + CUDA Version: 13.2 + Libcuda Path: /usr/lib64/libcuda.so.1 + Libnvrtc Path: /opt/cuda/lib64/libnvrtc.so.13 + Libcudart Path: /opt/cuda/lib64/libcudart.so.13 + +The following commands force KeOps to use the NVIDIA PyPI packages installed in +the active Python virtual environment: + +.. prompt:: bash $ + + pip install "cuda-toolkit[all]" + CUDA_PATH="$VIRTUAL_ENV/lib/python3.14/site-packages/nvidia/cu13" \ + PYKEOPS_VERBOSE=0 python -c "import pykeops; pykeops.config.cuda.print_all()" | head -n 9 + +This gives: + +.. code-block:: text + + ============================================================ + CUDA Support + ============================================================ + CUDA Support: Enabled ✅ + Number of GPUs Detected: 4 + CUDA Version: 13.0 + Libcuda Path: /usr/lib64/libcuda.so.1 + Libnvrtc Path: $VIRTUAL_ENV/lib/python3.14/site-packages/nvidia/cu13/lib/libnvrtc.so.13 + Libcudart Path: $VIRTUAL_ENV/lib/python3.14/site-packages/nvidia/cu13/lib/libcudart.so.13 + + +Compiler +........ -1. Install a compiler **system-wide**: for instance, on Debian-based Linux distributions, you can install g++ with apt and then use `update-alternatives `_ to choose a suitable compiler as default. Don't forget to pick compatible versions for both **gcc** and **g++**. +Recent Linux and macOS distributions usually provide suitable compilers. If +compilation fails, make sure that you are using a C++ compiler compatible with +the **C++11 revision**. Otherwise, formula compilation may fail in unexpected +ways. -2. Install a compiler **locally**: if you are using a conda environment, you can install a new instance of gcc and g++ by following the `documentation of conda `_. +1. Install a compiler **system-wide**: for instance, on Debian-based Linux + distributions, you can install g++ with apt and then use + `update-alternatives `_ + to choose a suitable compiler as default. -3. If you have a conda environment with CUDA toolkit and pyKeOps, the compiling test with ``pykeops.test_numpy_bindings()``will fail unless you also have a system-wide CUDA toolkit installation, due to missing ``cuda.h`` file. See `_, question "How can I compile CUDA (host or device) codes in my environment?" +2. Install a compiler **locally**: if you are using a conda environment, you can + install a new instance of gcc and g++ by following the + `Conda documentation `_. @@ -298,39 +501,44 @@ First of all, make sure that you are using a C++ compiler which is compatible wi Cache directory --------------- -If you experience problems with compilation, it may be a good idea to **flush the build folder** that KeOps uses as a cache for already-compiled formulas. To do this, just type: +If you experience compilation problems, it may be a good idea to **flush the +build folder** that KeOps uses as a cache for already-compiled formulas. To do +this, type: -.. code-block:: python +.. prompt:: python >>> import pykeops pykeops.clean_pykeops() -You can change the build folder by using the ``set_build_folder()`` function: +You can change the build folder with the ``set_build_folder()`` function: -.. code-block:: python +.. prompt:: python >>> import pykeops print(pykeops.get_build_folder()) # display current build_folder - pykeops.set_build_folder("/my/new/location") # change the build folder + pykeops.set_build_folder("/tmp/keops_cache_new_location") # change the build folder print(pykeops.get_build_folder()) # display new build_folder -Note that the command ``set_build_folder()`` without any argument will reset the location to the default one (``~/.keops/build`` on unix-like systems) +Calling ``set_build_folder()`` without any argument resets the location to the +default one (``~/.cache/keops`` on Unix-like systems). Verbosity level --------------- -You can deactivate all messages and warnings by setting the environment variable `PYKEOPS_VERBOSE` to 0. In a terminal, type: +The KeOps verbosity level is an integer equal to 0 (silent), 1 (default) or 2 +(verbose). You can deactivate all messages and warnings by setting the +``PYKEOPS_VERBOSE`` environment variable to 0. In a terminal, type: .. prompt:: bash $ - export PYKEOPS_VERBOSE=0 - python my_script_calling_pykeops.py + PYKEOPS_VERBOSE=0 python my_script_calling_pykeops.py -Alternatively, you can disable verbose compilation from your python script using the function ``pykeops.set_verbose()``. In a python shell, type: +Alternatively, you can disable verbose compilation from your Python script with +``pykeops.set_verbose()``: -.. code-block:: python +.. prompt:: python >>> import pykeops pykeops.set_verbose(0) # no output pykeops.set_verbose(1) # default verbosity level - pykeops.set_verbose(2) # maximum verbosity level \ No newline at end of file + pykeops.set_verbose(2) # maximum verbosity level diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index e3a5dc4e4..55376e32e 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -14,7 +14,15 @@ class CudaConfig: """ - Class for CUDA detection and configuration. + CUDA detection and configuration. + + The main state is stored in the ``_libcuda_info``, ``_libnvrtc_info`` and + ``_cudart_info`` dictionaries. These dictionaries are filled with the paths + to the corresponding libraries and headers, when found, and with their + ctypes handles when successfully loaded. Detection is performed by + ``_cuda_libraries_available``, which is called by ``set_use_cuda``. The + remaining configuration, including the CUDA version, include options and + preprocessing options, is set from the detection results. """ # CUDA constants From f728e688637d99ae6f58a6df0e348d564c2d9fc8 Mon Sep 17 00:00:00 2001 From: joanglaunes <34514333+joanglaunes@users.noreply.github.com> Date: Wed, 27 May 2026 21:25:02 +0200 Subject: [PATCH 82/98] fixed message --- pykeops/pykeops/numpy/test_install.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pykeops/pykeops/numpy/test_install.py b/pykeops/pykeops/numpy/test_install.py index 3d78bf1a5..4c8367644 100644 --- a/pykeops/pykeops/numpy/test_install.py +++ b/pykeops/pykeops/numpy/test_install.py @@ -25,7 +25,7 @@ def test_numpy_bindings(): if np.allclose(keops_res, expected_res): pyKeOps_Message( - "pyKeOps with torch bindings is working!", use_tag=False, level=1 + "pyKeOps with numpy bindings is working!", use_tag=False, level=1 ) return True else: From f16ba1f19827b3ab3b0af97a1dbcd69a2328a57f Mon Sep 17 00:00:00 2001 From: joanglaunes <34514333+joanglaunes@users.noreply.github.com> Date: Wed, 27 May 2026 21:31:51 +0200 Subject: [PATCH 83/98] fixed another message --- pykeops/pykeops/numpy/test_install.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pykeops/pykeops/numpy/test_install.py b/pykeops/pykeops/numpy/test_install.py index 4c8367644..e6ac4e805 100644 --- a/pykeops/pykeops/numpy/test_install.py +++ b/pykeops/pykeops/numpy/test_install.py @@ -21,7 +21,7 @@ def test_numpy_bindings(): try: keops_res = my_conv(x, y).flatten() except Exception as e: - raise ValueError(f"Error during computation: {e}", use_tag=False) + pyKeOps_Message(f"Error during computation: {e}", use_tag=False) if np.allclose(keops_res, expected_res): pyKeOps_Message( From 0ca18bf589a1bdad809b136d585d9990918994ab Mon Sep 17 00:00:00 2001 From: joanglaunes <34514333+joanglaunes@users.noreply.github.com> Date: Wed, 27 May 2026 21:37:36 +0200 Subject: [PATCH 84/98] better fix for message... --- pykeops/pykeops/numpy/test_install.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pykeops/pykeops/numpy/test_install.py b/pykeops/pykeops/numpy/test_install.py index e6ac4e805..e99d39adc 100644 --- a/pykeops/pykeops/numpy/test_install.py +++ b/pykeops/pykeops/numpy/test_install.py @@ -21,7 +21,7 @@ def test_numpy_bindings(): try: keops_res = my_conv(x, y).flatten() except Exception as e: - pyKeOps_Message(f"Error during computation: {e}", use_tag=False) + raise ValueError(f"Error during computation: {e}") if np.allclose(keops_res, expected_res): pyKeOps_Message( From 48b9b7f1fa05313e935d53bc983cebdcda3e8641 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Fri, 29 May 2026 15:28:35 +0200 Subject: [PATCH 85/98] add pip cuda toolkit support --- doc/python/installation.rst | 28 +- .../binders/nvrtc/Gpu_link_compile.py | 18 +- .../keopscore/binders/nvrtc/nvrtc_jit.cpp | 9 + keopscore/keopscore/config/Cuda.py | 334 +++++++++++------- keopscore/keopscore/config/KeOpsPath.py | 2 +- keopscore/keopscore/config/OpenMP.py | 89 +++-- keopscore/keopscore/config/Platform.py | 42 ++- keopscore/keopscore/config/__init__.py | 2 +- keopscore/keopscore/utils/gpu_utils.py | 59 ---- keopscore/keopscore/utils/path_utils.py | 60 +--- keopscore/keopscore/utils/system_utils.py | 3 + pykeops/setup.py | 2 + 12 files changed, 357 insertions(+), 291 deletions(-) delete mode 100644 keopscore/keopscore/utils/gpu_utils.py diff --git a/doc/python/installation.rst b/doc/python/installation.rst index 986c1e2f8..4157ab766 100644 --- a/doc/python/installation.rst +++ b/doc/python/installation.rst @@ -87,29 +87,25 @@ Using a Conda/Miniconda/Mamba environment Conda environments can be a convenient way to get a working configuration without the root permissions that may be needed to install system dependencies -such as the CUDA toolkit. They are not fully isolated from every other +such as the CUDA toolkit or OpenMP. They are not fully isolated from every other installation mechanism, though. If things do not behave as expected, check whether older packages are already installed somewhere else on your system. The dependencies for PyKeOps are listed :ref:`above `, -but the exact packages needed depend on your system. For instance, as of May -2026, on Ubuntu 24.04 LTS with an NVIDIA GPU and driver 580 installed, but -without ``nvidia-cuda-dev`` or ``nvidia-cuda-toolkit`` installed, the following -commands create a working environment: +but the exact packages needed depend on your system. For instance, the following +commands should create a working environment: .. prompt:: bash $ conda create --name keops_env python=3.14 conda activate keops_env - conda install conda-forge::openmp - conda install nvidia::cuda-toolkit==12.9.2 + conda install libgomp + conda install nvidia::cuda-toolkit pip install pykeops python -c "import pykeops; pykeops.test_numpy_bindings()" -On Ubuntu 24.04 LTS, pinning ``cuda-toolkit`` to version 12 may be necessary. -This pin may no longer be needed on later Ubuntu releases. On macOS @@ -427,16 +423,18 @@ the following order: 4. *`NVIDIA PyPI packages `_*: installed with the ``cuda-toolkit[all]`` module from PyPI. Some of these packages are also pulled - in by PyTorch, but the solution is still fragile. Indeed, CUDA 12 packages - have historically missed the ``crt`` header directory needed by PyKeOps, so - they do not work. CUDA 13 packages fix this layout, but may require a recent - distribution and can still lead to cryptic runtime errors on Ubuntu 24.04 LTS. + in by PyTorch, making this solution fragile as multiple versions of the cuda-toolkit can + co-exist in the same virtual environment. For instance, Torch may install + CUDA 12 packages while the user manually installs CUDA 13 packages... Using NVIDIA PyPI packages works on Arch Linux. The following command: .. prompt:: bash $ + python -m venv keops_venv + source keops_venv/bin/activate + pip install pykeops PYKEOPS_VERBOSE=0 python -c "import pykeops; pykeops.config.cuda.print_all()" | head -n 9 yields the detection of a system-wide CUDA installation under ``/opt/cuda``: @@ -457,8 +455,10 @@ The following commands force KeOps to use the NVIDIA PyPI packages installed in the active Python virtual environment: .. prompt:: bash $ + python -m venv keops_venv_cuda_pip + source keops_venv_cuda_pip/bin/activate - pip install "cuda-toolkit[all]" + pip install pykeops[cu13] CUDA_PATH="$VIRTUAL_ENV/lib/python3.14/site-packages/nvidia/cu13" \ PYKEOPS_VERBOSE=0 python -c "import pykeops; pykeops.config.cuda.print_all()" | head -n 9 diff --git a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py index 9e58683b4..cf1913e68 100644 --- a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py +++ b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py @@ -3,7 +3,6 @@ import keopscore.config from keopscore.binders.LinkCompile import LinkCompile -from keopscore.utils.gpu_utils import custom_cuda_include_fp16_path from keopscore.utils.messages import KeOps_Error, KeOps_Message from keopscore.utils.system_utils import KeOps_OS_Run @@ -31,13 +30,24 @@ def __init__(self): # generated by the JIT compiler, e.g. 7b9a611f7e_nvrtc.{ptx,cubin} self.low_level_code_file = os.path.join( keopscore.config.path.get_build_folder(), - self.gencode_filename + "_nvrtc." + keopscore.config.cuda.get_ir_type(), + self.gencode_filename + + "_nvrtc." + + keopscore.config.cuda.get_ir_type().lower(), ).encode("utf-8") # Load nvrtc_jit dll, that will be used to compile the cuda code in jit mode self.my_c_dll = ctypes.CDLL( keopscore.config.keops_jit_compile_name(type="target"), mode=os.RTLD_LAZY ) + self.my_c_dll.Compile.argtypes = [ + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_char_p, + ] + self.my_c_dll.Compile.restype = ctypes.c_int # file to check for existence to detect compilation is needed self.file_to_check = self.low_level_code_file @@ -57,7 +67,9 @@ def generate_code(self): ctypes.c_int(self.use_fast_math), ctypes.c_int(self.device_id), ctypes.create_string_buffer( - (custom_cuda_include_fp16_path() + os.path.sep).encode("utf-8") + ( + keopscore.config.cuda.custom_cuda_include_fp16_path() + os.path.sep + ).encode("utf-8") ), ) if res != keopscore.config.cuda.CUDA_SUCCESS: diff --git a/keopscore/keopscore/binders/nvrtc/nvrtc_jit.cpp b/keopscore/keopscore/binders/nvrtc/nvrtc_jit.cpp index e0215366d..717bb89a9 100644 --- a/keopscore/keopscore/binders/nvrtc/nvrtc_jit.cpp +++ b/keopscore/keopscore/binders/nvrtc/nvrtc_jit.cpp @@ -107,6 +107,15 @@ extern "C" int Compile(const char *target_file_name, const char *cu_code, } if (compileResult != NVRTC_SUCCESS) { + size_t logSize = 0; + nvrtcGetProgramLogSize(prog, &logSize); + if (logSize > 1) { + std::vector log(logSize); + nvrtcGetProgramLog(prog, log.data()); + std::cerr << "[KeOps] NVRTC compile log:\n" << log.data() << std::endl; + } + std::cerr << "[KeOps] nvrtcCompileProgram failed: " + << nvrtcGetErrorString(compileResult) << std::endl; return compileResult; } diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index 55376e32e..bae3f8f14 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -1,9 +1,11 @@ import ctypes +import glob import os from keopscore.utils.messages import print_envs from keopscore.utils.messages import enabled_dict, not_found_str -from keopscore.utils.messages import KeOps_Warning +from keopscore.utils.messages import KeOps_Error, KeOps_Warning +from keopscore.utils.file_utils import pack_header from keopscore.utils.path_utils import ( _first_matching_file, _ordered_search_roots, @@ -16,8 +18,8 @@ class CudaConfig: """ CUDA detection and configuration. - The main state is stored in the ``_libcuda_info``, ``_libnvrtc_info`` and - ``_cudart_info`` dictionaries. These dictionaries are filled with the paths + The main state is stored in the ``_libcuda_info``, ``_libnvrtc_info`` + and ``_headers_*_info`` dictionaries. These dictionaries are filled with the paths to the corresponding libraries and headers, when found, and with their ctypes handles when successfully loaded. Detection is performed by ``_cuda_libraries_available``, which is called by ``set_use_cuda``. The @@ -34,7 +36,7 @@ class CudaConfig: # Cuda detection variables _use_cuda = False _cuda_version = -1 - _ir_type = "" # "ptx" or "cubin" + _ir_type = "" # "PTX" or "CUBIN" _cuda_include_path = [""] # str or list of str _visible_devices = "" @@ -51,21 +53,33 @@ class CudaConfig: # Search location # # ------------------------ # + # default order of precedence for searching CUDA libraries and headers + order_precedence = ("env_vars", "conda", "pip", "system") + + # environment variables that may point to CUDA installations cuda_env_vars = [ - "CUDA_VISIBLE_DEVICES", "CUDA_PATH", "CUDA_HOME", "CUDA_ROOT", "CUDA_TOOLKIT_ROOT_DIR", + "CUDA_VISIBLE_DEVICES", ] - pip_suffixes = ( - os.path.join("nvidia", "cu13"), # for CUDA 13.X - # os.path.join("nvidia", "cuda_runtime"), # for CUDA 12.X but broken due to missing crt subfolder - # os.path.join("nvidia", "cuda_nvrtc") # for CUDA 12.X but broken due to missing crt subfolder - ) + # pip installation suffixes to look for when searching for CUDA libraries and headers + pip_suffixes = { + "cu12": [ + os.path.join("nvidia", "cuda_runtime"), + os.path.join("nvidia", "cuda_nvrtc"), + os.path.join("nvidia", "cuda_nvcc"), + os.path.join("nvidia", "cuda_cccl"), + ], + "cu13": [ + os.path.join("nvidia", "cu13"), + ], + } - system_suffixes = ( + # standard system suffixes to look for when searching for CUDA libraries and headers + system_prefixes = ( os.path.join(os.path.sep, "usr", "local", "cuda"), os.path.join(os.path.sep, "usr", "local"), os.path.join(os.path.sep, "opt", "cuda"), @@ -77,6 +91,7 @@ class CudaConfig: "lib64", "lib", os.path.join("lib", "x86_64-linux-gnu"), + "", ) include_suffixes = ( @@ -92,30 +107,49 @@ class CudaConfig: _libcuda_info = { "name": "cuda", - "lib_basename_candidate": ["libcuda.so.*", "libcuda.dylib", "cuda.lib"], - "header_basename": "cuda.h", + "lib_basename_candidate": ["libcuda.so.*"], "library": "", # to be filled later - "header": "", # to be filled later "ctype_handle": None, # to be filled later } _libnvrtc_info = { "name": "nvrtc", - "lib_basename_candidate": ["libnvrtc.so.*", "libnvrtc.dylib", "nvrtc.lib"], - "header_basename": "nvrtc.h", + "lib_basename_candidate": ["libnvrtc.so.*"], "library": "", # to be filled later - "header": "", # to be filled later "ctype_handle": None, # to be filled later } - _cudart_info = { - "name": "cudart", - "lib_basename_candidate": ["libcudart.so.*", "libcudart.dylib", "cudart.lib"], - "header_basename": None, # not needed + _libnvrtc_builtins_info = { + "name": "nvrtc-builtins", + "lib_basename_candidate": [ + "libnvrtc-builtins.so*", + "libnvrtc-builtins.alt.so*", + ], "library": "", # to be filled later - "header": None, # not needed "ctype_handle": None, # to be filled later } + _headers_nvrtc_info = { + "header_basename": "nvrtc.h", + "header": "", # to be filled later + } + _headers_cuda_info = { + "header_basename": "cuda.h", + "header": "", # to be filled later + } + _headers_crt_info = { + "header_basename": os.path.join("crt", "device_functions.h"), + "header": "", # to be filled later + } + _headers_nv_info = { + "header_basename": os.path.join("nv", "target"), + "header": "", # to be filled later + } + _headers_fp16_info = { + "header_basename": "cuda_fp16.h", + "header": "", # to be filled later + } - def __init__(self): + def __init__(self, platform): + + self.platform = platform self.set_visible_devices() @@ -125,7 +159,6 @@ def __init__(self): # If cuda is enabled, then we finalize the rest of the config if self.get_use_cuda(): - self.set_cuda_version() self.set_ir_type() self.set_cuda_include_path() self.set_linking_options() @@ -133,30 +166,25 @@ def __init__(self): self.set_preprocessing_options() self.set_include_options() - def find_install_path(self, lib_dict_info): + def find_library_path(self, lib_dict_info, where_to_search): """ - Locate a cuda and headers using an explicit ordered search. + Locate a cuda library using an explicit ordered search. Arguments: - lib_dict_info (dict): A dictionary containing at least the keys 'name', 'lib_basename_candidate', and 'header_basename' for the library to find. This allows the function to be used for finding libcuda, libcudart, or nvrtc by passing the appropriate info dict. + lib_dict_info (dict): A dictionary containing at least the keys 'name' and 'lib_basename_candidate' for the library to find. This allows the function to be used for finding libcuda or nvrtc by passing the appropriate info dict. + + where_to_search (dict): A dictionary containing the search locations, with keys corresponding to the search types (e.g., "env_vars", "conda", "system", "pip") and values containing the relevant paths or environment variable names. Returns: - result (dict): a copy of lib_dict_info completed with the ``library`` and ``header`` keys containing the absolute paths to the library file and include directory, or None if not found. + str: The absolute path to the library file if found, or None if not found. """ - # Start with the provided info, which may contain names and file patterns result = lib_dict_info.copy() candidate_roots = _ordered_search_roots( - env_vars=self.cuda_env_vars, - pip_suffixes=self.pip_suffixes, - conda_root="CONDA_PREFIX", - system_roots=self.system_suffixes, + **where_to_search, + order=self.order_precedence, ) - # ------------------------ # - # Search for library file # - # ------------------------ # - # First try to find the library file using the candidate roots and library suffixes result["library"] = _first_matching_file( @@ -164,33 +192,53 @@ def find_install_path(self, lib_dict_info): result["lib_basename_candidate"], ) if result["library"] is None: - result["library"] = _find_library_by_names((result["name"],)) + result["library"] = _find_library_by_names(result["name"]) + return result + + def find_header_path(self, header_dict_info, where_to_search): + """ + Locate a header file using an explicit ordered search. + + Arguments: + header_basename (str): The basename of the header to find. - # ------------------------ # - # Search for header files # - # ------------------------ # + where_to_search (dict): A dictionary containing the search locations, with keys corresponding to the search types (e.g., "env_vars", "conda", "system", "pip") and values containing the relevant paths or environment variable names. + + Returns: + str: The absolute path to the header file if found, or None if not found. + """ + + candidate_roots = _ordered_search_roots( + **where_to_search, + order=self.order_precedence, + ) - result["header"] = _first_matching_file( + header_dict_info["header"] = _first_matching_file( _path_candidates(candidate_roots, self.include_suffixes), - (result["header_basename"],), + header_dict_info["header_basename"], ) - return result + if not header_dict_info.get("header"): + return ( + False, + f"{header_dict_info['header_basename']} not found. Make sure the CUDA headers are installed and accessible.", + ) - def _find_and_load_libcuda(self): - """Locate, load, and initialize the CUDA driver library.""" - self._libcuda_info = self.find_install_path(self._libcuda_info) + return True, "" + + def _find_and_load_libcuda(self, where_to_search): + """ + Locate, load, and initialize the CUDA driver library. + Get the version of the CUDA driver and the number of visible GPUs, as well as some of their attributes. + Update the config state with this information. + """ + self._libcuda_info = self.find_library_path(self._libcuda_info, where_to_search) libcuda_path = self._libcuda_info["library"] if not libcuda_path: return ( False, "libcuda not found. Make sure the CUDA driver is installed and accessible.", ) - if not self._libcuda_info.get("header"): - return ( - False, - f"{self._libcuda_info['header_basename']} not found. Make sure the CUDA headers are installed and accessible.", - ) try: libcuda = ctypes.CDLL(libcuda_path, mode=ctypes.RTLD_GLOBAL) @@ -206,6 +254,14 @@ def _find_and_load_libcuda(self): "libcuda was detected, but driver API could not be initialized (flushing KeOps caches and/or rebooting the system may help).", ) + cuda_version = ctypes.c_int() + if libcuda.cuDriverGetVersion(ctypes.byref(cuda_version)) != self.CUDA_SUCCESS: + return ( + False, + "libcuda was detected and initialized, but failed to query CUDA driver API version.", + ) + self._cuda_version = int(cuda_version.value) + # If we successfully loaded libcuda and initialized it, store the handle in the config for potential future use self._libcuda_info["ctype_handle"] = libcuda @@ -237,20 +293,17 @@ def _find_and_load_libcuda(self): return True, "" - def _find_and_load_libnvrtc(self): + def _find_and_load_libnvrtc(self, where_to_search): """Locate and load the NVRTC runtime compilation library.""" - self._libnvrtc_info = self.find_install_path(self._libnvrtc_info) + self._libnvrtc_info = self.find_library_path( + self._libnvrtc_info, where_to_search + ) libnvrtc_path = self._libnvrtc_info["library"] if not libnvrtc_path: return ( False, "libnvrtc not found. Make sure the CUDA toolkit is installed and accessible.", ) - if not self._libnvrtc_info.get("header"): - return ( - False, - f"{self._libnvrtc_info['header_basename']} not found. Make sure the CUDA headers are installed and accessible.", - ) try: libnvrtc_handle = ctypes.CDLL(libnvrtc_path, mode=ctypes.RTLD_GLOBAL) @@ -265,50 +318,40 @@ def _find_and_load_libnvrtc(self): return True, "" - def _find_and_load_cudart(self): + def _find_and_load_libnvrtc_builtins(self, folder_to_search): """ - Attempt to find the CUDA version by loading the CUDA runtime library and querying its version. - Returns: - success (bool): True if the version was successfully determined, False otherwise. - error_msg (str): Contains error details if success==False, else "". + Preload of NVRTC builtins from the same folder as libnvrtc. + This check is non-blocking: warnings are emitted on failure, and CUDA + detection continues. """ - self._cudart_info = self.find_install_path(self._cudart_info) - libcudart_path = self._cudart_info["library"] - if not libcudart_path: + self._libnvrtc_builtins_info = self.find_library_path( + self._libnvrtc_builtins_info, {"system": (folder_to_search,)} + ) + libnvrtc_builtins_path = self._libnvrtc_builtins_info["library"] + if not libnvrtc_builtins_path: return ( - False, - "libcudart not found. Make sure the CUDA toolkit is installed and accessible.", + True, + f"NVRTC builtins library not found in {folder_to_search}. This may cause runtime compilation to fail in environments with multiple CUDA toolkit versions installed.", ) try: - libcudart_handle = ctypes.CDLL(libcudart_path) - except OSError as e: + handle = ctypes.CDLL(libnvrtc_builtins_path, mode=ctypes.RTLD_GLOBAL) + except OSError as err: return ( - False, - f"Failed to load '{os.path.basename(libcudart_path)}': {e}", + True, + f"Failed to load NVRTC builtins library '{os.path.basename(libnvrtc_builtins_path)}': {err}", ) - cuda_version = ctypes.c_int() - if ( - libcudart_handle.cudaRuntimeGetVersion(ctypes.byref(cuda_version)) - != self.CUDA_SUCCESS - ): - return ( - False, - "libcudart was found and loaded, but failed to get CUDA runtime version.", - ) - - # If we successfully loaded libcudart, store the handle in the config - self._cudart_info["ctype_handle"] = libcudart_handle - self._cuda_version = int(cuda_version.value) + # If we successfully loaded the library, store the handle in the config for potential future use + self._libnvrtc_builtins_info["ctype_handle"] = handle return True, "" def _cuda_libraries_available(self): """ - Check if libcuda (driver); libcudart (cudatoolkit) and nvrtc (cudatoolkit) libraries are available. + Check if libcuda (driver) and nvrtc (toolkit) libraries are available. This is where ```_cuda_include_path``` is set. Returns: @@ -323,23 +366,54 @@ def _cuda_libraries_available(self): ) return False + where_to_search = { + "env_vars": self.cuda_env_vars, + "conda": "CONDA_PREFIX", + "system": self.system_prefixes, + "pip": (), + } + # Libcuda (driver) loaded globally so it is available to KeOps shared objects. - success_cuda, err_cuda = self._find_and_load_libcuda() + # This is usually provided by the NVIDIA Driver (install system-wide) + success_cuda, err_cuda = self._find_and_load_libcuda(where_to_search) if not success_cuda: KeOps_Warning(f"{err_cuda}. Switching to CPU only.") return False + # Restrict the pip search to the suffixes corresponding to the detected CUDA version, if any. + where_to_search["pip"] = _path_candidates( + self.platform.get_python_package_roots(), + self.pip_suffixes.get(f"cu{self.get_cuda_version(out_type='major')}", ()), + ) + # libnvrtc as well, since it's required for the runtime compilation of CUDA code. - success_nvrtc, err_nvrtc = self._find_and_load_libnvrtc() + # This is usually provided by the CUDA Toolkit (install system-wide or in conda/pip). + success_nvrtc, err_nvrtc = self._find_and_load_libnvrtc(where_to_search) if not success_nvrtc: KeOps_Warning(f"{err_nvrtc}. Switching to CPU only.") return False - # Finally check cudart - success_cudart, err_cudart = self._find_and_load_cudart() - if not success_cudart: - KeOps_Warning(f"{err_cudart}. Switching to CPU only.") - return False + # Try load of NVRTC builtins from the same toolkit folder. + # This avoids runtime search path ambiguity in environments where + # multiple CUDA toolkit versions are installed. + _, warning_nvrtc = self._find_and_load_libnvrtc_builtins( + os.path.dirname(self._libnvrtc_info["library"]) + ) + if warning_nvrtc: + KeOps_Warning(warning_nvrtc, level=2) + + # Finally, check that we can find the CUDA toolkit headers + for header_dict_info in ( + self._headers_cuda_info, + self._headers_nvrtc_info, + self._headers_crt_info, + self._headers_nv_info, + self._headers_fp16_info, + ): + success, err = self.find_header_path(header_dict_info, where_to_search) + if not success: + KeOps_Warning(f"{err}. Switching to CPU only.") + return False return True @@ -374,7 +448,6 @@ def set_visible_devices(self): self._visible_devices = ( cuda_visible.replace(",", "_") if cuda_visible else "empty" ) - print(f"Parsed visible devices: {self._visible_devices}") def get_visible_devices(self): """Get the specific GPUs.""" @@ -445,23 +518,9 @@ def get_libnvrtc_path(self): def print_libnvrtc_path(self): print(f"Libnvrtc Path: {self.get_libnvrtc_path() or not_found_str}") - # Libcudart path - def set_libcudart_path(self): - """ - Return nothing if not using cuda - self.libcudart_path is already set in _cuda_libraries_available. - """ - pass - - def get_libcudart_path(self): - return self._cudart_info["library"] - - def print_libcudart_path(self): - print(f"Libcudart Path: {self.get_libcudart_path() or not_found_str}") - # CUDA Version def set_cuda_version(self): - """Set the CUDA version by querying the CUDA runtime library. Done in _find_and_load_cudart.""" + """Set in _find_and_load_libcuda.""" pass def get_cuda_version(self, out_type="single_value"): @@ -471,6 +530,8 @@ def get_cuda_version(self, out_type="single_value"): if out_type == "major,minor": return major, minor + if out_type == "major": + return major elif out_type == "string": return f"{major}.{minor}" else: @@ -487,11 +548,22 @@ def set_cuda_include_path(self): os.path.dirname(header) for header in ( self._libnvrtc_info.get("header"), - self._libcuda_info.get("header"), + self._headers_cuda_info.get("header"), + self._headers_fp16_info.get("header"), ) if header ] - self._cuda_include_path = list(set(include_dirs)) if include_dirs else "" + # Remove the crt in the header + include_dirs.extend( + [ + os.path.realpath(os.path.join(os.path.dirname(header), "..")) + for header in ( + self._headers_crt_info.get("header"), + self._headers_nv_info.get("header"), + ) + ] + ) + self._cuda_include_path = list(set(include_dirs)) if include_dirs else ("",) def get_cuda_include_path(self): """ @@ -505,14 +577,33 @@ def print_cuda_include_path(self): f"CUDA Include Path: {':'.join(self.get_cuda_include_path()) or not_found_str}" ) + def custom_cuda_include_fp16_path(self): + """ + Create (if needed) a packed standalone cuda_fp16.h in the KeOps build + folder for NVRTC compilation and return this folder. + """ + import keopscore.config + + build_folder = keopscore.config.path.get_build_folder() + fp16_header = "cuda_fp16.h" + fp16_header_path = os.path.join(build_folder, fp16_header) + if not os.path.isfile(fp16_header_path): + pack_header( + fp16_header, + os.path.dirname(self._headers_fp16_info["header"]), + build_folder, + ) + + return build_folder + # IR type def set_ir_type(self): """Set the IR type to be used for nvrtc compilation based on the CUDA version.""" if self.get_cuda_version() >= 11010: - self._ir_type = "cubin" + self._ir_type = "CUBIN" else: - self._ir_type = "ptx" + self._ir_type = "PTX" def get_ir_type(self): return self._ir_type @@ -545,14 +636,14 @@ def set_preprocessing_options(self): f"-DSHAREDMEMPERBLOCK{d}={self.get_SharedMemPerBlock()[d]}" ) - target_tag = "CUBIN" if self.get_ir_type == "cubin" else "PTX" + target_tag = "CUBIN" if self.get_ir_type() == "CUBIN" else "PTX" nvrtcGetTARGET = "nvrtcGet" + target_tag self.add_to_preprocessing_options(f"-DnvrtcGetTARGET={nvrtcGetTARGET}") nvrtcGetTARGETSize = nvrtcGetTARGET + "Size" self.add_to_preprocessing_options(f"-DnvrtcGetTARGETSize={nvrtcGetTARGETSize}") - arch_tag = '\\"sm\\"' if self.get_ir_type == "cubin" else '\\"compute\\"' + arch_tag = '\\"sm\\"' if self.get_ir_type() == "CUBIN" else '\\"compute\\"' self.add_to_preprocessing_options(f"-DARCHTAG={arch_tag}") def add_to_preprocessing_options(self, flags): @@ -596,7 +687,7 @@ def _get_device_attributes(self, libcuda, device_index): return ( 0, 0, - "failed to get device handle for device index {device_index}.", + f"failed to get device handle for device index {device_index}.", ) output = ctypes.c_int() @@ -611,7 +702,7 @@ def _get_device_attributes(self, libcuda, device_index): return ( 0, 0, - "failed to get max threads per block for device index {device_index}.", + f"failed to get max threads per block for device index {device_index}.", ) max_threads_per_block = output.value @@ -626,7 +717,7 @@ def _get_device_attributes(self, libcuda, device_index): return ( 0, 0, - "failed to get shared memory per block for device index {device_index}.", + f"failed to get shared memory per block for device index {device_index}.", ) shared_mem_per_block = output.value @@ -649,7 +740,7 @@ def print_all(self): # CUDA Support print("=" * 60) - print(f"CUDA Support") + print("CUDA Support") print("=" * 60) self.print_use_cuda() @@ -658,7 +749,6 @@ def print_all(self): self.print_cuda_version() self.print_libcuda_path() self.print_libnvrtc_path() - self.print_libcudart_path() self.print_cuda_include_path() self.print_preprocessing_options() @@ -684,5 +774,5 @@ def print_all(self): # openmp_info = OpenMPConfig(platform_info, cxx_compiler_info) # openmp_info.print_all() - cuda_info = CudaConfig() + cuda_info = CudaConfig(platform_info) cuda_info.print_all() diff --git a/keopscore/keopscore/config/KeOpsPath.py b/keopscore/keopscore/config/KeOpsPath.py index 6ad097f7a..52bd26009 100644 --- a/keopscore/keopscore/config/KeOpsPath.py +++ b/keopscore/keopscore/config/KeOpsPath.py @@ -234,7 +234,7 @@ def print_all(self): openmp_info = OpenMPConfig(platform_info, cxx_compiler_info) openmp_info.print_all() - cuda_info = CudaConfig() + cuda_info = CudaConfig(platform_info) cuda_info.print_all() keops_info = KeOpsPathConfig(platform_info, cuda_info) diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 1e866b798..7e7424ad2 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -37,20 +37,20 @@ class OpenMPConfig: "OpenMP_ROOT_DIR", ) - _openmp_system_roots = [ + system_prefixes = [ # self.get_brew_prefix() added later on, os.path.join(os.path.sep, "opt", "homebrew"), os.path.join(os.path.sep, "usr", "local"), os.path.join(os.path.sep, "usr"), ] - openmp_library_suffixes = ( + _library_suffixes = ( "lib", "lib64", os.path.join("opt", "libomp", "lib"), ) - _openmp_include_sufixes = ( + _include_sufixes = ( "include", os.path.join("opt", "libomp", "include"), ) @@ -78,7 +78,7 @@ def __init__(self, platform, cxx): # On MacOS: Add brew prefix to OpenMP search paths if it exists if self.platform.get_brew_prefix(): - self._openmp_system_roots.insert( + self.system_prefixes.insert( 0, self.platform.get_brew_prefix(), ) @@ -99,67 +99,62 @@ def find_install_path(self, lib_dict_info): result = lib_dict_info.copy() # First try to find OpenMP library using standard names via ctypes /cxx compiler. - if not self.platform.get_env_type().startswith("conda"): - header_path = get_include_file_abspath( - lib_dict_info["header_basename"], self.cxx.get_cxx_compiler() - ) + header_path = get_include_file_abspath( + lib_dict_info["header_basename"], self.cxx.get_cxx_compiler() + ) - #### - KeOps_Message("OpenMP header search using standard names:", level=2) + #### + KeOps_Message("OpenMP header search using standard names:", level=2) + KeOps_Message( + f" Trying header name: {lib_dict_info['header_basename']}", level=2 + ) + KeOps_Message(f" Found header path: {header_path or not_found_str}", level=2) + #### + + if not header_path: KeOps_Message( - f" Trying header name: {lib_dict_info['header_basename']}", level=2 + " Standard-name library probing skipped: omp.h was not found.", + level=2, ) + else: + result["header"] = header_path + + for name in lib_dict_info["name"]: + library_path = _find_library_by_names((name,)) + + #### + KeOps_Message("OpenMP library search using standard names:", level=2) + KeOps_Message(f" Trying library name: {name}", level=2) KeOps_Message( - f" Found header path: {header_path or not_found_str}", level=2 + f" Found library path: {library_path or not_found_str}", + level=2, ) #### - if not header_path: - KeOps_Message( - " Standard-name library probing skipped: omp.h was not found.", - level=2, - ) - else: - result["header"] = header_path - - for name in lib_dict_info["name"]: - library_path = _find_library_by_names((name,)) - - #### - KeOps_Message("OpenMP library search using standard names:", level=2) - KeOps_Message(f" Trying library name: {name}", level=2) - KeOps_Message( - f" Found library path: {library_path or not_found_str}", - level=2, - ) - #### - - if not library_path: - continue - - if header_path: - result["library"] = library_path - return result + if not library_path: + continue + + if header_path: + result["library"] = library_path + return result # If that fails, search for OpenMP headers and libraries in common locations. candidate_roots = _ordered_search_roots( env_vars=self.openmp_env_vars, - conda_root="CONDA_PREFIX", - system_roots=self._openmp_system_roots, + conda="CONDA_PREFIX", + system=self.system_prefixes, ) # libraries result["library"] = _first_matching_file( - _path_candidates(candidate_roots, self.openmp_library_suffixes), + _path_candidates(candidate_roots, self._library_suffixes), lib_dict_info["lib_basename_candidate"], ) #### KeOps_Message("OpenMP library search in common locations:", level=2) KeOps_Message(f" Candidate roots: {candidate_roots}", level=2) - KeOps_Message( - f" Library search suffixes: {self.openmp_library_suffixes}", level=2 - ) + KeOps_Message(f" Library search suffixes: {self._library_suffixes}", level=2) KeOps_Message( f" Found library path: {result['library'] or not_found_str}", level=2 ) @@ -167,16 +162,14 @@ def find_install_path(self, lib_dict_info): # headers result["header"] = _first_matching_file( - _path_candidates(candidate_roots, self._openmp_include_sufixes), + _path_candidates(candidate_roots, self._include_sufixes), lib_dict_info["header_basename"], ) #### KeOps_Message("OpenMP header search in common locations:", level=2) KeOps_Message(f" Candidate roots: {candidate_roots}", level=2) - KeOps_Message( - f" Header search suffixes: {self._openmp_include_sufixes}", level=2 - ) + KeOps_Message(f" Header search suffixes: {self._include_sufixes}", level=2) KeOps_Message( f" Found header path: {result['header'] or not_found_str}", level=2 ) diff --git a/keopscore/keopscore/config/Platform.py b/keopscore/keopscore/config/Platform.py index ab17145b7..aba21036b 100644 --- a/keopscore/keopscore/config/Platform.py +++ b/keopscore/keopscore/config/Platform.py @@ -1,8 +1,11 @@ import os import platform +import site import sys +import sysconfig -from ..utils.messages import print_envs +from keopscore.utils.path_utils import _unique_paths +from keopscore.utils.messages import not_found_str, print_envs from keopscore.utils.system_utils import KeOps_OS_Run @@ -17,6 +20,7 @@ class PlatformConfig: _uname = "" _python_version = "" _python_executable = "" + _python_package_roots = [] _env_type = "" _brew_prefix = "" @@ -35,6 +39,7 @@ def __init__(self): self.set_uname() self.set_python_version() self.set_python_executable() + self.set_python_package_roots() self.set_env_type() self.set_brew_prefix() @@ -110,6 +115,38 @@ def get_python_executable(self): def print_python_executable(self): print(f"Python Executable: {self.get_python_executable()}") + # Python package roots detection + def set_python_package_roots(self): + """Set the Python package roots where pip-installed wheels may live.""" + self._python_package_roots = self.detect_python_package_roots() + + def get_python_package_roots(self): + return self._python_package_roots + + def print_python_package_roots(self): + roots = self.get_python_package_roots() + rendered_roots = ":".join(str(path) for path in roots) + print(f"Python Package Roots: {rendered_roots or not_found_str}") + + @staticmethod + def detect_python_package_roots(): + """Return Python package roots where pip-installed wheels may live.""" + package_roots = [] + getters = [ + lambda: getattr(site, "getsitepackages", lambda: [])(), + lambda: [site.getusersitepackages()], + lambda: [sysconfig.get_path("purelib")], + lambda: [sysconfig.get_path("platlib")], + lambda: sys.path, + ] + for getter in getters: + try: + # Some site helpers are unavailable in embedded or non-standard Python builds. + package_roots.extend(getter() or []) + except Exception: + continue + return _unique_paths(package_roots) + # Environment Type Detection def set_env_type(self): """Set the environment type (conda, virtualenv, or system).""" @@ -150,7 +187,7 @@ def print_brew_prefix(self): @staticmethod def detect_env_type(): """Return whether Python runs in conda, virtualenv, or the system env.""" - if "CONDA_DEFAULT_ENV" in os.environ: + if "CONDA_DEFAULT_ENV" in os.environ and os.environ["CONDA_DEFAULT_ENV"]: return f"conda ({os.environ['CONDA_DEFAULT_ENV']})" if hasattr(sys, "real_prefix") or ( hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix @@ -172,6 +209,7 @@ def print_all(self): self.print_python_version() self.print_env_type() self.print_python_executable() + self.print_python_package_roots() self.print_brew_prefix() diff --git a/keopscore/keopscore/config/__init__.py b/keopscore/keopscore/config/__init__.py index 458501a46..20a456230 100644 --- a/keopscore/keopscore/config/__init__.py +++ b/keopscore/keopscore/config/__init__.py @@ -16,7 +16,7 @@ platform = PlatformConfig() cxx = CxxCompilerConfig(platform) openmp = OpenMPConfig(platform, cxx) -cuda = CudaConfig() +cuda = CudaConfig(platform) path = KeOpsPathConfig(platform, cuda) reduction = ReductionTuningConfig() diff --git a/keopscore/keopscore/utils/gpu_utils.py b/keopscore/keopscore/utils/gpu_utils.py deleted file mode 100644 index be1e82d63..000000000 --- a/keopscore/keopscore/utils/gpu_utils.py +++ /dev/null @@ -1,59 +0,0 @@ -import os - -import keopscore.config -from keopscore.utils.file_utils import pack_header -from keopscore.utils.messages import KeOps_Error -from keopscore.utils.path_utils import _first_matching_file -from keopscore.utils.system_utils import get_include_file_abspath - - -def orig_cuda_include_fp16_path(): - """ - We look for float 16 cuda headers cuda_fp16.h and cuda_fp16.hpp - based on cuda_path locations and return their directory - """ - - # First try to find the library file using the cuda includes - cuda_include_path = keopscore.config.cuda.get_cuda_include_path() - cuda_fp16_h_abspath = _first_matching_file( - cuda_include_path, - "cuda_fp16.h", - ) - cuda_fp16_hpp_abspath = _first_matching_file( - cuda_include_path, - "cuda_fp16.hpp", - ) - - if cuda_fp16_h_abspath and cuda_fp16_hpp_abspath: - return os.path.dirname(cuda_fp16_h_abspath) - - # Second try with compiler - cuda_fp16_h_abspath = get_include_file_abspath("cuda_fp16.h") - cuda_fp16_hpp_abspath = get_include_file_abspath("cuda_fp16.hpp") - if cuda_fp16_h_abspath and cuda_fp16_hpp_abspath: - path = os.path.dirname(cuda_fp16_h_abspath) - if path != os.path.dirname(cuda_fp16_hpp_abspath): - KeOps_Error("cuda_fp16.h and cuda_fp16.hpp are not in the same folder !") - return path - else: - KeOps_Error("cuda_fp16.h and cuda_fp16.hpp were not found") - - -def custom_cuda_include_fp16_path(): - """ - Here we will create (if not done already) a custom cuda_fp16.h file - to be included in nvrtc code compilation, and put it in the keops - build folder. - We need to create this custom cuda_fp16.h header because the original - cuda_fp16.h includes other cuda headers, and for some unknown reason, - providing all the recursively required headers to the nvrtc compiler - does not work. Hence we produce a packed stand-alone version of cuda_fp16.h - by replacing all #include statements by the corresponding headers contents. - """ - - build_folder = keopscore.config.path.get_build_folder() - fp16_header = "cuda_fp16.h" - fp16_header_path = os.path.join(build_folder, fp16_header) - if not os.path.isfile(fp16_header_path): - pack_header(fp16_header, orig_cuda_include_fp16_path(), build_folder) - return build_folder diff --git a/keopscore/keopscore/utils/path_utils.py b/keopscore/keopscore/utils/path_utils.py index 1da2babd8..2b12baa05 100644 --- a/keopscore/keopscore/utils/path_utils.py +++ b/keopscore/keopscore/utils/path_utils.py @@ -1,18 +1,18 @@ import os -import site import sys -import sysconfig from pathlib import Path def _unique_paths(paths): - """Return paths in first-seen order with duplicates and None values removed.""" + """Return existing paths in first-seen order with duplicates and None values removed.""" unique = [] seen = set() for path in paths: if path is None: continue path = Path(path) + if not path.exists() and not path.is_dir(): + continue key = str(path) if key not in seen: seen.add(key) @@ -34,56 +34,35 @@ def _path_candidates(roots, suffixes): return _unique_paths(candidates) -def _python_package_roots(): - """Return Python package roots where pip-installed wheels may live.""" - package_roots = [] - getters = [ - lambda: getattr(site, "getsitepackages", lambda: [])(), - lambda: [site.getusersitepackages()], - lambda: [sysconfig.get_path("purelib")], - lambda: [sysconfig.get_path("platlib")], - lambda: sys.path, - ] - for getter in getters: - try: - # Some site helpers are unavailable in embedded or non-standard Python builds. - package_roots.extend(getter() or []) - except Exception: - continue - return _unique_paths(package_roots) - - -def _ordered_search_roots( - env_vars=(), pip_suffixes=(), conda_root=None, system_roots=(), order=None -): +def _ordered_search_roots(env_vars=(), pip=(), conda=None, system=(), order=None): """Return normalized search roots in a caller-defined category order. Args: env_vars (tuple[str] | list[str]): Environment variable names whose values should be considered as root directories. - pip_suffixes (tuple[str] | list[str]): Relative suffixes appended to Python + pip (tuple[str] | list[str]): Relative suffixes appended to Python package roots discovered from the current interpreter. - conda_root (str | None): Name of an environment variable that points to a + conda (str | None): Name of an environment variable that points to a Conda root directory. - system_roots (tuple[str] | list[str]): Fallback root directories to append. + system (tuple[str] | list[str]): Fallback root directories to append. order (tuple[str] | list[str] | str | None): Ordered categories to apply. - Allowed items are ``env_vars``, ``pip_suffixes``, ``conda_root``, and - ``system_roots``. A comma-separated string is also accepted. + Allowed items are ``env_vars``, ``pip``, ``conda``, and + ``system``. A comma-separated string is also accepted. Returns: list[pathlib.Path]: Existing candidates, deduplicated in first-seen order. """ if order is None: - order = ("env_vars", "conda_root", "system_roots", "pip_suffixes") + order = ("env_vars", "conda", "system", "pip") else: isinstance(order, str) and (order := tuple(order.split(","))) for item in order: if item not in ( "env_vars", - "pip_suffixes", - "conda_root", - "system_roots", + "pip", + "conda", + "system", ): raise ValueError(f"Invalid order item: {item}") @@ -93,15 +72,14 @@ def _ordered_search_roots( if item == "env_vars" and env_vars: roots.extend(_env_roots(env_vars)) - if item == "pip_suffixes" and pip_suffixes: - # Pip wheels such as nvidia-cuda-runtime expose libraries under site-packages. - roots.extend(_path_candidates(_python_package_roots(), pip_suffixes)) + if item == "pip" and pip: + roots.extend(_unique_paths(pip)) - if item == "conda_root" and conda_root: - roots.extend(_env_roots((conda_root,))) + if item == "conda" and conda: + roots.extend(_env_roots((conda,))) - if item == "system_roots" and system_roots: - roots.extend(system_roots) + if item == "system" and system: + roots.extend(_unique_paths(system)) return _unique_paths(roots) diff --git a/keopscore/keopscore/utils/system_utils.py b/keopscore/keopscore/utils/system_utils.py index ca0ba4516..5ef5f7ad2 100644 --- a/keopscore/keopscore/utils/system_utils.py +++ b/keopscore/keopscore/utils/system_utils.py @@ -72,6 +72,9 @@ class LINKMAP(Structure): def _find_library_by_names(library_names): """Return the first library path resolved by ctypes for known library names.""" + if isinstance(library_names, str): + library_names = (library_names,) + for library_name in library_names: library_path = find_library(library_name) if library_path: diff --git a/pykeops/setup.py b/pykeops/setup.py index 166551a70..d19cadbbb 100644 --- a/pykeops/setup.py +++ b/pykeops/setup.py @@ -94,5 +94,7 @@ "pandas", ], "test:": ["pytest", "numpy", "torch"], + "cu12": ["cuda-toolkit[nvrtc,nvcc,cudart,cccl]==12.9.2"], + "cu13": ["cuda-toolkit[nvrtc,nvcc,cccl,cudart, crt]==13.*"], }, ) From 0a9a3fd01ab959154e5d44dfa517d84639dc2226 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Fri, 29 May 2026 15:49:19 +0200 Subject: [PATCH 86/98] test without custom headers --- .../binders/nvrtc/Gpu_link_compile.py | 10 +-- .../keopscore/binders/nvrtc/nvrtc_jit.cpp | 74 ++++++++++--------- 2 files changed, 43 insertions(+), 41 deletions(-) diff --git a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py index cf1913e68..bcbd72f42 100644 --- a/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py +++ b/keopscore/keopscore/binders/nvrtc/Gpu_link_compile.py @@ -60,17 +60,17 @@ def generate_code(self): self.write_code() # we execute the main dll, passing the code as argument, and the name of the low level code file to save the assembly instructions + cuda_include_paths = "\n".join( + path for path in keopscore.config.cuda.get_cuda_include_path() if path + ) + res = self.my_c_dll.Compile( ctypes.create_string_buffer(self.low_level_code_file), ctypes.create_string_buffer(self.code.encode("utf-8")), ctypes.c_int(self.use_half), ctypes.c_int(self.use_fast_math), ctypes.c_int(self.device_id), - ctypes.create_string_buffer( - ( - keopscore.config.cuda.custom_cuda_include_fp16_path() + os.path.sep - ).encode("utf-8") - ), + ctypes.create_string_buffer(cuda_include_paths.encode("utf-8")), ) if res != keopscore.config.cuda.CUDA_SUCCESS: KeOps_Error( diff --git a/keopscore/keopscore/binders/nvrtc/nvrtc_jit.cpp b/keopscore/keopscore/binders/nvrtc/nvrtc_jit.cpp index 717bb89a9..9b219e787 100644 --- a/keopscore/keopscore/binders/nvrtc/nvrtc_jit.cpp +++ b/keopscore/keopscore/binders/nvrtc/nvrtc_jit.cpp @@ -18,6 +18,7 @@ #include #include #include +#include #include #include @@ -36,22 +37,21 @@ extern "C" int Compile(const char *target_file_name, const char *cu_code, int use_half, int use_fast_math, int device_id, - const char *cuda_include_path) { + const char *cuda_include_paths) { nvrtcProgram prog; - int numHeaders; - const char *header_names[1]; - const char *header_sources[1]; - - if (use_half) { - numHeaders = 1; - std::ostringstream header_path; - header_path << cuda_include_path << "cuda_fp16.h"; - header_names[0] = "cuda_fp16.h"; - header_sources[0] = read_text_file(header_path.str().c_str()); - } else { - numHeaders = 0; + std::vector compile_options_storage; + std::vector compile_options; + + if (cuda_include_paths != nullptr && cuda_include_paths[0] != '\0') { + std::istringstream include_paths_stream(cuda_include_paths); + std::string include_path; + while (std::getline(include_paths_stream, include_path)) { + if (!include_path.empty()) { + compile_options_storage.emplace_back("--include-path=" + include_path); + } + } } // Get device id from Driver API @@ -71,39 +71,43 @@ extern "C" int Compile(const char *target_file_name, const char *cu_code, arch_flag << "-arch=" << ARCHTAG << "_" << deviceProp_major << deviceProp_minor; - char *arch_flag_char = new char[arch_flag.str().length()]; - arch_flag_char = strdup(arch_flag.str().c_str()); + compile_options_storage.push_back(arch_flag.str()); + if (use_fast_math) { + compile_options_storage.emplace_back("-use_fast_math"); + } + compile_options.reserve(compile_options_storage.size()); + for (const std::string &option : compile_options_storage) { + compile_options.push_back(option.c_str()); + } NVRTC_SAFE_CALL(nvrtcCreateProgram(&prog, // prog cu_code, // buffer NULL, // name - numHeaders, // numHeaders - header_sources, // headers - header_names // includeNames + 0, // numHeaders + NULL, // headers + NULL // includeNames )); - nvrtcResult compileResult; - if (use_fast_math) { - const char *opts[] = {arch_flag_char, "-use_fast_math"}; - compileResult = nvrtcCompileProgram(prog, // prog - 2, // numOptions - opts); // options - } else { - const char *opts[] = {arch_flag_char}; - compileResult = nvrtcCompileProgram(prog, // prog - 1, // numOptions - opts); // options - } + nvrtcResult compileResult = + nvrtcCompileProgram(prog, // prog + compile_options.size(), // numOptions + compile_options.data()); // options // following "if" block is when there is a mismatch between // the device compute capability and the cuda libs versions : typically // when the device is more recent than the lib, the -arch flag may fail to // compile. if (compileResult == NVRTC_ERROR_INVALID_OPTION) { - const char *new_opts[] = {"-use_fast_math"}; - compileResult = nvrtcCompileProgram(prog, // prog - 1, // numOptions - new_opts); // options + std::vector fallback_options; + fallback_options.reserve(compile_options_storage.size()); + for (const std::string &option : compile_options_storage) { + if (option.rfind("-arch=", 0) != 0) { + fallback_options.push_back(option.c_str()); + } + } + compileResult = nvrtcCompileProgram(prog, // prog + fallback_options.size(), // numOptions + fallback_options.data()); // options } if (compileResult != NVRTC_SUCCESS) { @@ -119,8 +123,6 @@ extern "C" int Compile(const char *target_file_name, const char *cu_code, return compileResult; } - delete[] arch_flag_char; - // Obtain PTX or CUBIN from the program. size_t targetSize; NVRTC_SAFE_CALL(nvrtcGetTARGETSize(prog, &targetSize)); From 9e13ff074bc92c501de20ee2157c42fbe35c8549 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Fri, 29 May 2026 16:16:48 +0200 Subject: [PATCH 87/98] fix OpenMP import on macOS --- keopscore/keopscore/config/Cuda.py | 81 ++++++++++------------------ keopscore/keopscore/config/OpenMP.py | 23 ++++++-- 2 files changed, 47 insertions(+), 57 deletions(-) diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index bae3f8f14..0f8feba30 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -5,7 +5,6 @@ from keopscore.utils.messages import print_envs from keopscore.utils.messages import enabled_dict, not_found_str from keopscore.utils.messages import KeOps_Error, KeOps_Warning -from keopscore.utils.file_utils import pack_header from keopscore.utils.path_utils import ( _first_matching_file, _ordered_search_roots, @@ -115,7 +114,6 @@ class CudaConfig: "name": "nvrtc", "lib_basename_candidate": ["libnvrtc.so.*"], "library": "", # to be filled later - "ctype_handle": None, # to be filled later } _libnvrtc_builtins_info = { "name": "nvrtc-builtins", @@ -124,7 +122,6 @@ class CudaConfig: "libnvrtc-builtins.alt.so*", ], "library": "", # to be filled later - "ctype_handle": None, # to be filled later } _headers_nvrtc_info = { "header_basename": "nvrtc.h", @@ -293,8 +290,8 @@ def _find_and_load_libcuda(self, where_to_search): return True, "" - def _find_and_load_libnvrtc(self, where_to_search): - """Locate and load the NVRTC runtime compilation library.""" + def _find_libnvrtc(self, where_to_search): + """Locate the NVRTC runtime compilation library.""" self._libnvrtc_info = self.find_library_path( self._libnvrtc_info, where_to_search ) @@ -305,25 +302,15 @@ def _find_and_load_libnvrtc(self, where_to_search): "libnvrtc not found. Make sure the CUDA toolkit is installed and accessible.", ) - try: - libnvrtc_handle = ctypes.CDLL(libnvrtc_path, mode=ctypes.RTLD_GLOBAL) - except OSError as e: - return ( - False, - f"Failed to load library '{os.path.basename(libnvrtc_path)}': {e}", - ) - - # If we successfully loaded libnvrtc, store the handle in the config - self._libnvrtc_info["ctype_handle"] = libnvrtc_handle - return True, "" - def _find_and_load_libnvrtc_builtins(self, folder_to_search): + def _find_libnvrtc_builtins(self, folder_to_search): """ - Preload of NVRTC builtins from the same folder as libnvrtc. + Locate NVRTC builtins from the same folder as libnvrtc. This check is non-blocking: warnings are emitted on failure, and CUDA - detection continues. + detection continues. We intentionally do not preload this library with + ctypes; runtime resolution is handled through linker rpath flags. """ self._libnvrtc_builtins_info = self.find_library_path( @@ -336,17 +323,6 @@ def _find_and_load_libnvrtc_builtins(self, folder_to_search): f"NVRTC builtins library not found in {folder_to_search}. This may cause runtime compilation to fail in environments with multiple CUDA toolkit versions installed.", ) - try: - handle = ctypes.CDLL(libnvrtc_builtins_path, mode=ctypes.RTLD_GLOBAL) - except OSError as err: - return ( - True, - f"Failed to load NVRTC builtins library '{os.path.basename(libnvrtc_builtins_path)}': {err}", - ) - - # If we successfully loaded the library, store the handle in the config for potential future use - self._libnvrtc_builtins_info["ctype_handle"] = handle - return True, "" def _cuda_libraries_available(self): @@ -388,15 +364,15 @@ def _cuda_libraries_available(self): # libnvrtc as well, since it's required for the runtime compilation of CUDA code. # This is usually provided by the CUDA Toolkit (install system-wide or in conda/pip). - success_nvrtc, err_nvrtc = self._find_and_load_libnvrtc(where_to_search) + success_nvrtc, err_nvrtc = self._find_libnvrtc(where_to_search) if not success_nvrtc: KeOps_Warning(f"{err_nvrtc}. Switching to CPU only.") return False - # Try load of NVRTC builtins from the same toolkit folder. - # This avoids runtime search path ambiguity in environments where + # Locate NVRTC builtins in the same toolkit folder. + # This helps set runtime search paths in environments where # multiple CUDA toolkit versions are installed. - _, warning_nvrtc = self._find_and_load_libnvrtc_builtins( + _, warning_nvrtc = self._find_libnvrtc_builtins( os.path.dirname(self._libnvrtc_info["library"]) ) if warning_nvrtc: @@ -577,27 +553,7 @@ def print_cuda_include_path(self): f"CUDA Include Path: {':'.join(self.get_cuda_include_path()) or not_found_str}" ) - def custom_cuda_include_fp16_path(self): - """ - Create (if needed) a packed standalone cuda_fp16.h in the KeOps build - folder for NVRTC compilation and return this folder. - """ - import keopscore.config - - build_folder = keopscore.config.path.get_build_folder() - fp16_header = "cuda_fp16.h" - fp16_header_path = os.path.join(build_folder, fp16_header) - if not os.path.isfile(fp16_header_path): - pack_header( - fp16_header, - os.path.dirname(self._headers_fp16_info["header"]), - build_folder, - ) - - return build_folder - # IR type - def set_ir_type(self): """Set the IR type to be used for nvrtc compilation based on the CUDA version.""" if self.get_cuda_version() >= 11010: @@ -668,6 +624,23 @@ def set_linking_options(self): lib_info["library"] if lib_info["library"] else f"-l{lib_info['name']}" ) + # Ensure runtime loader can resolve CUDA toolkit side dependencies + # (e.g. nvrtc-builtins) without requiring ctypes preloading. + rpath_dirs = [] + for lib_info in [ + self._libcuda_info, + self._libnvrtc_info, + self._libnvrtc_builtins_info, + ]: + lib_path = lib_info.get("library") + if lib_path: + lib_dir = os.path.dirname(lib_path) + if lib_dir and lib_dir not in rpath_dirs: + rpath_dirs.append(lib_dir) + + for lib_dir in rpath_dirs: + link_options.append(f"-Wl,-rpath,{lib_dir}") + self._linking_options = " ".join(link_options) def get_linking_options(self): diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 7e7424ad2..ff9c637da 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -322,9 +322,26 @@ def print_include_options(self): # C++ linking Options def set_linking_options(self): - self._linking_options += ( - f"-L{self.get_libomp_folder()}" if self.get_libomp_folder() else "" - ) + link_flags = [] + libomp_folder = self.get_libomp_folder() + libomp_path = self.get_libomp_path() + + if libomp_folder: + link_flags.append(f"-L{libomp_folder}") + + # Force-link OpenMP runtime to avoid unresolved symbols at dlopen time. + if libomp_path: + lib_basename = os.path.basename(libomp_path) + if lib_basename.startswith("lib"): + lib_name = lib_basename[3:].split(".")[0] + if lib_name: + link_flags.append(f"-l{lib_name}") + + # Ensure runtime loader can find libomp on macOS non-system paths. + if self.platform.get_platform() == "Darwin" and libomp_folder: + link_flags.append(f"-Wl,-rpath,{libomp_folder}") + + self._linking_options = " ".join(link_flags) def get_linking_options(self): return self._linking_options From 9b2f25b75d09cf242aafb84b988b72bd7a0a0260 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Fri, 29 May 2026 18:20:42 +0200 Subject: [PATCH 88/98] fix link to openmp lib on MacOS --- keopscore/keopscore/config/OpenMP.py | 5 ++-- keopscore/keopscore/utils/file_utils.py | 28 ----------------------- keopscore/keopscore/utils/system_utils.py | 22 ++++++++++++------ 3 files changed, 17 insertions(+), 38 deletions(-) diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index ff9c637da..639ffffde 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -62,8 +62,6 @@ class OpenMPConfig: "libomp.so*", "libgomp.dylib", "libgomp.so*", - "libm.dylib", - "libm.so*", ], "header_basename": "omp.h", "library": "", # to be filled later @@ -323,9 +321,10 @@ def print_include_options(self): # C++ linking Options def set_linking_options(self): link_flags = [] - libomp_folder = self.get_libomp_folder() libomp_path = self.get_libomp_path() + libomp_folder = libomp_path and os.path.dirname(libomp_path) + if libomp_folder: link_flags.append(f"-L{libomp_folder}") diff --git a/keopscore/keopscore/utils/file_utils.py b/keopscore/keopscore/utils/file_utils.py index a352bc37e..478dce8d6 100644 --- a/keopscore/keopscore/utils/file_utils.py +++ b/keopscore/keopscore/utils/file_utils.py @@ -12,31 +12,3 @@ def string_to_file(string, file_path): """Write a string to a text file.""" with open(file_path, "w", encoding="utf-8") as file: file.write(string) - - -def pack_header(filename, origin_folder, target_folder): - """ - Produce a stand-alone C/C++ header by inlining local quoted includes. - """ - code = file_to_string(os.path.join(origin_folder, filename)) - used_headers = [] - while True: - match = re.search('#include *"([^"]*)"', code) - if match is None: - break - header = match.groups()[0] - if header in used_headers: - code_to_insert = "" - else: - if not os.path.exists(os.path.join(origin_folder, header)): - header_found = os.path.basename(header) - elif os.path.exists(os.path.join(origin_folder, header)): - header_found = header - else: - raise FileNotFoundError( - f"Header file {header} not found in {origin_folder}." - ) - code_to_insert = file_to_string(os.path.join(origin_folder, header_found)) - used_headers.append(header) - code = code[: match.start()] + code_to_insert + code[match.end() :] - string_to_file(code, os.path.join(target_folder, filename)) diff --git a/keopscore/keopscore/utils/system_utils.py b/keopscore/keopscore/utils/system_utils.py index 5ef5f7ad2..a9c082123 100644 --- a/keopscore/keopscore/utils/system_utils.py +++ b/keopscore/keopscore/utils/system_utils.py @@ -59,15 +59,23 @@ class LINKMAP(Structure): try: dlinfo = libdl.dlinfo except AttributeError: - return "" - dlinfo.argtypes = c_void_p, c_int, c_void_p - dlinfo.restype = c_int + dlinfo = None + + if dlinfo is not None: + try: + dlinfo.argtypes = c_void_p, c_int, c_void_p + dlinfo.restype = c_int + + lmptr = c_void_p() + dlinfo(lib._handle, 2, byref(lmptr)) - lmptr = c_void_p() - dlinfo(lib._handle, 2, byref(lmptr)) + abspath = cast(lmptr, POINTER(LINKMAP)).contents.l_name + if abspath: + return abspath.decode("utf-8") + except Exception: + pass - abspath = cast(lmptr, POINTER(LINKMAP)).contents.l_name - return abspath.decode("utf-8") + return "" def _find_library_by_names(library_names): From 03c71526034c4b75059a1135972f873caec7dd27 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Thu, 4 Jun 2026 16:23:07 +0200 Subject: [PATCH 89/98] split strategy on MacOS openmp linkikng --- keopscore/keopscore/binders/LinkCompile.py | 3 ++ keopscore/keopscore/config/OpenMP.py | 18 +++++----- keopscore/keopscore/utils/Cache.py | 3 ++ pytest.sh | 41 ++++++++++++++++++++-- 4 files changed, 54 insertions(+), 11 deletions(-) diff --git a/keopscore/keopscore/binders/LinkCompile.py b/keopscore/keopscore/binders/LinkCompile.py index 9781f7d8c..ab79001f2 100644 --- a/keopscore/keopscore/binders/LinkCompile.py +++ b/keopscore/keopscore/binders/LinkCompile.py @@ -6,6 +6,9 @@ cpp_flag = keopscore.config.cxx.get_compile_options() cpp_flag += keopscore.config.cxx.get_linking_options() +cpp_flag += keopscore.config.openmp.get_compile_options() +cpp_flag += keopscore.config.openmp.get_include_options() +cpp_flag += keopscore.config.openmp.get_linking_options() cpp_flag += keopscore.config.cuda.get_include_options() cpp_flag += keopscore.config.cuda.get_linking_options() diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 639ffffde..e8f151ede 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -1,3 +1,4 @@ +import importlib.util import os import tempfile @@ -50,7 +51,7 @@ class OpenMPConfig: os.path.join("opt", "libomp", "lib"), ) - _include_sufixes = ( + _include_suffixes = ( "include", os.path.join("opt", "libomp", "include"), ) @@ -160,14 +161,14 @@ def find_install_path(self, lib_dict_info): # headers result["header"] = _first_matching_file( - _path_candidates(candidate_roots, self._include_sufixes), + _path_candidates(candidate_roots, self._include_suffixes), lib_dict_info["header_basename"], ) #### KeOps_Message("OpenMP header search in common locations:", level=2) KeOps_Message(f" Candidate roots: {candidate_roots}", level=2) - KeOps_Message(f" Header search suffixes: {self._include_sufixes}", level=2) + KeOps_Message(f" Header search suffixes: {self._include_suffixes}", level=2) KeOps_Message( f" Found header path: {result['header'] or not_found_str}", level=2 ) @@ -296,7 +297,6 @@ def set_compile_options(self): self._compile_options += " -Xpreprocessor" self._compile_options += " -fopenmp" - self._compile_options def get_compile_options(self): return self._compile_options.strip() @@ -328,22 +328,22 @@ def set_linking_options(self): if libomp_folder: link_flags.append(f"-L{libomp_folder}") - # Force-link OpenMP runtime to avoid unresolved symbols at dlopen time. - if libomp_path: + # Ensure runtime loader can find libomp on macOS non-system paths. + if self.platform.get_platform() == "Darwin" and libomp_folder and not importlib.util.find_spec("torch"): + # Force-link OpenMP runtime to avoid unresolved symbols at dlopen time. + lib_basename = os.path.basename(libomp_path) if lib_basename.startswith("lib"): lib_name = lib_basename[3:].split(".")[0] if lib_name: link_flags.append(f"-l{lib_name}") - # Ensure runtime loader can find libomp on macOS non-system paths. - if self.platform.get_platform() == "Darwin" and libomp_folder: link_flags.append(f"-Wl,-rpath,{libomp_folder}") self._linking_options = " ".join(link_flags) def get_linking_options(self): - return self._linking_options + return self._linking_options.rstrip() def print_linking_options(self): print(f"Linking Options: {self.get_linking_options()}") diff --git a/keopscore/keopscore/utils/Cache.py b/keopscore/keopscore/utils/Cache.py index 151235fdd..87ae47f36 100644 --- a/keopscore/keopscore/utils/Cache.py +++ b/keopscore/keopscore/utils/Cache.py @@ -9,6 +9,9 @@ env_param = ( lambda: keopscore.config.cxx.get_compile_options() + keopscore.config.cxx.get_linking_options() + + keopscore.config.openmp.get_compile_options() + + keopscore.config.openmp.get_include_options() + + keopscore.config.openmp.get_linking_options() + f" auto_factorize={keopscore.config.reduction.get_auto_factorize()}" + keopscore.config.cuda.get_preprocessing_options() + keopscore.config.cuda.get_include_options() diff --git a/pytest.sh b/pytest.sh index 2fa155bf2..98cc40586 100755 --- a/pytest.sh +++ b/pytest.sh @@ -5,6 +5,7 @@ set -euo pipefail readonly PYTHON_BIN="python3" readonly SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" readonly TEST_VENV="${SCRIPT_DIR}/.test_venv_pytest" +readonly NO_TORCH_TEST_VENV="${SCRIPT_DIR}/.test_venv_pytest_no_torch" readonly TEST_REQUIREMENTS=(pip) KEOPS_VERBOSE_LEVEL=-1 @@ -19,7 +20,7 @@ Usage: $0 [option...] -h Print the help -v <0|1|2> Verbosity level forwarded to KEOPS_VERBOSE and PYKEOPS_VERBOSE - --pip-constraint + --pip-constraint Constrain pip installs using the given constraints file. EOF } @@ -84,6 +85,17 @@ run_with_keops_verbose() { fi } +run_python_with_keops_verbose() { + local python_bin="$1" + local python_code="$2" + + if [[ "${KEOPS_VERBOSE_LEVEL}" -ne -1 ]]; then + KEOPS_VERBOSE="${KEOPS_VERBOSE_LEVEL}" PYKEOPS_VERBOSE="${KEOPS_VERBOSE_LEVEL}" "${python_bin}" -c "${python_code}" + else + "${python_bin}" -c "${python_code}" + fi +} + pip_install() { if [[ -n "${PIP_CONSTRAINT_FILE}" ]]; then run_with_keops_verbose "${PYTHON_BIN}" -m pip install --constraint "${PIP_CONSTRAINT_FILE}" "$@" @@ -116,7 +128,7 @@ run_python_outside_repo() { local python_code="$1" ( cd /tmp - run_with_keops_verbose "${PYTHON_BIN}" -c "${python_code}" + run_python_with_keops_verbose "${PYTHON_BIN}" "${python_code}" ) } @@ -130,6 +142,30 @@ run_pykeops_health_check() { run_python_outside_repo 'import pykeops; pykeops.check_health()' } +run_pykeops_no_torch_smoke_test() { + local no_torch_python="${NO_TORCH_TEST_VENV}/bin/python" + local no_torch_smoke_code="import importlib.util;" + no_torch_smoke_code+=" assert importlib.util.find_spec('torch') is None," + no_torch_smoke_code+=" 'torch must not be installed in no-torch smoke env';" + no_torch_smoke_code+=" import pykeops; pykeops.check_health();" + no_torch_smoke_code+=" assert pykeops.test_numpy_bindings()" + + echo "-- Running pykeops no-torch smoke test..." + "${PYTHON_BIN}" -m venv --clear "${NO_TORCH_TEST_VENV}" + + if [[ -n "${PIP_CONSTRAINT_FILE}" ]]; then + "${no_torch_python}" -m pip install --constraint "${PIP_CONSTRAINT_FILE}" -U pip + "${no_torch_python}" -m pip install --constraint "${PIP_CONSTRAINT_FILE}" -e "${SCRIPT_DIR}/keopscore" + "${no_torch_python}" -m pip install --constraint "${PIP_CONSTRAINT_FILE}" -e "${SCRIPT_DIR}/pykeops" + else + "${no_torch_python}" -m pip install -U pip + "${no_torch_python}" -m pip install -e "${SCRIPT_DIR}/keopscore" + "${no_torch_python}" -m pip install -e "${SCRIPT_DIR}/pykeops" + fi + + run_python_with_keops_verbose "${no_torch_python}" "${no_torch_smoke_code}" +} + run_test_suite() { local suite_name="$1" local suite_path="$2" @@ -140,6 +176,7 @@ run_test_suite() { main() { parse_options "$@" + run_pykeops_no_torch_smoke_test prepare_python_environment install_editable_package "keopscore" "${SCRIPT_DIR}/keopscore" install_editable_package "pykeops" "${SCRIPT_DIR}/pykeops[test]" From 88c6dcb49884e275eff261c006941e14e7f6b321 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Thu, 4 Jun 2026 16:24:37 +0200 Subject: [PATCH 90/98] lint --- keopscore/keopscore/config/OpenMP.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index e8f151ede..20d455e78 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -329,7 +329,11 @@ def set_linking_options(self): link_flags.append(f"-L{libomp_folder}") # Ensure runtime loader can find libomp on macOS non-system paths. - if self.platform.get_platform() == "Darwin" and libomp_folder and not importlib.util.find_spec("torch"): + if ( + self.platform.get_platform() == "Darwin" + and libomp_folder + and not importlib.util.find_spec("torch") + ): # Force-link OpenMP runtime to avoid unresolved symbols at dlopen time. lib_basename = os.path.basename(libomp_path) From 7929b4bb71f5aa1fc77f84a2b269c79b34bda41a Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Thu, 4 Jun 2026 16:27:51 +0200 Subject: [PATCH 91/98] remove shadowing in test --- pytest.sh | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pytest.sh b/pytest.sh index 98cc40586..9e1b8474b 100755 --- a/pytest.sh +++ b/pytest.sh @@ -163,7 +163,10 @@ run_pykeops_no_torch_smoke_test() { "${no_torch_python}" -m pip install -e "${SCRIPT_DIR}/pykeops" fi - run_python_with_keops_verbose "${no_torch_python}" "${no_torch_smoke_code}" + ( + cd /tmp + run_python_with_keops_verbose "${no_torch_python}" "${no_torch_smoke_code}" + ) } run_test_suite() { From 61ac93032c119a0248e34e37903bd1e0447b5df4 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Fri, 5 Jun 2026 23:30:22 +0200 Subject: [PATCH 92/98] prevent unconsistent CPUGPU configuration --- keopscore/keopscore/test/test_op.py | 6 +- .../test/test_block_sparse_reduction.py | 3 +- pykeops/pykeops/test/test_chunks.py | 4 +- pykeops/pykeops/test/test_chunks_ranges.py | 4 +- pykeops/pykeops/test/test_complex.py | 4 +- pykeops/pykeops/test/test_conv2d.py | 9 ++- pykeops/pykeops/test/test_finalchunks.py | 4 +- .../pykeops/test/test_finalchunks_ranges.py | 4 +- pykeops/pykeops/test/test_float16.py | 14 +++- pykeops/pykeops/test/test_gpu_cpu.py | 8 ++- pykeops/pykeops/test/test_lazytensor_clamp.py | 4 +- .../test/test_lazytensor_gaussian_batch.py | 4 +- .../test/test_lazytensor_gaussian_fromhost.py | 4 +- .../test/test_lazytensor_gaussian_inplace.py | 4 +- pykeops/pykeops/test/test_lazytensor_grad.py | 4 +- .../pykeops/test/test_lazytensor_tensordot.py | 4 +- pykeops/pykeops/test/test_shape_regression.py | 4 +- pykeops/pykeops/test/test_sinc.py | 4 +- pykeops/pykeops/test/test_torch.py | 20 +++--- pykeops/pykeops/torch/generic/generic_red.py | 10 ++- pytest.sh | 69 ++++++++++--------- 21 files changed, 126 insertions(+), 65 deletions(-) diff --git a/keopscore/keopscore/test/test_op.py b/keopscore/keopscore/test/test_op.py index 1aed2eb83..27f5b8536 100644 --- a/keopscore/keopscore/test/test_op.py +++ b/keopscore/keopscore/test/test_op.py @@ -8,6 +8,7 @@ import keopscore import keopscore.formulas +import pykeops.config as pykeopsconfig from keopscore.utils.messages import KeOps_Error from pykeops.torch import Genred @@ -63,8 +64,9 @@ def perform_test(op_str, tol=1e-4, dtype="float32", verbose=True): M = 300 N = 500 - # Choose the storage place for our data : CPU (host) or GPU (device) memory. - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + # Keep tensor placement aligned with the KeOps backend selected at runtime. + use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available + device = torch.device("cuda" if use_cuda else "cpu") if dtype == "float32": torchtype = torch.float32 diff --git a/pykeops/pykeops/test/test_block_sparse_reduction.py b/pykeops/pykeops/test/test_block_sparse_reduction.py index 90e1590d5..e9a1c1750 100644 --- a/pykeops/pykeops/test/test_block_sparse_reduction.py +++ b/pykeops/pykeops/test/test_block_sparse_reduction.py @@ -1,6 +1,7 @@ import time import numpy as np import torch +import pykeops.config as pykeopsconfig from pykeops.torch import LazyTensor # Import clustering functions from KeOps @@ -13,7 +14,7 @@ def test_block_sparse_reduction(): - use_cuda = torch.cuda.is_available() + use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available device = torch.device("cuda" if use_cuda else "cpu") dtype = torch.float32 diff --git a/pykeops/pykeops/test/test_chunks.py b/pykeops/pykeops/test/test_chunks.py index 0344b2345..7408e8245 100644 --- a/pykeops/pykeops/test/test_chunks.py +++ b/pykeops/pykeops/test/test_chunks.py @@ -1,5 +1,6 @@ import math import torch +import pykeops.config as pykeopsconfig from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -9,7 +10,8 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/pykeops/pykeops/test/test_chunks_ranges.py b/pykeops/pykeops/test/test_chunks_ranges.py index c48a03725..b0eafe409 100644 --- a/pykeops/pykeops/test/test_chunks_ranges.py +++ b/pykeops/pykeops/test/test_chunks_ranges.py @@ -1,5 +1,6 @@ import math import torch +import pykeops.config as pykeopsconfig from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -9,7 +10,8 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.rand(B1, B2, M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/pykeops/pykeops/test/test_complex.py b/pykeops/pykeops/test/test_complex.py index 9808ba70e..bc4570398 100644 --- a/pykeops/pykeops/test/test_complex.py +++ b/pykeops/pykeops/test/test_complex.py @@ -2,6 +2,7 @@ import math import torch +import pykeops.config as pykeopsconfig from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -11,7 +12,8 @@ M, N, D = 1000, 1000, 1 torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.rand(1, N, D, dtype=dtype_c, requires_grad=True, device=device_id) diff --git a/pykeops/pykeops/test/test_conv2d.py b/pykeops/pykeops/test/test_conv2d.py index 63bcbf62b..7ac2b4878 100644 --- a/pykeops/pykeops/test/test_conv2d.py +++ b/pykeops/pykeops/test/test_conv2d.py @@ -1,6 +1,7 @@ import math import unittest import torch +import pykeops.config as pykeopsconfig from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -11,7 +12,8 @@ torch.manual_seed(42) torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +device_id = "cuda" if use_cuda else "cpu" def fun(x, y, b, backend): @@ -33,7 +35,10 @@ def fun(x, y, b, backend): backends = ["keops2D", "torch"] -@unittest.skipUnless(torch.cuda.is_available(), "CUDA not available") +@unittest.skipUnless( + use_cuda, + "CUDA not available in torch or KeOps", +) class TestCase(unittest.TestCase): @classmethod def setUpClass(cls): diff --git a/pykeops/pykeops/test/test_finalchunks.py b/pykeops/pykeops/test/test_finalchunks.py index 4bcc05d0a..9665f2b73 100644 --- a/pykeops/pykeops/test/test_finalchunks.py +++ b/pykeops/pykeops/test/test_finalchunks.py @@ -1,5 +1,6 @@ import math import torch +import pykeops.config as pykeopsconfig from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -9,7 +10,8 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/pykeops/pykeops/test/test_finalchunks_ranges.py b/pykeops/pykeops/test/test_finalchunks_ranges.py index f42c08dae..14d97e0b0 100644 --- a/pykeops/pykeops/test/test_finalchunks_ranges.py +++ b/pykeops/pykeops/test/test_finalchunks_ranges.py @@ -1,5 +1,6 @@ import math import torch +import pykeops.config as pykeopsconfig from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -9,7 +10,8 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.rand(B1, B2, M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/pykeops/pykeops/test/test_float16.py b/pykeops/pykeops/test/test_float16.py index 1d2484a32..4552c273b 100644 --- a/pykeops/pykeops/test/test_float16.py +++ b/pykeops/pykeops/test/test_float16.py @@ -1,6 +1,7 @@ # Test for Clamp operation using LazyTensors import pytest import torch +import pykeops.config as pykeopsconfig from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -9,7 +10,8 @@ M, N, D = 5, 5, 1 torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.randn(M, 1, D, dtype=dtype, requires_grad=True, device=device_id) @@ -30,7 +32,10 @@ def fun(x, y, backend): class TestCase: out = [] - @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires a GPU") + @pytest.mark.skipif( + not use_cuda, + reason="Requires torch and KeOps CUDA support", + ) def test_float16_fw(self): for backend in ["torch", "keops"]: self.out.append(fun(x, y, backend).squeeze()) @@ -39,7 +44,10 @@ def test_float16_fw(self): self.out[0], self.out[1], atol=0.001, rtol=0.001, label="float16_fw" ) - @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires a GPU") + @pytest.mark.skipif( + not use_cuda, + reason="Requires torch and KeOps CUDA support", + ) def test_float16_bw(self): out_g = [] for k, backend in enumerate(["torch", "keops"]): diff --git a/pykeops/pykeops/test/test_gpu_cpu.py b/pykeops/pykeops/test/test_gpu_cpu.py index fe49cb7d0..89da4e426 100644 --- a/pykeops/pykeops/test/test_gpu_cpu.py +++ b/pykeops/pykeops/test/test_gpu_cpu.py @@ -1,8 +1,11 @@ import math import pytest import torch +import pykeops.config as pykeopsconfig from pykeops.torch import LazyTensor +use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available + M, N, D, DV = 1000, 1000, 3, 1 dtype = torch.float32 @@ -43,7 +46,10 @@ def test_torch_keops_cpu(self): out[0], out[1] ), f"torch vs keops_cpu mismatch: ||ref-test||_2={torch.norm(out[0] - out[1]).item():.6e} and ||ref||_2={torch.norm(out[0]).item():.6e}" - @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires a GPU") + @pytest.mark.skipif( + not use_cuda, + reason="Requires torch and KeOps CUDA support", + ) def test_torch_keops_gpu(self): out_gpu = fun(x, y, b, ["keops_gpu"]).squeeze() assert torch.allclose( diff --git a/pykeops/pykeops/test/test_lazytensor_clamp.py b/pykeops/pykeops/test/test_lazytensor_clamp.py index c9d962b03..6c0bcc652 100644 --- a/pykeops/pykeops/test/test_lazytensor_clamp.py +++ b/pykeops/pykeops/test/test_lazytensor_clamp.py @@ -1,11 +1,13 @@ import torch +import pykeops.config as pykeopsconfig from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose M, N, D = 1000, 1000, 3 torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.randn(M, 1, D, requires_grad=True, device=device_id) diff --git a/pykeops/pykeops/test/test_lazytensor_gaussian_batch.py b/pykeops/pykeops/test/test_lazytensor_gaussian_batch.py index 6f0a3779a..93f5cc86c 100644 --- a/pykeops/pykeops/test/test_lazytensor_gaussian_batch.py +++ b/pykeops/pykeops/test/test_lazytensor_gaussian_batch.py @@ -1,12 +1,14 @@ import math import torch +import pykeops.config as pykeopsconfig from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose B1, B2, M, N, D, DV = 3, 4, 20, 25, 3, 2 torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(1) x = torch.rand(1, B2, M, 1, D, device=device_id) / math.sqrt(D) diff --git a/pykeops/pykeops/test/test_lazytensor_gaussian_fromhost.py b/pykeops/pykeops/test/test_lazytensor_gaussian_fromhost.py index 1cc3b42e4..862f41cce 100644 --- a/pykeops/pykeops/test/test_lazytensor_gaussian_fromhost.py +++ b/pykeops/pykeops/test/test_lazytensor_gaussian_fromhost.py @@ -1,5 +1,6 @@ import math import torch +import pykeops.config as pykeopsconfig from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -9,7 +10,8 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = 0 if torch.cuda.is_available() else -1 +use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +device_id = 0 if use_cuda else -1 torch.manual_seed(0) x = torch.rand(M, 1, D, dtype=dtype) / math.sqrt(D) diff --git a/pykeops/pykeops/test/test_lazytensor_gaussian_inplace.py b/pykeops/pykeops/test/test_lazytensor_gaussian_inplace.py index 01d059764..f7d50c014 100644 --- a/pykeops/pykeops/test/test_lazytensor_gaussian_inplace.py +++ b/pykeops/pykeops/test/test_lazytensor_gaussian_inplace.py @@ -1,5 +1,6 @@ import math import torch +import pykeops.config as pykeopsconfig from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -9,7 +10,8 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.rand(M, 1, D, device=device_id, dtype=dtype) / math.sqrt(D) diff --git a/pykeops/pykeops/test/test_lazytensor_grad.py b/pykeops/pykeops/test/test_lazytensor_grad.py index 350742ed6..f83615fab 100644 --- a/pykeops/pykeops/test/test_lazytensor_grad.py +++ b/pykeops/pykeops/test/test_lazytensor_grad.py @@ -1,4 +1,5 @@ import torch +import pykeops.config as pykeopsconfig from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -8,7 +9,8 @@ dtype = torch.float32 torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) x = torch.rand(M, 1, D, requires_grad=True, device=device_id, dtype=dtype) diff --git a/pykeops/pykeops/test/test_lazytensor_tensordot.py b/pykeops/pykeops/test/test_lazytensor_tensordot.py index 2ca9ba8a1..a08d552f8 100644 --- a/pykeops/pykeops/test/test_lazytensor_tensordot.py +++ b/pykeops/pykeops/test/test_lazytensor_tensordot.py @@ -1,4 +1,5 @@ import torch +import pykeops.config as pykeopsconfig from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -6,7 +7,8 @@ # Matrix multiplication as a special case of Tensordot torch.backends.cuda.matmul.allow_tf32 = False -device_id = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) a = torch.randn(4 * 7, requires_grad=True, device=device_id, dtype=torch.float64) diff --git a/pykeops/pykeops/test/test_shape_regression.py b/pykeops/pykeops/test/test_shape_regression.py index a6e27e572..10f48cebd 100644 --- a/pykeops/pykeops/test/test_shape_regression.py +++ b/pykeops/pykeops/test/test_shape_regression.py @@ -1,6 +1,7 @@ import re import unittest import numpy as np +import pykeops.config as pykeopsconfig from pykeops.test import assert_np_allclose, assert_torch_allclose @@ -73,7 +74,8 @@ def test_numpy_vj_single_dim_formula_x_1d_rejected_cpp(self): class ShapeTorchTestCase(unittest.TestCase): def setUp(self): - device = "cuda" if torch.cuda.is_available() else "cpu" + use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available + device = "cuda" if use_cuda else "cpu" self.param = torch.tensor([0.4], dtype=torch.float32, device=device) self.param_scalar = torch.tensor(0.4, dtype=torch.float32, device=device) self.x = torch.tensor([[0.1], [0.2], [0.3]], dtype=torch.float32, device=device) diff --git a/pykeops/pykeops/test/test_sinc.py b/pykeops/pykeops/test/test_sinc.py index 996153ff3..0a79c3b6f 100644 --- a/pykeops/pykeops/test/test_sinc.py +++ b/pykeops/pykeops/test/test_sinc.py @@ -2,11 +2,13 @@ import pytest import torch +import pykeops.config as pykeopsconfig from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose -device = "cuda" if torch.cuda.is_available() else "cpu" +use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +device = "cuda" if use_cuda else "cpu" x = torch.rand(5, 1, dtype=torch.float64) * 2 * math.pi y = x.data.clone() diff --git a/pykeops/pykeops/test/test_torch.py b/pykeops/pykeops/test/test_torch.py index 7fd86ec25..e5ca563bc 100644 --- a/pykeops/pykeops/test/test_torch.py +++ b/pykeops/pykeops/test/test_torch.py @@ -19,6 +19,13 @@ import unittest +import pykeops +import pykeops.config +from pykeops.test import assert_torch_allclose + +HAS_TORCH = False +use_cuda = False + try: import torch @@ -26,19 +33,14 @@ torch.manual_seed(42) - use_cuda = torch.cuda.is_available() + use_cuda = torch.cuda.is_available() and pykeops.config.gpu_available + device = "cuda" if use_cuda else "cpu" + if use_cuda: - device = "cuda" torch.backends.cuda.matmul.allow_tf32 = False - else: - device = "cpu" except ImportError: - HAS_TORCH = False - -import pykeops -import pykeops.config -from pykeops.test import assert_torch_allclose + pass @unittest.skipUnless(HAS_TORCH, "torch not available") diff --git a/pykeops/pykeops/torch/generic/generic_red.py b/pykeops/pykeops/torch/generic/generic_red.py index 97e3eb898..edcabdfe2 100644 --- a/pykeops/pykeops/torch/generic/generic_red.py +++ b/pykeops/pykeops/torch/generic/generic_red.py @@ -45,7 +45,15 @@ def check_AD_supported(formula): def set_device(tagCPUGPU, tagHostDevice, device_id_request, *args): device_args = args[0].device - if tagCPUGPU == 1 & tagHostDevice == 1: + # Guard against unsupported configuration that can crash in native code: + # CPU backend cannot consume CUDA tensors. + if tagCPUGPU == 0 and tagHostDevice == 1: + raise ValueError( + "[KeOps] Incompatible configuration: CPU backend selected while input tensors are on CUDA. " + "Use CPU tensors or enable KeOps CUDA backend." + ) + + if tagCPUGPU == 1 and tagHostDevice == 1: for i in range(1, len(args)): if args[i].device.index != device_args.index: raise ValueError( diff --git a/pytest.sh b/pytest.sh index 9e1b8474b..56d39c315 100755 --- a/pytest.sh +++ b/pytest.sh @@ -85,17 +85,6 @@ run_with_keops_verbose() { fi } -run_python_with_keops_verbose() { - local python_bin="$1" - local python_code="$2" - - if [[ "${KEOPS_VERBOSE_LEVEL}" -ne -1 ]]; then - KEOPS_VERBOSE="${KEOPS_VERBOSE_LEVEL}" PYKEOPS_VERBOSE="${KEOPS_VERBOSE_LEVEL}" "${python_bin}" -c "${python_code}" - else - "${python_bin}" -c "${python_code}" - fi -} - pip_install() { if [[ -n "${PIP_CONSTRAINT_FILE}" ]]; then run_with_keops_verbose "${PYTHON_BIN}" -m pip install --constraint "${PIP_CONSTRAINT_FILE}" "$@" @@ -106,16 +95,29 @@ pip_install() { } prepare_python_environment() { + + local env_name="$1" + log_verbose "-- Preparing python environment for test..." - "${PYTHON_BIN}" -m venv --clear "${TEST_VENV}" + "${PYTHON_BIN}" -m venv --clear "${env_name}" # shellcheck disable=SC1091 - source "${TEST_VENV}/bin/activate" + source "${env_name}/bin/activate" log_verbose "---- Python version = $(${PYTHON_BIN} -V)" pip_install -U "${TEST_REQUIREMENTS[@]}" } +deactivate_python_environment() { + + log_verbose "-- Deactivate python environment..." + + # `deactivate` is provided by the activation script; there is no bin/deactivate file. + if declare -F deactivate >/dev/null; then + deactivate + fi +} + install_editable_package() { local package_name="$1" local package_path="$2" @@ -126,9 +128,10 @@ install_editable_package() { run_python_outside_repo() { local python_code="$1" + local python_bin="${2:-${PYTHON_BIN}}" ( cd /tmp - run_python_with_keops_verbose "${PYTHON_BIN}" "${python_code}" + run_with_keops_verbose "${python_bin}" -c "${python_code}" ) } @@ -147,26 +150,10 @@ run_pykeops_no_torch_smoke_test() { local no_torch_smoke_code="import importlib.util;" no_torch_smoke_code+=" assert importlib.util.find_spec('torch') is None," no_torch_smoke_code+=" 'torch must not be installed in no-torch smoke env';" - no_torch_smoke_code+=" import pykeops; pykeops.check_health();" + no_torch_smoke_code+=" import pykeops;" no_torch_smoke_code+=" assert pykeops.test_numpy_bindings()" - echo "-- Running pykeops no-torch smoke test..." - "${PYTHON_BIN}" -m venv --clear "${NO_TORCH_TEST_VENV}" - - if [[ -n "${PIP_CONSTRAINT_FILE}" ]]; then - "${no_torch_python}" -m pip install --constraint "${PIP_CONSTRAINT_FILE}" -U pip - "${no_torch_python}" -m pip install --constraint "${PIP_CONSTRAINT_FILE}" -e "${SCRIPT_DIR}/keopscore" - "${no_torch_python}" -m pip install --constraint "${PIP_CONSTRAINT_FILE}" -e "${SCRIPT_DIR}/pykeops" - else - "${no_torch_python}" -m pip install -U pip - "${no_torch_python}" -m pip install -e "${SCRIPT_DIR}/keopscore" - "${no_torch_python}" -m pip install -e "${SCRIPT_DIR}/pykeops" - fi - - ( - cd /tmp - run_python_with_keops_verbose "${no_torch_python}" "${no_torch_smoke_code}" - ) + run_python_outside_repo "${no_torch_smoke_code}" "${no_torch_python}" } run_test_suite() { @@ -179,14 +166,30 @@ run_test_suite() { main() { parse_options "$@" + + printf '%b' "****************************************************************************\n Start of pykeops no-torch smoke test\n****************************************************************************\n" + prepare_python_environment "${NO_TORCH_TEST_VENV}" + install_editable_package "keopscore" "${SCRIPT_DIR}/keopscore" + install_editable_package "pykeops" "${SCRIPT_DIR}/pykeops" + clean_pykeops_cache + + run_pykeops_health_check run_pykeops_no_torch_smoke_test - prepare_python_environment + + deactivate_python_environment + printf '%b' "****************************************************************************\n End of pykeops no-torch smoke test\n****************************************************************************\n\n\n\n\n\n" + + printf '%b' "****************************************************************************\n Start of pykeops tests\n****************************************************************************\n" + prepare_python_environment "${TEST_VENV}" install_editable_package "keopscore" "${SCRIPT_DIR}/keopscore" install_editable_package "pykeops" "${SCRIPT_DIR}/pykeops[test]" run_pykeops_health_check + clean_pykeops_cache run_test_suite "keopscore" "keopscore/keopscore/test/" run_test_suite "pykeops" "pykeops/pykeops/test/" + printf '%b' "****************************************************************************\n End of pykeops tests\n****************************************************************************\n\n\n\n\n\n" + } main "$@" From 746b63172bdce54b30df5ef9a2f28ca8f9dfcb28 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Sat, 6 Jun 2026 19:21:50 +0200 Subject: [PATCH 93/98] remove cirrus CI. Add pykeops.config.cuda.is_available() --- .cirrus.yml | 18 ------------------ keopscore/keopscore/config/Cuda.py | 3 +++ pykeops/pykeops/benchmarks/benchmark_KNN.py | 3 --- pykeops/pykeops/benchmarks/benchmark_utils.py | 3 ++- pykeops/pykeops/benchmarks/plot_accuracy.py | 3 ++- .../benchmarks/plot_benchmark_convolutions.py | 3 ++- .../plot_benchmark_grad1convolutions.py | 3 ++- .../plot_benchmark_high_dimension.py | 2 -- .../benchmarks/plot_benchmark_invkernel.py | 13 ++++++++++--- .../plot_benchmarks_convolutions_3D.py | 3 ++- pykeops/pykeops/common/get_options.py | 4 ++-- pykeops/pykeops/config.py | 3 ++- .../examples/numpy/plot_gpu_select_numpy.py | 4 ++-- .../examples/numpy/plot_grid_cluster_numpy.py | 4 ++-- .../numpy/plot_test_invkernel_numpy.py | 2 +- .../numpy/plot_test_invkernel_numpy_helper.py | 2 +- .../examples/pytorch/plot_advanced_formula.py | 5 ++++- .../pytorch/plot_anisotropic_kernels.py | 5 ++++- .../pytorch/plot_generic_syntax_pytorch.py | 4 +++- .../pytorch/plot_generic_syntax_pytorch_LSE.py | 5 ++++- .../plot_generic_syntax_pytorch_LSE_vect.py | 5 ++++- .../pytorch/plot_grid_cluster_pytorch.py | 3 ++- .../pytorch/plot_test_invkernel_torch.py | 16 +++++++++++----- .../plot_test_invkernel_torch_helper.py | 16 +++++++++++----- pykeops/pykeops/numpy/utils.py | 2 +- pykeops/pykeops/sandbox/test_gpu_cpu2.py | 4 ++-- .../test/test_block_sparse_reduction.py | 4 ++-- pykeops/pykeops/test/test_chunks.py | 4 ++-- pykeops/pykeops/test/test_chunks_ranges.py | 4 ++-- pykeops/pykeops/test/test_complex.py | 4 ++-- pykeops/pykeops/test/test_conv2d.py | 4 ++-- pykeops/pykeops/test/test_finalchunks.py | 4 ++-- .../pykeops/test/test_finalchunks_ranges.py | 4 ++-- pykeops/pykeops/test/test_float16.py | 4 ++-- pykeops/pykeops/test/test_gpu_cpu.py | 6 +++--- pykeops/pykeops/test/test_lazytensor_clamp.py | 4 ++-- .../test/test_lazytensor_gaussian_batch.py | 4 ++-- .../test/test_lazytensor_gaussian_fromhost.py | 4 ++-- .../test/test_lazytensor_gaussian_inplace.py | 4 ++-- pykeops/pykeops/test/test_lazytensor_grad.py | 4 ++-- .../pykeops/test/test_lazytensor_tensordot.py | 4 ++-- pykeops/pykeops/test/test_numpy.py | 6 +++--- pykeops/pykeops/test/test_shape_regression.py | 4 ++-- pykeops/pykeops/test/test_sinc.py | 4 ++-- pykeops/pykeops/test/test_torch.py | 10 +++++----- pykeops/pykeops/torch/test_install.py | 5 ++++- .../a_LazyTensors/plot_lazytensors_a.py | 6 +++--- .../a_LazyTensors/plot_lazytensors_b.py | 3 ++- .../a_LazyTensors/plot_lazytensors_c.py | 3 ++- .../pykeops/tutorials/backends/plot_scipy.py | 8 ++++---- .../gaussian_mixture/plot_gaussian_mixture.py | 4 +++- .../plot_RBF_interpolation_numpy.py | 2 +- .../plot_RBF_interpolation_torch.py | 3 ++- .../tutorials/kmeans/plot_kmeans_numpy.py | 2 +- .../tutorials/kmeans/plot_kmeans_torch.py | 3 ++- .../pykeops/tutorials/knn/plot_knn_mnist.py | 3 ++- .../pykeops/tutorials/knn/plot_knn_numpy.py | 4 ++-- .../pykeops/tutorials/knn/plot_knn_torch.py | 3 ++- 58 files changed, 152 insertions(+), 121 deletions(-) delete mode 100644 .cirrus.yml diff --git a/.cirrus.yml b/.cirrus.yml deleted file mode 100644 index 909977052..000000000 --- a/.cirrus.yml +++ /dev/null @@ -1,18 +0,0 @@ -linux_task: - name: Tests (linux) - container: - image: ubuntu:latest - install_script: - - apt-get update - - apt-get -y install python3 python3-venv python3-dev g++ - test_script: - - bash ./pytest.sh - -macos_task: - name: Tests (macOS) - macos_instance: - image: ghcr.io/cirruslabs/macos-runner:tahoe - script: - - brew install libomp - - python3 --version - - bash ./pytest.sh diff --git a/keopscore/keopscore/config/Cuda.py b/keopscore/keopscore/config/Cuda.py index 0f8feba30..d9716a7f2 100644 --- a/keopscore/keopscore/config/Cuda.py +++ b/keopscore/keopscore/config/Cuda.py @@ -401,6 +401,9 @@ def set_use_cuda(self): def get_use_cuda(self): return self._use_cuda + def is_available(self): + return self.get_use_cuda() + def print_use_cuda(self): print(f"CUDA Support: {enabled_dict[self.get_use_cuda() or False]}") diff --git a/pykeops/pykeops/benchmarks/benchmark_KNN.py b/pykeops/pykeops/benchmarks/benchmark_KNN.py index e4147ac4e..881e2b153 100644 --- a/pykeops/pykeops/benchmarks/benchmark_KNN.py +++ b/pykeops/pykeops/benchmarks/benchmark_KNN.py @@ -59,9 +59,6 @@ ) from dataset_utils import generate_samples -use_cuda = torch.cuda.is_available() - - ############################################## # We then specify the values of K that we will inspect: diff --git a/pykeops/pykeops/benchmarks/benchmark_utils.py b/pykeops/pykeops/benchmarks/benchmark_utils.py index 87e746a83..71b5c9b51 100644 --- a/pykeops/pykeops/benchmarks/benchmark_utils.py +++ b/pykeops/pykeops/benchmarks/benchmark_utils.py @@ -13,10 +13,11 @@ import numpy as np import torch +import pykeops.config # import jax -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() ################################################## # Utility functions: diff --git a/pykeops/pykeops/benchmarks/plot_accuracy.py b/pykeops/pykeops/benchmarks/plot_accuracy.py index 30db3093e..c765536ec 100644 --- a/pykeops/pykeops/benchmarks/plot_accuracy.py +++ b/pykeops/pykeops/benchmarks/plot_accuracy.py @@ -18,9 +18,10 @@ import numpy as np import torch +import pykeops.config from matplotlib import pyplot as plt -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() D = 3 diff --git a/pykeops/pykeops/benchmarks/plot_benchmark_convolutions.py b/pykeops/pykeops/benchmarks/plot_benchmark_convolutions.py index 04a88f253..ff26ab9d3 100644 --- a/pykeops/pykeops/benchmarks/plot_benchmark_convolutions.py +++ b/pykeops/pykeops/benchmarks/plot_benchmark_convolutions.py @@ -22,6 +22,7 @@ import numpy as np import timeit import matplotlib +import pykeops.config from matplotlib import pyplot as plt from pykeops.numpy.utils import np_kernel from pykeops.torch.utils import torch_kernel @@ -55,7 +56,7 @@ try: import torch - use_cuda = torch.cuda.is_available() + use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device = "cuda" if use_cuda else "cpu" torchtype = torch.float32 if dtype == "float32" else torch.float64 diff --git a/pykeops/pykeops/benchmarks/plot_benchmark_grad1convolutions.py b/pykeops/pykeops/benchmarks/plot_benchmark_grad1convolutions.py index a6922dc96..312d71f33 100644 --- a/pykeops/pykeops/benchmarks/plot_benchmark_grad1convolutions.py +++ b/pykeops/pykeops/benchmarks/plot_benchmark_grad1convolutions.py @@ -23,6 +23,7 @@ import numpy as np import timeit import matplotlib +import pykeops.config from matplotlib import pyplot as plt from pykeops.numpy.utils import grad_np_kernel, chain_rules from pykeops.torch.utils import torch_kernel @@ -61,7 +62,7 @@ try: import torch - use_cuda = torch.cuda.is_available() + use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device = "cuda" if use_cuda else "cpu" torchtype = torch.float32 if dtype == "float32" else torch.float64 diff --git a/pykeops/pykeops/benchmarks/plot_benchmark_high_dimension.py b/pykeops/pykeops/benchmarks/plot_benchmark_high_dimension.py index 67f07a515..cdba8c540 100644 --- a/pykeops/pykeops/benchmarks/plot_benchmark_high_dimension.py +++ b/pykeops/pykeops/benchmarks/plot_benchmark_high_dimension.py @@ -18,8 +18,6 @@ from benchmark_utils import random_normal, full_benchmark -use_cuda = torch.cuda.is_available() - ############################################## # Benchmark specifications: # diff --git a/pykeops/pykeops/benchmarks/plot_benchmark_invkernel.py b/pykeops/pykeops/benchmarks/plot_benchmark_invkernel.py index e5bdc89a4..6565d3165 100644 --- a/pykeops/pykeops/benchmarks/plot_benchmark_invkernel.py +++ b/pykeops/pykeops/benchmarks/plot_benchmark_invkernel.py @@ -33,6 +33,7 @@ import numpy as np import torch +import pykeops.config from matplotlib import pyplot as plt from scipy.sparse import diags @@ -45,7 +46,8 @@ from benchmark_utils import flatten, random_normal, unit_tensor, full_benchmark -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +torch_device = "cuda" if use_cuda else "cpu" if torch.__version__ >= "1.8": torchsolve = lambda A, B: torch.linalg.solve(A, B) @@ -72,7 +74,7 @@ # -def generate_samples(N, device="cuda", lang="torch", **kwargs): +def generate_samples(N, device=None, lang="torch", **kwargs): """Generates a point cloud x, a scalar signal b of size N and two regularization parameters. Args: @@ -83,6 +85,11 @@ def generate_samples(N, device="cuda", lang="torch", **kwargs): Returns: 3-uple of arrays: x, y, b """ + if device is None: + device = torch_device + if lang == "torch" and device == "cuda" and not use_cuda: + device = "cpu" + randn = random_normal(device=device, lang=lang) ones = unit_tensor(device=device, lang=lang) @@ -149,7 +156,7 @@ def Kinv_scipy(x, b, gamma, alpha, **kwargs): def Kinv_pytorch(x, b, gamma, alpha, **kwargs): - K_xx = alpha * torch.eye(x.shape[0], device=x.get_device()) + torch.exp( + K_xx = alpha * torch.eye(x.shape[0], device=x.device) + torch.exp( -gamma * sqdist_torch(x, x) ) res = torchsolve(K_xx, b) diff --git a/pykeops/pykeops/benchmarks/plot_benchmarks_convolutions_3D.py b/pykeops/pykeops/benchmarks/plot_benchmarks_convolutions_3D.py index 241eb5ea8..4c0dc816e 100644 --- a/pykeops/pykeops/benchmarks/plot_benchmarks_convolutions_3D.py +++ b/pykeops/pykeops/benchmarks/plot_benchmarks_convolutions_3D.py @@ -21,11 +21,12 @@ import numpy as np import torch +import pykeops.config from matplotlib import pyplot as plt from benchmark_utils import flatten, random_normal, full_benchmark -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() print( f"Running torch version {torch.__version__} with {'GPU' if use_cuda else 'CPU'}..." ) diff --git a/pykeops/pykeops/common/get_options.py b/pykeops/pykeops/common/get_options.py index ce4cb144b..0f1b9de24 100644 --- a/pykeops/pykeops/common/get_options.py +++ b/pykeops/pykeops/common/get_options.py @@ -53,7 +53,7 @@ def define_tag_backend(self, backend, variables): # auto : infer everything if backend == "auto": return ( - int(pykeopsconfig.gpu_available), + int(pykeopsconfig.cuda.is_available()), self._find_grid(), self._find_mem(variables), ) @@ -84,7 +84,7 @@ def define_backend(self, backend, variables): @staticmethod def _find_dev(): - return int(pykeopsconfig.gpu_available) + return int(pykeopsconfig.cuda.is_available()) @staticmethod def _find_mem(variables): diff --git a/pykeops/pykeops/config.py b/pykeops/pykeops/config.py index e9a2e4fe1..2e9bdee93 100644 --- a/pykeops/pykeops/config.py +++ b/pykeops/pykeops/config.py @@ -12,7 +12,8 @@ debug = keopscore.config.debug get_build_folder = path.get_build_folder -gpu_available = cuda.get_use_cuda() +# keep the old gpu_available variable for backward compatibility, but it is now recommended to use cuda.is_available() instead +gpu_available = cuda.is_available() # Initialize some variables: the values may be redefined later diff --git a/pykeops/pykeops/examples/numpy/plot_gpu_select_numpy.py b/pykeops/pykeops/examples/numpy/plot_gpu_select_numpy.py index ed6b8c279..d3100edf2 100644 --- a/pykeops/pykeops/examples/numpy/plot_gpu_select_numpy.py +++ b/pykeops/pykeops/examples/numpy/plot_gpu_select_numpy.py @@ -63,7 +63,7 @@ # And on our GPUs, with copies between # the Host and Device memories: # -if pykeops.config.gpu_available: +if pykeops.config.cuda.is_available(): for gpuid in gpuids: d = my_routine(x, y, a, p, backend="GPU", device_id=gpuid) print( @@ -96,7 +96,7 @@ #################################################################### # And on the GPUs, with copies between the Host and Device memories: # -if pykeops.config.gpu_available: +if pykeops.config.cuda.is_available(): for gpuid in gpuids: d = ((p - aj) ** 2 * (xi + yj).exp()).sum( axis=1, backend="GPU", device_id=gpuid diff --git a/pykeops/pykeops/examples/numpy/plot_grid_cluster_numpy.py b/pykeops/pykeops/examples/numpy/plot_grid_cluster_numpy.py index 048808156..db4c6f21f 100755 --- a/pykeops/pykeops/examples/numpy/plot_grid_cluster_numpy.py +++ b/pykeops/pykeops/examples/numpy/plot_grid_cluster_numpy.py @@ -27,7 +27,7 @@ ######################################################################## # Define our dataset: two point clouds on the unit square. # -M, N = (5000, 5000) if pykeops.config.gpu_available else (2000, 2000) +M, N = (5000, 5000) if pykeops.config.cuda.is_available() else (2000, 2000) t = np.linspace(0, 2 * np.pi, M + 1)[:-1] x = np.stack((0.4 + 0.4 * (t / 7) * np.cos(t), 0.5 + 0.3 * np.sin(t)), 1) @@ -189,7 +189,7 @@ backends = ( (["CPU", "GPU"] if M * N < 4e8 else ["GPU"]) - if pykeops.config.gpu_available + if pykeops.config.cuda.is_available() else ["CPU"] ) for backend in backends: diff --git a/pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy.py b/pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy.py index 4b460a452..264673511 100644 --- a/pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy.py +++ b/pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy.py @@ -27,7 +27,7 @@ # Define our dataset: # -N = 5000 if pykeops.config.gpu_available else 500 # Number of points +N = 5000 if pykeops.config.cuda.is_available() else 500 # Number of points D = 2 # Dimension of the ambient space Dv = 2 # Dimension of the vectors (= number of linear problems to solve) sigma = 0.1 # Radius of our RBF kernel diff --git a/pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy_helper.py b/pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy_helper.py index 4f18b855e..2a7556f85 100644 --- a/pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy_helper.py +++ b/pykeops/pykeops/examples/numpy/plot_test_invkernel_numpy_helper.py @@ -27,7 +27,7 @@ # Define our dataset: # -N = 5000 if pykeops.config.gpu_available else 500 # Number of points +N = 5000 if pykeops.config.cuda.is_available() else 500 # Number of points D = 2 # Dimension of the ambient space Dv = 2 # Dimension of the vectors (= number of linear problems to solve) sigma = 0.1 # Radius of our RBF kernel diff --git a/pykeops/pykeops/examples/pytorch/plot_advanced_formula.py b/pykeops/pykeops/examples/pytorch/plot_advanced_formula.py index 21327f745..0b4a4aeae 100644 --- a/pykeops/pykeops/examples/pytorch/plot_advanced_formula.py +++ b/pykeops/pykeops/examples/pytorch/plot_advanced_formula.py @@ -13,11 +13,14 @@ # First, the standard imports: import torch +import pykeops.config from pykeops.torch import Genred import matplotlib.pyplot as plt +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() + # Choose the storage place for our data : CPU (host) or GPU (device) memory. -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +device = torch.device("cuda" if use_cuda else "cpu") #################################################################### # Then, the definition of our dataset: diff --git a/pykeops/pykeops/examples/pytorch/plot_anisotropic_kernels.py b/pykeops/pykeops/examples/pytorch/plot_anisotropic_kernels.py index 26d0bc743..6dfcde53b 100644 --- a/pykeops/pykeops/examples/pytorch/plot_anisotropic_kernels.py +++ b/pykeops/pykeops/examples/pytorch/plot_anisotropic_kernels.py @@ -18,13 +18,16 @@ from matplotlib import pyplot as plt import matplotlib.cm as cm import torch +import pykeops.config from pykeops.torch import Vi, Vj, Pm, LazyTensor ############################################## # Dataset: +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() + # Choose the storage place for our data : CPU (host) or GPU (device) memory. -dtype = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor +dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor # Three points in the plane R^2 y = torch.tensor([[0.2, 0.7], [0.5, 0.3], [0.7, 0.5]]).type(dtype) diff --git a/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch.py b/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch.py index a7fd50a46..a5156358e 100644 --- a/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch.py +++ b/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch.py @@ -28,6 +28,7 @@ import matplotlib.pyplot as plt import torch +import pykeops.config from torch.autograd import grad from pykeops.torch import Genred @@ -39,7 +40,8 @@ N = 5000 # Choose the storage place for our data : CPU (host) or GPU (device) memory. -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device = torch.device("cuda" if use_cuda else "cpu") dtype = "float32" # Could be 'float32' or 'float64' torchtype = torch.float32 if dtype == "float32" else torch.float64 diff --git a/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch_LSE.py b/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch_LSE.py index 8b3a0dc0e..3ae45b873 100644 --- a/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch_LSE.py +++ b/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch_LSE.py @@ -28,6 +28,7 @@ import time import torch +import pykeops.config from matplotlib import pyplot as plt from torch.autograd import grad @@ -148,7 +149,9 @@ # # Of course, this will only work if you own a Gpu... -if torch.cuda.is_available(): +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() + +if use_cuda: # first transfer data on gpu pc, ac, xc, yc, ec = p.cuda(), a.cuda(), x.cuda(), y.cuda(), e.cuda() # then call the operations diff --git a/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch_LSE_vect.py b/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch_LSE_vect.py index 08c7670bc..3db79107c 100644 --- a/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch_LSE_vect.py +++ b/pykeops/pykeops/examples/pytorch/plot_generic_syntax_pytorch_LSE_vect.py @@ -30,6 +30,7 @@ import matplotlib.pyplot as plt import torch +import pykeops.config from torch.autograd import grad from pykeops.torch import Genred @@ -162,7 +163,9 @@ # # Of course, this will only work if you own a Gpu... -if torch.cuda.is_available(): +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() + +if use_cuda: # first transfer data on gpu pc, ac, xc, yc, bc, ec = p.cuda(), a.cuda(), x.cuda(), y.cuda(), b.cuda(), e.cuda() # then call the operations diff --git a/pykeops/pykeops/examples/pytorch/plot_grid_cluster_pytorch.py b/pykeops/pykeops/examples/pytorch/plot_grid_cluster_pytorch.py index 6865c501b..bcc6452e3 100644 --- a/pykeops/pykeops/examples/pytorch/plot_grid_cluster_pytorch.py +++ b/pykeops/pykeops/examples/pytorch/plot_grid_cluster_pytorch.py @@ -18,12 +18,13 @@ import numpy as np import torch +import pykeops.config from matplotlib import pyplot as plt from pykeops.torch import LazyTensor nump = lambda t: t.cpu().numpy() -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() # dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor device = torch.device("cuda" if use_cuda else "cpu") dtype = torch.float32 # Standard float type for both CPU/GPU diff --git a/pykeops/pykeops/examples/pytorch/plot_test_invkernel_torch.py b/pykeops/pykeops/examples/pytorch/plot_test_invkernel_torch.py index ac1be1424..83bffd0ca 100644 --- a/pykeops/pykeops/examples/pytorch/plot_test_invkernel_torch.py +++ b/pykeops/pykeops/examples/pytorch/plot_test_invkernel_torch.py @@ -19,6 +19,7 @@ import time import torch +import pykeops.config from matplotlib import pyplot as plt from pykeops.torch import KernelSolve @@ -32,16 +33,21 @@ # Define our dataset: # -N = 5000 if torch.cuda.is_available() else 500 # Number of points +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device = torch.device("cuda" if use_cuda else "cpu") + +N = 5000 if use_cuda else 500 # Number of points D = 2 # Dimension of the ambient space Dv = 2 # Dimension of the vectors (= number of linear problems to solve) sigma = 0.1 # Radius of our RBF kernel -x = torch.rand(N, D, requires_grad=True) -b = torch.rand(N, Dv) -g = torch.Tensor([0.5 / sigma**2]) # Parameter of the Gaussian RBF kernel +x = torch.rand(N, D, requires_grad=True, device=device) +b = torch.rand(N, Dv, device=device) +g = torch.tensor( + [0.5 / sigma**2], device=device +) # Parameter of the Gaussian RBF kernel -if torch.cuda.is_available(): +if use_cuda: sync = torch.cuda.synchronize else: diff --git a/pykeops/pykeops/examples/pytorch/plot_test_invkernel_torch_helper.py b/pykeops/pykeops/examples/pytorch/plot_test_invkernel_torch_helper.py index 6fc302f7c..dbd74010c 100644 --- a/pykeops/pykeops/examples/pytorch/plot_test_invkernel_torch_helper.py +++ b/pykeops/pykeops/examples/pytorch/plot_test_invkernel_torch_helper.py @@ -17,6 +17,7 @@ import time import torch +import pykeops.config from matplotlib import pyplot as plt from pykeops.torch import LazyTensor as keops @@ -26,14 +27,19 @@ # Define our dataset: # -N = 5000 if torch.cuda.is_available() else 500 # Number of points +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +device = torch.device("cuda" if use_cuda else "cpu") + +N = 5000 if use_cuda else 500 # Number of points D = 2 # Dimension of the ambient space Dv = 2 # Dimension of the vectors (= number of linear problems to solve) sigma = 0.1 # Radius of our RBF kernel -x = torch.rand(N, D, requires_grad=True) -b = torch.rand(N, Dv) -g = torch.Tensor([0.5 / sigma**2]) # Parameter of the Gaussian RBF kernel +x = torch.rand(N, D, requires_grad=True, device=device) +b = torch.rand(N, Dv, device=device) +g = torch.tensor( + [0.5 / sigma**2], device=device +) # Parameter of the Gaussian RBF kernel alpha = 0.01 # ridge regularization ############################################################################### @@ -59,7 +65,7 @@ # start = time.time() -K_xx = alpha * torch.eye(N) + torch.exp( +K_xx = alpha * torch.eye(N, device=device) + torch.exp( -torch.sum((x[:, None, :] - x[None, :, :]) ** 2, dim=2) / (2 * sigma**2) ) diff --git a/pykeops/pykeops/numpy/utils.py b/pykeops/pykeops/numpy/utils.py index d54518009..035ade1a8 100644 --- a/pykeops/pykeops/numpy/utils.py +++ b/pykeops/pykeops/numpy/utils.py @@ -236,7 +236,7 @@ def WarmUpGpu(): from pykeops.common.utils import pyKeOps_Message pyKeOps_Message("Warming up the Gpu (numpy bindings) !!!") - if pykeopsconfig.gpu_available: + if pykeopsconfig.cuda.is_available(): formula = "Exp(-oos2*SqDist(x,y))*b" aliases = [ "x = Vi(1)", # First arg : i-variable, of size 1 diff --git a/pykeops/pykeops/sandbox/test_gpu_cpu2.py b/pykeops/pykeops/sandbox/test_gpu_cpu2.py index 048808156..db4c6f21f 100644 --- a/pykeops/pykeops/sandbox/test_gpu_cpu2.py +++ b/pykeops/pykeops/sandbox/test_gpu_cpu2.py @@ -27,7 +27,7 @@ ######################################################################## # Define our dataset: two point clouds on the unit square. # -M, N = (5000, 5000) if pykeops.config.gpu_available else (2000, 2000) +M, N = (5000, 5000) if pykeops.config.cuda.is_available() else (2000, 2000) t = np.linspace(0, 2 * np.pi, M + 1)[:-1] x = np.stack((0.4 + 0.4 * (t / 7) * np.cos(t), 0.5 + 0.3 * np.sin(t)), 1) @@ -189,7 +189,7 @@ backends = ( (["CPU", "GPU"] if M * N < 4e8 else ["GPU"]) - if pykeops.config.gpu_available + if pykeops.config.cuda.is_available() else ["CPU"] ) for backend in backends: diff --git a/pykeops/pykeops/test/test_block_sparse_reduction.py b/pykeops/pykeops/test/test_block_sparse_reduction.py index e9a1c1750..defa32824 100644 --- a/pykeops/pykeops/test/test_block_sparse_reduction.py +++ b/pykeops/pykeops/test/test_block_sparse_reduction.py @@ -1,7 +1,7 @@ import time import numpy as np import torch -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.torch import LazyTensor # Import clustering functions from KeOps @@ -14,7 +14,7 @@ def test_block_sparse_reduction(): - use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available + use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") dtype = torch.float32 diff --git a/pykeops/pykeops/test/test_chunks.py b/pykeops/pykeops/test/test_chunks.py index 7408e8245..744f58f02 100644 --- a/pykeops/pykeops/test/test_chunks.py +++ b/pykeops/pykeops/test/test_chunks.py @@ -1,6 +1,6 @@ import math import torch -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -10,7 +10,7 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) diff --git a/pykeops/pykeops/test/test_chunks_ranges.py b/pykeops/pykeops/test/test_chunks_ranges.py index b0eafe409..f6170381f 100644 --- a/pykeops/pykeops/test/test_chunks_ranges.py +++ b/pykeops/pykeops/test/test_chunks_ranges.py @@ -1,6 +1,6 @@ import math import torch -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -10,7 +10,7 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) diff --git a/pykeops/pykeops/test/test_complex.py b/pykeops/pykeops/test/test_complex.py index bc4570398..ae2c917f2 100644 --- a/pykeops/pykeops/test/test_complex.py +++ b/pykeops/pykeops/test/test_complex.py @@ -2,7 +2,7 @@ import math import torch -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -12,7 +12,7 @@ M, N, D = 1000, 1000, 1 torch.backends.cuda.matmul.allow_tf32 = False -use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) diff --git a/pykeops/pykeops/test/test_conv2d.py b/pykeops/pykeops/test/test_conv2d.py index 7ac2b4878..74afd0d3d 100644 --- a/pykeops/pykeops/test/test_conv2d.py +++ b/pykeops/pykeops/test/test_conv2d.py @@ -1,7 +1,7 @@ import math import unittest import torch -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -12,7 +12,7 @@ torch.manual_seed(42) torch.backends.cuda.matmul.allow_tf32 = False -use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device_id = "cuda" if use_cuda else "cpu" diff --git a/pykeops/pykeops/test/test_finalchunks.py b/pykeops/pykeops/test/test_finalchunks.py index 9665f2b73..0bf164370 100644 --- a/pykeops/pykeops/test/test_finalchunks.py +++ b/pykeops/pykeops/test/test_finalchunks.py @@ -1,6 +1,6 @@ import math import torch -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -10,7 +10,7 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) diff --git a/pykeops/pykeops/test/test_finalchunks_ranges.py b/pykeops/pykeops/test/test_finalchunks_ranges.py index 14d97e0b0..d24939629 100644 --- a/pykeops/pykeops/test/test_finalchunks_ranges.py +++ b/pykeops/pykeops/test/test_finalchunks_ranges.py @@ -1,6 +1,6 @@ import math import torch -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -10,7 +10,7 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) diff --git a/pykeops/pykeops/test/test_float16.py b/pykeops/pykeops/test/test_float16.py index 4552c273b..4324d9424 100644 --- a/pykeops/pykeops/test/test_float16.py +++ b/pykeops/pykeops/test/test_float16.py @@ -1,7 +1,7 @@ # Test for Clamp operation using LazyTensors import pytest import torch -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -10,7 +10,7 @@ M, N, D = 5, 5, 1 torch.backends.cuda.matmul.allow_tf32 = False -use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) diff --git a/pykeops/pykeops/test/test_gpu_cpu.py b/pykeops/pykeops/test/test_gpu_cpu.py index 89da4e426..ada773e83 100644 --- a/pykeops/pykeops/test/test_gpu_cpu.py +++ b/pykeops/pykeops/test/test_gpu_cpu.py @@ -1,10 +1,10 @@ import math import pytest import torch -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.torch import LazyTensor -use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +use_cuda = pykeops.config.cuda.is_available() M, N, D, DV = 1000, 1000, 3, 1 @@ -48,7 +48,7 @@ def test_torch_keops_cpu(self): @pytest.mark.skipif( not use_cuda, - reason="Requires torch and KeOps CUDA support", + reason="Requires KeOps CUDA support", ) def test_torch_keops_gpu(self): out_gpu = fun(x, y, b, ["keops_gpu"]).squeeze() diff --git a/pykeops/pykeops/test/test_lazytensor_clamp.py b/pykeops/pykeops/test/test_lazytensor_clamp.py index 6c0bcc652..7aae6c2db 100644 --- a/pykeops/pykeops/test/test_lazytensor_clamp.py +++ b/pykeops/pykeops/test/test_lazytensor_clamp.py @@ -1,12 +1,12 @@ import torch -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose M, N, D = 1000, 1000, 3 torch.backends.cuda.matmul.allow_tf32 = False -use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) diff --git a/pykeops/pykeops/test/test_lazytensor_gaussian_batch.py b/pykeops/pykeops/test/test_lazytensor_gaussian_batch.py index 93f5cc86c..c6af5a374 100644 --- a/pykeops/pykeops/test/test_lazytensor_gaussian_batch.py +++ b/pykeops/pykeops/test/test_lazytensor_gaussian_batch.py @@ -1,13 +1,13 @@ import math import torch -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose B1, B2, M, N, D, DV = 3, 4, 20, 25, 3, 2 torch.backends.cuda.matmul.allow_tf32 = False -use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(1) diff --git a/pykeops/pykeops/test/test_lazytensor_gaussian_fromhost.py b/pykeops/pykeops/test/test_lazytensor_gaussian_fromhost.py index 862f41cce..0c0c078d7 100644 --- a/pykeops/pykeops/test/test_lazytensor_gaussian_fromhost.py +++ b/pykeops/pykeops/test/test_lazytensor_gaussian_fromhost.py @@ -1,6 +1,6 @@ import math import torch -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -10,7 +10,7 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device_id = 0 if use_cuda else -1 torch.manual_seed(0) diff --git a/pykeops/pykeops/test/test_lazytensor_gaussian_inplace.py b/pykeops/pykeops/test/test_lazytensor_gaussian_inplace.py index f7d50c014..eb75bc2b3 100644 --- a/pykeops/pykeops/test/test_lazytensor_gaussian_inplace.py +++ b/pykeops/pykeops/test/test_lazytensor_gaussian_inplace.py @@ -1,6 +1,6 @@ import math import torch -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -10,7 +10,7 @@ sum_scheme = "block_sum" torch.backends.cuda.matmul.allow_tf32 = False -use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) diff --git a/pykeops/pykeops/test/test_lazytensor_grad.py b/pykeops/pykeops/test/test_lazytensor_grad.py index f83615fab..2cf81bf31 100644 --- a/pykeops/pykeops/test/test_lazytensor_grad.py +++ b/pykeops/pykeops/test/test_lazytensor_grad.py @@ -1,5 +1,5 @@ import torch -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -9,7 +9,7 @@ dtype = torch.float32 torch.backends.cuda.matmul.allow_tf32 = False -use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) diff --git a/pykeops/pykeops/test/test_lazytensor_tensordot.py b/pykeops/pykeops/test/test_lazytensor_tensordot.py index a08d552f8..c428962f8 100644 --- a/pykeops/pykeops/test/test_lazytensor_tensordot.py +++ b/pykeops/pykeops/test/test_lazytensor_tensordot.py @@ -1,5 +1,5 @@ import torch -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose @@ -7,7 +7,7 @@ # Matrix multiplication as a special case of Tensordot torch.backends.cuda.matmul.allow_tf32 = False -use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device_id = "cuda" if use_cuda else "cpu" torch.manual_seed(0) diff --git a/pykeops/pykeops/test/test_numpy.py b/pykeops/pykeops/test/test_numpy.py index f1005a25d..28298c996 100644 --- a/pykeops/pykeops/test/test_numpy.py +++ b/pykeops/pykeops/test/test_numpy.py @@ -130,7 +130,7 @@ def test_generic_syntax_sum(self): formula = "Square(p-a)*Exp(x+y)" axis = 1 # 0 means summation over i, 1 means over j - if pykeops.config.gpu_available: + if pykeops.config.cuda.is_available(): backend_to_test = ["auto", "GPU_1D", "GPU_2D", "GPU"] else: backend_to_test = ["auto"] @@ -165,7 +165,7 @@ def test_generic_syntax_lse(self): aliases = ["p=Pm(0,1)", "a=Vj(1,1)", "x=Vi(2,3)", "y=Vj(3,3)"] formula = "Square(p-a)*Exp(-SqNorm2(x-y))" - if pykeops.config.gpu_available: + if pykeops.config.cuda.is_available(): backend_to_test = ["auto", "GPU_1D", "GPU_2D", "GPU"] else: backend_to_test = ["auto"] @@ -201,7 +201,7 @@ def test_generic_syntax_softmax(self): formula = "Square(p-a)*Exp(-SqNorm2(x-y))" formula_weights = "y" - if pykeops.config.gpu_available: + if pykeops.config.cuda.is_available(): backend_to_test = ["auto", "GPU_1D", "GPU_2D", "GPU"] else: backend_to_test = ["auto"] diff --git a/pykeops/pykeops/test/test_shape_regression.py b/pykeops/pykeops/test/test_shape_regression.py index 10f48cebd..87f46d5f5 100644 --- a/pykeops/pykeops/test/test_shape_regression.py +++ b/pykeops/pykeops/test/test_shape_regression.py @@ -1,7 +1,7 @@ import re import unittest import numpy as np -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.test import assert_np_allclose, assert_torch_allclose @@ -74,7 +74,7 @@ def test_numpy_vj_single_dim_formula_x_1d_rejected_cpp(self): class ShapeTorchTestCase(unittest.TestCase): def setUp(self): - use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available + use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device = "cuda" if use_cuda else "cpu" self.param = torch.tensor([0.4], dtype=torch.float32, device=device) self.param_scalar = torch.tensor(0.4, dtype=torch.float32, device=device) diff --git a/pykeops/pykeops/test/test_sinc.py b/pykeops/pykeops/test/test_sinc.py index 0a79c3b6f..9837504b4 100644 --- a/pykeops/pykeops/test/test_sinc.py +++ b/pykeops/pykeops/test/test_sinc.py @@ -2,12 +2,12 @@ import pytest import torch -import pykeops.config as pykeopsconfig +import pykeops.config from pykeops.torch import LazyTensor from pykeops.test import assert_torch_allclose -use_cuda = torch.cuda.is_available() and pykeopsconfig.gpu_available +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device = "cuda" if use_cuda else "cpu" x = torch.rand(5, 1, dtype=torch.float64) * 2 * math.pi diff --git a/pykeops/pykeops/test/test_torch.py b/pykeops/pykeops/test/test_torch.py index e5ca563bc..7283d1ff5 100644 --- a/pykeops/pykeops/test/test_torch.py +++ b/pykeops/pykeops/test/test_torch.py @@ -33,7 +33,7 @@ torch.manual_seed(42) - use_cuda = torch.cuda.is_available() and pykeops.config.gpu_available + use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() device = "cuda" if use_cuda else "cpu" if use_cuda: @@ -143,7 +143,7 @@ def test_generic_syntax_float(self): aliases = ["p=Pm(1)", "a=Vj(1)", "x=Vi(3)", "y=Vj(3)"] formula = "Square(p-a)*Exp(x+y)" - if pykeops.config.gpu_available: + if pykeops.config.cuda.is_available(): backend_to_test = ["auto", "GPU_1D", "GPU_2D", "GPU"] else: backend_to_test = ["auto"] @@ -172,7 +172,7 @@ def test_generic_syntax_double(self): aliases = ["p=Pm(1)", "a=Vj(1)", "x=Vi(3)", "y=Vj(3)"] formula = "Square(p-a)*Exp(x+y)" - if pykeops.config.gpu_available: + if pykeops.config.cuda.is_available(): backend_to_test = ["auto", "GPU_1D", "GPU_2D", "GPU"] else: backend_to_test = ["auto"] @@ -200,7 +200,7 @@ def test_generic_syntax_softmax(self): aliases = ["p=Pm(1)", "a=Vj(1)", "x=Vi(3)", "y=Vj(3)"] formula = "Square(p-a)*Exp(-SqNorm2(x-y))" formula_weights = "y" - if pykeops.config.gpu_available: + if pykeops.config.cuda.is_available(): backend_to_test = ["auto", "GPU_1D", "GPU_2D", "GPU"] else: backend_to_test = ["auto"] @@ -249,7 +249,7 @@ def test_generic_syntax_simple(self): formula = "Pow((X|Y),2) * ((Elem(P,0) * X) + (Elem(P,1) * Y))" - if pykeops.config.gpu_available: + if pykeops.config.cuda.is_available(): backend_to_test = ["auto", "GPU_1D", "GPU_2D", "GPU"] else: backend_to_test = ["auto"] diff --git a/pykeops/pykeops/torch/test_install.py b/pykeops/pykeops/torch/test_install.py index 2bc05931a..909054cc5 100644 --- a/pykeops/pykeops/torch/test_install.py +++ b/pykeops/pykeops/torch/test_install.py @@ -11,10 +11,13 @@ def test_torch_bindings(): """ Try to compile a simple KeOps formula using the PyTorch binder. """ + import pykeops + import pykeops.torch as pktorch + x = torch.arange(1, 10, dtype=torch.float32).view(-1, 3) y = torch.arange(3, 9, dtype=torch.float32).view(-1, 3) - import pykeops.torch as pktorch + use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() my_conv = pktorch.Genred(formula, var) diff --git a/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_a.py b/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_a.py index c0d5a1d7f..c4911111e 100644 --- a/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_a.py +++ b/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_a.py @@ -24,6 +24,8 @@ # clouds :math:`(x_i)_{i\in[1,M]}` and :math:`(y_j)_{j\in[1,N]}` in the unit square: import numpy as np +import torch +import pykeops.config M, N = 1000, 2000 x = np.random.rand(M, 2) @@ -55,9 +57,7 @@ # That's good! Going further, we can speed-up these computations # using the **CUDA routines** of the PyTorch library: -import torch - -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor x_i = tensor(x[:, None, :]) # (M, 1, 2) torch tensor diff --git a/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_b.py b/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_b.py index fe1ef1e57..b6de58335 100644 --- a/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_b.py +++ b/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_b.py @@ -27,10 +27,11 @@ # formulas: import torch +import pykeops.config from pykeops.torch import LazyTensor -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor M, N = (100000, 200000) if use_cuda else (1000, 2000) D = 3 diff --git a/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_c.py b/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_c.py index 4d52a2a0b..043c75aa1 100644 --- a/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_c.py +++ b/pykeops/pykeops/tutorials/a_LazyTensors/plot_lazytensors_c.py @@ -10,10 +10,11 @@ import time import torch +import pykeops.config from pykeops.torch import LazyTensor -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor ########################################################################### diff --git a/pykeops/pykeops/tutorials/backends/plot_scipy.py b/pykeops/pykeops/tutorials/backends/plot_scipy.py index 99249d74b..8ec0ab717 100644 --- a/pykeops/pykeops/tutorials/backends/plot_scipy.py +++ b/pykeops/pykeops/tutorials/backends/plot_scipy.py @@ -45,7 +45,7 @@ ################################################################### # Create a toy dataset, a spiral in 2D sampled with 10,000 points: -N = 10000 if pykeops.config.gpu_available else 1000 +N = 10000 if pykeops.config.cuda.is_available() else 1000 t = np.linspace(0, 2 * np.pi, N + 1)[:-1] x = np.stack((0.4 + 0.4 * (t / 7) * np.cos(t), 0.5 + 0.3 * np.sin(t)), 1) x = x + 0.01 * np.random.randn(*x.shape) @@ -221,7 +221,7 @@ # let's generate a large "noisy Swiss roll" with **1,000,000 points** in the unit cube: # -N = 1000000 if pykeops.config.gpu_available else 1000 +N = 1000000 if pykeops.config.cuda.is_available() else 1000 t = np.linspace(0, 2 * np.pi, N + 1)[:-1] x = np.stack( ( @@ -238,7 +238,7 @@ # To **display** our toy dataset with the (not-so-efficient) PyPlot library, # we pick **10,000 points** at random: -N_display = 10000 if pykeops.config.gpu_available else N +N_display = 10000 if pykeops.config.cuda.is_available() else N indices_display = np.random.randint(0, N, N_display) _, ax = plt.subplots(nrows=1, ncols=1, figsize=(8, 8), subplot_kw=dict(projection="3d")) @@ -317,7 +317,7 @@ # pattern can be encoded in a small boolean matrix **keep** computed through: sigma = ( - 0.01 if pykeops.config.gpu_available else 0.1 + 0.01 if pykeops.config.cuda.is_available() else 0.1 ) # Standard deviation of our Gaussian kernel # Compute a coarse Boolean mask: D = np.sum((x_centroids[:, None, :] - x_centroids[None, :, :]) ** 2, 2) diff --git a/pykeops/pykeops/tutorials/gaussian_mixture/plot_gaussian_mixture.py b/pykeops/pykeops/tutorials/gaussian_mixture/plot_gaussian_mixture.py index f90ad5e05..f082b115b 100644 --- a/pykeops/pykeops/tutorials/gaussian_mixture/plot_gaussian_mixture.py +++ b/pykeops/pykeops/tutorials/gaussian_mixture/plot_gaussian_mixture.py @@ -17,6 +17,7 @@ import matplotlib.cm as cm import numpy as np import torch +import pykeops.config from matplotlib import pyplot as plt from torch.nn import Module from torch.nn.functional import softmax, log_softmax @@ -28,7 +29,8 @@ # spiral in the unit square. # Choose the storage place for our data : CPU (host) or GPU (device) memory. -dtype = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() +dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor torch.manual_seed(0) N = 10000 # Number of samples t = torch.linspace(0, 2 * np.pi, N + 1)[:-1] diff --git a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py index cac02ea26..f6f58b042 100644 --- a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py +++ b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py @@ -40,7 +40,7 @@ dtype = "float64" -N = 10000 if pykeops.config.gpu_available else 1000 # Number of samples +N = 10000 if pykeops.config.cuda.is_available() else 1000 # Number of samples # Sampling locations: x = np.random.rand(N, 1).astype(dtype) diff --git a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py index 636ecf55a..4c9cde8f0 100644 --- a/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py +++ b/pykeops/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py @@ -30,13 +30,14 @@ import time import torch +import pykeops.config from matplotlib import pyplot as plt from pykeops.torch import LazyTensor ############################################################################################### # Generate some data: -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor N = 10000 if use_cuda else 1000 # Number of samples diff --git a/pykeops/pykeops/tutorials/kmeans/plot_kmeans_numpy.py b/pykeops/pykeops/tutorials/kmeans/plot_kmeans_numpy.py index 16a22b2b4..cd81c68d7 100644 --- a/pykeops/pykeops/tutorials/kmeans/plot_kmeans_numpy.py +++ b/pykeops/pykeops/tutorials/kmeans/plot_kmeans_numpy.py @@ -107,7 +107,7 @@ def KMeans(x, K=10, Niter=10, verbose=True): # ------------------------- # Second experiment with N=1,000,000 points in dimension D=100, with K=1,000 classes: -if pykeops.config.gpu_available: +if pykeops.config.cuda.is_available(): N, D, K = 1000000, 100, 1000 x = np.random.randn(N, D).astype(dtype) cl, c = KMeans(x, K) diff --git a/pykeops/pykeops/tutorials/kmeans/plot_kmeans_torch.py b/pykeops/pykeops/tutorials/kmeans/plot_kmeans_torch.py index a2043577d..f3b469cd0 100644 --- a/pykeops/pykeops/tutorials/kmeans/plot_kmeans_torch.py +++ b/pykeops/pykeops/tutorials/kmeans/plot_kmeans_torch.py @@ -24,10 +24,11 @@ import time import torch +import pykeops.config from matplotlib import pyplot as plt from pykeops.torch import LazyTensor -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() dtype = torch.float32 if use_cuda else torch.float64 device_id = "cuda" if use_cuda else "cpu" diff --git a/pykeops/pykeops/tutorials/knn/plot_knn_mnist.py b/pykeops/pykeops/tutorials/knn/plot_knn_mnist.py index 92caaa137..1d47c9d82 100644 --- a/pykeops/pykeops/tutorials/knn/plot_knn_mnist.py +++ b/pykeops/pykeops/tutorials/knn/plot_knn_mnist.py @@ -22,11 +22,12 @@ import time import torch +import pykeops.config from matplotlib import pyplot as plt from pykeops.torch import LazyTensor -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor ###################################################################### diff --git a/pykeops/pykeops/tutorials/knn/plot_knn_numpy.py b/pykeops/pykeops/tutorials/knn/plot_knn_numpy.py index a68fdcca2..74eec6690 100644 --- a/pykeops/pykeops/tutorials/knn/plot_knn_numpy.py +++ b/pykeops/pykeops/tutorials/knn/plot_knn_numpy.py @@ -32,7 +32,7 @@ # Dataset, in 2D: N, D = ( - 10000 if pykeops.config.gpu_available else 1000, + 10000 if pykeops.config.cuda.is_available() else 1000, 2, ) # Number of samples, dimension x = np.random.rand(N, D).astype(dtype) # Random samples on the unit square @@ -48,7 +48,7 @@ def fth(x): ############################# # Reference sampling grid, on the unit square: -M = 1000 if pykeops.config.gpu_available else 100 +M = 1000 if pykeops.config.cuda.is_available() else 100 tmp = np.linspace(0, 1, M).astype(dtype) g1, g2 = np.meshgrid(tmp, tmp) g = np.hstack((g1.reshape(-1, 1), g2.reshape(-1, 1))) diff --git a/pykeops/pykeops/tutorials/knn/plot_knn_torch.py b/pykeops/pykeops/tutorials/knn/plot_knn_torch.py index 34428b3e8..e878ef14a 100644 --- a/pykeops/pykeops/tutorials/knn/plot_knn_torch.py +++ b/pykeops/pykeops/tutorials/knn/plot_knn_torch.py @@ -22,11 +22,12 @@ import numpy as np import torch +import pykeops.config from matplotlib import pyplot as plt from pykeops.torch import LazyTensor -use_cuda = torch.cuda.is_available() +use_cuda = torch.cuda.is_available() and pykeops.config.cuda.is_available() dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor ###################################################################### From f35a321fccea464891fa3576df3bbab62e437207 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Tue, 9 Jun 2026 12:04:57 +0200 Subject: [PATCH 94/98] remove torch dependancies in test --- pykeops/pykeops/test/__init__.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/pykeops/pykeops/test/__init__.py b/pykeops/pykeops/test/__init__.py index 6fa62a837..7f70d7277 100644 --- a/pykeops/pykeops/test/__init__.py +++ b/pykeops/pykeops/test/__init__.py @@ -1,8 +1,7 @@ -import numpy as np -import torch +def assert_torch_allclose(actual, expected, *, label=None, **kwargs): + import torch -def assert_torch_allclose(actual, expected, *, label=None, **kwargs): # test dtype, convert long to float if needed if actual.dtype in [torch.int64, torch.int32]: @@ -20,6 +19,9 @@ def assert_torch_allclose(actual, expected, *, label=None, **kwargs): def assert_np_allclose(actual, expected, *, label=None, **kwargs): + + import numpy as np + ok = np.allclose(actual, expected, **kwargs) diff = float(np.linalg.norm(np.asarray(actual) - np.asarray(expected))) prefix = label if label is not None else "np.allclose failed" From 6730ba097f626c5e23dab163f491c65aa3ec6e08 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Joan=20Alexis=20Glaun=C3=A8s?= Date: Tue, 9 Jun 2026 17:59:02 +0200 Subject: [PATCH 95/98] fixed some issues on mac --- keopscore/keopscore/config/CxxCompiler.py | 20 ++++++++++---------- keopscore/keopscore/config/OpenMP.py | 2 +- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/keopscore/keopscore/config/CxxCompiler.py b/keopscore/keopscore/config/CxxCompiler.py index 25cc75d86..43f45a368 100644 --- a/keopscore/keopscore/config/CxxCompiler.py +++ b/keopscore/keopscore/config/CxxCompiler.py @@ -19,7 +19,7 @@ class CxxCompilerConfig: _linking_options = "" _disable_pragma_unrolls = True - _use_Apple_clang = False + _use_clang_on_macos = False cxx_envs = ["CXX", "CXXFLAGS"] @@ -29,7 +29,7 @@ def __init__(self, platform): self.platform = platform self.set_cxx_compiler() - self.set_use_Apple_clang() + self.set_use_clang_on_macos() self.set_cxx_env_flags() self.set_compile_options() self.set_linking_options() @@ -93,16 +93,16 @@ def print_cxx_compiler_version(self): f"C++ Compiler Version: {self.get_cxx_compiler_version() or not_found_str}" ) - def set_use_Apple_clang(self): - """Detect if using Apple Clang.""" - self._use_Apple_clang = ( - "Apple clang" in self.get_cxx_compiler_version() + def set_use_clang_on_macos(self): + """Detect if using clang on macOS.""" + self._use_clang_on_macos = ( + "clang" in self.get_cxx_compiler_version() and self.platform.get_platform() == "Darwin" if self.get_cxx_compiler_version() else False ) - def get_use_Apple_clang(self): - return self._use_Apple_clang + def get_use_clang_on_macos(self): + return self._use_clang_on_macos # Disable Pragma Unrolls def set_disable_pragma_unrolls(self): @@ -154,7 +154,7 @@ def set_linking_options(self): self.add_to_linking_options("-shared") # linker behavior (macOS) - if self.platform.get_platform() == "Darwin" and self.get_use_Apple_clang(): + if self.platform.get_platform() == "Darwin" and self.get_use_clang_on_macos(): self.add_to_linking_options("-undefined dynamic_lookup") # architecture (macOS ARM) @@ -191,7 +191,7 @@ def set_cxx_env_flags(self): if ( cxxflags and self.platform.get_platform() == "Darwin" - and self.get_use_Apple_clang() + and self.get_use_clang_on_macos() ): cxxflags = re.sub( r"(?:^|\s)-(?:march|mtune|mfpmath)(?:=\S+|\s+\S+)?", diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 20d455e78..814ae3314 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -293,7 +293,7 @@ def check_compiler_for_openmp(self): # C++ Compiler Options def set_compile_options(self): # Apple clang does not support -fopenmp directly; -Xpreprocessor is required. - if self.cxx.get_use_Apple_clang(): + if self.cxx.get_use_clang_on_macos(): self._compile_options += " -Xpreprocessor" self._compile_options += " -fopenmp" From b14e44f01c03b5745bb29079616bb14d33f48267 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Joan=20Alexis=20Glaun=C3=A8s?= Date: Tue, 9 Jun 2026 18:00:15 +0200 Subject: [PATCH 96/98] lint --- keopscore/keopscore/config/CxxCompiler.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/keopscore/keopscore/config/CxxCompiler.py b/keopscore/keopscore/config/CxxCompiler.py index 43f45a368..bbf4f4602 100644 --- a/keopscore/keopscore/config/CxxCompiler.py +++ b/keopscore/keopscore/config/CxxCompiler.py @@ -96,7 +96,8 @@ def print_cxx_compiler_version(self): def set_use_clang_on_macos(self): """Detect if using clang on macOS.""" self._use_clang_on_macos = ( - "clang" in self.get_cxx_compiler_version() and self.platform.get_platform() == "Darwin" + "clang" in self.get_cxx_compiler_version() + and self.platform.get_platform() == "Darwin" if self.get_cxx_compiler_version() else False ) From 2eda3a9788b45d95b081c6dc55e3b3ad9f536a7c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Joan=20Alexis=20Glaun=C3=A8s?= Date: Wed, 10 Jun 2026 00:37:37 +0200 Subject: [PATCH 97/98] another fix to avoid "duplicate rpath" warning --- keopscore/keopscore/config/OpenMP.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/keopscore/keopscore/config/OpenMP.py b/keopscore/keopscore/config/OpenMP.py index 814ae3314..dbc14235f 100644 --- a/keopscore/keopscore/config/OpenMP.py +++ b/keopscore/keopscore/config/OpenMP.py @@ -335,6 +335,8 @@ def set_linking_options(self): and not importlib.util.find_spec("torch") ): # Force-link OpenMP runtime to avoid unresolved symbols at dlopen time. + # N.B. (Joan) We added this because when libomp was installed via homebrew, we had these unresolved symbols at runtime. + # However I cannot reproduce the issue now, so it may be unnecessary. lib_basename = os.path.basename(libomp_path) if lib_basename.startswith("lib"): @@ -342,7 +344,10 @@ def set_linking_options(self): if lib_name: link_flags.append(f"-l{lib_name}") - link_flags.append(f"-Wl,-rpath,{libomp_folder}") + if ( + "homebrew" in libomp_folder + ): # only when libomp is installed via homebrew, because when it is installed via conda, it is useless and produces a "duplicate rpath" warning. + link_flags.append(f"-Wl,-rpath,{libomp_folder}") self._linking_options = " ".join(link_flags) From b979ac775e3799a3cc1139779ae46ea62c19f129 Mon Sep 17 00:00:00 2001 From: Benjamin Charlier Date: Fri, 12 Jun 2026 16:18:24 +0200 Subject: [PATCH 98/98] add some detail on cuda pip install --- doc/python/installation.rst | 51 ++++++++++++++++++------------------- pykeops/setup.py | 3 ++- 2 files changed, 27 insertions(+), 27 deletions(-) diff --git a/doc/python/installation.rst b/doc/python/installation.rst index 4157ab766..b1f3d7efa 100644 --- a/doc/python/installation.rst +++ b/doc/python/installation.rst @@ -15,7 +15,7 @@ Requirements Required dependencies: - **Python** (>= 3.8) with the **NumPy** package. -- A C++ compiler such as ``gcc`` or ``clang``, compatible with ``std=c++11``. +- A C++ compiler such as ``gcc`` (on Linux) or ``clang`` (on macOS). Optional, but highly recommended: @@ -70,7 +70,7 @@ On Google Colab Google provides free virtual machines, running on Ubuntu Linux, where KeOps runs out of the box. In a new -`Colab notebook `_, typing: +`Colab notebook `_, run: .. prompt:: python >>> @@ -93,7 +93,8 @@ whether older packages are already installed somewhere else on your system. The dependencies for PyKeOps are listed :ref:`above `, but the exact packages needed depend on your system. For instance, the following -commands should create a working environment: +commands should create a working environment on Ubuntu 24.04 equipped with +a GPU: .. prompt:: bash $ @@ -160,9 +161,9 @@ As an example, here are the steps that we follow to render this website on the cd ~/scratch/containers # Download the Docker image and store it as an immutable Singularity Image File: - # N.B.: Our image is pretty heavy (~7 Gb), so it is safer to create + # N.B.: Our image is fairly large (~7 GB), so it is safer to create # cache folders on the hard drive instead of relying on the RAM-only tmpfs: - # N.B.: This step may take 15mn to 60mn, so you may prefer to execute it on + # N.B.: This step may take 15 to 60 minutes, so you may prefer to execute it on # your local computer and then copy the resulting file `keops-full.sif` # to the cluster. # Alternatively, on the Jean Zay cluster, you may use the `prepost` partition @@ -178,12 +179,12 @@ As an example, here are the steps that we follow to render this website on the # to add our new environment to the cluster's container registry as explained here: # http://www.idris.fr/jean-zay/cpu/jean-zay-utilisation-singularity.html - # Then, create a separate home folder for this image. This is to ensure - # that we won't see any conflict between different versions of the KeOps binaries, + # Then, create a separate home folder for this image. This helps ensure + # that we do not encounter conflicts between different versions of the KeOps binaries, # stored in the ~/.cache folder of the virtual machine: mkdir -p ~/containers/singularity_homes/keops-full - # Ask the slurm scheduler to render our documentation. + # Ask the Slurm scheduler to render our documentation. sbatch keops-doc.batch @@ -258,7 +259,7 @@ And ``keops-doc.sh`` is an executable file that contains: # Render the website: make html - # Re-render the doc to remove compilation messages: + # Re-render the documentation to remove compilation messages: make clean make html @@ -295,7 +296,7 @@ Alternatively, you may: If you prefer not to install the packages, you can add ``/path/to/keops_cloned_repo/keopscore`` and ``/path/to/keops_cloned_repo/pykeops`` to your Python path. To do this once - and for all, add the paths to your ``~/.bashrc``: + permanently, add the paths to your ``~/.bashrc``: .. prompt:: bash $ @@ -392,7 +393,8 @@ CUDA toolkit detection Detecting the CUDA toolkit is not always straightforward because several CUDA versions may be present on the same system. For instance, PyTorch installations -may bring NVIDIA runtime packages into the active Python environment. +may bring a partial CUDA toolkit installation into the active Python environment (see +detail below). KeOps does not ship its own CUDA toolkit. It tries to detect a working toolkit with :class:`keopscore.config.CudaConfig ` in @@ -421,14 +423,11 @@ the following order: # Arch Linux and derivatives: yay -S cuda -4. *`NVIDIA PyPI packages `_*: installed - with the ``cuda-toolkit[all]`` module from PyPI. Some of these packages are also pulled - in by PyTorch, making this solution fragile as multiple versions of the cuda-toolkit can - co-exist in the same virtual environment. For instance, Torch may install - CUDA 12 packages while the user manually installs CUDA 13 packages... - - -Using NVIDIA PyPI packages works on Arch Linux. The following command: +4. *`NVIDIA PyPI packages `_*: needed package could be installed + with the ``pip install pykeops[cuda]`` or ``pip install pykeops[cu12]`` or or ``pip install pykeops[cu13]`` recipes. Beware, many cuda install can co-exist in the same virtual environment. + + +We illustrate here how the CUDA toolkit could be manually selected. On Arch linux, the following command: .. prompt:: bash $ python -m venv keops_venv @@ -451,8 +450,8 @@ yields the detection of a system-wide CUDA installation under ``/opt/cuda``: Libnvrtc Path: /opt/cuda/lib64/libnvrtc.so.13 Libcudart Path: /opt/cuda/lib64/libcudart.so.13 -The following commands force KeOps to use the NVIDIA PyPI packages installed in -the active Python virtual environment: +Now, forcing KeOps to use the NVIDIA PyPI packages installed in the active Python virtual + environment may be done by running: .. prompt:: bash $ python -m venv keops_venv_cuda_pip @@ -482,8 +481,8 @@ Compiler Recent Linux and macOS distributions usually provide suitable compilers. If compilation fails, make sure that you are using a C++ compiler compatible with -the **C++11 revision**. Otherwise, formula compilation may fail in unexpected -ways. +the **C++11 revision** (``std=c++11`` flag supported). Otherwise, formula compilation + may fail in unexpected ways. 1. Install a compiler **system-wide**: for instance, on Debian-based Linux distributions, you can install g++ with apt and then use @@ -539,6 +538,6 @@ Alternatively, you can disable verbose compilation from your Python script with .. prompt:: python >>> import pykeops - pykeops.set_verbose(0) # no output - pykeops.set_verbose(1) # default verbosity level - pykeops.set_verbose(2) # maximum verbosity level + pykeops.set_verbose(0) # no output + pykeops.set_verbose(1) # default verbosity level + pykeops.set_verbose(2) # maximum verbosity level diff --git a/pykeops/setup.py b/pykeops/setup.py index d19cadbbb..e82cd2c75 100644 --- a/pykeops/setup.py +++ b/pykeops/setup.py @@ -95,6 +95,7 @@ ], "test:": ["pytest", "numpy", "torch"], "cu12": ["cuda-toolkit[nvrtc,nvcc,cudart,cccl]==12.9.2"], - "cu13": ["cuda-toolkit[nvrtc,nvcc,cccl,cudart, crt]==13.*"], + "cu13": ["cuda-toolkit[nvrtc,nvcc,cccl,cudart,crt]==13.*"], + "cuda": ["cuda-toolkit[nvrtc,nvcc,cccl,cudart,crt]"], }, )