Enable parallel compilation with hiprtc (#1647)

* Add hiprtc driver

src/dynamic_loader.cpp

@@ -71,6 +71,16 @@ struct dynamic_loader_impl
     std::shared_ptr<tmp_dir> temp = nullptr;
 };
 
+fs::path dynamic_loader::path(void* address)
+{
+    fs::path p;
+    Dl_info info;
+    // Find the location of .so
+    if(dladdr(address, &info) != 0)
+        p = info.dli_fname;
+    return p;
+}
+
 dynamic_loader::dynamic_loader(const fs::path& p) : impl(std::make_shared<dynamic_loader_impl>(p))
 {
 }

src/include/migraphx/dynamic_loader.hpp

@@ -37,6 +37,12 @@ struct dynamic_loader_impl;
 
 struct dynamic_loader
 {
+    template <class T>
+    static fs::path path(T* address)
+    {
+        return path(reinterpret_cast<void*>(address));
+    }
+    static fs::path path(void* address);
     dynamic_loader() = default;
 
     dynamic_loader(const fs::path& p);

src/include/migraphx/msgpack.hpp

@@ -30,6 +30,7 @@
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
+void to_msgpack(const value& v, std::function<void(const char*, std::size_t)> writer);
 std::vector<char> to_msgpack(const value& v);
 value from_msgpack(const std::vector<char>& buffer);
 value from_msgpack(const char* buffer, std::size_t size);

src/include/migraphx/process.hpp

@@ -26,6 +26,7 @@
 
 #include <migraphx/config.hpp>
 #include <migraphx/filesystem.hpp>
+#include <functional>
 #include <string>
 #include <memory>
 
@@ -36,6 +37,7 @@ struct process_impl;
 
 struct process
 {
+    using writer = std::function<void(const char*, std::size_t)>;
     process(const std::string& cmd);
 
     // move constructor
@@ -49,6 +51,7 @@ struct process
     process& cwd(const fs::path& p);
 
     void exec();
+    void write(std::function<void(process::writer)> pipe_in);
 
     private:
     std::unique_ptr<process_impl> impl;

src/msgpack.cpp

@@ -172,6 +172,22 @@ struct vector_stream
     }
 };
 
+struct writer_stream
+{
+    std::function<void(const char*, std::size_t)> writer;
+    writer_stream& write(const char* b, std::size_t n)
+    {
+        writer(b, n);
+        return *this;
+    }
+};
+
+void to_msgpack(const value& v, std::function<void(const char*, std::size_t)> writer)
+{
+    writer_stream ws{std::move(writer)};
+    msgpack::pack(ws, v);
+}
+
 std::vector<char> to_msgpack(const value& v)
 {
     vector_stream vs;

src/process.cpp

@@ -38,27 +38,42 @@ std::function<void(const char*)> redirect_to(std::ostream& os)
     return [&](const char* x) { os << x; };
 }
 
-int exec(const std::string& cmd, const std::function<void(const char*)>& std_out)
+template <class F>
+int exec(const std::string& cmd, const char* type, F f)
 {
     int ec = 0;
     if(enabled(MIGRAPHX_TRACE_CMD_EXECUTE{}))
         std::cout << cmd << std::endl;
     auto closer = [&](FILE* stream) {
         auto status = pclose(stream);
-        ec          = WIFEXITED(status) ? 0 : WEXITSTATUS(status); // NOLINT
+        ec          = WIFEXITED(status) ? WEXITSTATUS(status) : 0; // NOLINT
     };
     {
         // TODO: Use execve instead of popen
-        std::unique_ptr<FILE, decltype(closer)> pipe(popen(cmd.c_str(), "r"), closer); // NOLINT
+        std::unique_ptr<FILE, decltype(closer)> pipe(popen(cmd.c_str(), type), closer); // NOLINT
         if(not pipe)
             MIGRAPHX_THROW("popen() failed: " + cmd);
-        std::array<char, 128> buffer;
-        while(fgets(buffer.data(), buffer.size(), pipe.get()) != nullptr)
-            std_out(buffer.data());
+        f(pipe.get());
     }
     return ec;
 }
 
+int exec(const std::string& cmd, const std::function<void(const char*)>& std_out)
+{
+    return exec(cmd, "r", [&](FILE* f) {
+        std::array<char, 128> buffer;
+        while(fgets(buffer.data(), buffer.size(), f) != nullptr)
+            std_out(buffer.data());
+    });
+}
+
+int exec(const std::string& cmd, std::function<void(process::writer)> std_in)
+{
+    return exec(cmd, "w", [&](FILE* f) {
+        std_in([&](const char* buffer, std::size_t n) { std::fwrite(buffer, 1, n, f); });
+    });
+}
+
 struct process_impl
 {
     std::string command{};
@@ -72,6 +87,15 @@ struct process_impl
         result += command;
         return result;
     }
