def main():
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_output
print('\nTEST: test_output')
count = 0
with ir.Context() as ctx, ir.Location.unknown():
# Loop over various sparse types: CSR, DCSR, CSC, DCSC.
levels = [[st.DimLevelType.dense, st.DimLevelType.compressed],
[st.DimLevelType.compressed, st.DimLevelType.compressed]]
orderings = [
ir.AffineMap.get_permutation([0, 1]),
ir.AffineMap.get_permutation([1, 0])
]
bitwidths = [8, 16, 32, 64]
for level in levels:
for ordering in orderings:
for bwidth in bitwidths:
attr = st.EncodingAttr.get(level, ordering, bwidth, bwidth)
compiler = sparse_compiler.SparseCompiler(options='')
build_compile_and_run_output(attr, support_lib, compiler)
count = count + 1
# CHECK: Passed 16 tests
print('Passed', count, 'tests')
def _run_test(support_lib, kernel):
"""Compiles, runs and checks results."""
compiler = sparse_compiler.SparseCompiler(
options='', opt_level=2, shared_libs=[support_lib])
module = ir.Module.parse(kernel)
engine = compiler.compile_and_jit(module)
# Set up numpy inputs and buffer for output.
a = np.array(
[[1.1, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 6.6, 0.0]],
np.float64)
b = np.array(
[[1.1, 0.0, 0.0, 2.8], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]],
np.float64)
mem_a = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(a)))
mem_b = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(b)))
# The sparse tensor output is a pointer to pointer of char.
out = ctypes.c_char(0)
mem_out = ctypes.pointer(ctypes.pointer(out))
# Invoke the kernel.
engine.invoke('main', mem_a, mem_b, mem_out)
# Retrieve and check the result.
rank, nse, shape, values, indices = test_tools.sparse_tensor_to_coo_tensor(
support_lib, mem_out[0], np.float64)
# CHECK: PASSED
if np.allclose(values, [2.2, 2.8, 6.6]) and np.allclose(
indices, [[0, 0], [0, 3], [2, 2]]):
print('PASSED')
else:
quit('FAILURE')
def main():
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: testSDDMMM
print('\nTEST: testSDDMMM')
with ir.Context() as ctx, ir.Location.unknown():
count = 0
# Loop over various ways to compile and annotate the SDDMM 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.
levels = [[st.DimLevelType.dense, st.DimLevelType.dense],
[st.DimLevelType.dense, st.DimLevelType.compressed],
[st.DimLevelType.compressed, st.DimLevelType.dense],
[st.DimLevelType.compressed, st.DimLevelType.compressed]]
orderings = [
ir.AffineMap.get_permutation([0, 1]),
ir.AffineMap.get_permutation([1, 0])
]
for level in levels:
for ordering in orderings:
for pwidth in [32]:
for iwidth in [32]:
for par in [0]:
for vec in [0, 1]:
for e in [True]:
vl = 1 if vec == 0 else 16
attr = st.EncodingAttr.get(
level, ordering, pwidth, iwidth)
opt = (f'parallelization-strategy={par} '
f'vectorization-strategy={vec} '
f'vl={vl} enable-simd-index32={e}')
compiler = sparse_compiler.SparseCompiler(
options=opt,
opt_level=0,
shared_libs=[support_lib])
build_compile_and_run_SDDMMM(
attr, compiler)
count = count + 1
# CHECK: Passed 16 tests
print('Passed ', count, 'tests')
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
# Disable direct sparse2sparse conversion, because it doubles the time!
# TODO: While direct s2s is far too slow for per-commit testing,
# we should have some framework ensure that we run this test with
# `s2s=0` on a regular basis, to ensure that it does continue to work.
s2s = 1
sparsification_options = (f'parallelization-strategy={par} '
f'vectorization-strategy={vec} '
f'vl={vl} '
f'enable-simd-index32={e} '
f's2s-strategy={s2s}')
compiler = sparse_compiler.SparseCompiler(
options=sparsification_options,
opt_level=0,
shared_libs=[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')