From a510e017e1c2f61cb04b9916fa3705d0f4b15af2 Mon Sep 17 00:00:00 2001
From: Adam Siemieniuk <adam.siemieniuk@intel.com>
Date: Thu, 25 Jul 2024 12:13:12 +0200
Subject: [PATCH] Pass XSMM runner lib to MLIR execution engine (#147)

Manually finds and passes XSMM runner library to the MLIR JIT engine
to resolve missing TPP xsmm_* symbols when executing from Python.

Works only for Linux currently.
---
 cmake/tpp-mlir.cmake                          |  6 +++-
 .../src/transformations/mlir/mlir_op.cpp      | 30 +++++++++++++++++--
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/cmake/tpp-mlir.cmake b/cmake/tpp-mlir.cmake
index 300ac99df3b734..5c658f7c6bd67e 100644
--- a/cmake/tpp-mlir.cmake
+++ b/cmake/tpp-mlir.cmake
@@ -27,7 +27,11 @@ if (TPP_MLIR_DIR)
             tpp_xsmm_runner_utils
         )
     function(add_tpp_mlir_includes target)
-        target_include_directories(${target} PRIVATE ${TPP_MLIR_DIR}/../include ${TPP_MLIR_DIR}/include)
+        target_include_directories(${target}
+                PRIVATE
+                    ${TPP_MLIR_DIR}/../include
+                    ${TPP_MLIR_DIR}/include
+                    ${TPP_MLIR_DIR}/../runtime)
     endfunction()
     function(add_tpp_mlir_libs target)
         target_link_directories(${target} PRIVATE ${TPP_MLIR_DIR}/lib)
diff --git a/src/common/transformations/src/transformations/mlir/mlir_op.cpp b/src/common/transformations/src/transformations/mlir/mlir_op.cpp
index eea1019f60db91..7bba79c2dffcec 100644
--- a/src/common/transformations/src/transformations/mlir/mlir_op.cpp
+++ b/src/common/transformations/src/transformations/mlir/mlir_op.cpp
@@ -4,16 +4,17 @@
 
 #include "mlir_op.hpp"
 
-#include <vector>
 #include <algorithm>
 #include <functional>
 #include <memory>
 #include <string>
+#include <vector>
 
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/Pass/PassManager.h"
 
 // TODO: Prune unused headers -- it's hard to understand needed ones
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/InitLLVM.h"
@@ -54,8 +55,16 @@
 #include "mlir/Target/LLVMIR/ModuleTranslation.h"
 
 #ifdef TPP_MLIR // If TPP is available
-#include "TPP/PassBundles.h"
-#include "TPP/Passes.h"
+#    if defined(__APPLE__) || defined(__linux__) || defined(__EMSCRIPTEN__)
+#        include <dlfcn.h>
+#    else
+#        error "Unsupported OS"
+#    endif
+
+#    include "PerfRunnerUtils.h"
+#    include "TPP/PassBundles.h"
+#    include "TPP/Passes.h"
+#    include "openvino/util/file_util.hpp"
 #endif
 
 namespace {
@@ -255,6 +264,19 @@ namespace mlir {
 
 using namespace ::mlir;
 
+static std::string get_xsmm_runner_path() {
+#if defined(__APPLE__) || defined(__linux__) || defined(__EMSCRIPTEN__)
+    // TODO: Add multiplatform shared library search support.
+    Dl_info info;
+    if (dladdr(reinterpret_cast<void*>(perf_start_timer), &info) == 0) {
+        llvm::errs() << "failed to find XSMM Runner library\n";
+        abort();
+    }
+    return ov::util::get_absolute_file_path(info.dli_fname);
+#else
+#    error "Unsupported OS"
+#endif
+}
 
 MLIREvaluate::MLIREvaluate(OwningOpRef<mlir::ModuleOp> _module) : module(std::move(_module)) {
     if (true) {
@@ -281,6 +303,8 @@ MLIREvaluate::MLIREvaluate(OwningOpRef<mlir::ModuleOp> _module) : module(std::mo
     engineOptions.transformer = optPipeline;  // opt level looks to be overriden in lowerToLLVMIR, but is still used
                                                 // in `create` independently
     engineOptions.llvmModuleBuilder = lowerToLLVMIR;
+    std::string xsmmLibPath = get_xsmm_runner_path();
+    engineOptions.sharedLibPaths = SmallVector<StringRef>{xsmmLibPath};
     auto maybeEngine = mlir::ExecutionEngine::create(module.get(), engineOptions);
     if (maybeEngine) {
         engine = std::move(maybeEngine.get());