1# RUN: SUPPORT_LIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext \
2# RUN:   %PYTHON %s | FileCheck %s
3
4import ctypes
5import numpy as np
6import os
7import sys
8
9from mlir import ir
10from mlir import runtime as rt
11
12from mlir.dialects import sparse_tensor as st
13from mlir.dialects import builtin
14from mlir.dialects import func
15from mlir.dialects.linalg.opdsl import lang as dsl
16
17_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
18sys.path.append(_SCRIPT_PATH)
19from tools import sparse_compiler
20
21@dsl.linalg_structured_op
22def matmul_dsl(
23    A=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.K),
24    B=dsl.TensorDef(dsl.T, dsl.S.K, dsl.S.N),
25    C=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.N, output=True)):
26    C[dsl.D.m, dsl.D.n] += A[dsl.D.m, dsl.D.k] * B[dsl.D.k, dsl.D.n]
27
28
29def build_SpMM(attr: st.EncodingAttr):
30    """Build SpMM kernel.
31
32  This method generates a linalg op with for matrix multiplication using
33  just the Python API. Effectively, a generic linalg op is constructed
34  that computes C(i,j) += A(i,k) * B(k,j) for annotated matrix A.
35  """
36    module = ir.Module.create()
37    f64 = ir.F64Type.get()
38    a = ir.RankedTensorType.get([3, 4], f64, attr)
39    b = ir.RankedTensorType.get([4, 2], f64)
40    c = ir.RankedTensorType.get([3, 2], f64)
41    arguments = [a, b, c]
42    with ir.InsertionPoint(module.body):
43
44        @func.FuncOp.from_py_func(*arguments)
45        def spMxM(*args):
46            return matmul_dsl(args[0], args[1], outs=[args[2]])
47
48    return module
49
50
51def boilerplate(attr: st.EncodingAttr):
52    """Returns boilerplate main method.
53
54  This method sets up a boilerplate main method that takes three tensors
55  (a, b, c), converts the first tensor a into s sparse tensor, and then
56  calls the sparse kernel for matrix multiplication. For convenience,
57  this part is purely done as string input.
58  """
59    return f"""
60func.func @main(%ad: tensor<3x4xf64>, %b: tensor<4x2xf64>, %c: tensor<3x2xf64>) -> tensor<3x2xf64>
61  attributes {{ llvm.emit_c_interface }} {{
62  %a = sparse_tensor.convert %ad : tensor<3x4xf64> to tensor<3x4xf64, {attr}>
63  %0 = call @spMxM(%a, %b, %c) : (tensor<3x4xf64, {attr}>,
64                                  tensor<4x2xf64>,
65                                  tensor<3x2xf64>) -> tensor<3x2xf64>
66  return %0 : tensor<3x2xf64>
67}}
68"""
69
70
71def build_compile_and_run_SpMM(attr: st.EncodingAttr, compiler):
72    # Build.
73    module = build_SpMM(attr)
74    func = str(module.operation.regions[0].blocks[0].operations[0].operation)
75    module = ir.Module.parse(func + boilerplate(attr))
76
77    # Compile.
78    engine = compiler.compile_and_jit(module)
79
80    # Set up numpy input and buffer for output.
81    a = np.array(
82        [[1.1, 0.0, 0.0, 1.4], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 3.3, 0.0]],
83        np.float64)
84    b = np.array([[1.0, 2.0], [4.0, 3.0], [5.0, 6.0], [8.0, 7.0]], np.float64)
85    c = np.zeros((3, 2), np.float64)
86
87    mem_a = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(a)))
88    mem_b = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(b)))
89    mem_c = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(c)))
90    # Allocate a MemRefDescriptor to receive the output tensor.
91    # The buffer itself is allocated inside the MLIR code generation.
92    ref_out = rt.make_nd_memref_descriptor(2, ctypes.c_double)()
93    mem_out = ctypes.pointer(ctypes.pointer(ref_out))
94
95    # Invoke the kernel and get numpy output.
96    # Built-in bufferization uses in-out buffers.
97    # TODO: replace with inplace comprehensive bufferization.
98    engine.invoke('main', mem_out, mem_a, mem_b, mem_c)
99
100    # Sanity check on computed result.
101    expected = np.matmul(a, b);
102    c = rt.ranked_memref_to_numpy(mem_out[0])
103    if np.allclose(c, expected):
104        pass
105    else:
106        quit(f'FAILURE')
107
108
109def main():
110    support_lib = os.getenv('SUPPORT_LIB')
111    assert support_lib is not None, 'SUPPORT_LIB is undefined'
112    if not os.path.exists(support_lib):
113        raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), support_lib)
114
115    # CHECK-LABEL: TEST: testSpMM
116    print('\nTEST: testSpMM')
117    with ir.Context() as ctx, ir.Location.unknown():
118        count = 0
119        # Loop over various ways to compile and annotate the SpMM kernel with
120        # a *single* sparse tensor. Note that we deliberate do not exhaustively
121        # search the full state space to reduce runtime of the test. It is
122        # straightforward to adapt the code below to explore more combinations.
123        par = 0
124        vec = 0
125        vl = 1
126        e = False
127        opt = (f'parallelization-strategy={par} '
128               f'vectorization-strategy={vec} '
129               f'vl={vl} enable-simd-index32={e}')
130        levels = [[st.DimLevelType.dense, st.DimLevelType.dense],
131                  [st.DimLevelType.dense, st.DimLevelType.compressed],
132                  [st.DimLevelType.compressed, st.DimLevelType.dense],
133                  [st.DimLevelType.compressed, st.DimLevelType.compressed]]
134        orderings = [
135            ir.AffineMap.get_permutation([0, 1]),
136            ir.AffineMap.get_permutation([1, 0])
137        ]
138        bitwidths = [0]
139        compiler = sparse_compiler.SparseCompiler(
140            options=opt, opt_level=0, shared_libs=[support_lib])
141        for level in levels:
142            for ordering in orderings:
143                for pwidth in bitwidths:
144                    for iwidth in bitwidths:
145                        attr = st.EncodingAttr.get(level, ordering, pwidth, iwidth)
146                        build_compile_and_run_SpMM(attr, compiler)
147                        count = count + 1
148        # CHECK: Passed 8 tests
149        print('Passed ', count, 'tests')
150
151
152if __name__ == '__main__':
153    main()
154