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
7
8import mlir.all_passes_registration
9
10from mlir import ir
11from mlir import runtime as rt
12from mlir import execution_engine
13from mlir import passmanager
14
15from mlir.dialects import sparse_tensor as st
16from mlir.dialects import builtin
17from mlir.dialects.linalg.opdsl import lang as dsl
18
19
20@dsl.linalg_structured_op
21def matmul_dsl(
22    A=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.K),
23    B=dsl.TensorDef(dsl.T, dsl.S.K, dsl.S.N),
24    C=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.N, output=True)):
25  C[dsl.D.m, dsl.D.n] += A[dsl.D.m, dsl.D.k] * B[dsl.D.k, dsl.D.n]
26
27
28def build_SpMM(attr: st.EncodingAttr):
29  """Build SpMM kernel.
30
31  This method generates a linalg op with for matrix multiplication using
32  just the Python API. Effectively, a generic linalg op is constructed
33  that computes C(i,j) += A(i,k) * B(k,j) for annotated matrix A.
34  """
35  module = ir.Module.create()
36  f64 = ir.F64Type.get()
37  a = ir.RankedTensorType.get([3, 4], f64, attr)
38  b = ir.RankedTensorType.get([4, 2], f64)
39  c = ir.RankedTensorType.get([3, 2], f64)
40  arguments = [a, b, c]
41  with ir.InsertionPoint(module.body):
42
43    @builtin.FuncOp.from_py_func(*arguments)
44    def spMxM(*args):
45      return matmul_dsl(args[0], args[1], outs=[args[2]])
46
47  return module
48
49
50def boilerplate(attr: st.EncodingAttr):
51  """Returns boilerplate main method.
52
53  This method sets up a boilerplate main method that takes three tensors
54  (a, b, c), converts the first tensor a into s sparse tensor, and then
55  calls the sparse kernel for matrix multiplication. For convenience,
56  this part is purely done as string input.
57  """
58  return f"""
59func @main(%ad: tensor<3x4xf64>, %b: tensor<4x2xf64>, %c: tensor<3x2xf64>) -> tensor<3x2xf64>
60  attributes {{ llvm.emit_c_interface }} {{
61  %a = sparse_tensor.convert %ad : tensor<3x4xf64> to tensor<3x4xf64, {attr}>
62  %0 = call @spMxM(%a, %b, %c) : (tensor<3x4xf64, {attr}>,
63                                  tensor<4x2xf64>,
64                                  tensor<3x2xf64>) -> tensor<3x2xf64>
65  return %0 : tensor<3x2xf64>
66}}
67"""
68
69
70def build_compile_and_run_SpMM(attr: st.EncodingAttr, support_lib: str,
71                               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  compiler(module)
79  engine = execution_engine.ExecutionEngine(
80      module, opt_level=0, shared_libs=[support_lib])
81
82  # Set up numpy input and buffer for output.
83  a = np.array(
84      [[1.1, 0.0, 0.0, 1.4], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 3.3, 0.0]],
85      np.float64)
86  b = np.array([[1.0, 2.0], [4.0, 3.0], [5.0, 6.0], [8.0, 7.0]], np.float64)
87  c = np.zeros((3, 2), np.float64)
88  out = np.zeros((3, 2), np.float64)
89
90  mem_a = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(a)))
91  mem_b = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(b)))
92  mem_c = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(c)))
93  mem_out = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(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
109class SparseCompiler:
110  """Sparse compiler passes."""
111
112  def __init__(self, options: str):
113    pipeline = (
114        f'builtin.func(linalg-generalize-named-ops,linalg-fuse-elementwise-ops),'
115        f'sparsification{{{options}}},'
116        f'sparse-tensor-conversion,'
117        f'builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),'
118        f'convert-scf-to-std,'
119        f'func-bufferize,'
120        f'tensor-constant-bufferize,'
121        f'builtin.func(tensor-bufferize,std-bufferize,finalizing-bufferize),'
122        f'convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},'
123        f'lower-affine,'
124        f'convert-memref-to-llvm,'
125        f'convert-std-to-llvm,'
126        f'reconcile-unrealized-casts')
127    self.pipeline = pipeline
128
129  def __call__(self, module: ir.Module):
130    passmanager.PassManager.parse(self.pipeline).run(module)
131
132
133def main():
134  support_lib = os.getenv('SUPPORT_LIB')
135  assert support_lib is not None, 'SUPPORT_LIB is undefined'
136  if not os.path.exists(support_lib):
137    raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), support_lib)
138
139  # CHECK-LABEL: TEST: testSpMM
140  print('\nTEST: testSpMM')
141  with ir.Context() as ctx, ir.Location.unknown():
142    count = 0
143    # Loop over various ways to compile and annotate the SpMM kernel with
144    # a *single* sparse tensor. Note that we deliberate do not exhaustively
145    # search the full state space to reduce runtime of the test. It is
146    # straightforward to adapt the code below to explore more combinations.
147    par = 0
148    vec = 0
149    vl = 1
150    e = False
151    opt = (f'parallelization-strategy={par} '
152           f'vectorization-strategy={vec} '
153           f'vl={vl} enable-simd-index32={e}')
154    levels = [[st.DimLevelType.dense, st.DimLevelType.dense],
155              [st.DimLevelType.dense, st.DimLevelType.compressed],
156              [st.DimLevelType.compressed, st.DimLevelType.dense],
157              [st.DimLevelType.compressed, st.DimLevelType.compressed]]
158    orderings = [
159        ir.AffineMap.get_permutation([0, 1]),
160        ir.AffineMap.get_permutation([1, 0])
161    ]
162    bitwidths = [0]
163    for level in levels:
164      for ordering in orderings:
165        for pwidth in bitwidths:
166          for iwidth in bitwidths:
167            attr = st.EncodingAttr.get(level, ordering, pwidth, iwidth)
168            compiler = SparseCompiler(options=opt)
169            build_compile_and_run_SpMM(attr, support_lib, compiler)
170            count = count + 1
171    # CHECK: Passed 8 tests
172    print('Passed ', count, 'tests')
173
174if __name__ == '__main__':
175  main()
176