I devised a simple program to explore solutions to a problem. Observing that it was very regular I expected to be able to easily take advantage of the GPU. But the performance was quite bad, so I came up with this minimal reproducible example.
import torch
def find_pairs(fn, y, n1, n2):
while torch.any(n2 > n1):
s = torch.sign(fn(n1) + fn(n2) - y)
yield s == 0, n1, n2
n1[s == -1] += 1
n2[(s != -1) & (n2 > n1)] -= 1;
Where fn represents any increasing function, n1 < n2 are integer numbers, it will yield solutions to fn(n1) + fn(n2) == y.
A simple example is to find Pythagorean triples
n = torch.arange(1, 2**13, dtype=torch.int64, device='cuda')
n1 = torch.ones_like(n)
n2 = torch.clone(n)
d = 0
for m, n1, n2 in find_pairs(lambda x: x**2, n**2, n1, n):
d += m & (n2 > n1)
torch.cuda.synchronize()
That takes around 8 seconds, if I run it on CPU it takes around 3.5 seconds.
Can anyone tell why it so slow, is there a way to make it to take full advantage of the GPU computation potential?