+
+    template <class... Ts>
+    void check_exec(Ts&&... xs) const
+    {
+        int ec = migraphx::exec(std::forward<Ts>(xs)...);
+        if(ec != 0)
+            MIGRAPHX_THROW("Command " + get_command() + " exited with status " +
+                           std::to_string(ec));
+    }
 };
 
 process::process(const std::string& cmd) : impl(std::make_unique<process_impl>())
@@ -95,12 +119,11 @@ process& process::cwd(const fs::path& p)
     return *this;
 }
 
-void process::exec()
+void process::exec() { impl->check_exec(impl->get_command(), redirect_to(std::cout)); }
+
+void process::write(std::function<void(process::writer)> pipe_in)
 {
-    auto ec = migraphx::exec(impl->get_command(), redirect_to(std::cout));
-    if(ec != 0)
-        MIGRAPHX_THROW("Command " + impl->get_command() + " exited with status " +
-                       std::to_string(ec));
+    impl->check_exec(impl->get_command(), std::move(pipe_in));
 }
 
 } // namespace MIGRAPHX_INLINE_NS

src/targets/gpu/CMakeLists.txt

@@ -240,6 +240,7 @@ target_link_libraries(migraphx_gpu PUBLIC migraphx MIOpen roc::rocblas)
 target_link_libraries(migraphx_gpu PRIVATE migraphx_device migraphx_kernels)
 
 add_subdirectory(driver)
