829 lines · python
1# RUN: env MLIR_RUNNER_UTILS=%mlir_runner_utils MLIR_C_RUNNER_UTILS=%mlir_c_runner_utils %PYTHON %s 2>&1 | FileCheck %s2# REQUIRES: host-supports-jit3import gc, sys, os, tempfile4from textwrap import dedent5from mlir.ir import *6from mlir.passmanager import *7from mlir.execution_engine import *8from mlir.runtime import *9 10try:11 from ml_dtypes import bfloat16, float8_e5m212 13 HAS_ML_DTYPES = True14except ModuleNotFoundError:15 HAS_ML_DTYPES = False16 17 18MLIR_RUNNER_UTILS = os.getenv(19 "MLIR_RUNNER_UTILS", "../../../../lib/libmlir_runner_utils.so"20)21MLIR_C_RUNNER_UTILS = os.getenv(22 "MLIR_C_RUNNER_UTILS", "../../../../lib/libmlir_c_runner_utils.so"23)24 25 26# Log everything to stderr and flush so that we have a unified stream to match27# errors/info emitted by MLIR to stderr.28def log(*args):29 print(*args, file=sys.stderr)30 sys.stderr.flush()31 32 33def run(f):34 log("\nTEST:", f.__name__)35 f()36 gc.collect()37 assert Context._get_live_count() == 038 39 40# Verify capsule interop.41# CHECK-LABEL: TEST: testCapsule42def testCapsule():43 with Context():44 module = Module.parse(45 r"""46llvm.func @none() {47 llvm.return48}49 """50 )51 execution_engine = ExecutionEngine(module)52 execution_engine_capsule = execution_engine._CAPIPtr53 # CHECK: mlir.execution_engine.ExecutionEngine._CAPIPtr54 log(repr(execution_engine_capsule))55 execution_engine._testing_release()56 execution_engine1 = ExecutionEngine._CAPICreate(execution_engine_capsule)57 # CHECK: _mlirExecutionEngine.ExecutionEngine58 log(repr(execution_engine1))59 60 61run(testCapsule)62 63 64# Test invalid ExecutionEngine creation65# CHECK-LABEL: TEST: testInvalidModule66def testInvalidModule():67 with Context():68 # Builtin function69 module = Module.parse(70 r"""71 func.func @foo() { return }72 """73 )74 # CHECK: Got RuntimeError: Failure while creating the ExecutionEngine.75 try:76 execution_engine = ExecutionEngine(module)77 except RuntimeError as e:78 log("Got RuntimeError: ", e)79 80 81run(testInvalidModule)82 83 84def lowerToLLVM(module):85 pm = PassManager.parse(86 "builtin.module(convert-complex-to-llvm,finalize-memref-to-llvm,convert-func-to-llvm,convert-arith-to-llvm,convert-cf-to-llvm,reconcile-unrealized-casts)"87 )88 pm.run(module.operation)89 return module90 91 92# Test simple ExecutionEngine execution93# CHECK-LABEL: TEST: testInvokeVoid94def testInvokeVoid():95 with Context():96 module = Module.parse(97 r"""98func.func @void() attributes { llvm.emit_c_interface } {99 return100}101 """102 )103 execution_engine = ExecutionEngine(lowerToLLVM(module))104 # Nothing to check other than no exception thrown here.105 execution_engine.invoke("void")106 107 108run(testInvokeVoid)109 110 111# Test argument passing and result with a simple float addition.112# CHECK-LABEL: TEST: testInvokeFloatAdd113def testInvokeFloatAdd():114 with Context():115 module = Module.parse(116 r"""117func.func @add(%arg0: f32, %arg1: f32) -> f32 attributes { llvm.emit_c_interface } {118 %add = arith.addf %arg0, %arg1 : f32119 return %add : f32120}121 """122 )123 execution_engine = ExecutionEngine(lowerToLLVM(module))124 # Prepare arguments: two input floats and one result.125 # Arguments must be passed as pointers.126 c_float_p = ctypes.c_float * 1127 arg0 = c_float_p(42.0)128 arg1 = c_float_p(2.0)129 res = c_float_p(-1.0)130 execution_engine.invoke("add", arg0, arg1, res)131 # CHECK: 42.0 + 2.0 = 44.0132 log("{0} + {1} = {2}".format(arg0[0], arg1[0], res[0]))133 134 135run(testInvokeFloatAdd)136 137 138# Test callback139# CHECK-LABEL: TEST: testBasicCallback140def testBasicCallback():141 # Define a callback function that takes a float and an integer and returns a float.142 @ctypes.CFUNCTYPE(ctypes.c_float, ctypes.c_float, ctypes.c_int)143 def callback(a, b):144 return a / 2 + b / 2145 146 with Context():147 # The module just forwards to a runtime function known as "some_callback_into_python".148 module = Module.parse(149 r"""150func.func @add(%arg0: f32, %arg1: i32) -> f32 attributes { llvm.emit_c_interface } {151 %resf = call @some_callback_into_python(%arg0, %arg1) : (f32, i32) -> (f32)152 return %resf : f32153}154func.func private @some_callback_into_python(f32, i32) -> f32 attributes { llvm.emit_c_interface }155 """156 )157 execution_engine = ExecutionEngine(lowerToLLVM(module))158 execution_engine.register_runtime("some_callback_into_python", callback)159 160 # Prepare arguments: two input floats and one result.161 # Arguments must be passed as pointers.162 c_float_p = ctypes.c_float * 1163 c_int_p = ctypes.c_int * 1164 arg0 = c_float_p(42.0)165 arg1 = c_int_p(2)166 res = c_float_p(-1.0)167 execution_engine.invoke("add", arg0, arg1, res)168 # CHECK: 42.0 + 2 = 44.0169 log("{0} + {1} = {2}".format(arg0[0], arg1[0], res[0] * 2))170 171 172run(testBasicCallback)173 174 175# Test callback with an unranked memref176# CHECK-LABEL: TEST: testUnrankedMemRefCallback177def testUnrankedMemRefCallback():178 # Define a callback function that takes an unranked memref, converts it to a numpy array and prints it.179 @ctypes.CFUNCTYPE(None, ctypes.POINTER(UnrankedMemRefDescriptor))180 def callback(a):181 arr = unranked_memref_to_numpy(a, np.float32)182 log("Inside callback: ")183 log(arr)184 185 with Context():186 # The module just forwards to a runtime function known as "some_callback_into_python".187 module = Module.parse(188 r"""189func.func @callback_memref(%arg0: memref<*xf32>) attributes { llvm.emit_c_interface } {190 call @some_callback_into_python(%arg0) : (memref<*xf32>) -> ()191 return192}193func.func private @some_callback_into_python(memref<*xf32>) -> () attributes { llvm.emit_c_interface }194"""195 )196 execution_engine = ExecutionEngine(lowerToLLVM(module))197 execution_engine.register_runtime("some_callback_into_python", callback)198 inp_arr = np.array([[1.0, 2.0], [3.0, 4.0]], np.float32)199 # CHECK: Inside callback:200 # CHECK{LITERAL}: [[1. 2.]201 # CHECK{LITERAL}: [3. 4.]]202 execution_engine.invoke(203 "callback_memref",204 ctypes.pointer(ctypes.pointer(get_unranked_memref_descriptor(inp_arr))),205 )206 inp_arr_1 = np.array([5, 6, 7], dtype=np.float32)207 strided_arr = np.lib.stride_tricks.as_strided(208 inp_arr_1, strides=(4, 0), shape=(3, 4)209 )210 # CHECK: Inside callback:211 # CHECK{LITERAL}: [[5. 5. 5. 5.]212 # CHECK{LITERAL}: [6. 6. 6. 6.]213 # CHECK{LITERAL}: [7. 7. 7. 7.]]214 execution_engine.invoke(215 "callback_memref",216 ctypes.pointer(ctypes.pointer(get_unranked_memref_descriptor(strided_arr))),217 )218 219 220run(testUnrankedMemRefCallback)221 222 223# Test callback with a ranked memref.224# CHECK-LABEL: TEST: testRankedMemRefCallback225def testRankedMemRefCallback():226 # Define a callback function that takes a ranked memref, converts it to a numpy array and prints it.227 @ctypes.CFUNCTYPE(228 None,229 ctypes.POINTER(230 make_nd_memref_descriptor(2, np.ctypeslib.as_ctypes_type(np.float32))231 ),232 )233 def callback(a):234 arr = ranked_memref_to_numpy(a)235 log("Inside Callback: ")236 log(arr)237 238 with Context():239 # The module just forwards to a runtime function known as "some_callback_into_python".240 module = Module.parse(241 r"""242func.func @callback_memref(%arg0: memref<2x2xf32>) attributes { llvm.emit_c_interface } {243 call @some_callback_into_python(%arg0) : (memref<2x2xf32>) -> ()244 return245}246func.func private @some_callback_into_python(memref<2x2xf32>) -> () attributes { llvm.emit_c_interface }247"""248 )249 execution_engine = ExecutionEngine(lowerToLLVM(module))250 execution_engine.register_runtime("some_callback_into_python", callback)251 inp_arr = np.array([[1.0, 5.0], [6.0, 7.0]], np.float32)252 # CHECK: Inside Callback:253 # CHECK{LITERAL}: [[1. 5.]254 # CHECK{LITERAL}: [6. 7.]]255 execution_engine.invoke(256 "callback_memref",257 ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(inp_arr))),258 )259 260 261run(testRankedMemRefCallback)262 263 264# Test callback with a ranked memref with non-zero offset.265# CHECK-LABEL: TEST: testRankedMemRefWithOffsetCallback266def testRankedMemRefWithOffsetCallback():267 # Define a callback function that takes a ranked memref, converts it to a numpy array and prints it.268 @ctypes.CFUNCTYPE(269 None,270 ctypes.POINTER(271 make_nd_memref_descriptor(1, np.ctypeslib.as_ctypes_type(np.float32))272 ),273 )274 def callback(a):275 arr = ranked_memref_to_numpy(a)276 log("Inside Callback: ")277 log(arr)278 279 with Context():280 # The module takes a subview of the argument memref and calls the callback with it281 module = Module.parse(282 r"""283func.func @callback_memref(%arg0: memref<5xf32>) attributes {llvm.emit_c_interface} {284 %base_buffer, %offset, %sizes, %strides = memref.extract_strided_metadata %arg0 : memref<5xf32> -> memref<f32>, index, index, index285 %reinterpret_cast = memref.reinterpret_cast %base_buffer to offset: [3], sizes: [2], strides: [1] : memref<f32> to memref<2xf32, strided<[1], offset: 3>>286 %cast = memref.cast %reinterpret_cast : memref<2xf32, strided<[1], offset: 3>> to memref<?xf32, strided<[?], offset: ?>>287 call @some_callback_into_python(%cast) : (memref<?xf32, strided<[?], offset: ?>>) -> ()288 return289}290func.func private @some_callback_into_python(memref<?xf32, strided<[?], offset: ?>>) attributes {llvm.emit_c_interface}291"""292 )293 execution_engine = ExecutionEngine(lowerToLLVM(module))294 execution_engine.register_runtime("some_callback_into_python", callback)295 inp_arr = np.array([0, 0, 0, 1, 2], np.float32)296 # CHECK: Inside Callback:297 # CHECK{LITERAL}: [1. 2.]298 execution_engine.invoke(299 "callback_memref",300 ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(inp_arr))),301 )302 303 304run(testRankedMemRefWithOffsetCallback)305 306 307# Test callback with an unranked memref with non-zero offset308# CHECK-LABEL: TEST: testUnrankedMemRefWithOffsetCallback309def testUnrankedMemRefWithOffsetCallback():310 # Define a callback function that takes an unranked memref, converts it to a numpy array and prints it.311 @ctypes.CFUNCTYPE(None, ctypes.POINTER(UnrankedMemRefDescriptor))312 def callback(a):313 arr = unranked_memref_to_numpy(a, np.float32)314 log("Inside callback: ")315 log(arr)316 317 with Context():318 # The module takes a subview of the argument memref, casts it to an unranked memref and319 # calls the callback with it.320 module = Module.parse(321 r"""322func.func @callback_memref(%arg0: memref<5xf32>) attributes {llvm.emit_c_interface} {323 %base_buffer, %offset, %sizes, %strides = memref.extract_strided_metadata %arg0 : memref<5xf32> -> memref<f32>, index, index, index324 %reinterpret_cast = memref.reinterpret_cast %base_buffer to offset: [3], sizes: [2], strides: [1] : memref<f32> to memref<2xf32, strided<[1], offset: 3>>325 %cast = memref.cast %reinterpret_cast : memref<2xf32, strided<[1], offset: 3>> to memref<*xf32>326 call @some_callback_into_python(%cast) : (memref<*xf32>) -> ()327 return328}329func.func private @some_callback_into_python(memref<*xf32>) attributes {llvm.emit_c_interface}330"""331 )332 execution_engine = ExecutionEngine(lowerToLLVM(module))333 execution_engine.register_runtime("some_callback_into_python", callback)334 inp_arr = np.array([1, 2, 3, 4, 5], np.float32)335 # CHECK: Inside callback:336 # CHECK{LITERAL}: [4. 5.]337 execution_engine.invoke(338 "callback_memref",339 ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(inp_arr))),340 )341 342 343run(testUnrankedMemRefWithOffsetCallback)344 345 346# Test addition of two memrefs.347# CHECK-LABEL: TEST: testMemrefAdd348def testMemrefAdd():349 with Context():350 module = Module.parse(351 """352 module {353 func.func @main(%arg0: memref<1xf32>, %arg1: memref<f32>, %arg2: memref<1xf32>) attributes { llvm.emit_c_interface } {354 %0 = arith.constant 0 : index355 %1 = memref.load %arg0[%0] : memref<1xf32>356 %2 = memref.load %arg1[] : memref<f32>357 %3 = arith.addf %1, %2 : f32358 memref.store %3, %arg2[%0] : memref<1xf32>359 return360 }361 } """362 )363 arg1 = np.array([32.5]).astype(np.float32)364 arg2 = np.array(6).astype(np.float32)365 res = np.array([0]).astype(np.float32)366 367 arg1_memref_ptr = ctypes.pointer(368 ctypes.pointer(get_ranked_memref_descriptor(arg1))369 )370 arg2_memref_ptr = ctypes.pointer(371 ctypes.pointer(get_ranked_memref_descriptor(arg2))372 )373 res_memref_ptr = ctypes.pointer(374 ctypes.pointer(get_ranked_memref_descriptor(res))375 )376 377 execution_engine = ExecutionEngine(lowerToLLVM(module))378 execution_engine.invoke(379 "main", arg1_memref_ptr, arg2_memref_ptr, res_memref_ptr380 )381 # CHECK: [32.5] + 6.0 = [38.5]382 log("{0} + {1} = {2}".format(arg1, arg2, res))383 384 385run(testMemrefAdd)386 387 388# Test addition of two f16 memrefs389# CHECK-LABEL: TEST: testF16MemrefAdd390def testF16MemrefAdd():391 with Context():392 module = Module.parse(393 """394 module {395 func.func @main(%arg0: memref<1xf16>,396 %arg1: memref<1xf16>,397 %arg2: memref<1xf16>) attributes { llvm.emit_c_interface } {398 %0 = arith.constant 0 : index399 %1 = memref.load %arg0[%0] : memref<1xf16>400 %2 = memref.load %arg1[%0] : memref<1xf16>401 %3 = arith.addf %1, %2 : f16402 memref.store %3, %arg2[%0] : memref<1xf16>403 return404 }405 } """406 )407 408 arg1 = np.array([11.0]).astype(np.float16)409 arg2 = np.array([22.0]).astype(np.float16)410 arg3 = np.array([0.0]).astype(np.float16)411 412 arg1_memref_ptr = ctypes.pointer(413 ctypes.pointer(get_ranked_memref_descriptor(arg1))414 )415 arg2_memref_ptr = ctypes.pointer(416 ctypes.pointer(get_ranked_memref_descriptor(arg2))417 )418 arg3_memref_ptr = ctypes.pointer(419 ctypes.pointer(get_ranked_memref_descriptor(arg3))420 )421 422 execution_engine = ExecutionEngine(lowerToLLVM(module))423 execution_engine.invoke(424 "main", arg1_memref_ptr, arg2_memref_ptr, arg3_memref_ptr425 )426 # CHECK: [11.] + [22.] = [33.]427 log("{0} + {1} = {2}".format(arg1, arg2, arg3))428 429 # test to-numpy utility430 # CHECK: [33.]431 npout = ranked_memref_to_numpy(arg3_memref_ptr[0])432 log(npout)433 434 435run(testF16MemrefAdd)436 437 438# Test addition of two complex memrefs439# CHECK-LABEL: TEST: testComplexMemrefAdd440def testComplexMemrefAdd():441 with Context():442 module = Module.parse(443 """444 module {445 func.func @main(%arg0: memref<1xcomplex<f64>>,446 %arg1: memref<1xcomplex<f64>>,447 %arg2: memref<1xcomplex<f64>>) attributes { llvm.emit_c_interface } {448 %0 = arith.constant 0 : index449 %1 = memref.load %arg0[%0] : memref<1xcomplex<f64>>450 %2 = memref.load %arg1[%0] : memref<1xcomplex<f64>>451 %3 = complex.add %1, %2 : complex<f64>452 memref.store %3, %arg2[%0] : memref<1xcomplex<f64>>453 return454 }455 } """456 )457 458 arg1 = np.array([1.0 + 2.0j]).astype(np.complex128)459 arg2 = np.array([3.0 + 4.0j]).astype(np.complex128)460 arg3 = np.array([0.0 + 0.0j]).astype(np.complex128)461 462 arg1_memref_ptr = ctypes.pointer(463 ctypes.pointer(get_ranked_memref_descriptor(arg1))464 )465 arg2_memref_ptr = ctypes.pointer(466 ctypes.pointer(get_ranked_memref_descriptor(arg2))467 )468 arg3_memref_ptr = ctypes.pointer(469 ctypes.pointer(get_ranked_memref_descriptor(arg3))470 )471 472 execution_engine = ExecutionEngine(lowerToLLVM(module))473 execution_engine.invoke(474 "main", arg1_memref_ptr, arg2_memref_ptr, arg3_memref_ptr475 )476 # CHECK: [1.+2.j] + [3.+4.j] = [4.+6.j]477 log("{0} + {1} = {2}".format(arg1, arg2, arg3))478 479 # test to-numpy utility480 # CHECK: [4.+6.j]481 npout = ranked_memref_to_numpy(arg3_memref_ptr[0])482 log(npout)483 484 485run(testComplexMemrefAdd)486 487 488# Test addition of two complex unranked memrefs489# CHECK-LABEL: TEST: testComplexUnrankedMemrefAdd490def testComplexUnrankedMemrefAdd():491 with Context():492 module = Module.parse(493 """494 module {495 func.func @main(%arg0: memref<*xcomplex<f32>>,496 %arg1: memref<*xcomplex<f32>>,497 %arg2: memref<*xcomplex<f32>>) attributes { llvm.emit_c_interface } {498 %A = memref.cast %arg0 : memref<*xcomplex<f32>> to memref<1xcomplex<f32>>499 %B = memref.cast %arg1 : memref<*xcomplex<f32>> to memref<1xcomplex<f32>>500 %C = memref.cast %arg2 : memref<*xcomplex<f32>> to memref<1xcomplex<f32>>501 %0 = arith.constant 0 : index502 %1 = memref.load %A[%0] : memref<1xcomplex<f32>>503 %2 = memref.load %B[%0] : memref<1xcomplex<f32>>504 %3 = complex.add %1, %2 : complex<f32>505 memref.store %3, %C[%0] : memref<1xcomplex<f32>>506 return507 }508 } """509 )510 511 arg1 = np.array([5.0 + 6.0j]).astype(np.complex64)512 arg2 = np.array([7.0 + 8.0j]).astype(np.complex64)513 arg3 = np.array([0.0 + 0.0j]).astype(np.complex64)514 515 arg1_memref_ptr = ctypes.pointer(516 ctypes.pointer(get_unranked_memref_descriptor(arg1))517 )518 arg2_memref_ptr = ctypes.pointer(519 ctypes.pointer(get_unranked_memref_descriptor(arg2))520 )521 arg3_memref_ptr = ctypes.pointer(522 ctypes.pointer(get_unranked_memref_descriptor(arg3))523 )524 525 execution_engine = ExecutionEngine(lowerToLLVM(module))526 execution_engine.invoke(527 "main", arg1_memref_ptr, arg2_memref_ptr, arg3_memref_ptr528 )529 # CHECK: [5.+6.j] + [7.+8.j] = [12.+14.j]530 log("{0} + {1} = {2}".format(arg1, arg2, arg3))531 532 # test to-numpy utility533 # CHECK: [12.+14.j]534 npout = unranked_memref_to_numpy(arg3_memref_ptr[0], np.dtype(np.complex64))535 log(npout)536 537 538run(testComplexUnrankedMemrefAdd)539 540 541# Test bf16 memrefs542# CHECK-LABEL: TEST: testBF16Memref543def testBF16Memref():544 with Context():545 module = Module.parse(546 """547 module {548 func.func @main(%arg0: memref<1xbf16>,549 %arg1: memref<1xbf16>) attributes { llvm.emit_c_interface } {550 %0 = arith.constant 0 : index551 %1 = memref.load %arg0[%0] : memref<1xbf16>552 memref.store %1, %arg1[%0] : memref<1xbf16>553 return554 }555 } """556 )557 558 arg1 = np.array([0.5]).astype(bfloat16)559 arg2 = np.array([0.0]).astype(bfloat16)560 561 arg1_memref_ptr = ctypes.pointer(562 ctypes.pointer(get_ranked_memref_descriptor(arg1))563 )564 arg2_memref_ptr = ctypes.pointer(565 ctypes.pointer(get_ranked_memref_descriptor(arg2))566 )567 568 execution_engine = ExecutionEngine(lowerToLLVM(module))569 execution_engine.invoke("main", arg1_memref_ptr, arg2_memref_ptr)570 571 # test to-numpy utility572 x = ranked_memref_to_numpy(arg2_memref_ptr[0])573 assert len(x) == 1574 assert x[0] == 0.5575 576 577if HAS_ML_DTYPES:578 run(testBF16Memref)579else:580 log("TEST: testBF16Memref")581 582 583# Test f8E5M2 memrefs584# CHECK-LABEL: TEST: testF8E5M2Memref585def testF8E5M2Memref():586 with Context():587 module = Module.parse(588 """589 module {590 func.func @main(%arg0: memref<1xf8E5M2>,591 %arg1: memref<1xf8E5M2>) attributes { llvm.emit_c_interface } {592 %0 = arith.constant 0 : index593 %1 = memref.load %arg0[%0] : memref<1xf8E5M2>594 memref.store %1, %arg1[%0] : memref<1xf8E5M2>595 return596 }597 } """598 )599 600 arg1 = np.array([0.5]).astype(float8_e5m2)601 arg2 = np.array([0.0]).astype(float8_e5m2)602 603 arg1_memref_ptr = ctypes.pointer(604 ctypes.pointer(get_ranked_memref_descriptor(arg1))605 )606 arg2_memref_ptr = ctypes.pointer(607 ctypes.pointer(get_ranked_memref_descriptor(arg2))608 )609 610 execution_engine = ExecutionEngine(lowerToLLVM(module))611 execution_engine.invoke("main", arg1_memref_ptr, arg2_memref_ptr)612 613 # test to-numpy utility614 x = ranked_memref_to_numpy(arg2_memref_ptr[0])615 assert len(x) == 1616 assert x[0] == 0.5617 618 619if HAS_ML_DTYPES:620 run(testF8E5M2Memref)621else:622 log("TEST: testF8E5M2Memref")623 624 625# Test addition of two 2d_memref626# CHECK-LABEL: TEST: testDynamicMemrefAdd2D627def testDynamicMemrefAdd2D():628 with Context():629 module = Module.parse(630 """631 module {632 func.func @memref_add_2d(%arg0: memref<2x2xf32>, %arg1: memref<?x?xf32>, %arg2: memref<2x2xf32>) attributes {llvm.emit_c_interface} {633 %c0 = arith.constant 0 : index634 %c2 = arith.constant 2 : index635 %c1 = arith.constant 1 : index636 cf.br ^bb1(%c0 : index)637 ^bb1(%0: index): // 2 preds: ^bb0, ^bb5638 %1 = arith.cmpi slt, %0, %c2 : index639 cf.cond_br %1, ^bb2, ^bb6640 ^bb2: // pred: ^bb1641 %c0_0 = arith.constant 0 : index642 %c2_1 = arith.constant 2 : index643 %c1_2 = arith.constant 1 : index644 cf.br ^bb3(%c0_0 : index)645 ^bb3(%2: index): // 2 preds: ^bb2, ^bb4646 %3 = arith.cmpi slt, %2, %c2_1 : index647 cf.cond_br %3, ^bb4, ^bb5648 ^bb4: // pred: ^bb3649 %4 = memref.load %arg0[%0, %2] : memref<2x2xf32>650 %5 = memref.load %arg1[%0, %2] : memref<?x?xf32>651 %6 = arith.addf %4, %5 : f32652 memref.store %6, %arg2[%0, %2] : memref<2x2xf32>653 %7 = arith.addi %2, %c1_2 : index654 cf.br ^bb3(%7 : index)655 ^bb5: // pred: ^bb3656 %8 = arith.addi %0, %c1 : index657 cf.br ^bb1(%8 : index)658 ^bb6: // pred: ^bb1659 return660 }661 }662 """663 )664 arg1 = np.random.randn(2, 2).astype(np.float32)665 arg2 = np.random.randn(2, 2).astype(np.float32)666 res = np.random.randn(2, 2).astype(np.float32)667 668 arg1_memref_ptr = ctypes.pointer(669 ctypes.pointer(get_ranked_memref_descriptor(arg1))670 )671 arg2_memref_ptr = ctypes.pointer(672 ctypes.pointer(get_ranked_memref_descriptor(arg2))673 )674 res_memref_ptr = ctypes.pointer(675 ctypes.pointer(get_ranked_memref_descriptor(res))676 )677 678 execution_engine = ExecutionEngine(lowerToLLVM(module))679 execution_engine.invoke(680 "memref_add_2d", arg1_memref_ptr, arg2_memref_ptr, res_memref_ptr681 )682 # CHECK: True683 log(np.allclose(arg1 + arg2, res))684 685 686run(testDynamicMemrefAdd2D)687 688 689# Test loading of shared libraries.690# CHECK-LABEL: TEST: testSharedLibLoad691def testSharedLibLoad():692 with Context():693 module = Module.parse(694 """695 module {696 func.func @main(%arg0: memref<1xf32>) attributes { llvm.emit_c_interface } {697 %c0 = arith.constant 0 : index698 %cst42 = arith.constant 42.0 : f32699 memref.store %cst42, %arg0[%c0] : memref<1xf32>700 %u_memref = memref.cast %arg0 : memref<1xf32> to memref<*xf32>701 call @printMemrefF32(%u_memref) : (memref<*xf32>) -> ()702 return703 }704 func.func private @printMemrefF32(memref<*xf32>) attributes { llvm.emit_c_interface }705 } """706 )707 arg0 = np.array([0.0]).astype(np.float32)708 709 arg0_memref_ptr = ctypes.pointer(710 ctypes.pointer(get_ranked_memref_descriptor(arg0))711 )712 713 if sys.platform == "win32":714 shared_libs = [715 "../../../../bin/mlir_runner_utils.dll",716 "../../../../bin/mlir_c_runner_utils.dll",717 ]718 elif sys.platform == "darwin":719 shared_libs = [720 "../../../../lib/libmlir_runner_utils.dylib",721 "../../../../lib/libmlir_c_runner_utils.dylib",722 ]723 else:724 shared_libs = [725 MLIR_RUNNER_UTILS,726 MLIR_C_RUNNER_UTILS,727 ]728 729 execution_engine = ExecutionEngine(730 lowerToLLVM(module), opt_level=3, shared_libs=shared_libs731 )732 execution_engine.invoke("main", arg0_memref_ptr)733 # CHECK: Unranked Memref734 # CHECK-NEXT: [42]735 736 737run(testSharedLibLoad)738 739 740# Test that nano time clock is available.741# CHECK-LABEL: TEST: testNanoTime742def testNanoTime():743 with Context():744 module = Module.parse(745 """746 module {747 func.func @main() attributes { llvm.emit_c_interface } {748 %now = call @nanoTime() : () -> i64749 %memref = memref.alloca() : memref<1xi64>750 %c0 = arith.constant 0 : index751 memref.store %now, %memref[%c0] : memref<1xi64>752 %u_memref = memref.cast %memref : memref<1xi64> to memref<*xi64>753 call @printMemrefI64(%u_memref) : (memref<*xi64>) -> ()754 return755 }756 func.func private @nanoTime() -> i64 attributes { llvm.emit_c_interface }757 func.func private @printMemrefI64(memref<*xi64>) attributes { llvm.emit_c_interface }758 }"""759 )760 761 if sys.platform == "win32":762 shared_libs = [763 "../../../../bin/mlir_runner_utils.dll",764 "../../../../bin/mlir_c_runner_utils.dll",765 ]766 else:767 shared_libs = [768 MLIR_RUNNER_UTILS,769 MLIR_C_RUNNER_UTILS,770 ]771 772 execution_engine = ExecutionEngine(773 lowerToLLVM(module), opt_level=3, shared_libs=shared_libs774 )775 execution_engine.invoke("main")776 # CHECK: Unranked Memref777 # CHECK: [{{.*}}]778 779 780run(testNanoTime)781 782 783# Test that nano time clock is available.784# CHECK-LABEL: TEST: testDumpToObjectFile785def testDumpToObjectFile():786 fd, object_path = tempfile.mkstemp(suffix=".o")787 788 try:789 with Context():790 module = Module.parse(791 dedent(792 """793 func.func private @printF32(f32)794 func.func @main(%arg0: f32) attributes { llvm.emit_c_interface } {795 call @printF32(%arg0) : (f32) -> ()796 return797 }798 """799 )800 )801 802 execution_engine = ExecutionEngine(803 lowerToLLVM(module),804 opt_level=3,805 # Loading MLIR_C_RUNNER_UTILS is necessary even though we don't actually run the code (i.e., call printF32)806 # because RTDyldObjectLinkingLayer::emit will try to resolve symbols before dumping807 # (see the jitLinkForORC call at the bottom there).808 shared_libs=[MLIR_C_RUNNER_UTILS],809 )810 811 # CHECK: Object file exists: True812 print(f"Object file exists: {os.path.exists(object_path)}")813 # CHECK: Object file is empty: True814 print(f"Object file is empty: {os.path.getsize(object_path) == 0}")815 816 execution_engine.dump_to_object_file(object_path)817 818 # CHECK: Object file exists: True819 print(f"Object file exists: {os.path.exists(object_path)}")820 # CHECK: Object file is empty: False821 print(f"Object file is empty: {os.path.getsize(object_path) == 0}")822 823 finally:824 os.close(fd)825 os.remove(object_path)826 827 828run(testDumpToObjectFile)829