#!/usr/bin/env python

from __future__ import print_function

import numpy

import time

from scipy.sparse import csr_matrix

from itertools import chain

from collections import OrderedDict

import kernel_tuner

def tune():

with open('spmv.cu', 'r') as f:

kernel_string = f.read()

nrows = numpy.int32(128*1024)

ncols = 64*1024

nnz = int(nrows*ncols*0.001)

#problem_size = (nrows, 1)

problem_size = nrows

#generate sparse matrix in CSR

rows = numpy.asarray([0]+sorted(numpy.random.rand(nrows-1)*nnz)+[nnz]).astype(numpy.int32)

cols = (numpy.random.rand(nnz)*ncols).astype(numpy.int32)

vals = numpy.random.randn(nnz).astype(numpy.float32)

#input and output vector (y = matrix * x)

x = numpy.random.randn(ncols).astype(numpy.float32)

y = numpy.zeros(nrows).astype(numpy.float32)

args = [y, rows, cols, vals, x, nrows]

tune_params = OrderedDict()

tune_params["block_size_x"] = [32*i for i in range(1,33)]

tune_params["threads_per_row"] = [1, 32]

tune_params["read_only"] = [0, 1]

grid_div_x = ["block_size_x/threads_per_row"]

#compute reference answer using scipy.sparse

row_ind = list(chain.from_iterable([[i] * (rows[i+1]-rows[i]) for i in range(nrows)]))

matrix = csr_matrix((vals, (row_ind, cols)), shape=(nrows, ncols))

start = time.time()

expected_y = matrix.dot(x)

end = time.time()

print("computing reference using scipy.sparse took: " + str(start-end / 1000.0) + " ms.")

answer = [expected_y, None, None, None, None, None]

return kernel_tuner.tune_kernel("spmv_kernel", kernel_string,

problem_size, args, tune_params,

grid_div_x=grid_div_x, verbose=True, answer=answer, atol=1e-4)

if __name__ == "__main__":

tune()