mlir/test/Integration/Dialect/SparseTensor/taco/test_SpMV.py

   1 # RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
   2
   3 import numpy as np
   4 import os
   5 import sys
   6 import tempfile
   7
   8 _SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
   9 sys.path.append(_SCRIPT_PATH)
  10 from tools import mlir_pytaco_api as pt
  11
  12 ###### This PyTACO part is taken from the TACO open-source project. ######
  13 # See http://tensor-compiler.org/docs/scientific_computing/index.html.
  14
  15 compressed = pt.compressed
  16 dense = pt.dense
  17
  18 # Define formats for storing the sparse matrix and dense vectors.
  19 csr = pt.format([dense, compressed])
  20 dv = pt.format([dense])
  21
  22 # Load a sparse matrix stored in the matrix market format) and store it
  23 # as a CSR matrix.  The matrix in this test is a reduced version of the data
  24 # downloaded from here:
  25 # https://www.cise.ufl.edu/research/sparse/MM/Boeing/pwtk.tar.gz
  26 # In order to run the program using the matrix above, you can download the
  27 # matrix and replace this path to the actual path to the file.
  28 A = pt.read(os.path.join(_SCRIPT_PATH, "data/pwtk.mtx"), csr)
  29
  30 # These two lines have been modified from the original program to use static
  31 # data to support result comparison.
  32 x = pt.from_array(np.full((A.shape[1],), 1, dtype=np.float64))
  33 z = pt.from_array(np.full((A.shape[0],), 2, dtype=np.float64))
  34
  35 # Declare the result to be a dense vector
  36 y = pt.tensor([A.shape[0]], dv)
  37
  38 # Declare index vars
  39 i, j = pt.get_index_vars(2)
  40
  41 # Define the SpMV computation
  42 y[i] = A[i, j] * x[j] + z[i]
  43
  44 ##########################################################################
  45
  46 # CHECK: Compare result True
  47 # Perform the SpMV computation and write the result to file
  48 with tempfile.TemporaryDirectory() as test_dir:
  49   actual_file = os.path.join(test_dir, "y.tns")
  50   pt.write(actual_file, y)
  51   actual = np.loadtxt(actual_file, np.float64)
  52   expected = np.loadtxt(
  53       os.path.join(_SCRIPT_PATH, "data/gold_y.tns"), np.float64)
  54   print(f"Compare result {np.allclose(actual, expected, rtol=0.01)}")