tensor-compiler / taco

The Tensor Algebra Compiler (taco) computes sparse tensor expressions on CPUs and GPUs
http://tensor-compiler.org
Other
1.26k stars 189 forks source link

Memory leak with `pytaco` #546

Open f-dangel opened 1 year ago

f-dangel commented 1 year ago

Hi,

I ran into a memory leak when using pytaco (I am using commit 2b8ece4c230a5f0f0a74bc6f48e28edfb6c1c95e, the current master at the time of writing).

The following script demonstrates the memory leak. I have two vectors a, b and repeatedly compute their outer product C[i, j] = a[i] * b[j].

"""Repeatedly compute the outer product of two random vectors."""

import os
import random

import psutil
import pytaco

dim = 2000
iterations = 201
print_every = 40

a = pytaco.tensor([dim], pytaco.dense)
b = pytaco.tensor([dim], pytaco.dense)

for i in range(dim):
    a.insert((i,), random.random())
    b.insert((i,), random.random())

a.pack()
b.pack()

i, j = pytaco.get_index_vars(2)

for n in range(iterations):
    C = pytaco.tensor([dim, dim], pytaco.dense)
    C[i, j] = a[i] * b[j]

    C.compile()
    C.assemble()
    C.compute()

    if n % print_every == 0:
        process = psutil.Process(os.getpid())
        print(f"Iter. {n}, memory (bytes): {process.memory_info().rss:.2e}")

Here is the output, which shows that memory increases at each iteration:

Iter. 0, memory (bytes): 2.84e+08
Iter. 40, memory (bytes): 9.25e+08
Iter. 80, memory (bytes): 1.56e+09
Iter. 120, memory (bytes): 2.21e+09
Iter. 160, memory (bytes): 2.84e+09
Iter. 200, memory (bytes): 3.46e+09

I tried adding a del C, and a import gc; gc.collect() at the end of each iteration without success. Therefore I suspect the leak is in the C++ layer.

Best, Felix