diff --git a/mlir/lib/ExecutionEngine/RunnerUtils.cpp b/mlir/lib/ExecutionEngine/RunnerUtils.cpp index 83ead7c575a2..3a5f4713c673 100644 --- a/mlir/lib/ExecutionEngine/RunnerUtils.cpp +++ b/mlir/lib/ExecutionEngine/RunnerUtils.cpp @@ -14,6 +14,7 @@ //===----------------------------------------------------------------------===// #include "mlir/ExecutionEngine/RunnerUtils.h" +#include extern "C" void _mlir_ciface_print_memref_shape_i8(UnrankedMemRefType *M) { @@ -75,6 +76,14 @@ extern "C" void _mlir_ciface_print_memref_f64(UnrankedMemRefType *M) { impl::printMemRef(*M); } +extern "C" int64_t _mlir_ciface_nano_time() { + auto now = std::chrono::high_resolution_clock::now(); + auto duration = now.time_since_epoch(); + auto nanoseconds = + std::chrono::duration_cast(duration); + return nanoseconds.count(); +} + extern "C" void print_memref_i32(int64_t rank, void *ptr) { UnrankedMemRefType descriptor = {rank, ptr}; _mlir_ciface_print_memref_i32(&descriptor); diff --git a/mlir/test/python/execution_engine.py b/mlir/test/python/execution_engine.py index ffb0de9b69bd..3cb116b04ae3 100644 --- a/mlir/test/python/execution_engine.py +++ b/mlir/test/python/execution_engine.py @@ -358,3 +358,37 @@ def testSharedLibLoad(): run(testSharedLibLoad) + + +# Test that nano time clock is available. +# CHECK-LABEL: TEST: testNanoTime +def testNanoTime(): + with Context(): + module = Module.parse(""" + module { + func @main() attributes { llvm.emit_c_interface } { + %now = call @nano_time() : () -> i64 + %memref = memref.alloca() : memref<1xi64> + %c0 = arith.constant 0 : index + memref.store %now, %memref[%c0] : memref<1xi64> + %u_memref = memref.cast %memref : memref<1xi64> to memref<*xi64> + call @print_memref_i64(%u_memref) : (memref<*xi64>) -> () + return + } + func private @nano_time() -> i64 attributes { llvm.emit_c_interface } + func private @print_memref_i64(memref<*xi64>) attributes { llvm.emit_c_interface } + }""") + + execution_engine = ExecutionEngine( + lowerToLLVM(module), + opt_level=3, + shared_libs=[ + "../../../../lib/libmlir_runner_utils.so", + "../../../../lib/libmlir_c_runner_utils.so" + ]) + execution_engine.invoke("main") + # CHECK: Unranked Memref + # CHECK: [{{.*}}] + + +run(testNanoTime)