Merge pull request #1130 from lmoriche/master

Add support for code object v3

hipamd/CMakeLists.txt

@@ -271,6 +271,21 @@ if(HIP_PLATFORM STREQUAL "hcc")
         target_link_libraries(hiprtc PUBLIC stdc++fs)
     endif()
 
+
+    if(HIP_PLATFORM STREQUAL "hcc")
+        find_package(amd_comgr REQUIRED CONFIG
+            PATHS
+                /opt/rocm/
+            PATH_SUFFIXES
+                cmake/amd_comgr
+                lib/cmake/amd_comgr
+        )
+        MESSAGE(STATUS "Code Object Manager found at ${amd_comgr_DIR}.")
+    endif()
+
+    target_link_libraries(hip_hcc PRIVATE amd_comgr)
+    target_link_libraries(hip_hcc_static PRIVATE amd_comgr)
+
     string(REPLACE " " ";" HCC_CXX_FLAGS_LIST ${HCC_CXX_FLAGS})
     foreach(TARGET hip_hcc hip_hcc_static)
         target_include_directories(${TARGET} SYSTEM INTERFACE $<INSTALL_INTERFACE:$<INSTALL_PREFIX>/include>;${HSA_PATH}/include)

hipamd/src/hip_module.cpp

@@ -55,6 +55,18 @@ THE SOFTWARE.
 using namespace ELFIO;
 using namespace std;
 
