95 lines · python
1# RUN: env SUPPORT_LIB=%mlir_c_runner_utils \2# RUN: %PYTHON %s | FileCheck %s3 4import ctypes5import os6import sys7import tempfile8 9from mlir import ir10from mlir import runtime as rt11from mlir.dialects import builtin12from mlir.dialects import sparse_tensor as st13import numpy as np14 15_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))16sys.path.append(_SCRIPT_PATH)17from tools import sparsifier18 19 20def boilerplate():21 """Returns boilerplate main method."""22 return """23#Dense = #sparse_tensor.encoding<{24 map = (i, j) -> (i: dense, j: dense)25}>26 27#map = affine_map<(d0, d1) -> (d0, d1)>28func.func @add(%st_0 : tensor<3x4xf64, #Dense>,29 %st_1 : tensor<3x4xf64, #Dense>) attributes { llvm.emit_c_interface } {30 %out_st = tensor.empty() : tensor<3x4xf64, #Dense>31 %res = linalg.generic {indexing_maps = [#map, #map, #map],32 iterator_types = ["parallel", "parallel"]}33 ins(%st_0, %st_1 : tensor<3x4xf64, #Dense>, tensor<3x4xf64, #Dense>)34 outs(%out_st : tensor<3x4xf64, #Dense>) {35 ^bb0(%in_0: f64, %in_1: f64, %out: f64):36 %2 = sparse_tensor.binary %in_0, %in_1 : f64, f64 to f6437 overlap = {38 ^bb0(%arg1: f64, %arg2: f64):39 %3 = arith.addf %arg1, %arg2 : f6440 sparse_tensor.yield %3 : f6441 }42 left = {43 ^bb0(%arg1: f64):44 sparse_tensor.yield %arg1 : f6445 }46 right = {47 ^bb0(%arg1: f64):48 sparse_tensor.yield %arg1 : f6449 }50 linalg.yield %2 : f6451 } -> tensor<3x4xf64, #Dense>52 sparse_tensor.print %res : tensor<3x4xf64, #Dense>53 return54}55"""56 57 58def main():59 support_lib = os.getenv("SUPPORT_LIB")60 assert support_lib is not None, "SUPPORT_LIB is undefined"61 if not os.path.exists(support_lib):62 raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), support_lib)63 64 # CHECK-LABEL: TEST: all dense65 # CHECK: ---- Sparse Tensor ----66 # CHECK: nse = 1267 # CHECK: dim = ( 3, 4 )68 # CHECK: lvl = ( 3, 4 )69 # CHECK: values : ( 1, 1, 0, 1, 0, 6, 2, 3, 0, 0, 0, 2 )70 # CHECK: ----71 print("\nTEST: all dense")72 with ir.Context() as ctx, ir.Location.unknown():73 compiler = sparsifier.Sparsifier(74 extras="sparse-assembler,",75 options="enable-runtime-library=false",76 opt_level=2,77 shared_libs=[support_lib],78 )79 module = ir.Module.parse(boilerplate())80 engine = compiler.compile_and_jit(module)81 print(module)82 83 a = np.array([1, 0, 0, 1, 0, 2, 2, 0, 0, 0, 0, 1], dtype=np.float64)84 b = np.array([0, 1, 0, 0, 0, 4, 0, 3, 0, 0, 0, 1], dtype=np.float64)85 mem_a = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(a)))86 mem_b = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(b)))87 88 # Invoke the kernel and get numpy output.89 # Built-in bufferization uses in-out buffers.90 engine.invoke("add", mem_a, mem_b)91 92 93if __name__ == "__main__":94 main()95