103 lines
3.6 KiB
Python
103 lines
3.6 KiB
Python
from data_stat import Stat, Format, MatrixType
|
|
|
|
import torch, scipy
|
|
import numpy as np
|
|
import argparse
|
|
import time
|
|
import json
|
|
import sys, os
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('matrix_type', type=str,
|
|
choices=[t.name.lower() for t in MatrixType],
|
|
help='the type of matrix')
|
|
parser.add_argument('format', type=str,
|
|
choices=[fmt.name.lower() for fmt in Format],
|
|
help='the sparse format to use')
|
|
parser.add_argument('iterations', type=int, help='the number of iterations of multiplication to perform')
|
|
parser.add_argument('-m', '--matrix_file', type=str,
|
|
help='the input matrix (.mtx) file')
|
|
parser.add_argument('-ss', '--synthetic_size', type=int,
|
|
help='the synthetic matrix parameters size (rows)')
|
|
parser.add_argument('-sd', '--synthetic_density', type=float,
|
|
help='the synthetic matrix density')
|
|
args = parser.parse_args()
|
|
args.matrix_type = MatrixType[args.matrix_type.upper()]
|
|
args.format = Format[args.format.upper()]
|
|
|
|
device = 'cpu'
|
|
|
|
if args.matrix_type == MatrixType.SUITESPARSE:
|
|
if args.matrix_file is None:
|
|
exit("Matrix file not specified!")
|
|
|
|
matrix = scipy.io.mmread(args.matrix_file)
|
|
matrix = torch.sparse_coo_tensor(
|
|
np.vstack((matrix.row, matrix.col)),
|
|
matrix.data, matrix.shape,
|
|
device=device, dtype=torch.float32)
|
|
elif args.matrix_type == MatrixType.SYNTHETIC:
|
|
if args.synthetic_size is None and args.synthetic_density is None:
|
|
exit("Synthetic matrix parameters not specified!")
|
|
|
|
nnz = int((args.synthetic_size ** 2) * args.synthetic_density)
|
|
row_indices = torch.randint(0, args.synthetic_size, (nnz,))
|
|
col_indices = torch.randint(0, args.synthetic_size, (nnz,))
|
|
indices = torch.stack([row_indices, col_indices])
|
|
values = torch.randn(nnz)
|
|
|
|
matrix = torch.sparse_coo_tensor(
|
|
indices, values,
|
|
size=(args.synthetic_size, args.synthetic_size),
|
|
device=device, dtype=torch.float32)
|
|
else:
|
|
exit("Unrecognized matrix type!")
|
|
|
|
if args.format == Format.CSR:
|
|
matrix = matrix.to_sparse_csr().type(torch.float32)
|
|
elif args.format == Format.COO:
|
|
pass
|
|
else:
|
|
exit("Unrecognized format!")
|
|
|
|
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)
|
|
#torch.sparse.mm(matrix, vector.unsqueeze(-1)).squeeze(-1)
|
|
#print(i)
|
|
end = time.time()
|
|
|
|
result = dict()
|
|
|
|
result[Stat.MATRIX_TYPE.name] = args.matrix_type.value
|
|
print(f"Matrix: {result[Stat.MATRIX_TYPE.name]}", file=sys.stderr)
|
|
|
|
if args.matrix_type == MatrixType.SUITESPARSE:
|
|
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_FORMAT.name] = args.format.value
|
|
print(f"Matrix: {result[Stat.MATRIX_FORMAT.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))
|