From aaea92e1cd8f802a99b1c3d3f49b684276578037 Mon Sep 17 00:00:00 2001
From: Denys Shabalin <shabalin@google.com>
Date: Tue, 9 Nov 2021 17:03:51 +0100
Subject: [PATCH] [mlir] Reintroduce nano time to execution_engine

Prior change had a broken test that wasn't run by accident.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D113488
---
 mlir/lib/ExecutionEngine/RunnerUtils.cpp |  9 +++++++
 mlir/test/python/execution_engine.py     | 34 ++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

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 <chrono>
 
 extern "C" void
 _mlir_ciface_print_memref_shape_i8(UnrankedMemRefType<int8_t> *M) {
@@ -75,6 +76,14 @@ extern "C" void _mlir_ciface_print_memref_f64(UnrankedMemRefType<double> *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<std::chrono::nanoseconds>(duration);
+  return nanoseconds.count();
+}
+
 extern "C" void print_memref_i32(int64_t rank, void *ptr) {
   UnrankedMemRefType<int32_t> 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)