76647fce13
This is a major step along the way towards the new STEA design. While a great deal of this patch is simple renaming, there are several significant changes as well. I've done my best to ensure that this patch retains the previous behavior and error-conditions, even though those are at odds with the eventual intended semantics of the `dimToLvl` mapping. Since the majority of the compiler does not yet support non-permutations, I've also added explicit assertions in places that previously had implicitly assumed it was dealing with permutations. Reviewed By: aartbik Differential Revision: https://reviews.llvm.org/D151505
104 lines
3.1 KiB
Python
104 lines
3.1 KiB
Python
# RUN: env SUPPORT_LIB=%mlir_c_runner_utils \
|
|
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
import ctypes
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
|
|
from mlir import ir
|
|
from mlir import runtime as rt
|
|
|
|
from mlir.dialects import builtin
|
|
from mlir.dialects import sparse_tensor as st
|
|
|
|
_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
|
|
sys.path.append(_SCRIPT_PATH)
|
|
from tools import sparse_compiler
|
|
|
|
# TODO: move more into actual IR building.
|
|
def boilerplate(attr: st.EncodingAttr):
|
|
"""Returns boilerplate main method."""
|
|
return f"""
|
|
func.func @main(%p : !llvm.ptr<i8>) -> () attributes {{ llvm.emit_c_interface }} {{
|
|
%d = arith.constant sparse<[[0, 0], [1, 1], [0, 9], [9, 0], [4, 4]],
|
|
[1.0, 2.0, 3.0, 4.0, 5.0]> : tensor<10x10xf64>
|
|
%a = sparse_tensor.convert %d : tensor<10x10xf64> to tensor<10x10xf64, {attr}>
|
|
sparse_tensor.out %a, %p : tensor<10x10xf64, {attr}>, !llvm.ptr<i8>
|
|
return
|
|
}}
|
|
"""
|
|
|
|
|
|
def expected():
|
|
"""Returns expected contents of output.
|
|
|
|
Regardless of the dimension ordering, compression, and bitwidths that are
|
|
used in the sparse tensor, the output is always lexicographically sorted
|
|
by natural index order.
|
|
"""
|
|
return f"""; extended FROSTT format
|
|
2 5
|
|
10 10
|
|
1 1 1
|
|
1 10 3
|
|
2 2 2
|
|
5 5 5
|
|
10 1 4
|
|
"""
|
|
|
|
|
|
def build_compile_and_run_output(attr: st.EncodingAttr, compiler):
|
|
# Build and Compile.
|
|
module = ir.Module.parse(boilerplate(attr))
|
|
engine = compiler.compile_and_jit(module)
|
|
|
|
# Invoke the kernel and compare output.
|
|
with tempfile.TemporaryDirectory() as test_dir:
|
|
out = os.path.join(test_dir, "out.tns")
|
|
buf = out.encode("utf-8")
|
|
mem_a = ctypes.pointer(ctypes.pointer(ctypes.create_string_buffer(buf)))
|
|
engine.invoke("main", mem_a)
|
|
|
|
actual = open(out).read()
|
|
if actual != expected():
|
|
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: 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]
|
|
compiler = sparse_compiler.SparseCompiler(
|
|
options="", opt_level=2, shared_libs=[support_lib]
|
|
)
|
|
for level in levels:
|
|
for ordering in orderings:
|
|
for bwidth in bitwidths:
|
|
attr = st.EncodingAttr.get(level, ordering, bwidth, bwidth)
|
|
build_compile_and_run_output(attr, compiler)
|
|
count = count + 1
|
|
|
|
# CHECK: Passed 16 tests
|
|
print("Passed", count, "tests")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|