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