1# RUN: env SUPPORT_LIB=%mlir_c_runner_utils \ 2# RUN: %PYTHON %s | FileCheck %s 3 4import ctypes 5import os 6import sys 7import tempfile 8 9from mlir import ir 10from mlir import runtime as rt 11from mlir.dialects import builtin 12from mlir.dialects import sparse_tensor as st 13import numpy as np 14 15_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__)) 16sys.path.append(_SCRIPT_PATH) 17from tools import sparsifier 18 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 f64 37 overlap = { 38 ^bb0(%arg1: f64, %arg2: f64): 39 %3 = arith.addf %arg1, %arg2 : f64 40 sparse_tensor.yield %3 : f64 41 } 42 left = { 43 ^bb0(%arg1: f64): 44 sparse_tensor.yield %arg1 : f64 45 } 46 right = { 47 ^bb0(%arg1: f64): 48 sparse_tensor.yield %arg1 : f64 49 } 50 linalg.yield %2 : f64 51 } -> tensor<3x4xf64, #Dense> 52 sparse_tensor.print %res : tensor<3x4xf64, #Dense> 53 return 54} 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 dense 65 # CHECK: ---- Sparse Tensor ---- 66 # CHECK: nse = 12 67 # 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