Add TORCH_CHECK_INDEX in convert_indices_from_coo_to_csr_cpu #138068
Conversation
|
This PR needs a
|
|
@Kh4L - have you tried using the |
Yes, I gave it a try I think the cost of adding this index check is relatively small compared to the time it could save developers in the long run |
|
@Kh4L - I agree, I'd have hoped that |
|
If this doesn't work, I suggest we modify this PR to use https://pytorch.org/docs/main/generated/torch.sparse.check_sparse_tensor_invariants.html#torch.sparse.check_sparse_tensor_invariants.is_enabled as a guard. That way this can be turned on for debugging, but doesn't affect performance by default. |
|
@cpuhrsch I understand but is it really ok to have python crashing with errors such as |
|
@cpuhrsch import torch
import time
import numpy as np
num_nonzeros = 2 ** 14
dense_size = (2**16, 2**14)
row_indices = torch.randint(0, dense_size[0], (num_nonzeros,))
col_indices = torch.randint(0, dense_size[1], (num_nonzeros,))
coo_indices = torch.stack((row_indices, col_indices))
sparse_coo = torch.sparse_coo_tensor(
coo_indices,
torch.ones(coo_indices.size(1), dtype=bool),
dense_size,
)
def benchmark_to_sparse_csr():
return sparse_coo.to_sparse_csr()
num_runs = 100
times = []
for _ in range(num_runs):
start_time = time.time()
benchmark_to_sparse_csr()
end_time = time.time()
times.append(end_time - start_time)
print(f"Mean execution time: {np.mean(times):.6f} seconds")
print(f"Standard deviation: {np.std(times):.6f} seconds")before the change (without index check): after the change (with index check): |
|
@Kh4L - when the Tensor is big it's fine, but maybe when it's small it's not fine anymore for some applications. By supporting the guard the user can choose what works best for them. |
The
to_sparse_csrCPU implementationconvert_indices_from_coo_to_csr_cpudoesn't validate the COO indices, which may lead to illegal memory access.This PR fixes that by adding checks on the index before accessing the data.