merge from master

src/hip_hcc_internal.h

@@ -1022,6 +1022,27 @@ inline std::ostream& operator<<(std::ostream& os, const ihipCtx_t* c) {
 namespace hip_internal {
 hipError_t memcpyAsync(void* dst, const void* src, size_t sizeBytes, hipMemcpyKind kind,
                        hipStream_t stream);
+
+hipError_t ihipHostMalloc(TlsData *tls, void** ptr, size_t sizeBytes, unsigned int flags);
+
+hipError_t ihipHostFree(TlsData *tls, void* ptr);
+
+};
+
+#define MAX_COOPERATIVE_GPUs 255
+
+// do not change these two structs without changing the device library
+struct mg_sync {
+    uint w0;
+    uint w1;
+};
+
+struct mg_info {
+    struct mg_sync *mgs;
+    uint grid_id;
+    uint num_grids;
+    ulong prev_sum;
+    ulong all_sum;
 };
 
 //---

src/hip_memory.cpp

@@ -141,6 +141,103 @@ void* allocAndSharePtr(const char* msg, size_t sizeBytes, ihipCtx_t* ctx, bool s
     return ptr;
 }
 
+hipError_t ihipHostMalloc(TlsData *tls, void** ptr, size_t sizeBytes, unsigned int flags) {
+    hipError_t hip_status = hipSuccess;
+
+    if (HIP_SYNC_HOST_ALLOC) {
+        hipDeviceSynchronize();
+    }
+
+    auto ctx = ihipGetTlsDefaultCtx();
+    if ((ctx == nullptr) || (ptr == nullptr)) {
+        hip_status = hipErrorInvalidValue;
+    }
+    else if (sizeBytes == 0) {
+        hip_status = hipSuccess;
+        // TODO - should size of 0 return err or be siliently ignored?
+    } else {
+        unsigned trueFlags = flags;
+        if (flags == hipHostMallocDefault) {
+            // HCC/ROCM provide a modern system with unified memory and should set both of these
+            // flags by default:
+            trueFlags = hipHostMallocMapped | hipHostMallocPortable;
+        }
+
+
+        const unsigned supportedFlags = hipHostMallocPortable | hipHostMallocMapped |
+                                        hipHostMallocWriteCombined | hipHostMallocCoherent |
+                                        hipHostMallocNonCoherent;
+
+
+        const unsigned coherencyFlags = hipHostMallocCoherent | hipHostMallocNonCoherent;
+
+        if ((flags & ~supportedFlags) || ((flags & coherencyFlags) == coherencyFlags)) {
+            *ptr = nullptr;
+            // can't specify unsupported flags, can't specify both Coherent + NonCoherent
+            hip_status = hipErrorInvalidValue;
+        } else {
+            auto device = ctx->getWriteableDevice();
+#if (__hcc_workweek__ >= 19115)
+            //Avoid mapping host pinned memory to all devices by HCC
+            unsigned amFlags = amHostUnmapped;
+#else
+            unsigned amFlags = 0;
+#endif
+            if (flags & hipHostMallocCoherent) {
+                amFlags |= amHostCoherent;
+            } else if (flags & hipHostMallocNonCoherent) {
+                amFlags |= amHostNonCoherent;
+            } else {
+                // depends on env variables:
+                amFlags |= HIP_HOST_COHERENT ? amHostCoherent : amHostNonCoherent;
+            }
+
+
+            *ptr = hip_internal::allocAndSharePtr(
+                (amFlags & amHostCoherent) ? "finegrained_host" : "pinned_host", sizeBytes, ctx,
+                true  /*shareWithAll*/, amFlags, flags, 0);
+
+            if (sizeBytes && (*ptr == NULL)) {
+                hip_status = hipErrorMemoryAllocation;
+            }
+        }
+    }
+
+    if (HIP_SYNC_HOST_ALLOC) {
+        hipDeviceSynchronize();
+    }
+    return hip_status;
+}
+
+hipError_t ihipHostFree(TlsData *tls, void* ptr) {
+
+    // Synchronize to ensure all work has finished.
+    ihipGetTlsDefaultCtx()->locked_waitAllStreams();  // ignores non-blocking streams, this waits
+                                                      // for all activity to finish.
+
+    hipError_t hipStatus = hipErrorInvalidValue;
+    if (ptr) {
+        hc::accelerator acc;
+#if (__hcc_workweek__ >= 17332)
+        hc::AmPointerInfo amPointerInfo(NULL, NULL, NULL, 0, acc, 0, 0);
+#else
+        hc::AmPointerInfo amPointerInfo(NULL, NULL, 0, acc, 0, 0);
+#endif
+        am_status_t status = hc::am_memtracker_getinfo(&amPointerInfo, ptr);
+        if (status == AM_SUCCESS) {
+            if (amPointerInfo._hostPointer == ptr) {
+                hc::am_free(ptr);
+                hipStatus = hipSuccess;
+            }
+        }
+    } else {
+        // free NULL pointer succeeds and is common technique to initialize runtime
+        hipStatus = hipSuccess;
+    }
+
+    return hipStatus;
+}
+
 
 }  // end namespace hip_internal
 
@@ -301,79 +398,12 @@ hipError_t hipExtMallocWithFlags(void** ptr, size_t sizeBytes, unsigned int flag
     return ihipLogStatus(hip_status);
 }
 