+struct amd_kernel_code_v3_t {
+  uint32_t group_segment_fixed_size;
+  uint32_t private_segment_fixed_size;
+  uint8_t reserved0[8];
+  int64_t kernel_code_entry_byte_offset;
+  uint8_t reserved1[24];
+  uint32_t compute_pgm_rsrc1;
+  uint32_t compute_pgm_rsrc2;
+  uint16_t kernel_code_properties;
+  uint8_t reserved2[6];
+};
+
 // calculate MD5 checksum
 inline std::string checksum(size_t size, const char *source) {
     // FNV-1a hashing, 64-bit version
@@ -206,10 +218,20 @@ hipError_t ihipModuleLaunchKernel(hipFunction_t f, uint32_t globalWorkSizeX,
         aql.grid_size_x = globalWorkSizeX;
         aql.grid_size_y = globalWorkSizeY;
         aql.grid_size_z = globalWorkSizeZ;
-        aql.group_segment_size =
-            f->_header->workgroup_group_segment_byte_size + sharedMemBytes;
-        aql.private_segment_size =
-            f->_header->workitem_private_segment_byte_size;
+        bool is_code_object_v3 = f->_name.find(".kd") != std::string::npos;
+        if (is_code_object_v3) {
+            const auto* header =
+                reinterpret_cast<const amd_kernel_code_v3_t*>(f->_header);
+            aql.group_segment_size =
+                header->group_segment_fixed_size + sharedMemBytes;
+            aql.private_segment_size =
+                header->private_segment_fixed_size;
+        } else {
+            aql.group_segment_size =
+                f->_header->workgroup_group_segment_byte_size + sharedMemBytes;
+            aql.private_segment_size =
+                f->_header->workitem_private_segment_byte_size;
+        }
         aql.kernel_object = f->_object;
         aql.setup = 3 << HSA_KERNEL_DISPATCH_PACKET_SETUP_DIMENSIONS;
         aql.header =
@@ -462,6 +484,12 @@ hipError_t ihipModuleGetFunction(hipFunction_t* func, hipModule_t hmod, const ch
 
     auto kernel = find_kernel_by_name(hmod->executable, name, agent);
 
+    if (kernel.handle == 0u) {
+        std::string name_str(name);
+        name_str.append(".kd");
+        kernel = find_kernel_by_name(hmod->executable, name_str.c_str(), agent);
+    }
+
     if (kernel.handle == 0u) return hipErrorNotFound;
 
     // TODO: refactor the whole ihipThisThat, which is a mess and yields the
@@ -486,7 +514,11 @@ hipError_t hipModuleGetFunctionEx(hipFunction_t* hfunc, hipModule_t hmod,
 }
 
 namespace {
-hipFuncAttributes make_function_attributes(const amd_kernel_code_t& header) {
+const amd_kernel_code_v3_t *header_v3(const ihipModuleSymbol_t& kd) {
+  return reinterpret_cast<const amd_kernel_code_v3_t*>(kd._header);
+}
+
+hipFuncAttributes make_function_attributes(const ihipModuleSymbol_t& kd) {
     hipFuncAttributes r{};
 
     hipDeviceProp_t prop{};
@@ -495,16 +527,31 @@ hipFuncAttributes make_function_attributes(const amd_kernel_code_t& header) {
     //       available per CU, therefore we hardcode it to 64 KiRegisters.
     prop.regsPerBlock = prop.regsPerBlock ? prop.regsPerBlock : 64 * 1024;
 
-    r.localSizeBytes = header.workitem_private_segment_byte_size;
-    r.sharedSizeBytes = header.workgroup_group_segment_byte_size;
+    bool is_code_object_v3 = kd._name.find(".kd") != std::string::npos;
+    if (is_code_object_v3) {
+        r.localSizeBytes = header_v3(kd)->private_segment_fixed_size;
+        r.sharedSizeBytes = header_v3(kd)->group_segment_fixed_size;
+    } else {
+        r.localSizeBytes = kd._header->workitem_private_segment_byte_size;
+        r.sharedSizeBytes = kd._header->workgroup_group_segment_byte_size;
+    }
     r.maxDynamicSharedSizeBytes = prop.sharedMemPerBlock - r.sharedSizeBytes;
-    r.numRegs = header.workitem_vgpr_count;
+    if (is_code_object_v3) {
+        r.numRegs = ((header_v3(kd)->compute_pgm_rsrc1 & 0x3F) + 1) << 2;
+    } else {
+        r.numRegs = kd._header->workitem_vgpr_count;
+    }
     r.maxThreadsPerBlock = r.numRegs ?
         std::min(prop.maxThreadsPerBlock, prop.regsPerBlock / r.numRegs) :
         prop.maxThreadsPerBlock;
-    r.binaryVersion =
-        header.amd_machine_version_major * 10 +
-        header.amd_machine_version_minor;
+    if (is_code_object_v3) {
+        r.binaryVersion = 0; // FIXME: should it be the ISA version or code
+                             //        object format version?
+    } else {
+        r.binaryVersion =
+            kd._header->amd_machine_version_major * 10 +
+            kd._header->amd_machine_version_minor;
+    }
     r.ptxVersion = prop.major * 10 + prop.minor; // HIP currently presents itself as PTX 3.0.
 
     return r;
@@ -520,11 +567,10 @@ hipError_t hipFuncGetAttributes(hipFuncAttributes* attr, const void* func)
 
     auto agent = this_agent();
     auto kd = get_program_state().kernel_descriptor(reinterpret_cast<uintptr_t>(func), agent);
-    const auto header = kd->_header;
 
-    if (!header) throw runtime_error{"Ill-formed Kernel_descriptor."};
+    if (!kd->_header) throw runtime_error{"Ill-formed Kernel_descriptor."};
 
-    *attr = make_function_attributes(*header);
+    *attr = make_function_attributes(*kd);
 
     return hipSuccess;
 }
@@ -555,11 +601,9 @@ hipError_t ihipModuleLoadData(hipModule_t* module, const void* image) {
     (*module)->executable = get_program_state().load_executable(
                                             content.data(), content.size(), (*module)->executable,
                                             this_agent());
-    istringstream elf{content};
-    ELFIO::elfio reader;
-    if (reader.load(elf)) {
-        program_state_impl::read_kernarg_metadata(reader, (*module)->kernargs);
-    }
+
+    std::vector<char> blob(content.cbegin(), content.cend());
+    program_state_impl::read_kernarg_metadata(blob, (*module)->kernargs);
 
     // compute the hash of the code object
     (*module)->hash = checksum(content.length(), content.data());

hipamd/src/program_state.inl

@@ -17,6 +17,7 @@
 #include <hsa/hsa.h>
 #include <hsa/hsa_ext_amd.h>
 #include <hsa/hsa_ven_amd_loader.h>
+#include <amd_comgr.h>
 
 #include <link.h>
 
@@ -540,9 +541,13 @@ public:
 
         std::call_once(functions[agent].first, [this](hsa_agent_t aa) {
             for (auto&& function : get_function_names()) {
-                const auto it = get_kernels(aa).find(function.second);
+                auto it = get_kernels(aa).find(function.second);
 
-                if (it == get_kernels(aa).cend()) continue;
+                if (it == get_kernels(aa).cend()) {
+                    it = get_kernels(aa).find(function.second + ".kd");
+                    if (it == get_kernels(aa).cend())
+                        continue;
+                }
 
                 for (auto&& kernel_symbol : it->second) {
                     functions[aa].second.emplace(
@@ -556,92 +561,172 @@ public:
     }
 
     static
-    std::size_t parse_args(
-            const std::string& metadata,
-            std::size_t f,
-            std::size_t l,
+    std::string metadata_to_string(const amd_comgr_metadata_node_t& md) {
+        std::string str;
+        size_t size;
+
+        if (amd_comgr_get_metadata_string(md, &size, NULL)
+            == AMD_COMGR_STATUS_SUCCESS) {
+            str.resize(size - 1);
+            amd_comgr_get_metadata_string(md, &size, &str[0]);
+        }
+        return str;
+    }
+
+    static
+    void parse_args(
+            const amd_comgr_metadata_node_t& args_md,
+            bool is_code_object_v3,
             std::vector<std::pair<std::size_t, std::size_t>>& size_align) {
-        if (f == l) return f;
-        if (!size_align.empty()) return l;
+        size_t arg_count = 0;
+        if (amd_comgr_get_metadata_list_size(args_md, &arg_count)
+            != AMD_COMGR_STATUS_SUCCESS)
+            return;
 
-        do {
-            static constexpr size_t size_sz{5};
-            f = metadata.find("Size:", f) + size_sz;
+        for (size_t i = 0; i < arg_count; ++i) {
+            amd_comgr_metadata_node_t arg_md;
 
-            if (l <= f) return f;
+            if (amd_comgr_index_list_metadata(args_md, i, &arg_md)
+                != AMD_COMGR_STATUS_SUCCESS)
+                return;
 
-            auto size = std::strtoul(&metadata[f], nullptr, 10);
+            amd_comgr_metadata_node_t arg_size_md;
+            if (amd_comgr_metadata_lookup(arg_md,
+                                          is_code_object_v3 ? ".size" : "Size",
+                                          &arg_size_md)
+                != AMD_COMGR_STATUS_SUCCESS)
+                return;
 
-            static constexpr size_t align_sz{6};
-            f = metadata.find("Align:", f) + align_sz;
+            size_t arg_size = std::stoul(metadata_to_string(arg_size_md));
 
-            char* l{};
-            auto align = std::strtoul(&metadata[f], &l, 10);
+            if (amd_comgr_destroy_metadata(arg_size_md)
+                != AMD_COMGR_STATUS_SUCCESS)
+                return;
 
-            f += (l - &metadata[f]) + 1;
+            size_t arg_align;
 
-            size_align.emplace_back(size, align);
-        } while (true);
+            if (is_code_object_v3) {
+                amd_comgr_metadata_node_t arg_offset_md;
+                if (amd_comgr_metadata_lookup(arg_md, ".offset", &arg_offset_md)
+                    != AMD_COMGR_STATUS_SUCCESS)
+                    return;
+
+                size_t arg_offset
+                    = std::stoul(metadata_to_string(arg_offset_md));
+
+                if (amd_comgr_destroy_metadata(arg_offset_md)
+                    != AMD_COMGR_STATUS_SUCCESS)
+                    return;
+
+                arg_align = 1;
+                while (arg_offset && (arg_offset & 1) == 0) {
+                    arg_offset >>= 1;
+                    arg_align <<= 1;
+                }
+            } else {
+                amd_comgr_metadata_node_t arg_align_md;
+                if (amd_comgr_metadata_lookup(arg_md, "Align", &arg_align_md)
+                    != AMD_COMGR_STATUS_SUCCESS)
+                    return;
+
+                arg_align = std::stoul(metadata_to_string(arg_align_md));
+
+                if (amd_comgr_destroy_metadata(arg_align_md)
+                    != AMD_COMGR_STATUS_SUCCESS)
+                    return;
+            }
+
+            size_align.emplace_back(arg_size, arg_align);
+
+            if (amd_comgr_destroy_metadata(arg_md)
+                != AMD_COMGR_STATUS_SUCCESS)
+                return;
+        }
     }
 
     static
     void read_kernarg_metadata(
-            ELFIO::elfio& reader,
+            const std::vector<char>& blob,
             std::unordered_map<
             std::string,
             std::vector<std::pair<std::size_t, std::size_t>>>& kernargs) {
-        // TODO: this is inefficient.
-        auto it = find_section_if(reader, [](const ELFIO::section* x) {
-                return x->get_type() == SHT_NOTE;
-                });
+        amd_comgr_data_t dataIn;
+        amd_comgr_status_t status;
 
-        if (!it) return;
+        if (amd_comgr_create_data(AMD_COMGR_DATA_KIND_RELOCATABLE, &dataIn)
+            != AMD_COMGR_STATUS_SUCCESS)
+            return;
 
-        const ELFIO::note_section_accessor acc{reader, it};
-        for (decltype(acc.get_notes_num()) i = 0; i != acc.get_notes_num(); ++i) {
-            ELFIO::Elf_Word type{};
-            std::string name{};
-            void* desc{};
-            ELFIO::Elf_Word desc_size{};
+        if (amd_comgr_set_data(dataIn, blob.size(), blob.data())
+            != AMD_COMGR_STATUS_SUCCESS)
+            return;
 
-            acc.get_note(i, type, name, desc, desc_size);
+        amd_comgr_metadata_node_t metadata;
+        if (amd_comgr_get_data_metadata(dataIn, &metadata)
+            != AMD_COMGR_STATUS_SUCCESS)
+            return;
 
-            if (name != "AMD") continue; // TODO: switch to using NT_AMD_AMDGPU_HSA_METADATA.
-
-            std::string tmp{
-                static_cast<char*>(desc), static_cast<char*>(desc) + desc_size};
-
-            auto dx = tmp.find("Kernels:");
-
-            if (dx == std::string::npos) continue;
-
-            static constexpr decltype(tmp.size()) kernels_sz{8};
-            dx += kernels_sz;
-
-            do {
-                dx = tmp.find("Name:", dx);
-
-                if (dx == std::string::npos) break;
-
-                static constexpr decltype(tmp.size()) name_sz{5};
-                dx = tmp.find_first_not_of(" '", dx + name_sz);
-
-                auto fn = tmp.substr(dx, tmp.find_first_of("'\n", dx) - dx);
-                dx += fn.size();
-
-                auto dx1 = tmp.find("CodeProps", dx);
-                dx = tmp.find("Args:", dx);
-
-                if (dx1 < dx) {
-                    dx = dx1;
-                    continue;
-                }
-                if (dx == std::string::npos) break;
-
-                static constexpr decltype(tmp.size()) args_sz{5};
-                dx = parse_args(tmp, dx + args_sz, dx1, kernargs[fn]);
-            } while (true);
+        bool is_code_object_v3 = false;
+        amd_comgr_metadata_node_t kernels_md;
+        if (amd_comgr_metadata_lookup(metadata, "Kernels", &kernels_md)
+            != AMD_COMGR_STATUS_SUCCESS) {
+            if (amd_comgr_metadata_lookup(metadata,
+                                          "amdhsa.kernels",
+                                          &kernels_md)
+                != AMD_COMGR_STATUS_SUCCESS)
+                return;
+            is_code_object_v3 = true;
         }
+
+        size_t kernel_count = 0;
+        if (amd_comgr_get_metadata_list_size(kernels_md, &kernel_count)
+            != AMD_COMGR_STATUS_SUCCESS)
+            return;
+
+        for (size_t i = 0; i < kernel_count; i++) {
+            amd_comgr_metadata_node_t kernel_md;
+
+            if (amd_comgr_index_list_metadata(kernels_md, i, &kernel_md)
+                != AMD_COMGR_STATUS_SUCCESS)
+                continue;
+
+            amd_comgr_metadata_node_t name_md;
+            if (amd_comgr_metadata_lookup(kernel_md,
+                                          is_code_object_v3 ? ".name" : "Name",
+                                          &name_md)
+                != AMD_COMGR_STATUS_SUCCESS)
+                continue;
+
+            std::string kernel_name_str = metadata_to_string(name_md);
+
+            if (amd_comgr_destroy_metadata(name_md)
+                != AMD_COMGR_STATUS_SUCCESS)
+                continue;
+
+            if (is_code_object_v3)
+                kernel_name_str.append(".kd");
+
+
+            amd_comgr_metadata_node_t args_md;
+            if (amd_comgr_metadata_lookup(kernel_md,
+                                          is_code_object_v3 ? ".args" : "Args",
+                                          &args_md)
+                != AMD_COMGR_STATUS_SUCCESS)
+                continue;
+
+            parse_args(args_md, is_code_object_v3, kernargs[kernel_name_str]);
+
+            if (amd_comgr_destroy_metadata(args_md) != AMD_COMGR_STATUS_SUCCESS
+                || amd_comgr_destroy_metadata(kernel_md)
+                   != AMD_COMGR_STATUS_SUCCESS)
+                continue;
+        }
+
+        if (amd_comgr_destroy_metadata(kernels_md) != AMD_COMGR_STATUS_SUCCESS
+            || amd_comgr_destroy_metadata(metadata) != AMD_COMGR_STATUS_SUCCESS)
+            return;
+
+        amd_comgr_release_data(dataIn);
     }
 
     const std::unordered_map<std::string, 
@@ -651,13 +736,7 @@ public:
             for (auto&& name_and_isa_blobs : get_code_object_blobs()) {
                 for (auto&& isa_blobs : name_and_isa_blobs.second) {
                     for (auto&& blob : isa_blobs.second) {
-                        std::stringstream tmp{std::string{blob.cbegin(), blob.cend()}};
-
-                        ELFIO::elfio reader;
-
-                        if (!reader.load(tmp)) continue;
-
-                        read_kernarg_metadata(reader, kernargs.second);
+                        read_kernarg_metadata(blob, kernargs.second);
                     }
                 }
             }
@@ -711,8 +790,11 @@ public:
 
         auto it1 = get_kernargs().find(it->second);
         if (it1 == get_kernargs().end()) {
-            hip_throw(std::runtime_error{
-                      "Missing metadata for __global__ function: " + it->second});
+            it1 = get_kernargs().find(it->second + ".kd");
+            if (it1 == get_kernargs().end()) {
+                hip_throw(std::runtime_error{
+                          "Missing metadata for __global__ function: " + it->second});
+            }
         }
 
         return it1->second;