[mlir][sparse] Adding a stress test

[wrengr]

Summary

Addresses https://bugs.llvm.org/show_bug.cgi?id=52410
Depends on D114192

[mehdi_amini]

I am not sure that mlir/test/python/... (and mlir/test/python/dialects/...) is intended for end-to-end integration tests. This likely would better live something in the integration testing directory, but I don't know if we have precedents to use python for this kind of testing at the moment.

Also: how long does this take to run?

[aartbik]

In D114118#3138961, @mehdi_amini wrote:

I don't know if we have precedents to use python for this kind of testing at the moment.

I am afraid I set the wrong example with "test_SpMM.py" in this directory already. But I agree that we should migrate this to a better place since this does not test python functionality, but using the python dsl to test sparse functionality ;-)

Do you have suggestions for a better place? Would test/Integration/Python or something like that make sense. Or perhaps test/Integration/Dialect/SparseTensor/Python to make it more clear that we are testing sparse dialect through Python DSL.

[wrengr]

In D114118#3138961, @mehdi_amini wrote:

Also: how long does this take to run?

Running on my local machine and writing the modules out, it takes:

real    0m3.369s
user    0m2.027s
sys     0m0.536s

If not writing the modules out (as it wouldn't be for the lit+FileCheck runs), it's a bit quicker:

real    0m2.703s
user    0m2.093s
sys     0m0.301s

I intentionally set the shape/rank small enough that it runs so quickly. Whereas if I increase the rank by 1, then it takes much longer; without writing the modules out it comes to:

real    0m23.446s
user    0m22.019s
sys     0m1.041s

Hence why I use rank 4 rather than the rank 7 suggested in https://bugs.llvm.org/show_bug.cgi?id=52410

[wrengr]

Removing unused imports

[mehdi_amini]

Thanks! Timing is reasonable, any place suggested by Aart in the integration directory looks fine.

[wrengr]

rebasing and splitting off D114192

[wrengr]

Moved the file per D114192. (Apparently this never got uploaded when that differential got split off)

[wrengr]

Rebase

[aartbik]

mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py

15

I would put empty lines between the different "kinds" of imports

22

empty line before this separator

184

empty line after separator

192

We now have the

@run
def testXXX

style in one file, and plain main() in the other. Let's unify how our tests are run.

240

is this some hidden feature to write the mlir to file for debugging, if so document with sample invocation; otherwise, I would remove this debugging code from the actual test

[wrengr]

addressing comments

[wrengr]

Forgot to remove the unused run() decorator

[aartbik]

almost there!

mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py

161

we should move this also in its own file and share between all our examples, just like the TODO you added above

220

space after comma

236

one last request, the current version generates

 %0 = sparse_tensor.convert %arg0 : tensor<2x3x4x5xf64> to tensor<2x3x4x5xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense", "dense" ], dimOrdering = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, pointerBitWidth = 0, indexBitWidth = 0 }>>
  %1 = sparse_tensor.convert %0 : tensor<2x3x4x5xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense", "dense" ], dimOrdering = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, pointerBitWidth = 0, indexBitWidth = 0 }>> to tensor<2x3x4x5xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense", "dense" ], dimOrdering = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3, d2)>, pointerBitWidth = 0, indexBitWidth = 0 }>>
....
  %384 = sparse_tensor.convert %383 : tensor<2x3x4x5xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed", "compressed" ], dimOrdering = affine_map<(d0, d1, d2, d3) -> (d3, d2, d1, d0)>, pointerBitWidth = 0, indexBitWidth = 0 }>> to tensor<2x3x4x5xf64>
  return %384 : tensor<2x3x4x5xf64>

which means that all, except the last one could be freed. So just to set the right example, let's emit sparse_tensor.release on all the intermediate sparse tensors

[wrengr]

Releasing intermediate tensors.

mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py

161

Will do; though I think it ought to be in a separate CL, just to help keep the impact of each CL clearer (and hence also help track down bugs, as needed)

[aartbik]

mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py

161

Oh yes, I meant as a TODO, not in this revision.

[wrengr]

addressing comment (adding a todo). And rebasing

mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py

-# RUN: SUPPORT_LIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
+# RUN: SUPPORT_LIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext \
+# RUN:   %PYTHON %s | FileCheck %s
 
 import ctypes
 import numpy as np
 import os
 
 import mlir.all_passes_registration
 
 from mlir import ir
 from mlir import runtime as rt
 from mlir import execution_engine
 from mlir import passmanager
 
 from mlir.dialects import sparse_tensor as st
 from mlir.dialects import builtin
 from mlir.dialects.linalg.opdsl import lang as dsl
 
 
