implementation of external intrinsics with processing & replacment header content

.gitignore

@@ -1131,4 +1131,4 @@ secring.*
 *build*
 
 .vscode/
-lib/core/builtin_headers/okl_intrinsic_*
+lib/core/intrinsics/okl_intrinsic_*

include/oklt/core/transpiler_session/user_input.h

@@ -4,6 +4,7 @@
 
 #include <filesystem>
 #include <map>
+#include <optional>
 #include <vector>
 
 namespace oklt {
@@ -19,6 +20,7 @@ struct UserInput {
     std::vector<std::filesystem::path> includeDirectories;  ///< The include directories.
     std::vector<std::string> defines;                       ///< The defined macroses.
     std::string hash;                                       ///< OKL hash
+    std::optional<std::filesystem::path> userIntrincis;     ///< OKL user external intrincis folder
 };
 
 }  // namespace oklt

lib/CMakeLists.txt

@@ -182,8 +182,10 @@ set (OCCA_TRANSPILER_SOURCES
     core/rewriter/rewriter_proxy.cpp
     core/rewriter/rewriter_fabric.cpp
 
-    core/builtin_headers/intrinsic_impl.cpp
-    core/builtin_headers/intrinsic_impl.h
+    core/intrinsics/builtin_intrinsics.h
+    core/intrinsics/builtin_intrinsics.cpp
+    core/intrinsics/external_intrinsics.h
+    core/intrinsics/external_intrinsics.cpp
 
     core/handler_manager/handler_manager.cpp
     core/handler_manager/handler_map.cpp
@@ -277,22 +279,22 @@ target_link_libraries(occa-transpiler
 )
 
 embed_resource_txt(${ROOT_DIR}/lib/resources/okl_intrinsic_cuda.h
-    ${ROOT_DIR}/lib/core/builtin_headers/okl_intrinsic_cuda.h
+    ${ROOT_DIR}/lib/core/intrinsics/okl_intrinsic_cuda.h
     INTRINSIC_CUDA
 )
 
 embed_resource_txt(${ROOT_DIR}/lib/resources/okl_intrinsic_dpcpp.h
-    ${ROOT_DIR}/lib/core/builtin_headers/okl_intrinsic_dpcpp.h
+    ${ROOT_DIR}/lib/core/intrinsics/okl_intrinsic_dpcpp.h
     INTRINSIC_DPCPP
 )
 
 embed_resource_txt(${ROOT_DIR}/lib/resources/okl_intrinsic_hip.h
-    ${ROOT_DIR}/lib/core/builtin_headers/okl_intrinsic_hip.h
+    ${ROOT_DIR}/lib/core/intrinsics/okl_intrinsic_hip.h
     INTRINSIC_HIP
 )
 
 embed_resource_txt(${ROOT_DIR}/lib/resources/okl_intrinsic_host.h
-    ${ROOT_DIR}/lib/core/builtin_headers/okl_intrinsic_host.h
+    ${ROOT_DIR}/lib/core/intrinsics/okl_intrinsic_host.h
     INTRINSIC_HOST
 )
 

lib/attributes/backend/launcher.cpp

@@ -201,7 +201,8 @@ void collectLoops(OklLoopInfo& loopInfo, std::list<OklLoopInfo*>& out) {
 }
 #endif
 
-std::pair<LoopMetaData, LoopMetaData> splitTileAttr(OklLoopInfo& loopInfo, const oklt::Rewriter& r) {
+std::pair<LoopMetaData, LoopMetaData> splitTileAttr(OklLoopInfo& loopInfo,
+                                                    const oklt::Rewriter& r) {
     auto sz = util::parseStrTo<size_t>(loopInfo.tileSize);
 
     // Prepare first loop
@@ -358,7 +359,6 @@ HandleResult handleLauncherTranslationUnit(SessionStage& s, const TranslationUni
 
     SPDLOG_DEBUG("Handle translation unit");
 
-    //    s.getRewriter().InsertTextBefore(loc, "#include " + includeOCCA + "\n\n");
     auto& backendDeps = s.tryEmplaceUserCtx<HeaderDepsInfo>().backendHeaders;
     backendDeps.clear();
     backendDeps.emplace_back("#include " + std::string(includeOCCA) + "\n\n");
@@ -374,8 +374,8 @@ HandleResult handleLauncherKernelAttribute(SessionStage& s,
     auto& rewriter = s.getRewriter();
 
     if (!sema.getParsingKernelInfo()) {
-        return tl::make_unexpected(Error{OkltPipelineErrorCode::INTERNAL_ERROR_KERNEL_INFO_NULL,
-                                         "handleKernelAttribute"});
+        return tl::make_unexpected(
+            Error{OkltPipelineErrorCode::INTERNAL_ERROR_KERNEL_INFO_NULL, "handleKernelAttribute"});
     }
 
     auto kernelInfo = *sema.getParsingKernelInfo();

lib/attributes/utils/replace_attribute.cpp

@@ -1,6 +1,5 @@
 #include "attributes/utils/replace_attribute.h"
 #include "attributes/attribute_names.h"
-#include "core/builtin_headers/intrinsic_impl.h"
 #include "core/transpiler_session/header_info.h"
 #include "core/transpiler_session/session_stage.h"
 #include "core/utils/var_decl.h"

lib/core/builtin_headers/intrinsic_impl.cpp → lib/core/intrinsics/builtin_intrinsics.cpp

@@ -1,8 +1,8 @@
-#include "core/builtin_headers/intrinsic_impl.h"
-#include "core/builtin_headers/okl_intrinsic_cuda.h"
-#include "core/builtin_headers/okl_intrinsic_dpcpp.h"
-#include "core/builtin_headers/okl_intrinsic_hip.h"
-#include "core/builtin_headers/okl_intrinsic_host.h"
+#include "core/intrinsics/builtin_intrinsics.h"
+#include "core/intrinsics/okl_intrinsic_cuda.h"
+#include "core/intrinsics/okl_intrinsic_dpcpp.h"
+#include "core/intrinsics/okl_intrinsic_hip.h"
+#include "core/intrinsics/okl_intrinsic_host.h"
 
 #include <clang/Frontend/CompilerInstance.h>
 #include "core/transpiler_session/transpiler_session.h"

lib/core/builtin_headers/intrinsic_impl.h → lib/core/intrinsics/builtin_intrinsics.h

@@ -14,6 +14,9 @@ constexpr const char INTRINSIC_INCLUDE_FILENAME[] = "okl_intrinsic.h";
 void addInstrinsicStub(TranspilerSession &session,
                        clang::CompilerInstance &compiler);
 
+void addExternalIntrincis(TranspilerSession &session,
+                          clang::CompilerInstance &compiler);
+
 std::vector<std::string> embedInstrinsic(std::string &input,
                                          TargetBackend backend);
 

lib/core/intrinsics/external_intrinsics.cpp

@@ -0,0 +1,123 @@
+#include "core/intrinsics/external_intrinsics.h"
+
+#include <clang/Frontend/CompilerInstance.h>
+#include "core/transpiler_session/transpiler_session.h"
+#include "util/string_utils.hpp"
+
+#include <algorithm>
+
+namespace oklt {
+
+using namespace llvm;
+namespace fs = std::filesystem;
+
+tl::expected<fs::path, std::string>
+getIntrincisImplSourcePath(TargetBackend backend,
+                       const fs::path &intrincisPath)
+{
+    switch (backend) {
+        case TargetBackend::CUDA:
+            return intrincisPath / "cuda";
+        case TargetBackend::HIP:
+            return intrincisPath / "hip";
+        case TargetBackend::DPCPP:
+            return intrincisPath / "dpcpp";
+        case TargetBackend::OPENMP:
+            return intrincisPath / "openmp";
+        case TargetBackend::SERIAL:
+            return intrincisPath / "serial";
+        default:
+            return tl::make_unexpected("User intrinsic does not implement target backend");
+     }
+}
+
+bool isExternalInstrincisInclude(TranspilerSession &session,
+                                 const std::string &fileName)
+{
+     auto maybeUserIntrinsic = session.getInput().userIntrincis;
+     if(!maybeUserIntrinsic) {
+        return false;
+     }
+     auto instrinsic = maybeUserIntrinsic.value();
+     auto folderPrefix = instrinsic.filename().string();
+     return util::starts_with(fileName, folderPrefix);
+}
+
+
+tl::expected<std::unique_ptr<llvm::MemoryBuffer>, std::string>
+getExternalIntrinsicSource(TargetBackend backend,
+                           const fs::path &intrinsicPath,
+                           clang::SourceManager& sm)
+{
+
+     auto implPathResult = getIntrincisImplSourcePath(backend, intrinsicPath);
+     if(!implPathResult) {
+        return tl::make_unexpected(implPathResult.error());
+     }
+
+     auto sourceFolder = implPathResult.value();
+     if(!std::filesystem::exists(sourceFolder)) {
+        return tl::make_unexpected("Intrinsic implementation folder does not exist");
+     }
+
+     std::vector<fs::path> files(fs::directory_iterator(sourceFolder), {});
+     if(files.empty()) {
+        return tl::make_unexpected("Intrinsic implementation files is missing");
+     }
+
+     auto it = std::find_if(files.cbegin(), files.cend(), [](const fs::path &p) -> bool {
+         return p.extension().string() == std::string(".h");
+     });
+
+     if(it == files.cend()) {
+        std::string error = "Can't' find implementation file with path: " + sourceFolder.string();
+        return tl::make_unexpected(error);
+     }
+
+     auto &fm = sm.getFileManager();
+     auto backendFilePath = it->string();
+     auto maybeReplacedFile = fm.getFile(backendFilePath);
+     if(!maybeReplacedFile) {
+        std::string error = "Can't open file: " + backendFilePath;
+        return tl::make_unexpected(error);
+     }
+     auto maybeBuffer = fm.getBufferForFile(maybeReplacedFile.get());
+     if (!maybeBuffer) {
+        std::string error = "Can't get memory buffer for: " + backendFilePath;
+        return tl::make_unexpected(error);
+     }
+     return std::move(maybeBuffer.get());
+}
+
+void overrideExternalIntrinsic(TranspilerSession &session,
+                               const std::string &includedFileName,
+                               clang::OptionalFileEntryRef includedFile,
+                               clang::SourceManager &sourceManager)
+{
+     auto maybeIntrinsicPath = session.getInput().userIntrincis;
+     if(maybeIntrinsicPath && isExternalInstrincisInclude(session, includedFileName)) {
+        auto intrinsicPath = maybeIntrinsicPath.value();
+        auto infoResult = getExternalIntrinsicSource(session.getInput().backend, intrinsicPath, sourceManager);
+        if(!infoResult) {
+            session.pushError(std::error_code(), infoResult.error());
+            return;
+        }
+        auto buffer = std::move(infoResult.value());
+        auto &fm = sourceManager.getFileManager();
+        if(includedFile) {
+            auto fileRef = includedFile;
+            const auto &fileEntry = fileRef->getFileEntry();
+            sourceManager.overrideFileContents(&fileEntry, std::move(buffer));
+        } else {
+            //INFO: case when the file can be found by relative path
+            //      it happens when the include path is relative to WORKING DIR path
+            auto maybeFileRef = fm.getFileRef(includedFileName);
+            if(maybeFileRef) {
+                auto foundFileRef = maybeFileRef.get();
+                sourceManager.overrideFileContents(foundFileRef, std::move(buffer));
+            }
+        }
+     }
+}
+
+}  // namespace oklt

lib/core/intrinsics/external_intrinsics.h

@@ -0,0 +1,22 @@
+#pragma once
+#include <string>
+#include <clang/Basic/FileEntry.h>
+
+namespace clang {
+class CompilerInstance;
+class SourceManager;
+}
+
+namespace oklt {
+
+class TranspilerSession;
+
+//bool isExternalInstrincisInclude(TranspilerSession &session,
+//                                 const std::string &fileName);
+
+void overrideExternalIntrinsic(TranspilerSession &session,
+                               const std::string &includedFileName,
+                               clang::OptionalFileEntryRef includedFile,
+                               clang::SourceManager &sourceManager);
+
+}  // namespace oklt

lib/core/transpiler_session/code_generator.cpp

@@ -6,7 +6,7 @@
 #include "core/transpiler_session/transpilation_node.h"
 #include "core/transpiler_session/transpiler_session.h"
 
-#include "core/builtin_headers/intrinsic_impl.h"
+#include "core/intrinsics/builtin_intrinsics.h"
 
 #include "core/handler_manager/handler_manager.h"
 
@@ -94,14 +94,17 @@ TransformedFiles gatherTransformedFiles(SessionStage& stage) {
     // to preserve them for possible laucher generator
     auto clone = stage.getSession().getStagedHeaders();
     inputs.fileMap.merge(clone);
-    inputs.fileMap["okl_kernel.cpp"] = stage.getRewriterResultForMainFile();std::ostringstream oss;
-oss << std::this_thread::get_id() << std::endl;
-printf("%s\n", oss.str().c_str());
+    inputs.fileMap["okl_kernel.cpp"] = stage.getRewriterResultForMainFile();
+
+    //    std::ostringstream oss;
+    //    oss << std::this_thread::get_id() << std::endl;
+    //    printf("%s\n", oss.str().c_str());
+
     return inputs;
 }
 
 tl::expected<std::string, Error> preprocesseInputs(SessionStage& stage,
-                                                    const TransformedFiles& inputs) {
+                                                   const TransformedFiles& inputs) {
     auto invocation = std::make_shared<CompilerInvocation>();
 
     auto& ppOutOpt = invocation->getPreprocessorOutputOpts();
@@ -147,17 +150,15 @@ tl::expected<std::string, Error> preprocesseInputs(SessionStage& stage,
     return preprocessedAndFused.value();
 }
 
-std::string restoreSystemAndBackendHeaders(
-    TargetBackend backend,
-    std::string& input,
-    const HeaderDepsInfo& deps)
-{
+std::string restoreSystemAndBackendHeaders(TargetBackend backend,
+                                           std::string& input,
+                                           const HeaderDepsInfo& deps) {
     // insert backend specific headers and namespaces
     for (auto it = deps.backendNss.rbegin(); it < deps.backendNss.rend(); ++it) {
         input.insert(0, *it);
     }
 
-    if(deps.useOklIntrinsic) {
+    if (deps.useOklIntrinsic) {
         auto intrinsicHeaders = embedInstrinsic(input, backend);
 
         for (auto it = intrinsicHeaders.rbegin(); it < intrinsicHeaders.rend(); ++it) {
@@ -190,9 +191,8 @@ tl::expected<std::string, Error> fuseIncludeDeps(SessionStage& stage, const Head
         return preprocessedResult;
     }
 
-    auto finalTranspiledKernel = restoreSystemAndBackendHeaders(stage.getBackend(),
-                                                                preprocessedResult.value(),
-                                                                deps);
+    auto finalTranspiledKernel =
+        restoreSystemAndBackendHeaders(stage.getBackend(), preprocessedResult.value(), deps);
     return finalTranspiledKernel;
 }
 }  // namespace

lib/core/transpiler_session/header_info.cpp

@@ -1,12 +1,19 @@
 #include "core/transpiler_session/header_info.h"
-#include "core/builtin_headers/intrinsic_impl.h"
+#include "core/intrinsics/builtin_intrinsics.h"
+#include "core/intrinsics/external_intrinsics.h"
+#include "core/transpiler_session/transpiler_session.h"
 
 namespace {}
 namespace oklt {
-InclusionDirectiveCallback::InclusionDirectiveCallback(HeaderDepsInfo& deps_,
-                                                       const clang::SourceManager& sm_)
+
+using namespace llvm;
+
+InclusionDirectiveCallback::InclusionDirectiveCallback(TranspilerSession& session,
+                                                       HeaderDepsInfo& deps_,
+                                                       clang::SourceManager& sm_)
     : deps(deps_),
-      sm(sm_) {}
+      sm(sm_),
+      _session(session) {}
 
 void InclusionDirectiveCallback::InclusionDirective(clang::SourceLocation hashLoc,
                                                     const clang::Token& includeTok,
@@ -18,7 +25,7 @@ void InclusionDirectiveCallback::InclusionDirective(clang::SourceLocation hashLo
                                                     clang::StringRef relativePath,
                                                     const clang::Module* imported,
                                                     clang::SrcMgr::CharacteristicKind fileType) {
-    if(!deps.useOklIntrinsic) {
+    if (!deps.useOklIntrinsic) {
         deps.useOklIntrinsic = fileName == INTRINSIC_INCLUDE_FILENAME;
     }
 
@@ -27,10 +34,13 @@ void InclusionDirectiveCallback::InclusionDirective(clang::SourceLocation hashLo
         return;
     }
 
+    auto fileNameStr = fileName.str();
+    overrideExternalIntrinsic(_session, fileNameStr, file, sm);
+
     deps.topLevelDeps.push_back(HeaderDep{
         .hashLoc = hashLoc,
         .includeTok = includeTok,
-        .fileName = fileName.str(),
+        .fileName = fileNameStr,
         .isAngled = isAngled,
         .filenameRange = filenameRange,
         .file = file,
@@ -41,5 +51,4 @@ void InclusionDirectiveCallback::InclusionDirective(clang::SourceLocation hashLo
 
     });
 }
-
 }  // namespace oklt