def test_sparse_coo_and_csr(m, n, k, nnz, test_count): start = Event(enable_timing=True) stop = Event(enable_timing=True) coo, csr = gen_sparse_coo_and_csr((m, k), nnz) mat = torch.randn((k, n), dtype=torch.double) times = [] for _ in range(test_count): start.record() coo.matmul(mat) stop.record() times.append(start.elapsed_time(stop)) coo_mean_time = sum(times) / len(times) times = [] for _ in range(test_count): start.record() csr.matmul(mat) stop.record() times.append(start.elapsed_time(stop)) csr_mean_time = sum(times) / len(times) return coo_mean_time, csr_mean_time
def test_sparse_csr(m, n, k, nnz, test_count): start_timer = Event(enable_timing=True) stop_timer = Event(enable_timing=True) csr = gen_sparse_csr((m, k), nnz) mat = torch.randn(k, n, dtype=torch.double) times = [] for _ in range(test_count): start_timer.record() csr.matmul(mat) stop_timer.record() times.append(start_timer.elapsed_time(stop_timer)) return sum(times) / len(times)