-def run(f):
-  print('\nTEST:', f.__name__)
-  f()
-  return f
-
-
 @dsl.linalg_structured_op
 def matmul_dsl(
     A=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.K),
     B=dsl.TensorDef(dsl.T, dsl.S.K, dsl.S.N),
     C=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.N, output=True)):
   C[dsl.D.m, dsl.D.n] += A[dsl.D.m, dsl.D.k] * B[dsl.D.k, dsl.D.n]
 
 
pipeline = (
         f'convert-std-to-llvm,'
         f'reconcile-unrealized-casts')
     self.pipeline = pipeline
 
   def __call__(self, module: ir.Module):
     passmanager.PassManager.parse(self.pipeline).run(module)
 
 
-# CHECK-LABEL: TEST: testSpMM
-# CHECK: Passed 8 tests
-@run
-def testSpMM():
-  # Obtain path to runtime support library.
+def main():
   support_lib = os.getenv('SUPPORT_LIB')
-  assert os.path.exists(support_lib), f'{support_lib} does not exist'
+  assert support_lib is not None, 'SUPPORT_LIB is undefined'
+  if not os.path.exists(support_lib):
+    raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), support_lib)
 
+  # CHECK-LABEL: TEST: testSpMM
+  print('\nTEST: testSpMM')
   with ir.Context() as ctx, ir.Location.unknown():
     count = 0
     # Loop over various ways to compile and annotate the SpMM kernel with
     # a *single* sparse tensor. Note that we deliberate do not exhaustively
     # search the full state space to reduce runtime of the test. It is
     # straightforward to adapt the code below to explore more combinations.
     par = 0
     vec = 0
with ir.Context() as ctx, ir.Location.unknown():
     for level in levels:
       for ordering in orderings:
         for pwidth in bitwidths:
           for iwidth in bitwidths:
             attr = st.EncodingAttr.get(level, ordering, pwidth, iwidth)
             compiler = SparseCompiler(options=opt)
             build_compile_and_run_SpMM(attr, support_lib, compiler)
             count = count + 1
+    # CHECK: Passed 8 tests
     print('Passed ', count, 'tests')
+
+if __name__ == '__main__':
+  main()

mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py