+add_subdirectory(hiprtc)
 
 rocm_install_targets(
     TARGETS migraphx_gpu migraphx_device compile_for_gpu

src/targets/gpu/compile_hip.cpp

@@ -32,6 +32,13 @@
 #ifdef MIGRAPHX_USE_HIPRTC
 #include <hip/hiprtc.h>
 #include <migraphx/manage_ptr.hpp>
+#include <migraphx/value.hpp>
+#include <migraphx/tmp_dir.hpp>
+#include <migraphx/dynamic_loader.hpp>
+#include <migraphx/process.hpp>
+#include <migraphx/msgpack.hpp>
+#include <migraphx/serialize.hpp>
+#include <migraphx/file_buffer.hpp>
 #else
 #include <migraphx/compile_src.hpp>
 #include <migraphx/process.hpp>
@@ -63,6 +70,7 @@ void hiprtc_check_error(hiprtcResult err, const std::string& msg, const std::str
         throw make_exception(ctx, hiprtc_error(err, msg));
 }
 
+// NOLINTNEXTLINE
 #define MIGRAPHX_HIPRTC(...) \
     hiprtc_check_error(__VA_ARGS__, #__VA_ARGS__, MIGRAPHX_MAKE_SOURCE_CTX())
 
@@ -110,21 +118,19 @@ struct hiprtc_program
     std::string cpp_src  = "";
     std::string cpp_name = "";
 
-    hiprtc_program(const std::vector<src_file>& srcs)
+    hiprtc_program(std::vector<hiprtc_src_file> srcs)
     {
         for(auto&& src : srcs)
         {
-            std::string content{src.content.first, src.content.second};
-            std::string path = src.path.string();
-            if(src.path.extension().string() == ".cpp")
+            if(ends_with(src.path, ".cpp"))
             {
-                cpp_src  = std::move(content);
-                cpp_name = std::move(path);
+                cpp_src  = std::move(src.content);
+                cpp_name = std::move(src.path);
             }
             else
             {
-                headers.push_back(std::move(content));
-                include_names.push_back(std::move(path));
+                headers.push_back(std::move(src.content));
+                include_names.push_back(std::move(src.path));
             }
         }
         prog = hiprtc_program_create(cpp_src.c_str(),
@@ -134,7 +140,7 @@ struct hiprtc_program
                                      include_names.data());
     }
 
-    void compile(const std::vector<std::string>& options)
+    void compile(const std::vector<std::string>& options) const
     {
         if(enabled(MIGRAPHX_TRACE_HIPRTC{}))
             std::cout << "hiprtc " << join_strings(options, " ") << " " << cpp_name << std::endl;
@@ -175,10 +181,11 @@ struct hiprtc_program
     }
 };
 
-std::vector<std::vector<char>>
-compile_hip_src(const std::vector<src_file>& srcs, std::string params, const std::string& arch)
+std::vector<std::vector<char>> compile_hip_src_with_hiprtc(std::vector<hiprtc_src_file> srcs,
+                                                           std::string params,
+                                                           const std::string& arch)
 {
-    hiprtc_program prog(srcs);
+    hiprtc_program prog(std::move(srcs));
     auto options = split_string(params, ' ');
     options.push_back("-DMIGRAPHX_USE_HIPRTC=1");
     // remove following three compilation flags for HIPRTC once fixes from hipRTC are available in
@@ -205,8 +212,41 @@ compile_hip_src(const std::vector<src_file>& srcs, std::string params, const std
     return {prog.get_code_obj()};
 }
 
+std::vector<std::vector<char>>
+compile_hip_src(const std::vector<src_file>& srcs, std::string params, const std::string& arch)
+{
+    std::vector<hiprtc_src_file> hsrcs{srcs.begin(), srcs.end()};
+    auto p      = dynamic_loader::path(&compile_hip_src_with_hiprtc);
+    auto driver = p.parent_path().parent_path() / "bin" / "migraphx-hiprtc-driver";
+
+    if(fs::exists(driver))
+    {
+        value v;
+        v["srcs"]   = to_value(hsrcs);
+        v["params"] = to_value(params);
+        v["arch"]   = to_value(arch);
+
+        tmp_dir td{};
+        auto out = td.path / "output";
+
+        process(driver.string() + " " + out.string()).write([&](auto writer) {
+            to_msgpack(v, writer);
+        });
+        if(fs::exists(out))
+            return {read_buffer(out.string())};
+    }
+    return compile_hip_src_with_hiprtc(std::move(hsrcs), std::move(params), arch);
+}
+
 #else // MIGRAPHX_USE_HIPRTC
 
+std::vector<std::vector<char>> compile_hip_src_with_hiprtc(std::vector<hiprtc_src_file>, // NOLINT
+                                                           std::string,                  // NOLINT
+                                                           const std::string&)
+{
+    MIGRAPHX_THROW("Not using hiprtc");
+}
+
 bool is_hip_clang_compiler()
 {
     static const auto result = ends_with(MIGRAPHX_STRINGIZE(MIGRAPHX_HIP_COMPILER), "clang++");

src/targets/gpu/hiprtc/CMakeLists.txt

+#####################################################################################
+# The MIT License (MIT)
+#
+# Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+#####################################################################################
+
+add_executable(migraphx-hiprtc-driver
+    main.cpp
+)
+rocm_clang_tidy_check(migraphx-hiprtc-driver)
+target_link_libraries(migraphx-hiprtc-driver PRIVATE migraphx_gpu)
+add_dependencies(migraphx_all_targets migraphx-hiprtc-driver)
+rocm_install_targets(
+    TARGETS migraphx-hiprtc-driver
+)

src/targets/gpu/hiprtc/main.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#include <migraphx/gpu/compile_hip.hpp>
+#include <migraphx/serialize.hpp>
+#include <migraphx/value.hpp>
+#include <migraphx/msgpack.hpp>
+#include <migraphx/file_buffer.hpp>
+#include <migraphx/ranges.hpp>
+#include <iostream>
+#include <cstring>
+
+std::vector<char> read_stdin()
+{
+    std::vector<char> result;
+
+    std::array<char, 1024> buffer;
+    std::size_t len = 0;
+    while((len = std::fread(buffer.data(), 1, buffer.size(), stdin)) > 0)
+    {
+        if(std::ferror(stdin) != 0 and std::feof(stdin) == 0)
+            MIGRAPHX_THROW(std::strerror(errno));
+
+        result.insert(result.end(), buffer.data(), buffer.data() + len);
+    }
+    return result;
+}
+
+int main(int argc, char const* argv[])
+{
+    if(argc < 2 or migraphx::contains({"-h", "--help", "-v", "--version"}, std::string(argv[1])))
+    {
+        std::cout << "USAGE:" << std::endl;
+        std::cout << "    ";
+        std::cout << "Used internally by migraphx to compile hip programs out-of-process."
+                  << std::endl;
+        std::exit(0);
+    }
+    std::string output_name = argv[1];
+
+    auto v = migraphx::from_msgpack(read_stdin());
+    std::vector<migraphx::gpu::hiprtc_src_file> srcs;
+    migraphx::from_value(v.at("srcs"), srcs);
+    auto out = migraphx::gpu::compile_hip_src_with_hiprtc(
+        std::move(srcs), v.at("params").to<std::string>(), v.at("arch").to<std::string>());
+    if(not out.empty())
+        migraphx::write_buffer(output_name, out.front());
+}

src/targets/gpu/include/migraphx/gpu/compile_hip.hpp

@@ -27,6 +27,7 @@
 #include <migraphx/config.hpp>
 #include <migraphx/filesystem.hpp>
 #include <migraphx/compile_src.hpp>
+#include <migraphx/functional.hpp>
 #include <string>
 #include <utility>
 #include <vector>
@@ -35,6 +36,26 @@ namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 namespace gpu {
 
+struct hiprtc_src_file
+{
+    hiprtc_src_file() = default;
+    hiprtc_src_file(const src_file& s)
+        : path(s.path.string()), content(s.content.first, s.content.second)
+    {
+    }
+    std::string path;
+    std::string content;
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.path, "path"), f(self.content, "content"));
+    }
+};
+
+std::vector<std::vector<char>> compile_hip_src_with_hiprtc(std::vector<hiprtc_src_file> srcs,
+                                                           std::string params,
+                                                           const std::string& arch);
+
 std::vector<std::vector<char>>
 compile_hip_src(const std::vector<src_file>& srcs, std::string params, const std::string& arch);