-hipError_t ihipHostMalloc(TlsData *tls, void** ptr, size_t sizeBytes, unsigned int flags) {
-    hipError_t hip_status = hipSuccess;
-
-    if (HIP_SYNC_HOST_ALLOC) {
-        hipDeviceSynchronize();
-    }
-
-    auto ctx = ihipGetTlsDefaultCtx();
-    if ((ctx == nullptr) || (ptr == nullptr)) {
-        hip_status = hipErrorInvalidValue;
-    }
-    else if (sizeBytes == 0) {
-        hip_status = hipSuccess;
-        // TODO - should size of 0 return err or be siliently ignored?
-    } else {
-        unsigned trueFlags = flags;
-        if (flags == hipHostMallocDefault) {
-            // HCC/ROCM provide a modern system with unified memory and should set both of these
-            // flags by default:
-            trueFlags = hipHostMallocMapped | hipHostMallocPortable;
-        }
-
-
-        const unsigned supportedFlags = hipHostMallocPortable | hipHostMallocMapped |
-                                        hipHostMallocWriteCombined | hipHostMallocCoherent |
-                                        hipHostMallocNonCoherent;
-
-
-        const unsigned coherencyFlags = hipHostMallocCoherent | hipHostMallocNonCoherent;
-
-        if ((flags & ~supportedFlags) || ((flags & coherencyFlags) == coherencyFlags)) {
-            *ptr = nullptr;
-            // can't specify unsupported flags, can't specify both Coherent + NonCoherent
-            hip_status = hipErrorInvalidValue;
-        } else {
-            auto device = ctx->getWriteableDevice();
-#if (__hcc_workweek__ >= 19115)
-            //Avoid mapping host pinned memory to all devices by HCC
-            unsigned amFlags = amHostUnmapped;
-#else
-            unsigned amFlags = 0;
-#endif
-            if (flags & hipHostMallocCoherent) {
-                amFlags |= amHostCoherent;
-            } else if (flags & hipHostMallocNonCoherent) {
-                amFlags |= amHostNonCoherent;
-            } else {
-                // depends on env variables:
-                amFlags |= HIP_HOST_COHERENT ? amHostCoherent : amHostNonCoherent;
-            }
-
-
-            *ptr = hip_internal::allocAndSharePtr(
-                (amFlags & amHostCoherent) ? "finegrained_host" : "pinned_host", sizeBytes, ctx,
-                true  /*shareWithAll*/, amFlags, flags, 0);
-
-            if (sizeBytes && (*ptr == NULL)) {
-                hip_status = hipErrorMemoryAllocation;
-            }
-        }
-    }
-
-    if (HIP_SYNC_HOST_ALLOC) {
-        hipDeviceSynchronize();
-    }
-    return hip_status;
-}
 
 hipError_t hipHostMalloc(void** ptr, size_t sizeBytes, unsigned int flags) {
     HIP_INIT_SPECIAL_API(hipHostMalloc, (TRACE_MEM), ptr, sizeBytes, flags);
     HIP_SET_DEVICE();
     hipError_t hip_status = hipSuccess;
-    hip_status = ihipHostMalloc(tls, ptr, sizeBytes, flags);
+    hip_status = hip_internal::ihipHostMalloc(tls, ptr, sizeBytes, flags);
     return ihipLogStatus(hip_status);
 }
 
@@ -384,7 +414,7 @@ hipError_t hipMallocManaged(void** devPtr, size_t size, unsigned int flags) {
     if(flags != hipMemAttachGlobal)
         hip_status = hipErrorInvalidValue;
     else
-        hip_status = ihipHostMalloc(tls, devPtr, size, hipHostMallocDefault);
+        hip_status = hip_internal::ihipHostMalloc(tls, devPtr, size, hipHostMallocDefault);
     return ihipLogStatus(hip_status);
 }
 
@@ -1935,15 +1965,15 @@ hipError_t hipMemset2DAsync(void* dst, size_t pitch, int value, size_t width, si
     return ihipLogStatus(e);
 };
 
