File: test_SpMM.py

package info (click to toggle)
swiftlang 6.0.3-2
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 2,519,992 kB
  • sloc: cpp: 9,107,863; ansic: 2,040,022; asm: 1,135,751; python: 296,500; objc: 82,456; f90: 60,502; lisp: 34,951; pascal: 19,946; sh: 18,133; perl: 7,482; ml: 4,937; javascript: 4,117; makefile: 3,840; awk: 3,535; xml: 914; fortran: 619; cs: 573; ruby: 573
file content (35 lines) | stat: -rw-r--r-- 1,061 bytes parent folder | download | duplicates (2) 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
 
# RUN: env SUPPORTLIB=%mlir_c_runner_utils %PYTHON %s | FileCheck %s

import filecmp
import numpy as np
import os
import sys
import tempfile

_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
sys.path.append(_SCRIPT_PATH)

from tools import mlir_pytaco_api as pt
from tools import testing_utils as utils

# Define the CSR format.
csr = pt.format([pt.dense, pt.compressed], [0, 1])

# Read matrices A and B from file, infer size of output matrix C.
A = pt.read(os.path.join(_SCRIPT_PATH, "data/A.mtx"), csr)
B = pt.read(os.path.join(_SCRIPT_PATH, "data/B.mtx"), csr)
C = pt.tensor([A.shape[0], B.shape[1]], csr)

# Define the kernel.
i, j, k = pt.get_index_vars(3)
C[i, j] = A[i, k] * B[k, j]

# Force evaluation of the kernel by writing out C.
with tempfile.TemporaryDirectory() as test_dir:
    golden_file = os.path.join(_SCRIPT_PATH, "data/gold_C.tns")
    out_file = os.path.join(test_dir, "C.tns")
    pt.write(out_file, C)
    #
    # CHECK: Compare result True
    #
    print(f"Compare result {utils.compare_sparse_tns(golden_file, out_file)}")