161 lines · python
1# RUN: env SUPPORT_LIB=%mlir_c_runner_utils \2# RUN: %PYTHON %s | FileCheck %s3 4import ctypes5import numpy as np6import os7import sys8 9from mlir import ir10from mlir import runtime as rt11 12from mlir.dialects import sparse_tensor as st13from mlir.dialects import builtin14from mlir.dialects import func15from mlir.dialects.linalg.opdsl import lang as dsl16 17_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))18sys.path.append(_SCRIPT_PATH)19from tools import sparsifier20 21 22@dsl.linalg_structured_op23def matmul_dsl(24 A=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.K),25 B=dsl.TensorDef(dsl.T, dsl.S.K, dsl.S.N),26 C=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.N, output=True),27):28 C[dsl.D.m, dsl.D.n] += A[dsl.D.m, dsl.D.k] * B[dsl.D.k, dsl.D.n]29 30 31def build_SpMM(attr: st.EncodingAttr):32 """Build SpMM kernel.33 34 This method generates a linalg op with for matrix multiplication using35 just the Python API. Effectively, a generic linalg op is constructed36 that computes C(i,j) += A(i,k) * B(k,j) for annotated matrix A.37 """38 module = ir.Module.create()39 f64 = ir.F64Type.get()40 a = ir.RankedTensorType.get([3, 4], f64, attr)41 b = ir.RankedTensorType.get([4, 2], f64)42 c = ir.RankedTensorType.get([3, 2], f64)43 arguments = [a, b, c]44 with ir.InsertionPoint(module.body):45 46 @func.FuncOp.from_py_func(*arguments)47 def spMxM(*args):48 return matmul_dsl(args[0], args[1], outs=[args[2]])49 50 return module51 52 53def boilerplate(attr: st.EncodingAttr):54 """Returns boilerplate main method.55 56 This method sets up a boilerplate main method that takes three tensors57 (a, b, c), converts the first tensor a into s sparse tensor, and then58 calls the sparse kernel for matrix multiplication. For convenience,59 this part is purely done as string input.60 """61 return f"""62func.func @main(%ad: tensor<3x4xf64>, %b: tensor<4x2xf64>, %c: tensor<3x2xf64>) -> tensor<3x2xf64>63 attributes {{ llvm.emit_c_interface }} {{64 %a = sparse_tensor.convert %ad : tensor<3x4xf64> to tensor<3x4xf64, {attr}>65 %0 = call @spMxM(%a, %b, %c) : (tensor<3x4xf64, {attr}>,66 tensor<4x2xf64>,67 tensor<3x2xf64>) -> tensor<3x2xf64>68 return %0 : tensor<3x2xf64>69}}70"""71 72 73def build_compile_and_run_SpMM(attr: st.EncodingAttr, compiler):74 # Build.75 module = build_SpMM(attr)76 func = str(module.operation.regions[0].blocks[0].operations[0].operation)77 module = ir.Module.parse(func + boilerplate(attr))78 79 # Compile.80 engine = compiler.compile_and_jit(module)81 82 # Set up numpy input and buffer for output.83 a = np.array(84 [[1.1, 0.0, 0.0, 1.4], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 3.3, 0.0]], np.float6485 )86 b = np.array([[1.0, 2.0], [4.0, 3.0], [5.0, 6.0], [8.0, 7.0]], np.float64)87 c = np.zeros((3, 2), np.float64)88 89 mem_a = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(a)))90 mem_b = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(b)))91 mem_c = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(c)))92 # Allocate a MemRefDescriptor to receive the output tensor.93 # The buffer itself is allocated inside the MLIR code generation.94 ref_out = rt.make_nd_memref_descriptor(2, ctypes.c_double)()95 mem_out = ctypes.pointer(ctypes.pointer(ref_out))96 97 # Invoke the kernel and get numpy output.98 # Built-in bufferization uses in-out buffers.99 engine.invoke("main", mem_out, mem_a, mem_b, mem_c)100 101 # Sanity check on computed result.102 expected = np.matmul(a, b)103 c = rt.ranked_memref_to_numpy(mem_out[0])104 if np.allclose(c, expected):105 pass106 else:107 quit(f"FAILURE")108 109 110def main():111 support_lib = os.getenv("SUPPORT_LIB")112 assert support_lib is not None, "SUPPORT_LIB is undefined"113 if not os.path.exists(support_lib):114 raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), support_lib)115 116 # CHECK-LABEL: TEST: testSpMM117 print("\nTEST: testSpMM")118 count = 0119 with ir.Context() as ctx, ir.Location.unknown():120 # Loop over various ways to compile and annotate the SpMM kernel with121 # a *single* sparse tensor. Note that we deliberate do not exhaustively122 # search the full state space to reduce runtime of the test. It is123 # straightforward to adapt the code below to explore more combinations.124 # For these simple orderings, dim2lvl and lvl2dim are the same.125 vl = 1126 e = False127 opt = f"parallelization-strategy=none"128 builder = st.EncodingAttr.build_level_type129 fmt = st.LevelFormat130 prop = st.LevelProperty131 levels = [132 [builder(fmt.compressed, [prop.non_unique]), builder(fmt.singleton)],133 [builder(fmt.dense), builder(fmt.dense)],134 [builder(fmt.dense), builder(fmt.compressed)],135 [builder(fmt.compressed), builder(fmt.dense)],136 [builder(fmt.compressed), builder(fmt.compressed)],137 ]138 orderings = [139 ir.AffineMap.get_permutation([0, 1]),140 ir.AffineMap.get_permutation([1, 0]),141 ]142 bitwidths = [0]143 compiler = sparsifier.Sparsifier(144 extras="", options=opt, opt_level=0, shared_libs=[support_lib]145 )146 for level in levels:147 for ordering in orderings:148 for pwidth in bitwidths:149 for iwidth in bitwidths:150 attr = st.EncodingAttr.get(151 level, ordering, ordering, pwidth, iwidth152 )153 build_compile_and_run_SpMM(attr, compiler)154 count = count + 1155 # CHECK: Passed 10 tests156 print("Passed ", count, "tests")157 158 159if __name__ == "__main__":160 main()161