-hipError_t hipMemsetD8(hipDeviceptr_t dst, unsigned char value, size_t sizeBytes) {
-    HIP_INIT_SPECIAL_API(hipMemsetD8, (TRACE_MCMD), dst, value, sizeBytes);
+hipError_t hipMemsetD8(hipDeviceptr_t dst, unsigned char value, size_t count) {
+    HIP_INIT_SPECIAL_API(hipMemsetD8, (TRACE_MCMD), dst, value, count);
 
     hipError_t e = hipSuccess;
 
     hipStream_t stream = hipStreamNull;
     stream = ihipSyncAndResolveStream(stream);
     if (stream) {
-        e = ihipMemset(dst, value, sizeBytes, stream, ihipMemsetDataTypeChar);
+        e = ihipMemset(dst, value, count, stream, ihipMemsetDataTypeChar);
         stream->locked_wait();
     } else {
         e = hipErrorInvalidValue;
@@ -1951,23 +1981,23 @@ hipError_t hipMemsetD8(hipDeviceptr_t dst, unsigned char value, size_t sizeBytes
     return ihipLogStatus(e);
 }
 
-hipError_t hipMemsetD8Async(hipDeviceptr_t dst, unsigned char value, size_t sizeBytes , hipStream_t stream ) {
-    HIP_INIT_SPECIAL_API(hipMemsetD8Async, (TRACE_MCMD), dst, value, sizeBytes, stream);
+hipError_t hipMemsetD8Async(hipDeviceptr_t dst, unsigned char value, size_t count , hipStream_t stream ) {
+    HIP_INIT_SPECIAL_API(hipMemsetD8Async, (TRACE_MCMD), dst, value, count, stream);
 
     stream = ihipSyncAndResolveStream(stream);
     if (stream) {
-        return ihipLogStatus(ihipMemset(dst, value, sizeBytes, stream, ihipMemsetDataTypeChar));
+        return ihipLogStatus(ihipMemset(dst, value, count, stream, ihipMemsetDataTypeChar));
     } else {
         return ihipLogStatus(hipErrorInvalidValue);
     }
 }
 
-hipError_t hipMemsetD16(hipDeviceptr_t dst, unsigned short value, size_t sizeBytes){
-    HIP_INIT_SPECIAL_API(hipMemsetD16, (TRACE_MCMD), dst, value, sizeBytes);
+hipError_t hipMemsetD16(hipDeviceptr_t dst, unsigned short value, size_t count){
+    HIP_INIT_SPECIAL_API(hipMemsetD16, (TRACE_MCMD), dst, value, count);
     hipError_t e = hipSuccess;
     hipStream_t stream = ihipSyncAndResolveStream(hipStreamNull);
     if (stream) {
-        e = ihipMemset(dst, value, sizeBytes, stream, ihipMemsetDataTypeShort);
+        e = ihipMemset(dst, value, count, stream, ihipMemsetDataTypeShort);
         if(hipSuccess == e)
             stream->locked_wait();
     } else {
@@ -1976,12 +2006,12 @@ hipError_t hipMemsetD16(hipDeviceptr_t dst, unsigned short value, size_t sizeByt
     return ihipLogStatus(e);
 }
 
-hipError_t hipMemsetD16Async(hipDeviceptr_t dst, unsigned short value, size_t sizeBytes, hipStream_t stream ){
-    HIP_INIT_SPECIAL_API(hipMemsetD16Async, (TRACE_MCMD), dst, value, sizeBytes, stream);
+hipError_t hipMemsetD16Async(hipDeviceptr_t dst, unsigned short value, size_t count, hipStream_t stream ){
+    HIP_INIT_SPECIAL_API(hipMemsetD16Async, (TRACE_MCMD), dst, value, count, stream);
 
     stream = ihipSyncAndResolveStream(stream);
     if (stream) {
-        return ihipLogStatus(ihipMemset(dst, value, sizeBytes, stream, ihipMemsetDataTypeShort));
+        return ihipLogStatus(ihipMemset(dst, value, count, stream, ihipMemsetDataTypeShort));
     } else {
         return ihipLogStatus(hipErrorInvalidValue);
     }
@@ -2146,30 +2176,8 @@ hipError_t hipFree(void* ptr) {
 hipError_t hipHostFree(void* ptr) {
     HIP_INIT_SPECIAL_API(hipHostFree, (TRACE_MEM), ptr);
 
-    // Synchronize to ensure all work has finished.
-    ihipGetTlsDefaultCtx()->locked_waitAllStreams();  // ignores non-blocking streams, this waits
-                                                      // for all activity to finish.
-
-
-    hipError_t hipStatus = hipErrorInvalidValue;
-    if (ptr) {
-        hc::accelerator acc;
-#if (__hcc_workweek__ >= 17332)
-        hc::AmPointerInfo amPointerInfo(NULL, NULL, NULL, 0, acc, 0, 0);
-#else
-        hc::AmPointerInfo amPointerInfo(NULL, NULL, 0, acc, 0, 0);
-#endif
-        am_status_t status = hc::am_memtracker_getinfo(&amPointerInfo, ptr);
-        if (status == AM_SUCCESS) {
-            if (amPointerInfo._hostPointer == ptr) {
-                hc::am_free(ptr);
-                hipStatus = hipSuccess;
-            }
-        }
-    } else {
-        // free NULL pointer succeeds and is common technique to initialize runtime
-        hipStatus = hipSuccess;
-    }
+    hipError_t hipStatus = hipSuccess;
+    hipStatus = hip_internal::ihipHostFree(tls, ptr);
 
     return ihipLogStatus(hipStatus);
 };

src/hip_module.cpp

@@ -109,6 +109,7 @@ struct ihipModuleSymbol_t {
     amd_kernel_code_t const* _header{};
     string _name;  // TODO - review for performance cost.  Name is just used for debug.
     vector<pair<size_t, size_t>> _kernarg_layout{};
+    bool _is_code_object_v3{};
 };
 
 template <>
@@ -137,7 +138,8 @@ hipError_t ihipModuleLaunchKernel(TlsData *tls, hipFunction_t f, uint32_t global
                                   uint32_t localWorkSizeX, uint32_t localWorkSizeY,
                                   uint32_t localWorkSizeZ, size_t sharedMemBytes,
                                   hipStream_t hStream, void** kernelParams, void** extra,
-                                  hipEvent_t startEvent, hipEvent_t stopEvent, uint32_t flags, bool isStreamLocked = 0) {
+                                  hipEvent_t startEvent, hipEvent_t stopEvent, uint32_t flags, bool isStreamLocked = 0,
+                                  void** impCoopParams = 0) {
     using namespace hip_impl;
 
     auto ctx = ihipGetTlsDefaultCtx();
@@ -181,10 +183,17 @@ hipError_t ihipModuleLaunchKernel(TlsData *tls, hipFunction_t f, uint32_t global
             return hipErrorInvalidValue;
         }
 
-        // Insert 48-bytes at the end for implicit kernel arguments and fill with value zero.
+        // Insert 56-bytes at the end for implicit kernel arguments and fill with value zero.
         size_t padSize = (~kernargs.size() + 1) & (HIP_IMPLICIT_KERNARG_ALIGNMENT - 1);
         kernargs.insert(kernargs.end(), padSize + HIP_IMPLICIT_KERNARG_SIZE, 0);
 
+        if (impCoopParams) {
+            const auto p{static_cast<const char*>(*impCoopParams)};
+            // The sixth index is for multi-grid synchronization
+            kernargs.insert((kernargs.cend() - padSize - HIP_IMPLICIT_KERNARG_SIZE) + 6 * HIP_IMPLICIT_KERNARG_ALIGNMENT,
+                            p, p + HIP_IMPLICIT_KERNARG_ALIGNMENT);
+        }
+
         /*
           Kernel argument preparation.
         */
@@ -208,8 +217,7 @@ hipError_t ihipModuleLaunchKernel(TlsData *tls, hipFunction_t f, uint32_t global
         aql.grid_size_x = globalWorkSizeX;
         aql.grid_size_y = globalWorkSizeY;
         aql.grid_size_z = globalWorkSizeZ;
-        bool is_code_object_v3 = f->_name.find(".kd") != std::string::npos;
-        if (is_code_object_v3) {
+        if (f->_is_code_object_v3) {
             const auto* header =
                 reinterpret_cast<const amd_kernel_code_v3_t*>(f->_header);
             aql.group_segment_size =
@@ -449,6 +457,10 @@ hipError_t hipLaunchCooperativeKernel(const void* f, dim3 gridDim,
         return ihipLogStatus(hipErrorLaunchFailure);
     }
 
+    size_t impCoopArg = 1;
+    void* impCoopParams[1];
+    impCoopParams[0] = &impCoopArg;
+
     // launch the main kernel
     result = ihipModuleLaunchKernel(tls, kd,
             gridDim.x * blockDimX.x,
@@ -456,7 +468,7 @@ hipError_t hipLaunchCooperativeKernel(const void* f, dim3 gridDim,
             gridDim.z * blockDimX.z,
             blockDimX.x, blockDimX.y, blockDimX.z,
             sharedMemBytes, stream, kernelParams, nullptr, nullptr,
-            nullptr, 0, true);
+            nullptr, 0, true, impCoopParams);
 
     stream->criticalData().unlock();
 #if (__hcc_workweek__ >= 19213)
@@ -472,7 +484,7 @@ hipError_t hipLaunchCooperativeKernelMultiDevice(hipLaunchParams* launchParamsLi
     HIP_INIT_API(hipLaunchCooperativeKernelMultiDevice, launchParamsList, numDevices, flags);
     hipError_t result;
 
-    if (numDevices > g_deviceCnt || launchParamsList == nullptr) {
+    if (numDevices > g_deviceCnt || launchParamsList == nullptr || numDevices > MAX_COOPERATIVE_GPUs) {
         return ihipLogStatus(hipErrorInvalidValue);
     }
 
@@ -523,6 +535,32 @@ hipError_t hipLaunchCooperativeKernelMultiDevice(hipLaunchParams* launchParamsLi
                 kargs.getHandle());
     }
 
+    mg_sync *mg_sync_ptr = 0;
+    mg_info *mg_info_ptr[MAX_COOPERATIVE_GPUs] = {0};
+
+    result = hip_internal::ihipHostMalloc(tls, (void **)&mg_sync_ptr, sizeof(mg_sync), hipHostMallocDefault);
+    if (result != hipSuccess) {
+        return ihipLogStatus(hipErrorInvalidValue);
+    }
+    mg_sync_ptr->w0 = 0;
+    mg_sync_ptr->w1 = 0;
+
+    uint all_sum = 0;
+    for (int i = 0; i < numDevices; ++i) {
+        result = hip_internal::ihipHostMalloc(tls, (void **)&mg_info_ptr[i], sizeof(mg_info), hipHostMallocDefault);
+        if (result != hipSuccess) {
+            hip_internal::ihipHostFree(tls, mg_sync_ptr);
+            for (int j = 0; j < i; ++j) {
+                hip_internal::ihipHostFree(tls, mg_info_ptr[j]);
+            }
+            return ihipLogStatus(hipErrorInvalidValue);
+        }
+        // calculate the sum of sizes of all grids
+        const hipLaunchParams& lp = launchParamsList[i];
+        all_sum += lp.blockDim.x * lp.blockDim.y * lp.blockDim.z *
+                   lp.gridDim.x  * lp.gridDim.y  * lp.gridDim.z;
+    }
+
     // lock all streams before launching the blit kernels for initializing the GWS and main kernels to each device
     for (int i = 0; i < numDevices; ++i) {
         LockedAccessor_StreamCrit_t streamCrit(launchParamsList[i].stream->criticalData(), false);
@@ -531,7 +569,7 @@ hipError_t hipLaunchCooperativeKernelMultiDevice(hipLaunchParams* launchParamsLi
 #endif
     }
 
-    // launch the init_gws kernel to initialize the GWS followed by launching the main kernels for each device
+    // launch the init_gws kernel to initialize the GWS for each device
     for (int i = 0; i < numDevices; ++i) {
         const hipLaunchParams& lp = launchParamsList[i];
 
@@ -549,8 +587,32 @@ hipError_t hipLaunchCooperativeKernelMultiDevice(hipLaunchParams* launchParamsLi
                 launchParamsList[j].stream->criticalData()._av.release_locked_hsa_queue();
 #endif
             }
+            hip_internal::ihipHostFree(tls, mg_sync_ptr);
+            for (int j = 0; j < numDevices; ++j) {
+                hip_internal::ihipHostFree(tls, mg_info_ptr[j]);
+            }
+
             return ihipLogStatus(hipErrorLaunchFailure);
         }
+    }
+
+    void* impCoopParams[1];
+    ulong prev_sum = 0;
+    // launch the main kernels for each device
+    for (int i = 0; i < numDevices; ++i) {
+        const hipLaunchParams& lp = launchParamsList[i];
+
+        //initialize and setup the implicit kernel argument for multi-grid sync
+        mg_info_ptr[i]->mgs       = mg_sync_ptr;
+        mg_info_ptr[i]->grid_id   = i;
+        mg_info_ptr[i]->num_grids = numDevices;
+        mg_info_ptr[i]->all_sum   = all_sum;
+        mg_info_ptr[i]->prev_sum  = prev_sum;
+        prev_sum += lp.blockDim.x * lp.blockDim.y * lp.blockDim.z *
+                    lp.gridDim.x  * lp.gridDim.y  * lp.gridDim.z;
+
+
+        impCoopParams[0] = &mg_info_ptr[i];
 
         result = ihipModuleLaunchKernel(tls, kds[i],
                 lp.gridDim.x * lp.blockDim.x,
@@ -559,7 +621,23 @@ hipError_t hipLaunchCooperativeKernelMultiDevice(hipLaunchParams* launchParamsLi
                 lp.blockDim.x, lp.blockDim.y,
                 lp.blockDim.z, lp.sharedMem,
                 lp.stream, lp.args, nullptr, nullptr, nullptr, 0,
-                true);
+                true, impCoopParams);
+
+        if (result != hipSuccess) {
+            for (int j = 0; j < numDevices; ++j) {
+                launchParamsList[j].stream->criticalData().unlock();
+#if (__hcc_workweek__ >= 19213)
+                launchParamsList[j].stream->criticalData()._av.release_locked_hsa_queue();
+#endif
+            }
+            hip_internal::ihipHostFree(tls, mg_sync_ptr);
+            for (int j = 0; j < numDevices; ++j) {
+                hip_internal::ihipHostFree(tls, mg_info_ptr[j]);
+            }
+
+            return ihipLogStatus(hipErrorLaunchFailure);
+        }
+
     }
 
     // unlock all streams
@@ -573,6 +651,11 @@ hipError_t hipLaunchCooperativeKernelMultiDevice(hipLaunchParams* launchParamsLi
     free(gwsKds);
     free(kds);
 
+    hip_internal::ihipHostFree(tls, mg_sync_ptr);
+    for (int j = 0; j < numDevices; ++j) {
+        hip_internal::ihipHostFree(tls, mg_info_ptr[j]);
+    }
+
     return ihipLogStatus(result);
 }
 
@@ -977,31 +1060,24 @@ hipFuncAttributes make_function_attributes(TlsData *tls, const ihipModuleSymbol_
     //       available per CU, therefore we hardcode it to 64 KiRegisters.
     prop.regsPerBlock = prop.regsPerBlock ? prop.regsPerBlock : 64 * 1024;
 
-    bool is_code_object_v3 = kd._name.find(".kd") != std::string::npos;
-    if (is_code_object_v3) {
+    if (kd._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;
-    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;
-    if (is_code_object_v3) {
         r.binaryVersion = 0; // FIXME: should it be the ISA version or code
                              //        object format version?
     } else {
+        r.localSizeBytes = kd._header->workitem_private_segment_byte_size;
+        r.sharedSizeBytes = kd._header->workgroup_group_segment_byte_size;
+        r.numRegs = kd._header->workitem_vgpr_count;
         r.binaryVersion =
             kd._header->amd_machine_version_major * 10 +
             kd._header->amd_machine_version_minor;
     }
+    r.maxDynamicSharedSizeBytes = prop.sharedMemPerBlock - r.sharedSizeBytes;
+    r.maxThreadsPerBlock = r.numRegs ?
+        std::min(prop.maxThreadsPerBlock, prop.regsPerBlock / r.numRegs) :
+        prop.maxThreadsPerBlock;
     r.ptxVersion = prop.major * 10 + prop.minor; // HIP currently presents itself as PTX 3.0.
 
     return r;
@@ -1099,8 +1175,7 @@ hipError_t ihipModuleLoadData(TlsData *tls, hipModule_t* module, const void* ima
                                             content.data(), content.size(), (*module)->executable,
                                             this_agent());
 
-    std::vector<char> blob(content.cbegin(), content.cend());
-    program_state_impl::read_kernarg_metadata(blob, (*module)->kernargs);
+    program_state_impl::read_kernarg_metadata(content, (*module)->kernargs);
 
     // compute the hash of the code object
     (*module)->hash = checksum(content.length(), content.data());
@@ -1152,8 +1227,7 @@ hipError_t hipModuleGetTexRef(textureReference** texRef, hipModule_t hmod, const
 
 void getGprsLdsUsage(hipFunction_t f, size_t* usedVGPRS, size_t* usedSGPRS, size_t* usedLDS)
 {
-    bool is_code_object_v3 = f->_name.find(".kd") != std::string::npos;
-    if (is_code_object_v3) {
+    if (f->_is_code_object_v3) {
         const auto header = reinterpret_cast<const amd_kernel_code_v3_t*>(f->_header);
         // GRANULATED_WAVEFRONT_VGPR_COUNT is specified in 0:5 bits of COMPUTE_PGM_RSRC1
         // the granularity for gfx6-gfx9 is max(0, ceil(vgprs_used / 4) - 1)
@@ -1174,9 +1248,9 @@ void getGprsLdsUsage(hipFunction_t f, size_t* usedVGPRS, size_t* usedSGPRS, size
     }
 }
 
-hipError_t ihipOccupancyMaxPotentialBlockSize(TlsData *tls, uint32_t* gridSize, uint32_t* blockSize,
-                                              hipFunction_t f, size_t dynSharedMemPerBlk,
-                                              uint32_t blockSizeLimit)
+hipError_t ihipOccupancyMaxPotentialBlockSize(TlsData *tls, int* gridSize, int* blockSize,
+                                              hipFunction_t f, size_t dynamicSMemSize,
+                                              int blockSizeLimit)
 {
     using namespace hip_impl;
 
@@ -1257,7 +1331,7 @@ hipError_t ihipOccupancyMaxPotentialBlockSize(TlsData *tls, uint32_t* gridSize,
         }
         else {
             size_t availableSharedMemPerCU = prop.maxSharedMemoryPerMultiProcessor;
-            size_t workgroupPerCU = availableSharedMemPerCU / (usedLDS + dynSharedMemPerBlk);
+            size_t workgroupPerCU = availableSharedMemPerCU / (usedLDS + dynamicSMemSize);
             wavefrontsLDS = min(workgroupPerCU, maxWorkgroupPerCU) * wavefrontsPerWG;
         }
 
@@ -1286,18 +1360,19 @@ hipError_t ihipOccupancyMaxPotentialBlockSize(TlsData *tls, uint32_t* gridSize,
     return hipSuccess;
 }
 
-hipError_t hipOccupancyMaxPotentialBlockSize(uint32_t* gridSize, uint32_t* blockSize,
-                                             hipFunction_t f, size_t dynSharedMemPerBlk,
-                                             uint32_t blockSizeLimit)
+hipError_t hipOccupancyMaxPotentialBlockSize(int* gridSize, int* blockSize,
+                                             const void* f, size_t dynamicSMemSize,
+                                             int blockSizeLimit)
 {
-    HIP_INIT_API(hipOccupancyMaxPotentialBlockSize, gridSize, blockSize, f, dynSharedMemPerBlk, blockSizeLimit);
-
+    HIP_INIT_API(hipOccupancyMaxPotentialBlockSize, gridSize, blockSize, f, dynamicSMemSize, blockSizeLimit);
+    auto F = hip_impl::get_program_state().kernel_descriptor((std::uintptr_t)(f),
+                                                   hip_impl::target_agent(0));
     return ihipLogStatus(ihipOccupancyMaxPotentialBlockSize(tls,
-        gridSize, blockSize, f, dynSharedMemPerBlk, blockSizeLimit));
+        gridSize, blockSize, F, dynamicSMemSize, blockSizeLimit));
 }
 
 hipError_t ihipOccupancyMaxActiveBlocksPerMultiprocessor(
-   TlsData *tls, uint32_t* numBlocks, hipFunction_t f, uint32_t blockSize, size_t dynSharedMemPerBlk)
+   TlsData *tls, int* numBlocks, hipFunction_t f, int blockSize, size_t dynamicSMemSize)
 {
     using namespace hip_impl;
 
@@ -1326,45 +1401,50 @@ hipError_t ihipOccupancyMaxActiveBlocksPerMultiprocessor(
     size_t numWavefronts = (blockSize + wavefrontSize - 1) / wavefrontSize;
 
     size_t availableVGPRs = (prop.regsPerBlock / wavefrontSize / simdPerCU);
-    size_t vgprs_alu_occupancy = simdPerCU * std::min(maxWavesPerSimd, availableVGPRs / usedVGPRS);
+    size_t vgprs_alu_occupancy = simdPerCU * (usedVGPRS == 0 ? maxWavesPerSimd
+        : std::min(maxWavesPerSimd, availableVGPRs / usedVGPRS));
 
     // Calculate blocks occupancy per CU based on VGPR usage
     *numBlocks = vgprs_alu_occupancy / numWavefronts;
 
     const size_t availableSGPRs = (prop.gcnArch < 800) ? 512 : 800;
-    size_t sgprs_alu_occupancy = simdPerCU * ((usedSGPRS == 0) ? maxWavesPerSimd
+    size_t sgprs_alu_occupancy = simdPerCU * (usedSGPRS == 0 ? maxWavesPerSimd
         : std::min(maxWavesPerSimd, availableSGPRs / usedSGPRS));
 
     // Calculate blocks occupancy per CU based on SGPR usage
-    *numBlocks = std::min(*numBlocks, (uint32_t) (sgprs_alu_occupancy / numWavefronts)); 
+    *numBlocks = std::min(*numBlocks, (int) (sgprs_alu_occupancy / numWavefronts));
 
-    size_t total_used_lds = usedLDS + dynSharedMemPerBlk;
+    size_t total_used_lds = usedLDS + dynamicSMemSize;
     if (total_used_lds != 0) {
       // Calculate LDS occupacy per CU. lds_per_cu / (static_lsd + dynamic_lds)
       size_t lds_occupancy = prop.maxSharedMemoryPerMultiProcessor / total_used_lds;
-      *numBlocks = std::min(*numBlocks, (uint32_t) lds_occupancy);
+      *numBlocks = std::min(*numBlocks, (int) lds_occupancy);
     }
 
     return hipSuccess;
 }
 
 hipError_t hipOccupancyMaxActiveBlocksPerMultiprocessor(
-   uint32_t* numBlocks, hipFunction_t f, uint32_t blockSize, size_t dynSharedMemPerBlk)
+   int* numBlocks, const void* f, int blockSize, size_t dynamicSMemSize)
 {
-    HIP_INIT_API(hipOccupancyMaxActiveBlocksPerMultiprocessor, numBlocks, f, blockSize, dynSharedMemPerBlk);
+    HIP_INIT_API(hipOccupancyMaxActiveBlocksPerMultiprocessor, numBlocks, f, blockSize, dynamicSMemSize);
 
+    auto F = hip_impl::get_program_state().kernel_descriptor((std::uintptr_t)(f),
+                                                   hip_impl::target_agent(0));
     return ihipLogStatus(ihipOccupancyMaxActiveBlocksPerMultiprocessor(
-        tls, numBlocks, f, blockSize, dynSharedMemPerBlk));
+        tls, numBlocks, F, blockSize, dynamicSMemSize));
 }
 
 hipError_t hipOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(
-   uint32_t* numBlocks, hipFunction_t f, uint32_t  blockSize, size_t dynSharedMemPerBlk,
+   int* numBlocks, const void* f, int  blockSize, size_t dynamicSMemSize,
    unsigned int flags)
 {
-    HIP_INIT_API(hipOccupancyMaxActiveBlocksPerMultiprocessorWithFlags, numBlocks, f, blockSize, dynSharedMemPerBlk, flags);
+    HIP_INIT_API(hipOccupancyMaxActiveBlocksPerMultiprocessorWithFlags, numBlocks, f, blockSize, dynamicSMemSize, flags);
 
+    auto F = hip_impl::get_program_state().kernel_descriptor((std::uintptr_t)(f),
+                                                   hip_impl::target_agent(0));
     return ihipLogStatus(ihipOccupancyMaxActiveBlocksPerMultiprocessor(
-        tls, numBlocks, f, blockSize, dynSharedMemPerBlk));
+        tls, numBlocks, F, blockSize, dynamicSMemSize));
 }
 
 hipError_t hipLaunchKernel(

src/program_state.inl

@@ -89,9 +89,10 @@ struct Symbol {
 
 class Kernel_descriptor {
     std::uint64_t kernel_object_{};
-    amd_kernel_code_t const* kernel_header_{nullptr};
-    std::string name_{};
+    amd_kernel_code_t const* header_{};
+    std::string name_;
     std::vector<std::pair<std::size_t, std::size_t>> kernarg_layout_{};
+    bool is_code_object_v3_{};
 public:
     Kernel_descriptor() = default;
     Kernel_descriptor(
@@ -101,7 +102,8 @@ public:
         :
         kernel_object_{kernel_object},
         name_{name},
-        kernarg_layout_{std::move(kernarg_layout)}
+        kernarg_layout_{std::move(kernarg_layout)},
+        is_code_object_v3_{name.find(".kd") != std::string::npos}
     {
         bool supported{false};
         std::uint16_t min_v{UINT16_MAX};
@@ -123,7 +125,7 @@ public:
 
         r = tbl.hsa_ven_amd_loader_query_host_address(
             reinterpret_cast<void*>(kernel_object_),
-            reinterpret_cast<const void**>(&kernel_header_));
+            reinterpret_cast<const void**>(&header_));
 
         if (r != HSA_STATUS_SUCCESS) return;
     }
@@ -149,7 +151,7 @@ public:
             std::string,
                 std::unordered_map<
                     hsa_isa_t,
-                    std::vector<std::vector<char>>>>> code_object_blobs;
+                    std::vector<std::string>>>> code_object_blobs;
 
     std::pair<
         std::once_flag,
@@ -213,7 +215,7 @@ public:
         std::string,
             std::unordered_map<
                 hsa_isa_t,
-                std::vector<std::vector<char>>>>& get_code_object_blobs() {
+                std::vector<std::string>>>& get_code_object_blobs() {
 
         std::call_once(code_object_blobs.first, [this]() {
             dl_iterate_phdr([](dl_phdr_info* info, std::size_t, void* p) {
@@ -584,6 +586,68 @@ public:
         return functions[agent].second;
     }
 
+    static
+    std::size_t parse_args_v2(
+            const std::string& metadata,
+            std::size_t f,
+            std::size_t l,
+            std::vector<std::pair<std::size_t, std::size_t>>& size_align) {
+        if (f == l) return f;
+        if (!size_align.empty()) return l;
+
+        do {
+            static constexpr size_t size_sz{5};
+            f = metadata.find("Size:", f) + size_sz;
+
+            if (l <= f) return f;
+
+            auto size = std::strtoul(&metadata[f], nullptr, 10);
+
+            static constexpr size_t align_sz{6};
+            f = metadata.find("Align:", f) + align_sz;
+
+            char* l{};
+            auto align = std::strtoul(&metadata[f], &l, 10);
+
+            f += (l - &metadata[f]) + 1;
+
+            size_align.emplace_back(size, align);
+        } while (true);
+    }
+
+    static
+    void read_kernarg_metadata_v2(
+        const std::string& kernels_md,
+        std::size_t dx,
+        std::unordered_map<
+            std::string,
+            std::vector<std::pair<std::size_t, std::size_t>>>& kernargs) {
+        do {
+            dx = kernels_md.find("Name:", dx);
+
+            if (dx == std::string::npos) break;
+
+            static constexpr decltype(kernels_md.size()) name_sz{5};
+            dx = kernels_md.find_first_not_of(" '", dx + name_sz);
+
+            auto fn =
+                kernels_md.substr(dx, kernels_md.find_first_of("'\n", dx) - dx);
+            dx += fn.size();
+
+            auto dx1 = kernels_md.find("CodeProps", dx);
+            dx = kernels_md.find("Args:", dx);
+
+            if (dx1 < dx) {
+                dx = dx1;
+                continue;
+            }
+            if (dx == std::string::npos) break;
+
+            static constexpr decltype(kernels_md.size()) args_sz{5};
+            dx = parse_args_v2(kernels_md, dx + args_sz, dx1, kernargs[fn]);
+        } while (true);
+    }
+
     static
     std::string metadata_to_string(const amd_comgr_metadata_node_t& md) {
         std::string str;
@@ -598,9 +662,8 @@ public:
     }
 
     static
-    void parse_args(
+    void parse_args_v3(
             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) {
         size_t arg_count = 0;
         if (amd_comgr_get_metadata_list_size(args_md, &arg_count)
@@ -615,9 +678,7 @@ public:
                 return;
 
             amd_comgr_metadata_node_t arg_size_md;
-            if (amd_comgr_metadata_lookup(arg_md,
-                                          is_code_object_v3 ? ".size" : "Size",
-                                          &arg_size_md)
+            if (amd_comgr_metadata_lookup(arg_md, ".size", &arg_size_md)
                 != AMD_COMGR_STATUS_SUCCESS)
                 return;
 
@@ -629,35 +690,21 @@ public:
 
             size_t arg_align;
 
-            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;
+            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));
+            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;
+            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;
+            arg_align = 1;
+            while (arg_offset && (arg_offset & 1) == 0) {
+                arg_offset >>= 1;
+                arg_align <<= 1;
             }
 
             size_align.emplace_back(arg_size, arg_align);
@@ -669,11 +716,11 @@ public:
     }
 
     static
-    void read_kernarg_metadata(
-            const std::vector<char>& blob,
+    void read_kernarg_metadata_v3(
+            const std::string& blob,
             std::unordered_map<
-            std::string,
-            std::vector<std::pair<std::size_t, std::size_t>>>& kernargs) {
+                std::string,
+                std::vector<std::pair<std::size_t, std::size_t>>>& kernargs) {
         amd_comgr_data_t dataIn;
         amd_comgr_status_t status;
 
@@ -690,7 +737,6 @@ public:
             != AMD_COMGR_STATUS_SUCCESS)
             return;
 
-        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) {
@@ -699,7 +745,6 @@ public:
                                           &kernels_md)
                 != AMD_COMGR_STATUS_SUCCESS)
                 return;
-            is_code_object_v3 = true;
         }
 
         size_t kernel_count = 0;
@@ -715,9 +760,7 @@ public:
                 continue;
 
             amd_comgr_metadata_node_t name_md;
-            if (amd_comgr_metadata_lookup(kernel_md,
-                                          is_code_object_v3 ? ".name" : "Name",
-                                          &name_md)
+            if (amd_comgr_metadata_lookup(kernel_md, ".name", &name_md)
                 != AMD_COMGR_STATUS_SUCCESS)
                 continue;
 
@@ -727,21 +770,15 @@ public:
                 != 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)
+            if (amd_comgr_metadata_lookup(kernel_md, ".args", &args_md)
                 != AMD_COMGR_STATUS_SUCCESS)
                 continue;
 
             auto foundKernel = kernargs.find(kernel_name_str);
             // parse arguments for a given kernel only once
             if (foundKernel == kernargs.end()) {
-                parse_args(args_md, is_code_object_v3, kernargs[kernel_name_str]);
+                parse_args_v3(args_md, kernargs[kernel_name_str]);
             }
 
             if (amd_comgr_destroy_metadata(args_md) != AMD_COMGR_STATUS_SUCCESS
@@ -757,7 +794,52 @@ public:
         amd_comgr_release_data(dataIn);
     }
 
-    const std::unordered_map<std::string, 
+    static
+    void read_kernarg_metadata(
+        const std::string& blob,
+        std::unordered_map<
+            std::string,
+            std::vector<std::pair<std::size_t, std::size_t>>>& kernargs)
+    {
+        std::istringstream istr{blob};
+        ELFIO::elfio reader;
+
+        if (!reader.load(istr)) return;
+
+        // TODO: this is inefficient.
+        auto it = find_section_if(reader, [](const ELFIO::section* x) {
+            return x->get_type() == SHT_NOTE;
+        });
+
+        if (!it) return;
+
+        const ELFIO::note_section_accessor acc{reader, it};
+        auto n{acc.get_notes_num()};
+        while (n--) {
+            ELFIO::Elf_Word type{};
+            std::string name{};
+            void* desc{};
+            ELFIO::Elf_Word desc_size{};
+
+            acc.get_note(n, type, name, desc, desc_size);
+
+            if (name == "AMDGPU") {
+                return read_kernarg_metadata_v3(blob, kernargs);
+            }
+            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;
+
+            return read_kernarg_metadata_v2(tmp, dx + 8u, kernargs); // Skip "Kernels:".
+        }
+    }
+
+    const std::unordered_map<std::string,
         std::vector<std::pair<std::size_t, std::size_t>>>& get_kernargs() {
 
         std::call_once(kernargs.first, [this]() {