[Mlir-commits] [mlir] 4748cc6 - [mlir][sparse] Adding a stress test
wren romano
llvmlistbot at llvm.org
Fri Dec 3 14:59:46 PST 2021
Author: wren romano
Date: 2021-12-03T14:59:39-08:00
New Revision: 4748cc69314ad1ffd85fd6f0265d64fbbeba4430
URL: https://github.com/llvm/llvm-project/commit/4748cc69314ad1ffd85fd6f0265d64fbbeba4430
DIFF: https://github.com/llvm/llvm-project/commit/4748cc69314ad1ffd85fd6f0265d64fbbeba4430.diff
LOG: [mlir][sparse] Adding a stress test
Addresses https://bugs.llvm.org/show_bug.cgi?id=52410
Depends on D114192
Reviewed By: aartbik, mehdi_amini
Differential Revision: https://reviews.llvm.org/D114118
Added:
mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py
Modified:
mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py
Removed:
################################################################################
diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py
index 38ca333c788b1..b5887263cc1cb 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py
@@ -1,4 +1,5 @@
-# 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
@@ -16,12 +17,6 @@
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),
@@ -135,14 +130,14 @@ def __call__(self, module: ir.Module):
passmanager.PassManager.parse(self.pipeline).run(module)
-# CHECK-LABEL: TEST: testSpMM
-# CHECK: Passed 8 tests
- at 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
@@ -173,4 +168,8 @@ def testSpMM():
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()
diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py
new file mode 100644
index 0000000000000..55e64668d4955
--- /dev/null
+++ b/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py
@@ -0,0 +1,271 @@
+# 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()
More information about the Mlir-commits
mailing list