ampere_research/pytorch/spmv.py
2024-12-09 10:57:15 -05:00

55 lines
1.6 KiB
Python

from data_stat import Stat
import torch, scipy
import numpy as np
import argparse
import time
import json
import sys, os
parser = argparse.ArgumentParser()
parser.add_argument('matrix_file', help='the input matrix (.mtx) file')
parser.add_argument('iterations', type=int, help='the number of iterations of multiplication to perform')
args = parser.parse_args()
device = 'cpu'
matrix = scipy.io.mmread(args.matrix_file)
matrix = torch.sparse_coo_tensor(
np.vstack((matrix.row, matrix.col)),
matrix.data, matrix.shape,
device=device
).to_sparse_csr().type(torch.float)
vector = torch.rand(matrix.shape[1], device=device)
print(matrix, file=sys.stderr)
print(vector, file=sys.stderr)
start = time.time()
for i in range(0, args.iterations):
torch.mv(matrix, vector)
end = time.time()
result = dict()
result[Stat.MATRIX_FILE.name] = os.path.splitext(os.path.basename(args.matrix_file))[0]
print(f"Matrix: {result[Stat.MATRIX_FILE.name]}", file=sys.stderr)
result[Stat.MATRIX_SHAPE.name] = matrix.shape
print(f"Shape: {result[Stat.MATRIX_SHAPE.name]}", file=sys.stderr)
result[Stat.MATRIX_SIZE.name] = matrix.shape[0] * matrix.shape[1]
print(f"Size: {result[Stat.MATRIX_SIZE.name]}", file=sys.stderr)
result[Stat.MATRIX_NNZ.name] = matrix.values().shape[0]
print(f"NNZ: {result[Stat.MATRIX_NNZ.name]}", file=sys.stderr)
result[Stat.MATRIX_DENSITY.name] = matrix.values().shape[0] / (matrix.shape[0] * matrix.shape[1])
print(f"Density: {result[Stat.MATRIX_DENSITY.name]}", file=sys.stderr)
result[Stat.TIME_S.name] = end - start
print(f"Time: {result[Stat.TIME_S.name]} seconds", file=sys.stderr)
print(json.dumps(result))