259 lines · python
1# RUN: env SUPPORT_LIB=%mlir_c_runner_utils \2# RUN: %PYTHON %s | FileCheck %s3 4import ctypes5import errno6import itertools7import os8import sys9 10from typing import List, Callable11 12import numpy as np13 14from mlir import ir15from mlir import runtime as rt16 17from mlir.dialects import bufferization18from mlir.dialects import builtin19from mlir.dialects import func20from mlir.dialects import sparse_tensor as st21 22_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))23sys.path.append(_SCRIPT_PATH)24from tools import sparsifier25 26# ===----------------------------------------------------------------------=== #27 28 29class TypeConverter:30 """Converter between NumPy types and MLIR types."""31 32 def __init__(self, context: ir.Context):33 # Note 1: these are numpy "scalar types" (i.e., the values of34 # np.sctypeDict) not numpy "dtypes" (i.e., the np.dtype class).35 #36 # Note 2: we must construct the MLIR types in the same context as the37 # types that'll be passed to irtype_to_sctype() or irtype_to_dtype();38 # otherwise, those methods will raise a KeyError.39 types_list = [40 (np.float64, ir.F64Type.get(context=context)),41 (np.float32, ir.F32Type.get(context=context)),42 (np.int64, ir.IntegerType.get_signless(64, context=context)),43 (np.int32, ir.IntegerType.get_signless(32, context=context)),44 (np.int16, ir.IntegerType.get_signless(16, context=context)),45 (np.int8, ir.IntegerType.get_signless(8, context=context)),46 ]47 self._sc2ir = dict(types_list)48 self._ir2sc = dict(((ir, sc) for sc, ir in types_list))49 50 def dtype_to_irtype(self, dtype: np.dtype) -> ir.Type:51 """Returns the MLIR equivalent of a NumPy dtype."""52 try:53 return self.sctype_to_irtype(dtype.type)54 except KeyError as e:55 raise KeyError(f"Unknown dtype: {dtype}") from e56 57 def sctype_to_irtype(self, sctype) -> ir.Type:58 """Returns the MLIR equivalent of a NumPy scalar type."""59 if sctype in self._sc2ir:60 return self._sc2ir[sctype]61 else:62 raise KeyError(f"Unknown sctype: {sctype}")63 64 def irtype_to_dtype(self, tp: ir.Type) -> np.dtype:65 """Returns the NumPy dtype equivalent of an MLIR type."""66 return np.dtype(self.irtype_to_sctype(tp))67 68 def irtype_to_sctype(self, tp: ir.Type):69 """Returns the NumPy scalar-type equivalent of an MLIR type."""70 if tp in self._ir2sc:71 return self._ir2sc[tp]72 else:73 raise KeyError(f"Unknown ir.Type: {tp}")74 75 def get_RankedTensorType_of_nparray(76 self, nparray: np.ndarray77 ) -> ir.RankedTensorType:78 """Returns the ir.RankedTensorType of a NumPy array. Note that NumPy79 arrays can only be converted to/from dense tensors, not sparse tensors."""80 return ir.RankedTensorType.get(81 nparray.shape, self.dtype_to_irtype(nparray.dtype)82 )83 84 85# ===----------------------------------------------------------------------=== #86 87 88class StressTest:89 def __init__(self, tyconv: TypeConverter):90 self._tyconv = tyconv91 self._roundtripTp = None92 self._module = None93 self._engine = None94 95 def _assertEqualsRoundtripTp(self, tp: ir.RankedTensorType):96 assert self._roundtripTp is not None, "StressTest: uninitialized roundtrip type"97 if tp != self._roundtripTp:98 raise AssertionError(99 f"Type is not equal to the roundtrip type.\n"100 f"\tExpected: {self._roundtripTp}\n"101 f"\tFound: {tp}\n"102 )103 104 def build(self, types: List[ir.Type]):105 """Builds the ir.Module. The module has only the @main function,106 which will convert the input through the list of types and then back107 to the initial type. The roundtrip type must be a dense tensor."""108 assert self._module is None, "StressTest: must not call build() repeatedly"109 self._module = ir.Module.create()110 with ir.InsertionPoint(self._module.body):111 tp0 = types.pop(0)112 self._roundtripTp = tp0113 types.append(tp0)114 funcTp = ir.FunctionType.get(inputs=[tp0], results=[tp0])115 funcOp = func.FuncOp(name="main", type=funcTp)116 funcOp.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()117 with ir.InsertionPoint(funcOp.add_entry_block()):118 arg0 = funcOp.entry_block.arguments[0]119 self._assertEqualsRoundtripTp(arg0.type)120 v = st.ConvertOp(types.pop(0), arg0)121 for tp in types:122 w = st.ConvertOp(tp, v)123 # Release intermediate tensors before they fall out of scope.124 bufferization.DeallocTensorOp(v.result)125 v = w126 self._assertEqualsRoundtripTp(v.result.type)127 func.ReturnOp(v)128 return self129 130 def writeTo(self, filename):131 """Write the ir.Module to the given file. If the file already exists,132 then raises an error. If the filename is None, then is a no-op."""133 assert (134 self._module is not None135 ), "StressTest: must call build() before writeTo()"136 if filename is None:137 # Silent no-op, for convenience.138 return self139 if os.path.exists(filename):140 raise FileExistsError(errno.EEXIST, os.strerror(errno.EEXIST), filename)141 with open(filename, "w") as f:142 f.write(str(self._module))143 return self144 145 def compile(self, compiler):146 """Compile the ir.Module."""147 assert (148 self._module is not None149 ), "StressTest: must call build() before compile()"150 assert self._engine is None, "StressTest: must not call compile() repeatedly"151 self._engine = compiler.compile_and_jit(self._module)152 return self153 154 def run(self, np_arg0: np.ndarray) -> np.ndarray:155 """Runs the test on the given numpy array, and returns the resulting156 numpy array."""157 assert self._engine is not None, "StressTest: must call compile() before run()"158 self._assertEqualsRoundtripTp(159 self._tyconv.get_RankedTensorType_of_nparray(np_arg0)160 )161 np_out = np.zeros(np_arg0.shape, dtype=np_arg0.dtype)162 self._assertEqualsRoundtripTp(163 self._tyconv.get_RankedTensorType_of_nparray(np_out)164 )165 mem_arg0 = ctypes.pointer(166 ctypes.pointer(rt.get_ranked_memref_descriptor(np_arg0))167 )168 mem_out = ctypes.pointer(169 ctypes.pointer(rt.get_ranked_memref_descriptor(np_out))170 )171 self._engine.invoke("main", mem_out, mem_arg0)172 return rt.ranked_memref_to_numpy(mem_out[0])173 174 175# ===----------------------------------------------------------------------=== #176 177 178def main():179 """180 USAGE: python3 test_stress.py [raw_module.mlir [compiled_module.mlir]]181 182 The environment variable SUPPORT_LIB must be set to point to the183 libmlir_c_runner_utils shared library. There are two optional184 arguments, for debugging purposes. The first argument specifies where185 to write out the raw/generated ir.Module. The second argument specifies186 where to write out the compiled version of that ir.Module.187 """188 support_lib = os.getenv("SUPPORT_LIB")189 assert support_lib is not None, "SUPPORT_LIB is undefined"190 if not os.path.exists(support_lib):191 raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), support_lib)192 193 # CHECK-LABEL: TEST: test_stress194 print("\nTEST: test_stress")195 with ir.Context() as ctx, ir.Location.unknown():196 sparsification_options = f"parallelization-strategy=none "197 compiler = sparsifier.Sparsifier(198 extras="",199 options=sparsification_options,200 opt_level=0,201 shared_libs=[support_lib],202 )203 f64 = ir.F64Type.get()204 # Be careful about increasing this because205 # len(types) = 1 + len(level_choices)^rank * rank! * len(bitwidths)^2206 shape = range(2, 3)207 rank = len(shape)208 # All combinations.209 dense_lvl = st.EncodingAttr.build_level_type(st.LevelFormat.dense)210 sparse_lvl = st.EncodingAttr.build_level_type(st.LevelFormat.compressed)211 levels = list(212 itertools.product(*itertools.repeat([dense_lvl, sparse_lvl], rank))213 )214 # All permutations.215 orderings = list(216 map(ir.AffineMap.get_permutation, itertools.permutations(range(rank)))217 )218 bitwidths = [0]219 # The first type must be a dense tensor for numpy conversion to work.220 types = [ir.RankedTensorType.get(shape, f64)]221 for level in levels:222 for ordering in orderings:223 for pwidth in bitwidths:224 for iwidth in bitwidths:225 attr = st.EncodingAttr.get(226 level, ordering, None, pwidth, iwidth227 )228 types.append(ir.RankedTensorType.get(shape, f64, attr))229 #230 # For exhaustiveness we should have one or more StressTest, such231 # that their paths cover all 2*n*(n-1) directed pairwise combinations232 # of the `types` set. However, since n is already superexponential,233 # such exhaustiveness would be prohibitive for a test that runs on234 # every commit. So for now we'll just pick one particular path that235 # at least hits all n elements of the `types` set.236 #237 tyconv = TypeConverter(ctx)238 size = 1239 for d in shape:240 size *= d241 np_arg0 = np.arange(size, dtype=tyconv.irtype_to_dtype(f64)).reshape(*shape)242 np_out = (243 StressTest(tyconv)244 .build(types)245 .writeTo(sys.argv[1] if len(sys.argv) > 1 else None)246 .compile(compiler)247 .writeTo(sys.argv[2] if len(sys.argv) > 2 else None)248 .run(np_arg0)249 )250 # CHECK: Passed251 if np.allclose(np_out, np_arg0):252 print("Passed")253 else:254 sys.exit("FAILURE")255 256 257if __name__ == "__main__":258 main()259