From 4962d3e107333a3de8efbed49f32a4ccb2861eec Mon Sep 17 00:00:00 2001
From: Dave Airlie <airlied@redhat.com>
Date: Wed, 13 May 2020 10:43:56 +1000
Subject: [PATCH] gallivm: add cache interface to mcjit

MCJIT uses an ObjectCache object to implement the cache,
this creates and instances of it and adds it to the MCJIT
instances, it stores the cached object for later use by
the outer layers.

Reviewed-by: Roland Scheidegger <sroland@vmware.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/5049>
---
 src/gallium/auxiliary/gallivm/lp_bld_init.c   |  1 +
 src/gallium/auxiliary/gallivm/lp_bld_misc.cpp | 46 ++++++++++++++++++-
 src/gallium/auxiliary/gallivm/lp_bld_misc.h   |  3 ++
 3 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/src/gallium/auxiliary/gallivm/lp_bld_init.c b/src/gallium/auxiliary/gallivm/lp_bld_init.c
index 43d02083287..6ebc9ebf25a 100644
--- a/src/gallium/auxiliary/gallivm/lp_bld_init.c
+++ b/src/gallium/auxiliary/gallivm/lp_bld_init.c
@@ -212,6 +212,7 @@ gallivm_free_ir(struct gallivm_state *gallivm)
    }
 
    if (gallivm->cache) {
+      lp_free_objcache(gallivm->cache->jit_obj_cache);
       free(gallivm->cache->data);
    }
    FREE(gallivm->module_name);
diff --git a/src/gallium/auxiliary/gallivm/lp_bld_misc.cpp b/src/gallium/auxiliary/gallivm/lp_bld_misc.cpp
index 186965f32bb..6f63ab50ded 100644
--- a/src/gallium/auxiliary/gallivm/lp_bld_misc.cpp
+++ b/src/gallium/auxiliary/gallivm/lp_bld_misc.cpp
@@ -62,7 +62,7 @@
 #include <llvm/Support/CommandLine.h>
 #include <llvm/Support/Host.h>
 #include <llvm/Support/PrettyStackTrace.h>
-
+#include <llvm/ExecutionEngine/ObjectCache.h>
 #include <llvm/Support/TargetSelect.h>
 
 #if LLVM_VERSION_MAJOR < 11
@@ -289,6 +289,36 @@ class ShaderMemoryManager : public DelegatingJITMemoryManager {
       }
 };
 
+class LPObjectCache : public llvm::ObjectCache {
+private:
+   bool has_object;
+   struct lp_cached_code *cache_out;
+public:
+   LPObjectCache(struct lp_cached_code *cache) {
+      cache_out = cache;
+      has_object = false;
+   }
+
+   ~LPObjectCache() {
+   }
+   void notifyObjectCompiled(const llvm::Module *M, llvm::MemoryBufferRef Obj) {
+      const std::string ModuleID = M->getModuleIdentifier();
+      if (has_object)
+         fprintf(stderr, "CACHE ALREADY HAS MODULE OBJECT\n");
+      has_object = true;
+      cache_out->data_size = Obj.getBufferSize();
+      cache_out->data = malloc(cache_out->data_size);
+      memcpy(cache_out->data, Obj.getBufferStart(), cache_out->data_size);
+   }
+
+   virtual std::unique_ptr<llvm::MemoryBuffer> getObject(const llvm::Module *M) {
+      if (cache_out->data_size) {
+         return llvm::MemoryBuffer::getMemBuffer(llvm::StringRef((const char *)cache_out->data, cache_out->data_size));
+      }
+      return NULL;
+   }
+
+};
 
 /**
  * Same as LLVMCreateJITCompilerForModule, but:
@@ -502,6 +532,13 @@ lp_build_create_jit_compiler_for_module(LLVMExecutionEngineRef *OutJIT,
    ExecutionEngine *JIT;
 
    JIT = builder.create();
+
+   if (cache_out) {
+      LPObjectCache *objcache = new LPObjectCache(cache_out);
+      JIT->setObjectCache(objcache);
+      cache_out->jit_obj_cache = (void *)objcache;
+   }
+
 #if LLVM_USE_INTEL_JITEVENTS
    JITEventListener *JEL = JITEventListener::createIntelJITEventListener();
    JIT->RegisterJITEventListener(JEL);
@@ -541,6 +578,13 @@ lp_free_memory_manager(LLVMMCJITMemoryManagerRef memorymgr)
    delete reinterpret_cast<BaseMemoryManager*>(memorymgr);
 }
 
+extern "C" void
+lp_free_objcache(void *objcache_ptr)
+{
+   LPObjectCache *objcache = (LPObjectCache *)objcache_ptr;
+   delete objcache;
+}
+
 extern "C" LLVMValueRef
 lp_get_called_value(LLVMValueRef call)
 {
diff --git a/src/gallium/auxiliary/gallivm/lp_bld_misc.h b/src/gallium/auxiliary/gallivm/lp_bld_misc.h
index f3be195554b..f2a15f19e47 100644
--- a/src/gallium/auxiliary/gallivm/lp_bld_misc.h
+++ b/src/gallium/auxiliary/gallivm/lp_bld_misc.h
@@ -50,6 +50,7 @@ struct lp_cached_code {
    void *data;
    size_t data_size;
    bool dont_cache;
+   void *jit_obj_cache;
 };
 
 struct lp_generated_code;
@@ -88,6 +89,8 @@ lp_get_called_value(LLVMValueRef call);
 extern bool
 lp_is_function(LLVMValueRef v);
 
+void
+lp_free_objcache(void *objcache);
 #ifdef __cplusplus
 }
 #endif
-- 
2.30.2