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

@@ -19,6 +19,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::vector<std::filesystem::path> userIntrinsics;      ///< 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/diag/diag_consumer.cpp

@@ -26,7 +26,8 @@ DiagConsumer::DiagConsumer(SessionStage& session)
       clang::DiagnosticConsumer(){};
 
 void DiagConsumer::HandleDiagnostic(DiagnosticsEngine::Level DiagLevel, const Diagnostic& Info) {
-    if (!_includeDiag.test_and_set()) {
+    if (!_includeDiag) {
+        _includeDiag = 1;
         return;
     }
 
@@ -44,13 +45,13 @@ bool DiagConsumer::IncludeInDiagnosticCounts() const {
 
     // Accept only Warning, Error and Fatal
     if (diagLevel < DiagnosticsEngine::Level::Warning) {
-        const_cast<std::atomic_flag&>(_includeDiag).clear();
+        _includeDiag = 0;
         return false;
     }
 
     for (auto& ptr : getDiagDiagHandleInstances()) {
         if (ptr->_id == info.getID() && ptr->HandleDiagnostic(_session, diagLevel, info)) {
-            const_cast<std::atomic_flag&>(_includeDiag).clear();
+            _includeDiag = 0;
             return false;
         }
     }

lib/core/diag/diag_consumer.h

@@ -22,7 +22,7 @@ class DiagConsumer : public clang::DiagnosticConsumer {
 
    protected:
     SessionStage& _session;
-    std::atomic_flag _includeDiag = true;
+    mutable int _includeDiag = 0;
 };
 
 }  // namespace oklt

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

lib/core/intrinsics/external_intrinsics.cpp

@@ -0,0 +1,158 @@
+#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>
+#include <optional>
+
+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 isExternalIntrinsicInclude(TranspilerSession& session, const std::string& fileName) {
+    const auto& userIntrinsic = session.getInput().userIntrinsics;
+    if (userIntrinsic.empty()) {
+        return false;
+    }
+    auto normalizedName = fileName;
+    if (util::startsWith(normalizedName, "./")) {
+        normalizedName = normalizedName.substr(2);
+    }
+    for (const auto& intrinsic : userIntrinsic) {
+        auto folderPrefix = intrinsic.filename().string();
+        if (util::startsWith(normalizedName, folderPrefix)) {
+            return true;
+        }
+    }
+    return false;
+}
+
+std::optional<fs::path> getExternalInstrincisInclude(TranspilerSession& session,
+                                                     const std::string& fileName) {
+    const auto& userIntrinsic = session.getInput().userIntrinsics;
+    if (userIntrinsic.empty()) {
+        return std::nullopt;
+    }
+    for (const auto& intrinsic : userIntrinsic) {
+        auto folderPrefix = intrinsic.filename().string();
+        if (util::startsWith(fileName, folderPrefix)) {
+            return intrinsic;
+        }
+    }
+    return std::nullopt;
+}
+
+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());
+}
+
+bool overrideExternalIntrinsic(TranspilerSession& session,
+                               const std::string& includedFileName,
+                               clang::OptionalFileEntryRef includedFile,
+                               clang::SourceManager& sourceManager) {
+    const auto& userIntrinsics = session.getInput().userIntrinsics;
+    if (!userIntrinsics.empty()) {
+        auto maybeIntrinsicPath = getExternalInstrincisInclude(session, includedFileName);
+        if (!maybeIntrinsicPath) {
+            return false;
+        }
+        auto intrinsicPath = maybeIntrinsicPath.value();
+        auto infoResult =
+            getExternalIntrinsicSource(session.getInput().backend, intrinsicPath, sourceManager);
+        if (!infoResult) {
+            session.pushError(std::error_code(), infoResult.error());
+            return false;
+        }
+        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));
+            }
+        }
+        return true;
+    }
+    return false;
+}
+
+void nullyExternalIntrinsics(TransformedFiles& inputs, TranspilerSession& session) {
+    const auto& intrinsics = session.getInput().userIntrinsics;
+    if (intrinsics.empty()) {
+        return;
+    }
+    for (auto& mappedFile : inputs.fileMap) {
+        if (isExternalIntrinsicInclude(session, mappedFile.first)) {
+            mappedFile.second.clear();
+        }
+    }
+}
+
+}  // namespace oklt

lib/core/intrinsics/external_intrinsics.h

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