1# RUN: SUPPORT_LIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext \
2# RUN:   %PYTHON %s | FileCheck %s
3
4import ctypes
5import errno
6import itertools
7import os
8import sys
9
10from typing import List, Callable
11
12import numpy as np
13
14from mlir import ir
15from mlir import runtime as rt
16from mlir.execution_engine import ExecutionEngine
17
18from mlir.dialects import builtin
19from mlir.dialects import func
20from mlir.dialects import sparse_tensor as st
21
22_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
23sys.path.append(_SCRIPT_PATH)
24from tools import sparse_compiler
25
26# ===----------------------------------------------------------------------=== #
27
28# TODO: move this boilerplate to its own module, so it can be used by
29# other tests and programs.
30class TypeConverter:
31  """Converter between NumPy types and MLIR types."""
32
33  def __init__(self, context: ir.Context):
34    # Note 1: these are numpy "scalar types" (i.e., the values of
35    # np.sctypeDict) not numpy "dtypes" (i.e., the np.dtype class).
36    #
37    # Note 2: we must construct the MLIR types in the same context as the
38    # types that'll be passed to irtype_to_sctype() or irtype_to_dtype();
39    # otherwise, those methods will raise a KeyError.
40    types_list = [
41      (np.float64, ir.F64Type.get(context=context)),
42      (np.float32, ir.F32Type.get(context=context)),
43      (np.int64, ir.IntegerType.get_signless(64, context=context)),
44      (np.int32, ir.IntegerType.get_signless(32, context=context)),
45      (np.int16, ir.IntegerType.get_signless(16, context=context)),
46      (np.int8, ir.IntegerType.get_signless(8, context=context)),
47    ]
48    self._sc2ir = dict(types_list)
49    self._ir2sc = dict(( (ir,sc) for sc,ir in types_list ))
50
51  def dtype_to_irtype(self, dtype: np.dtype) -> ir.Type:
52    """Returns the MLIR equivalent of a NumPy dtype."""
53    try:
54      return self.sctype_to_irtype(dtype.type)
55    except KeyError as e:
56      raise KeyError(f'Unknown dtype: {dtype}') from e
57
58  def sctype_to_irtype(self, sctype) -> ir.Type:
59    """Returns the MLIR equivalent of a NumPy scalar type."""
60    if sctype in self._sc2ir:
61      return self._sc2ir[sctype]
62    else:
63      raise KeyError(f'Unknown sctype: {sctype}')
64
65  def irtype_to_dtype(self, tp: ir.Type) -> np.dtype:
66    """Returns the NumPy dtype equivalent of an MLIR type."""
67    return np.dtype(self.irtype_to_sctype(tp))
68
69  def irtype_to_sctype(self, tp: ir.Type):
70    """Returns the NumPy scalar-type equivalent of an MLIR type."""
71    if tp in self._ir2sc:
72      return self._ir2sc[tp]
73    else:
74      raise KeyError(f'Unknown ir.Type: {tp}')
75
76  def get_RankedTensorType_of_nparray(self, nparray: np.ndarray) -> ir.RankedTensorType:
77    """Returns the ir.RankedTensorType of a NumPy array.  Note that NumPy
78    arrays can only be converted to/from dense tensors, not sparse tensors."""
79    # TODO: handle strides as well?
80    return ir.RankedTensorType.get(nparray.shape,
81                                   self.dtype_to_irtype(nparray.dtype))
82
83# ===----------------------------------------------------------------------=== #
84
85class StressTest:
86  def __init__(self, tyconv: TypeConverter):
87    self._tyconv = tyconv
88    self._roundtripTp = None
89    self._module = None
90    self._engine = None
91
92  def _assertEqualsRoundtripTp(self, tp: ir.RankedTensorType):
93    assert self._roundtripTp is not None, \
94        'StressTest: uninitialized roundtrip type'
95    if tp != self._roundtripTp:
96      raise AssertionError(
97          f"Type is not equal to the roundtrip type.\n"
98          f"\tExpected: {self._roundtripTp}\n"
99          f"\tFound:    {tp}\n")
100
101  def build(self, types: List[ir.Type]):
102    """Builds the ir.Module.  The module has only the @main function,
103    which will convert the input through the list of types and then back
104    to the initial type.  The roundtrip type must be a dense tensor."""
105    assert self._module is None, 'StressTest: must not call build() repeatedly'
106    self._module = ir.Module.create()
107    with ir.InsertionPoint(self._module.body):
108      tp0 = types.pop(0)
109      self._roundtripTp = tp0
110      # TODO: assert dense? assert element type is recognised by the TypeConverter?
111      types.append(tp0)
112      funcTp = ir.FunctionType.get(inputs=[tp0], results=[tp0])
113      funcOp = func.FuncOp(name='main', type=funcTp)
114      funcOp.attributes['llvm.emit_c_interface'] = ir.UnitAttr.get()
115      with ir.InsertionPoint(funcOp.add_entry_block()):
116        arg0 = funcOp.entry_block.arguments[0]
117        self._assertEqualsRoundtripTp(arg0.type)
118        v = st.ConvertOp(types.pop(0), arg0)
119        for tp in types:
120          w = st.ConvertOp(tp, v)
121          # Release intermediate tensors before they fall out of scope.
122          st.ReleaseOp(v.result)
123          v = w
124        self._assertEqualsRoundtripTp(v.result.type)
125        func.ReturnOp(v)
126    return self
127
128  def writeTo(self, filename):
129    """Write the ir.Module to the given file.  If the file already exists,
130    then raises an error.  If the filename is None, then is a no-op."""
131    assert self._module is not None, \
132        'StressTest: must call build() before writeTo()'
133    if filename is None:
134      # Silent no-op, for convenience.
135      return self
136    if os.path.exists(filename):
137      raise FileExistsError(errno.EEXIST, os.strerror(errno.EEXIST), filename)
138    with open(filename, 'w') as f:
139      f.write(str(self._module))
140    return self
141
142  def compile(self, compiler, support_lib: str):
143    """Compile the ir.Module."""
144    assert self._module is not None, \
145        'StressTest: must call build() before compile()'
146    assert self._engine is None, \
147        'StressTest: must not call compile() repeatedly'
148    compiler(self._module)
149    self._engine = ExecutionEngine(
150        self._module, opt_level=0, shared_libs=[support_lib])
151    return self
152
153  def run(self, np_arg0: np.ndarray) -> np.ndarray:
154    """Runs the test on the given numpy array, and returns the resulting
155    numpy array."""
156    assert self._engine is not None, \
157        'StressTest: must call compile() before run()'
158    self._assertEqualsRoundtripTp(
159        self._tyconv.get_RankedTensorType_of_nparray(np_arg0))
160    np_out = np.zeros(np_arg0.shape, dtype=np_arg0.dtype)
161    self._assertEqualsRoundtripTp(
162        self._tyconv.get_RankedTensorType_of_nparray(np_out))
163    mem_arg0 = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(np_arg0)))
164    mem_out = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(np_out)))
165    self._engine.invoke('main', mem_out, mem_arg0)
166    return rt.ranked_memref_to_numpy(mem_out[0])
167
168# ===----------------------------------------------------------------------=== #
169
170def main():
171  """
172  USAGE: python3 test_stress.py [raw_module.mlir [compiled_module.mlir]]
173
174  The environment variable SUPPORT_LIB must be set to point to the
175  libmlir_c_runner_utils shared library.  There are two optional
176  arguments, for debugging purposes.  The first argument specifies where
177  to write out the raw/generated ir.Module.  The second argument specifies
178  where to write out the compiled version of that ir.Module.
179  """
180  support_lib = os.getenv('SUPPORT_LIB')
181  assert support_lib is not None, 'SUPPORT_LIB is undefined'
182  if not os.path.exists(support_lib):
183    raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), support_lib)
184
185  # CHECK-LABEL: TEST: test_stress
186  print("\nTEST: test_stress")
187  with ir.Context() as ctx, ir.Location.unknown():
188    par = 0
189    vec = 0
190    vl = 1
191    e = False
192    sparsification_options = (
193        f'parallelization-strategy={par} '
194        f'vectorization-strategy={vec} '
195        f'vl={vl} '
196        f'enable-simd-index32={e}')
197    compiler = sparse_compiler.SparseCompiler(options=sparsification_options)
198    f64 = ir.F64Type.get()
199    # Be careful about increasing this because
200    #     len(types) = 1 + 2^rank * rank! * len(bitwidths)^2
201    shape = range(2, 6)
202    rank = len(shape)
203    # All combinations.
204    levels = list(itertools.product(*itertools.repeat(
205      [st.DimLevelType.dense, st.DimLevelType.compressed], rank)))
206    # All permutations.
207    orderings = list(map(ir.AffineMap.get_permutation,
208      itertools.permutations(range(rank))))
209    bitwidths = [0]
210    # The first type must be a dense tensor for numpy conversion to work.
211    types = [ir.RankedTensorType.get(shape, f64)]
212    for level in levels:
213      for ordering in orderings:
214        for pwidth in bitwidths:
215          for iwidth in bitwidths:
216            attr = st.EncodingAttr.get(level, ordering, pwidth, iwidth)
217            types.append(ir.RankedTensorType.get(shape, f64, attr))
218    #
219    # For exhaustiveness we should have one or more StressTest, such
220    # that their paths cover all 2*n*(n-1) directed pairwise combinations
221    # of the `types` set.  However, since n is already superexponential,
222    # such exhaustiveness would be prohibitive for a test that runs on
223    # every commit.  So for now we'll just pick one particular path that
224    # at least hits all n elements of the `types` set.
225    #
226    tyconv = TypeConverter(ctx)
227    size = 1
228    for d in shape:
229      size *= d
230    np_arg0 = np.arange(size, dtype=tyconv.irtype_to_dtype(f64)).reshape(*shape)
231    np_out = (
232        StressTest(tyconv).build(types).writeTo(
233            sys.argv[1] if len(sys.argv) > 1 else None).compile(
234                compiler, support_lib).writeTo(
235                    sys.argv[2] if len(sys.argv) > 2 else None).run(np_arg0))
236    # CHECK: Passed
237    if np.allclose(np_out, np_arg0):
238      print('Passed')
239    else:
240      sys.exit('FAILURE')
241
242if __name__ == '__main__':
243  main()
244