408 lines · python
1# RUN: %PYTHON %s 2>&1 | FileCheck %s2 3import ctypes4import sys5from mlir.ir import *6from mlir.dialects import builtin7from mlir.dialects import func8from mlir.dialects import linalg9from mlir.passmanager import *10from mlir.execution_engine import *11 12from mlir.dialects.linalg.opdsl.lang import *13 14 15# Log everything to stderr and flush so that we have a unified stream to match16# errors/info emitted by MLIR to stderr.17def log(*args):18 print(*args, file=sys.stderr)19 sys.stderr.flush()20 21 22fill_boiler = """23func.func @main() -> i32 attributes {llvm.emit_c_interface} {24 %O0 = memref.alloc() : memref<i32>25 %O1 = memref.alloc() : memref<16xi32>26 %O2 = memref.alloc() : memref<4x16xi32>27 28 %val0 = arith.constant 1 : i3229 %val1 = arith.constant 2 : i3230 %val2 = arith.constant 3 : i3231 32 call @fill_0d_on_buffers(%val0, %O0) : (i32, memref<i32>) -> ()33 call @fill_1d_on_buffers(%val1, %O1) : (i32, memref<16xi32>) -> ()34 call @fill_2d_on_buffers(%val2, %O2) : (i32, memref<4x16xi32>) -> ()35 36 %c0 = arith.constant 0 : index37 %res0 = memref.load %O0[] : memref<i32>38 %c8 = arith.constant 8 : index39 %res1 = memref.load %O1[%c8] : memref<16xi32>40 %c2 = arith.constant 2 : index41 %res2 = memref.load %O2[%c2, %c8] : memref<4x16xi32>42 43 %0 = arith.addi %res0, %res1 : i3244 %1 = arith.addi %0, %res2 : i3245 46 // TODO: FFI-based solution to allow testing and printing with python code.47 return %1 : i3248}49"""50 51fill_rng_boiler = """52func.func @main() -> i32 attributes {llvm.emit_c_interface} {53 %O = memref.alloc() : memref<4x16xi32>54 %min = arith.constant -1000.0 : f6455 %max = arith.constant 1000.0 : f6456 %seed = arith.constant 42 : i3257 58 call @fill_rng_on_buffers(%min, %max, %seed, %O) :59 (f64, f64, i32, memref<4x16xi32>) -> ()60 61 %c0 = arith.constant 0 : index62 %0 = memref.load %O[%c0, %c0] : memref<4x16xi32>63 64 // TODO: FFI-based solution to allow testing and printing with python code.65 return %0 : i3266}67"""68 69conv_boiler = """70func.func @main() -> i32 attributes {llvm.emit_c_interface} {71 %v0 = arith.constant 0 : i3272 %v1 = arith.constant 1.0 : f6473 %v2 = arith.constant 2.0 : f6474 75 %input = memref.alloc() : memref<1x4x16x1xf64>76 %filter = memref.alloc() : memref<2x2x1xf64>77 %output = memref.alloc() : memref<1x2x4x1xi32>78 linalg.fill ins(%v1 : f64) outs(%input : memref<1x4x16x1xf64>)79 linalg.fill ins(%v2 : f64) outs(%filter : memref<2x2x1xf64>)80 linalg.fill ins(%v0 : i32) outs(%output : memref<1x2x4x1xi32>)81 82 call @conv_on_buffers(%input, %filter, %output) :83 (memref<1x4x16x1xf64>, memref<2x2x1xf64>, memref<1x2x4x1xi32>) -> ()84 85 %c0 = arith.constant 0 : index86 %0 = memref.load %output[%c0, %c0, %c0, %c0] : memref<1x2x4x1xi32>87 88 // TODO: FFI-based solution to allow testing and printing with python code.89 return %0 : i3290}91"""92 93pooling_boiler = """94func.func @main() -> i32 attributes {llvm.emit_c_interface} {95 %v0 = arith.constant 0 : i3296 %v42 = arith.constant 42.0 : f6497 %v77 = arith.constant 77.0 : f6498 %v-13 = arith.constant -13.0 : f6499 %v1 = arith.constant 1.0 : f64100 101 %input = memref.alloc() : memref<1x4x16x1xf64>102 %shape = memref.alloc() : memref<2x2xf64>103 %output = memref.alloc() : memref<1x2x4x1xi32>104 linalg.fill ins(%v1 : f64) outs(%input : memref<1x4x16x1xf64>)105 linalg.fill ins(%v1 : f64) outs(%shape : memref<2x2xf64>)106 linalg.fill ins(%v0 : i32) outs(%output : memref<1x2x4x1xi32>)107 108 %c0 = arith.constant 0 : index109 %c1 = arith.constant 1 : index110 %c2 = arith.constant 2 : index111 memref.store %v42, %input[%c0, %c0, %c0, %c0] : memref<1x4x16x1xf64>112 memref.store %v77, %input[%c0, %c0, %c1, %c0] : memref<1x4x16x1xf64>113 memref.store %v-13, %input[%c0, %c1, %c0, %c0] : memref<1x4x16x1xf64>114 115 call @pooling_on_buffers(%input, %shape, %output) :116 (memref<1x4x16x1xf64>, memref<2x2xf64>, memref<1x2x4x1xi32>) -> ()117 118 %0 = memref.load %output[%c0, %c0, %c0, %c0] : memref<1x2x4x1xi32>119 120 // TODO: FFI-based solution to allow testing and printing with python code.121 return %0 : i32122}123"""124 125 126def transform(module, boilerplate):127 # TODO: Allow cloning functions from one module to another.128 # Atm we have to resort to string concatenation.129 ops = module.operation.regions[0].blocks[0].operations130 mod = Module.parse("\n".join([str(op) for op in ops]) + boilerplate)131 132 pm = PassManager("builtin.module")133 pm.add("func.func(convert-linalg-to-loops)")134 pm.add("func.func(lower-affine)")135 pm.add("func.func(convert-math-to-llvm)")136 pm.add("func.func(convert-scf-to-cf)")137 pm.add("func.func(arith-expand)")138 pm.add("func.func(memref-expand)")139 pm.add("convert-vector-to-llvm")140 pm.add("finalize-memref-to-llvm")141 pm.add("convert-func-to-llvm")142 pm.add("convert-arith-to-llvm")143 pm.add("convert-cf-to-llvm")144 pm.add("reconcile-unrealized-casts")145 pm.run(mod.operation)146 return mod147 148 149def test_fill_builtin():150 with Context() as ctx, Location.unknown():151 module = Module.create()152 i32 = IntegerType.get_signless(32)153 with InsertionPoint(module.body):154 155 @func.FuncOp.from_py_func(i32, MemRefType.get([], i32))156 def fill_0d_on_buffers(value, out):157 linalg.fill(value, outs=[out])158 159 @func.FuncOp.from_py_func(i32, MemRefType.get([16], i32))160 def fill_1d_on_buffers(value, out):161 linalg.fill(value, outs=[out])162 163 @func.FuncOp.from_py_func(i32, MemRefType.get([4, 16], i32))164 def fill_2d_on_buffers(value, out):165 linalg.fill(value, outs=[out])166 167 execution_engine = ExecutionEngine(transform(module, fill_boiler))168 169 # TODO: FFI-based solution to allow testing and printing with python code.170 # Prepare arguments: one result i32.171 # Arguments must be passed as pointers.172 c_int_p = ctypes.c_int * 1173 res = c_int_p(-1)174 execution_engine.invoke("main", res)175 176 log("RESULT: ", res[0])177 # CHECK: RESULT: 6178 179 180test_fill_builtin()181 182 183def test_fill_generic():184 with Context() as ctx, Location.unknown():185 module = Module.create()186 i32 = IntegerType.get_signless(32)187 with InsertionPoint(module.body):188 189 @func.FuncOp.from_py_func(i32, MemRefType.get([], i32))190 def fill_0d_on_buffers(value, out):191 linalg.fill(value, outs=[out], emit_generic=True)192 193 @func.FuncOp.from_py_func(i32, MemRefType.get([16], i32))194 def fill_1d_on_buffers(value, out):195 linalg.fill(value, outs=[out], emit_generic=True)196 197 @func.FuncOp.from_py_func(i32, MemRefType.get([4, 16], i32))198 def fill_2d_on_buffers(value, out):199 linalg.fill(value, outs=[out], emit_generic=True)200 201 execution_engine = ExecutionEngine(transform(module, fill_boiler))202 203 # TODO: FFI-based solution to allow testing and printing with python code.204 # Prepare arguments: one result i32.205 # Arguments must be passed as pointers.206 c_int_p = ctypes.c_int * 1207 res = c_int_p(-1)208 execution_engine.invoke("main", res)209 210 log("RESULT: ", res[0])211 # CHECK: RESULT: 6212 213 214test_fill_generic()215 216 217def test_fill_rng_builtin():218 with Context() as ctx, Location.unknown():219 module = Module.create()220 f64 = F64Type.get()221 i32 = IntegerType.get_signless(32)222 with InsertionPoint(module.body):223 224 @func.FuncOp.from_py_func(f64, f64, i32, MemRefType.get((4, 16), i32))225 def fill_rng_on_buffers(min, max, seed, out):226 linalg.fill_rng_2d(min, max, seed, outs=[out])227 228 execution_engine = ExecutionEngine(transform(module, fill_rng_boiler))229 230 # TODO: FFI-based solution to allow testing and printing with python code.231 # Prepare arguments: one result i32.232 # Arguments must be passed as pointers.233 c_int_p = ctypes.c_int * 1234 res = c_int_p(-1)235 execution_engine.invoke("main", res)236 237 log("RESULT: ", res[0])238 # CHECK: RESULT: -480239 240 241test_fill_rng_builtin()242 243 244def test_fill_rng_generic():245 with Context() as ctx, Location.unknown():246 module = Module.create()247 f64 = F64Type.get()248 i32 = IntegerType.get_signless(32)249 with InsertionPoint(module.body):250 251 @func.FuncOp.from_py_func(f64, f64, i32, MemRefType.get((4, 16), i32))252 def fill_rng_on_buffers(min, max, seed, out):253 linalg.fill_rng_2d(min, max, seed, outs=[out], emit_generic=True)254 255 execution_engine = ExecutionEngine(transform(module, fill_rng_boiler))256 257 # TODO: FFI-based solution to allow testing and printing with python code.258 # Prepare arguments: one result i32.259 # Arguments must be passed as pointers.260 c_int_p = ctypes.c_int * 1261 res = c_int_p(-1)262 execution_engine.invoke("main", res)263 264 log("RESULT: ", res[0])265 # CHECK: RESULT: -480266 267 268test_fill_rng_generic()269 270 271def test_max_pooling_builtin():272 with Context() as ctx, Location.unknown():273 module = Module.create()274 f64 = F64Type.get()275 i32 = IntegerType.get_signless(32)276 with InsertionPoint(module.body):277 278 @func.FuncOp.from_py_func(279 MemRefType.get((1, 4, 16, 1), f64),280 MemRefType.get((2, 2), f64),281 MemRefType.get((1, 2, 4, 1), i32),282 )283 def pooling_on_buffers(input, shape, output):284 linalg.pooling_nhwc_max(285 input, shape, outs=[output], strides=[2, 4], dilations=[1, 2]286 )287 288 execution_engine = ExecutionEngine(transform(module, pooling_boiler))289 290 # TODO: FFI-based solution to allow testing and printing with python code.291 # Prepare arguments: one result i32.292 # Arguments must be passed as pointers.293 c_int_p = ctypes.c_int * 1294 res = c_int_p(-1)295 execution_engine.invoke("main", res)296 297 log("RESULT: ", res[0])298 # 77 is not selected due to the dilation 2 in the second dimension.299 # CHECK: RESULT: 42300 301 302test_max_pooling_builtin()303 304 305def test_max_pooling_generic():306 with Context() as ctx, Location.unknown():307 module = Module.create()308 f64 = F64Type.get()309 i32 = IntegerType.get_signless(32)310 with InsertionPoint(module.body):311 312 @func.FuncOp.from_py_func(313 MemRefType.get((1, 4, 16, 1), f64),314 MemRefType.get((2, 2), f64),315 MemRefType.get((1, 2, 4, 1), i32),316 )317 def pooling_on_buffers(input, shape, output):318 linalg.pooling_nhwc_max(319 input,320 shape,321 outs=[output],322 strides=[2, 4],323 dilations=[1, 2],324 emit_generic=True,325 )326 327 execution_engine = ExecutionEngine(transform(module, pooling_boiler))328 329 # TODO: FFI-based solution to allow testing and printing with python code.330 # Prepare arguments: one result i32.331 # Arguments must be passed as pointers.332 c_int_p = ctypes.c_int * 1333 res = c_int_p(-1)334 execution_engine.invoke("main", res)335 336 log("RESULT: ", res[0])337 # 77 is not selected due to the dilation 2 in the second dimension.338 # CHECK: RESULT: 42339 340 341test_max_pooling_generic()342 343 344def test_min_pooling_builtin():345 with Context() as ctx, Location.unknown():346 module = Module.create()347 f64 = F64Type.get()348 i32 = IntegerType.get_signless(32)349 with InsertionPoint(module.body):350 351 @func.FuncOp.from_py_func(352 MemRefType.get((1, 4, 16, 1), f64),353 MemRefType.get((2, 2), f64),354 MemRefType.get((1, 2, 4, 1), i32),355 )356 # Set the strides and use the default dilations.357 def pooling_on_buffers(input, shape, output):358 linalg.pooling_nhwc_min(input, shape, outs=[output], strides=[2, 4])359 360 execution_engine = ExecutionEngine(transform(module, pooling_boiler))361 362 # TODO: FFI-based solution to allow testing and printing with python code.363 # Prepare arguments: one result i32.364 # Arguments must be passed as pointers.365 c_int_p = ctypes.c_int * 1366 res = c_int_p(-1)367 execution_engine.invoke("main", res)368 369 log("RESULT: ", res[0])370 # CHECK: RESULT: -13371 372 373test_min_pooling_builtin()374 375 376def test_min_pooling_generic():377 with Context() as ctx, Location.unknown():378 module = Module.create()379 f64 = F64Type.get()380 i32 = IntegerType.get_signless(32)381 with InsertionPoint(module.body):382 383 @func.FuncOp.from_py_func(384 MemRefType.get((1, 4, 16, 1), f64),385 MemRefType.get((2, 2), f64),386 MemRefType.get((1, 2, 4, 1), i32),387 )388 # Set the strides and use the default dilations.389 def pooling_on_buffers(input, shape, output):390 linalg.pooling_nhwc_min(391 input, shape, outs=[output], strides=[2, 4], emit_generic=True392 )393 394 execution_engine = ExecutionEngine(transform(module, pooling_boiler))395 396 # TODO: FFI-based solution to allow testing and printing with python code.397 # Prepare arguments: one result i32.398 # Arguments must be passed as pointers.399 c_int_p = ctypes.c_int * 1400 res = c_int_p(-1)401 execution_engine.invoke("main", res)402 403 log("RESULT: ", res[0])404 # CHECK: RESULT: -13405 406 407test_min_pooling_generic()408