From 312c51406da68aa512641d444cb7e369b2c4f1cf Mon Sep 17 00:00:00 2001 From: Aart Bik Date: Thu, 9 Dec 2021 15:47:21 -0800 Subject: [PATCH] [mlir][sparse] python driven test for SDDMM explores various sparsity combinations of the SDMM kernel and verifies that the computed result is the same for all cases Reviewed By: bixia Differential Revision: https://reviews.llvm.org/D115476 --- .../Dialect/SparseTensor/python/test_SDDMM.py | 188 +++++++++++++++++++++ .../Dialect/SparseTensor/python/test_SpMM.py | 1 + 2 files changed, 189 insertions(+) create mode 100644 mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py new file mode 100644 index 0000000..feeedcc --- /dev/null +++ b/mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py @@ -0,0 +1,188 @@ +# 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 + + +@dsl.linalg_structured_op +def sddmm_dsl( + A=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.K), + B=dsl.TensorDef(dsl.T, dsl.S.K, dsl.S.N), + S=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.N), + C=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.N, output=True)): + C[dsl.D.m, + dsl.D.n] += S[dsl.D.m, dsl.D.n] * A[dsl.D.m, dsl.D.k] * B[dsl.D.k, dsl.D.n] + + +def build_SDDMM(attr: st.EncodingAttr): + """Build SDDMM kernel. + + This method generates a linalg op with for matrix multiplication using + just the Python API. Effectively, a generic linalg op is constructed + that computes C(i,j) += S(i,j) SUM_k A(i,k) B(k,j) for sparse S. + """ + module = ir.Module.create() + f64 = ir.F64Type.get() + a = ir.RankedTensorType.get([8, 8], f64) + b = ir.RankedTensorType.get([8, 8], f64) + c = ir.RankedTensorType.get([8, 8], f64) + s = ir.RankedTensorType.get([8, 8], f64, attr) + arguments = [a, b, s, c] + with ir.InsertionPoint(module.body): + + @builtin.FuncOp.from_py_func(*arguments) + def sddmm(*args): + return sddmm_dsl(args[0], args[1], args[2], outs=[args[3]]) + + return module + + +def boilerplate(attr: st.EncodingAttr): + """Returns boilerplate code for main driver.""" + return f""" +func @main(%a: tensor<8x8xf64>, + %b: tensor<8x8xf64>, + %c: tensor<8x8xf64>) -> tensor<8x8xf64> attributes {{ llvm.emit_c_interface }} {{ + %t = arith.constant sparse<[[0,0], [0,2], [4,1]], [1.0, 2.0, 3.0]> : tensor<8x8xf64> + %s = sparse_tensor.convert %t : tensor<8x8xf64> to tensor<8x8xf64, {attr}> + %0 = call @sddmm(%a, %b, %s, %c) : (tensor<8x8xf64>, + tensor<8x8xf64>, + tensor<8x8xf64, {attr}>, + tensor<8x8xf64>) -> tensor<8x8xf64> + return %0 : tensor<8x8xf64> +}} +""" + + +def build_compile_and_run_SDDMMM(attr: st.EncodingAttr, opt: str, + support_lib: str, compiler): + # Build. + module = build_SDDMM(attr) + func = str(module.operation.regions[0].blocks[0].operations[0].operation) + module = ir.Module.parse(func + boilerplate(attr)) + + # Compile. + compiler(module) + engine = execution_engine.ExecutionEngine( + module, opt_level=0, shared_libs=[support_lib]) + + # Set up numpy input and buffer for output. + a = np.array([[1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1, 8.1], + [1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2, 8.2], + [1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3, 8.3], + [1.4, 2.4, 3.4, 4.4, 5.4, 6.4, 7.4, 8.4], + [1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5], + [1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6], + [1.7, 2.7, 3.7, 4.7, 5.7, 6.7, 7.7, 8.7], + [1.8, 2.8, 3.8, 4.8, 5.8, 6.8, 7.8, 8.8]], np.float64) + b = np.ones((8, 8), np.float64) + c = np.zeros((8, 8), 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))) + mem_c = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(c))) + + # Allocate a MemRefDescriptor to receive the output tensor. + # The buffer itself is allocated inside the MLIR code generation. + ref_out = rt.make_nd_memref_descriptor(2, ctypes.c_double)() + mem_out = ctypes.pointer(ctypes.pointer(ref_out)) + + # Invoke the kernel and get numpy output. + # Built-in bufferization uses in-out buffers. + # TODO: replace with inplace comprehensive bufferization. + engine.invoke('main', mem_out, mem_a, mem_b, mem_c) + + # Sanity check on computed result. Only a few elements + # are sampled from the full dense matrix multiplication. + full_matmul = np.matmul(a, b) + expected = np.zeros((8, 8), np.float64) + expected[0, 0] = 1.0 * full_matmul[0, 0] + expected[0, 2] = 2.0 * full_matmul[0, 2] + expected[4, 1] = 3.0 * full_matmul[4, 1] + c = rt.ranked_memref_to_numpy(mem_out[0]) + if np.allclose(c, expected): + pass + else: + quit(f'FAILURE') + + +class SparseCompiler: + """Sparse compiler passes.""" + + def __init__(self, options: str): + pipeline = ( + f'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') + self.pipeline = pipeline + + def __call__(self, module: ir.Module): + passmanager.PassManager.parse(self.pipeline).run(module) + + +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 = SparseCompiler(options=opt) + build_compile_and_run_SDDMMM(attr, opt, support_lib, compiler) + count = count + 1 + # CHECK: Passed 16 tests + print('Passed ', count, 'tests') + + +if __name__ == '__main__': + main() diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py index b6a1eaa..76ff846 100644 --- a/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py +++ b/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py @@ -173,5 +173,6 @@ def main(): # CHECK: Passed 8 tests print('Passed ', count, 'tests') + if __name__ == '__main__': main() -- 2.7.4