search

getkeops / keops

KErnel OPerationS, on CPUs and GPUs, with autodiff and without memory overflows
https://www.kernel-operations.io
MIT License
1.03k stars 65 forks source link

Error: incompatible dimensions in VectApply with chunking on Pm variable #294

Closed Giodiro closed 1 year ago

Giodiro commented 1 year ago

Hi again, Another crash which I think should have an easy fix, in version 2.1.1 of KeOps.

The minimal example is again a RBF kernel, but with higher dimensions so that some kind of chunking kicks in

import torch
from pykeops.torch import Genred

X1 = torch.randn(20, 100).cuda().requires_grad_()
X2 = torch.randn(5, 100).cuda().requires_grad_()
v = torch.randn(5, 120).cuda().requires_grad_()
sigma = torch.randn(100).cuda().requires_grad_()

formula = 'Exp(SqDist((x1 / s), (x2 / s)) * IntInv(-2)) * v'
aliases = [
    'x1 = Vi(%d)' % (X1.shape[1]),
    'x2 = Vj(%d)' % (X2.shape[1]),
    'v = Vj(%d)' % (v.shape[1]),
    's = Pm(%d)' % (sigma.shape[0])
]
other_vars = [sigma]

fn = Genred(formula, aliases, reduction_op="Sum", axis=1)
out = fn(X1, X2, v, *other_vars, backend="GPU_1D")
print(out.shape)

This fails when calling fn with the following traceback

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/pykeops/torch/generic/generic_red.py", line 627, in __call__
    out = GenredAutograd.apply(
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/pykeops/torch/generic/generic_red.py", line 78, in forward
    myconv = keops_binder["nvrtc" if tagCPUGPU else "cpp"](
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/keopscore/utils/Cache.py", line 68, in __call__
    obj = self.cls(*args)
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/pykeops/common/keops_io/LoadKeOps_nvrtc.py", line 15, in __init__
    super().__init__(*args, fast_init=fast_init)
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/pykeops/common/keops_io/LoadKeOps.py", line 18, in __init__
    self.init(*args)
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/pykeops/common/keops_io/LoadKeOps.py", line 126, in init
    ) = get_keops_dll(
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/keopscore/utils/Cache.py", line 27, in __call__
    self.library[str_id] = self.fun(*args)
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/keopscore/get_keops_dll.py", line 124, in get_keops_dll_impl
    res = map_reduce_obj.get_dll_and_params()
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/keopscore/binders/LinkCompile.py", line 101, in get_dll_and_params
    self.generate_code()
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/keopscore/binders/nvrtc/Gpu_link_compile.py", line 63, in generate_code
    self.get_code()
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py", line 167, in get_code
    chunk_sub_routine = do_chunk_sub(
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/keopscore/mapreduce/gpu/GpuReduc1D_chunks.py", line 104, in do_chunk_sub
    {fun_chunked_curr(fout_tmp_chunk, chktable)}
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/keopscore/formulas/Operation.py", line 82, in __call__
    string += child(arg, table)
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/keopscore/formulas/Operation.py", line 82, in __call__
    string += child(arg, table)
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/keopscore/formulas/Operation.py", line 85, in __call__
    string += self.Op(out, table, *args)
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/keopscore/formulas/VectorizedScalarOp.py", line 26, in Op
    return VectApply(self.ScalarOp, out, *args)
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/keopscore/utils/code_gen_utils.py", line 385, in VectApply
    KeOps_Error("incompatible dimensions in VectApply")
  File "/anaconda3/envs/torch_111/lib/python3.9/site-packages/keopscore/utils/misc_utils.py", line 28, in KeOps_Error
    raise ValueError(message)
ValueError: [KeOps] Error : incompatible dimensions in VectApply (error at line 385 in file /anaconda3/envs/torch_111/lib/python3.9/site-packages/keopscore/utils/code_gen_utils.py)

while it works when sigma is 1D. Digging a little bit, with 1D sigma we have dims=[64, 64, 1] and dimloop=64 and the check passes. With 100D sigma instead dims=[64, 64, 100] and dimloop=100, hence the check fails. I tried to get rid of the if statement and nothing else crashes, so hopefully it's just a matter of adding another edge case to it.

Also do let me know if the preferred way of writing kernel-vector multiplications has changed with the evolution of your library, and I'm using code that is somewhat outdated. I have not yet adopted the LazyTensor interface as I quite liked the string formulas, but maybe this error would not have happened with LazyTensors?

Thanks again, Giacomo

joanglaunes commented 1 year ago

Hello @Giodiro , Sorry again for the late reply, I had done the fix some time ago but forgot to merge. Can you check that it works now ? About your last question, there should be absolutely no problem in using the Genred syntax. The LazyTensor interface always redirects to Genred in the end.

Giodiro commented 1 year ago

Hi @joanglaunes Working perfectly now, looks great! Thanks for the fix :)