mlir/test/Integration/Dialect/SparseTensor/taco/test_MTTKRP.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/data_analytics/index.html.
  14
  15 compressed = pt.compressed
  16 dense = pt.dense
  17
  18 # Define formats for storing the sparse tensor and dense matrices.
  19 csf = pt.format([compressed, compressed, compressed])
  20 rm = pt.format([dense, dense])
  21
  22 # Load a sparse three-dimensional tensor from file (stored in the FROSTT
  23 # format) and store it as a compressed sparse fiber tensor. We use a small
  24 # tensor for the purpose of testing. To run the program using the data from
  25 # the real application, please download the data from:
  26 # http://frostt.io/tensors/nell-2/
  27 B = pt.read(os.path.join(_SCRIPT_PATH, "data/nell-2.tns"), csf)
  28
  29 # These two lines have been modified from the original program to use static
  30 # data to support result comparison.
  31 C = pt.from_array(np.full((B.shape[1], 25), 1, dtype=np.float64))
  32 D = pt.from_array(np.full((B.shape[2], 25), 2, dtype=np.float64))
  33
  34 # Declare the result to be a dense matrix.
  35 A = pt.tensor([B.shape[0], 25], rm)
  36
  37 # Declare index vars.
  38 i, j, k, l = pt.get_index_vars(4)
  39
  40 # Define the MTTKRP computation.
  41 A[i, j] = B[i, k, l] * D[l, j] * C[k, j]
  42
  43 ##########################################################################
  44
  45 # CHECK: Compare result True
  46 # Perform the MTTKRP computation and write the result to file.
  47 with tempfile.TemporaryDirectory() as test_dir:
  48   actual_file = os.path.join(test_dir, "A.tns")
  49   pt.write(actual_file, A)
  50   actual = np.loadtxt(actual_file, np.float64)
  51   expected = np.loadtxt(
  52       os.path.join(_SCRIPT_PATH, "data/gold_A.tns"), np.float64)
  53   print(f"Compare result {np.allclose(actual, expected, rtol=0.01)}")