+# RUN: SUPPORT_LIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext \
+# RUN:   %PYTHON %s | FileCheck %s
+
+import ctypes
+import errno
+import itertools
+import os
+import sys
+from typing import List, Callable
+
+import numpy as np
+
+import mlir.all_passes_registration
+
+from mlir import ir
+from mlir import runtime as rt
+from mlir.execution_engine import ExecutionEngine
+from mlir.passmanager import PassManager
+
+from mlir.dialects import builtin
+from mlir.dialects import std
+from mlir.dialects import sparse_tensor as st
+
+# ===----------------------------------------------------------------------=== #
+
+# TODO: move this boilerplate to its own module, so it can be used by
+# other tests and programs.
+class TypeConverter:
+  """Converter between NumPy types and MLIR types."""
+
+  def __init__(self, context: ir.Context):
+    # Note 1: these are numpy "scalar types" (i.e., the values of
+    # np.sctypeDict) not numpy "dtypes" (i.e., the np.dtype class).
+    #
+    # Note 2: we must construct the MLIR types in the same context as the
+    # types that'll be passed to irtype_to_sctype() or irtype_to_dtype();
+    # otherwise, those methods will raise a KeyError.
+    types_list = [
+      (np.float64, ir.F64Type.get(context=context)),
+      (np.float32, ir.F32Type.get(context=context)),
+      (np.int64, ir.IntegerType.get_signless(64, context=context)),
+      (np.int32, ir.IntegerType.get_signless(32, context=context)),
+      (np.int16, ir.IntegerType.get_signless(16, context=context)),
+      (np.int8, ir.IntegerType.get_signless(8, context=context)),
+    ]
+    self._sc2ir = dict(types_list)
+    self._ir2sc = dict(( (ir,sc) for sc,ir in types_list ))
+
+  def dtype_to_irtype(self, dtype: np.dtype) -> ir.Type:
+    """Returns the MLIR equivalent of a NumPy dtype."""
+    try:
+      return self.sctype_to_irtype(dtype.type)
+    except KeyError as e:
+      raise KeyError(f'Unknown dtype: {dtype}') from e
+
+  def sctype_to_irtype(self, sctype) -> ir.Type:
+    """Returns the MLIR equivalent of a NumPy scalar type."""
+    if sctype in self._sc2ir:
+      return self._sc2ir[sctype]
+    else:
+      raise KeyError(f'Unknown sctype: {sctype}')
+
+  def irtype_to_dtype(self, tp: ir.Type) -> np.dtype:
+    """Returns the NumPy dtype equivalent of an MLIR type."""
+    return np.dtype(self.irtype_to_sctype(tp))
+
+  def irtype_to_sctype(self, tp: ir.Type):
+    """Returns the NumPy scalar-type equivalent of an MLIR type."""
+    if tp in self._ir2sc:
+      return self._ir2sc[tp]
+    else:
+      raise KeyError(f'Unknown ir.Type: {tp}')
+
+  def get_RankedTensorType_of_nparray(self, nparray: np.ndarray) -> ir.RankedTensorType:
+    """Returns the ir.RankedTensorType of a NumPy array.  Note that NumPy
+    arrays can only be converted to/from dense tensors, not sparse tensors."""
+    # TODO: handle strides as well?
+    return ir.RankedTensorType.get(nparray.shape,
+                                   self.dtype_to_irtype(nparray.dtype))
+
+# ===----------------------------------------------------------------------=== #
+
+class StressTest:
+  def __init__(self, tyconv: TypeConverter):
+    self._tyconv = tyconv
+    self._roundtripTp = None
+    self._module = None
+    self._engine = None
+
+  def _assertEqualsRoundtripTp(self, tp: ir.RankedTensorType):
+    assert self._roundtripTp is not None, \
+        'StressTest: uninitialized roundtrip type'
+    if tp != self._roundtripTp:
+      raise AssertionError(
+          f"Type is not equal to the roundtrip type.\n"
+          f"\tExpected: {self._roundtripTp}\n"
+          f"\tFound:    {tp}\n")
+
+  def build(self, types: List[ir.Type]):
+    """Builds the ir.Module.  The module has only the @main function,
+    which will convert the input through the list of types and then back
+    to the initial type.  The roundtrip type must be a dense tensor."""
+    assert self._module is None, 'StressTest: must not call build() repeatedly'
+    self._module = ir.Module.create()
+    with ir.InsertionPoint(self._module.body):
+      tp0 = types.pop(0)
+      self._roundtripTp = tp0
+      # TODO: assert dense? assert element type is recognised by the TypeConverter?
+      types.append(tp0)
+      funcTp = ir.FunctionType.get(inputs=[tp0], results=[tp0])
+      funcOp = builtin.FuncOp(name='main', type=funcTp)
+      funcOp.attributes['llvm.emit_c_interface'] = ir.UnitAttr.get()
+      with ir.InsertionPoint(funcOp.add_entry_block()):
+        arg0 = funcOp.entry_block.arguments[0]
+        self._assertEqualsRoundtripTp(arg0.type)
+        v = st.ConvertOp(types.pop(0), arg0)
+        for tp in types:
+          w = st.ConvertOp(tp, v)
+          # Release intermediate tensors before they fall out of scope.
+          st.ReleaseOp(v.result)
+          v = w
+        self._assertEqualsRoundtripTp(v.result.type)
+        std.ReturnOp(v)
+    return self
+
+  def writeTo(self, filename):
+    """Write the ir.Module to the given file.  If the file already exists,
+    then raises an error.  If the filename is None, then is a no-op."""
+    assert self._module is not None, \
+        'StressTest: must call build() before writeTo()'
+    if filename is None:
+      # Silent no-op, for convenience.
+      return self
+    if os.path.exists(filename):
+      raise FileExistsError(errno.EEXIST, os.strerror(errno.EEXIST), filename)
+    with open(filename, 'w') as f:
+      f.write(str(self._module))
+    return self
+
+  def compile(self, compiler: Callable[[ir.Module], ExecutionEngine]):
+    """Compile the ir.Module."""
+    assert self._module is not None, \
+        'StressTest: must call build() before compile()'
+    assert self._engine is None, \
+        'StressTest: must not call compile() repeatedly'
+    self._engine = compiler(self._module)
+    return self
+
+  def run(self, np_arg0: np.ndarray) -> np.ndarray:
+    """Runs the test on the given numpy array, and returns the resulting
+    numpy array."""
+    assert self._engine is not None, \
+        'StressTest: must call compile() before run()'
+    self._assertEqualsRoundtripTp(
+        self._tyconv.get_RankedTensorType_of_nparray(np_arg0))
+    np_out = np.zeros(np_arg0.shape, dtype=np_arg0.dtype)
+    self._assertEqualsRoundtripTp(
+        self._tyconv.get_RankedTensorType_of_nparray(np_out))
+    mem_arg0 = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(np_arg0)))
+    mem_out = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(np_out)))
+    self._engine.invoke('main', mem_out, mem_arg0)
+    return rt.ranked_memref_to_numpy(mem_out[0])
+
+# ===----------------------------------------------------------------------=== #
+
+# TODO: move this boilerplate to its own module, so it can be used by
+# other tests and programs.
+class SparseCompiler:
+  """Sparse compiler passes."""
+
+  def __init__(self, sparsification_options: str, support_lib: str):
+    self._support_lib = support_lib
+    self._pipeline = (
+        f'builtin.func(linalg-generalize-named-ops,linalg-fuse-elementwise-ops),'
+        f'sparsification{{{sparsification_options}}},'
+        f'sparse-tensor-conversion,'
+        f'builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),'
+        f'convert-scf-to-std,'
+        f'func-bufferize,'
+        f'tensor-constant-bufferize,'
+        f'builtin.func(tensor-bufferize,std-bufferize,finalizing-bufferize),'
+        f'convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},'
+        f'lower-affine,'
+        f'convert-memref-to-llvm,'
+        f'convert-std-to-llvm,'
+        f'reconcile-unrealized-casts')
+    # Must be in the scope of a `with ir.Context():`
+    self._passmanager = PassManager.parse(self._pipeline)
+
+  def __call__(self, module: ir.Module) -> ExecutionEngine:
+    self._passmanager.run(module)
+    return ExecutionEngine(module, opt_level=0, shared_libs=[self._support_lib])
+
+# ===----------------------------------------------------------------------=== #
+
+def main():
+  """
+  USAGE: python3 test_stress.py [raw_module.mlir [compiled_module.mlir]]
+
+  The environment variable SUPPORT_LIB must be set to point to the
+  libmlir_c_runner_utils shared library.  There are two optional
+  arguments, for debugging purposes.  The first argument specifies where
+  to write out the raw/generated ir.Module.  The second argument specifies
+  where to write out the compiled version of that ir.Module.
+  """
+  support_lib = os.getenv('SUPPORT_LIB')
+  assert support_lib is not None, 'SUPPORT_LIB is undefined'
+  if not os.path.exists(support_lib):
+    raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), support_lib)
+
+  # CHECK-LABEL: TEST: test_stress
+  print("\nTEST: test_stress")
+  with ir.Context() as ctx, ir.Location.unknown():
+    par = 0
+    vec = 0
+    vl = 1
+    e = False
+    sparsification_options = (
+        f'parallelization-strategy={par} '
+        f'vectorization-strategy={vec} '
+        f'vl={vl} '
+        f'enable-simd-index32={e}')
+    compiler = SparseCompiler(sparsification_options, support_lib)
+    f64 = ir.F64Type.get()
+    # Be careful about increasing this because
+    #     len(types) = 1 + 2^rank * rank! * len(bitwidths)^2
+    shape = range(2, 6)
+    rank = len(shape)
+    # All combinations.
+    levels = list(itertools.product(*itertools.repeat(
+      [st.DimLevelType.dense, st.DimLevelType.compressed], rank)))
+    # All permutations.
+    orderings = list(map(ir.AffineMap.get_permutation,
+      itertools.permutations(range(rank))))
+    bitwidths = [0]
+    # The first type must be a dense tensor for numpy conversion to work.
+    types = [ir.RankedTensorType.get(shape, f64)]
+    for level in levels:
+      for ordering in orderings:
+        for pwidth in bitwidths:
+          for iwidth in bitwidths:
+            attr = st.EncodingAttr.get(level, ordering, pwidth, iwidth)
+            types.append(ir.RankedTensorType.get(shape, f64, attr))
+    #
+    # For exhaustiveness we should have one or more StressTest, such
+    # that their paths cover all 2*n*(n-1) directed pairwise combinations
+    # of the `types` set.  However, since n is already superexponential,
+    # such exhaustiveness would be prohibitive for a test that runs on
+    # every commit.  So for now we'll just pick one particular path that
+    # at least hits all n elements of the `types` set.
+    #
+    tyconv = TypeConverter(ctx)
+    size = 1
+    for d in shape:
+      size *= d
+    np_arg0 = np.arange(size, dtype=tyconv.irtype_to_dtype(f64)).reshape(*shape)
+    np_out = (
+        StressTest(tyconv)
+        .build(types)
+        .writeTo(sys.argv[1] if len(sys.argv) > 1 else None)
+        .compile(compiler)
+        .writeTo(sys.argv[2] if len(sys.argv) > 2 else None)
+        .run(np_arg0))
+    # CHECK: Passed
+    if np.allclose(np_out, np_arg0):
+      print('Passed')
+    else:
+      sys.exit('FAILURE')
+
+if __name__ == '__main__':
+  main()