[opt] Incorrect global id (gid) computation via hipExtModuleLaunchKernel for non-divisible global sizes when using -fno-offload-uniform-block

[0oyyo0]

Environment

• This issue has been reproduced on ROCm 6.0 and ROCm 7.0
• Kernel launch API: hipExtModuleLaunchKernel
• Total thread size not divisible by block size (tail workgroup exists)
• Add the -fno-offload-uniform-block option during compilation.

Problem Description
When launching a kernel via hipExtModuleLaunchKernel with:
totalThreads % blockSize != 0
the kernel exhibits incorrect global id (gid) computation:

• Some lower gid values are executed twice
• This only happens when a tail workgroup is present(i.e., global size is not divisible by block size)
• When the global size is divisible by block size, the problem disappears.

demo

• Reproducing the problem
  hipcc --offload-device-only kernel.hip -o kernel.co -fno-offload-uniform-block #-mcode-object-version=5
hipcc main.cpp -o main_test
./main_test

• kernel.hip

#include <hip/hip_runtime.h>
#include <stdint.h>

extern "C" __global__
void write_gid_kernel(uint64_t* out)
{
    uint64_t gid = blockIdx.x * blockDim.x + threadIdx.x;
    printf("gid = %lu\n", gid);
    out[gid] = gid;
}

• main.cpp

#include <hip/hip_runtime.h>
#include <hip/hip_ext.h>
#include <cstdio>
#include <cstdlib>
#include <vector>


int main()
{
    hipModule_t module;
    hipFunction_t kernel;

    hipModuleLoad(&module, "kernel.co");
    hipModuleGetFunction(&kernel, module, "write_gid_kernel");

    // const int block_size    = 256;
    // const int total_threads = 1000;

    const int block_size    = 512;
    const int total_threads = 822528; // 822528 % 512 = 256 

    printf("logical threads = %d\n", total_threads);
    printf("block size      = %d\n", block_size);

    uint64_t* d_out;
    hipMalloc(&d_out, total_threads * sizeof(uint64_t));

    void* args[] = {
        (void*)&d_out
    };

    
    hipExtModuleLaunchKernel(
        kernel,
        total_threads, 1, 1,   
        block_size, 1, 1, 
        0,
        nullptr,
        args,
        nullptr
    );

    hipDeviceSynchronize();

    std::vector<uint64_t> h_out(total_threads, 0);
    hipMemcpy(
        h_out.data(), d_out,
        total_threads * sizeof(uint64_t),
        hipMemcpyDeviceToHost);

    hipFree(d_out);
    hipModuleUnload(module);
    return 0;
}

Observed Behavior
When running the above demo,

• gid in the range:[411136, 411391] is printed twice

• gid in the final tail region: [822272, totalThreads) is not printed at all

• 822528/512=1606; 822528 % 512 = 256; 1606*256=411136; 411136+256-1=411391

• It is suspected that, for the final remainder workgroup(256 threads), gid computation uses 256 * blockIdx + lane_id, leading to incorrect global IDs.

Expected Behavior
Each gid in [0, totalThreads) should be printed exactly once

[schung-amd]

Hi @0oyyo0,

gid in the range:[411136, 411391] is printed twice
gid in the final tail region: [822272, totalThreads) is not printed at all
822528/512=1606; 822528 % 512 = 256; 1606*256=411136; 411136+256-1=411391

This matches your gid computation; since you've specified non-uniform blocks, the last block to be launched has blockDim.x = 256 instead of 512. uint64_t gid = blockIdx.x * blockDim.x + threadIdx.x assumes a uniform block size, so since the last block is half of the size of the other blocks the gid computation here will return a value roughly half what it should be, as you've observed.

[0oyyo0]

Hi @0oyyo0,

gid in the range:[411136, 411391] is printed twice
gid in the final tail region: [822272, totalThreads) is not printed at all
822528/512=1606; 822528 % 512 = 256; 1606*256=411136; 411136+256-1=411391

This matches your gid computation; since you've specified non-uniform blocks, the last block to be launched has blockDim.x = 256 instead of 512. uint64_t gid = blockIdx.x * blockDim.x + threadIdx.x assumes a uniform block size, so since the last block is half of the size of the other blocks the gid computation here will return a value roughly half what it should be, as you've observed.

Thank you for your reply. As you said, the actual size of the last block is 256. However, logically, even if the size of the last block is 256, when calculating the gid, it should be calculated by multiplying the preceding normal size blockSize(512) by blockIdx, is it an error to use 256 multiplied by blockIdx here? Or is the operation of calculating gid using blockIdx.x * blockDim.x + threadIdx.x when using the -fno-offload-uniform-block option incorrect?

[schung-amd]

it should be calculated by multiplying the preceding normal size blockSize(512) by blockIdx, is it an error to use 256 multiplied by blockIdx here? Or is the operation of calculating gid using blockIdx.x * blockDim.x + threadIdx.x when using the -fno-offload-uniform-block option incorrect?

Right, that computation is correct when the block size is fixed but not necessarily with a non-uniform block size. 512 is the correct value to multiply by in this reproducer but since blockIdx.x is a local value to the block, the last block will use 256. You can verify this by printing out the value of blockIdx.x, you'll see that the duplicated gids you're observing happen alongside blockIdx.x being 256 rather than 512.

I don't think we provide an API to get the global thread ID in this case, so you might have to hardcode or template a value for the intended block size here.