PTXAS says registers not enough for m64n256k32 integer wgmma after setting regcount from 128 to 200 by setmaxnreg.inc

Accelerated Computing CUDA CUDA Programming and Performance
1

When compiling the following code for sm_90a architecture, PTXAS complains that number of registers is not enough for wgmma.mma_async(need 154, but actual 128), which is very strange since there’s already a setmaxnreg instruction which increase the regcount from 128 to 200.

__launch_bounds__(512, 1)
__global__ void kernel(const char *input_s8, const char *weight_s8, int *output,
  unsigned infeaturesdiv128) {

  extern __shared__ char buffer[];
  char *buf = &buffer[0];
  asm ("cvta.to.shared.u64 %0,%0;\n":"+l"(buf));
  //blockDim.x == 128, blockDim.y == 4, blockDim.z == 1, gridDim.z == 8
  //barrier initialization omitted

  if (threadIdx.y < 2) {
    /* consumer: 2 warpgroups */
    asm volatile ("setmaxnreg.inc.sync.aligned.u32 200;\n":::"memory");
    auto mdenc = [](size_t x)->size_t { return (x & 0x3FFFF) >> 4; };
    const size_t mdbase_inp = (mdenc(1024) << 16) | (mdenc(128) << 32);
    const size_t mdbase_wt = (mdenc(128) << 16) | (mdenc(1024) << 32);
    int acc[128] = {};
    asm volatile ("wgmma.fence.sync.aligned;\n":::"memory");
    int buffer_id = 0;
    auto commit_k128 = [&]()->void {
      size_t baseaddr = reinterpret_cast<size_t>(buf) + buffer_id * 384 * 128;
      size_t input_addr = baseaddr + threadIdx.y * 8192;
      size_t md_input = mdbase_inp | mdenc(input_addr);
      size_t wt_addr = baseaddr + 16384;
      size_t md_weight1 = mdbase_wt | mdenc(wt_addr);
      for (int i = 0; i < 4; ++i) {
        asm ("wgmma.mma_async.sync.aligned.m64n256k32.s32.s8.s8 "
          "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16,%17,%18,"
          "%19,%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31,%32,%33,%34,%35,"
          "%36,%37,%38,%39,%40,%41,%42,%43,%44,%45,%46,%47,%48,%49,%50,%51,%52,"
          "%53,%54,%55,%56,%57,%58,%59,%60,%61,%62,%63,%64,%65,%66,%67,%68,%69,"
          "%70,%71,%72,%73,%74,%75,%76,%77,%78,%79,%80,%81,%82,%83,%84,%85,%86,"
          "%87,%88,%89,%90,%91,%92,%93,%94,%95,%96,%97,%98,%99,%100,%101,%102,"
          "%103,%104,%105,%106,%107,%108,%109,%110,%111,%112,%113,%114,%115,"
          "%116,%117,%118,%119,%120,%121,%122,%123,%124,%125,%126,%127},%128,%129,1;\n"
          :"+r"(acc[0]),"+r"(acc[1]),"+r"(acc[2]),"+r"(acc[3]),"+r"(acc[4]),
          "+r"(acc[5]),"+r"(acc[6]),"+r"(acc[7]),"+r"(acc[8]),"+r"(acc[9]),
          "+r"(acc[10]),"+r"(acc[11]),"+r"(acc[12]),"+r"(acc[13]),"+r"(acc[14]),
          "+r"(acc[15]),"+r"(acc[16]),"+r"(acc[17]),"+r"(acc[18]),"+r"(acc[19]),
          "+r"(acc[20]),"+r"(acc[21]),"+r"(acc[22]),"+r"(acc[23]),"+r"(acc[24]),
          "+r"(acc[25]),"+r"(acc[26]),"+r"(acc[27]),"+r"(acc[28]),"+r"(acc[29]),
          "+r"(acc[30]),"+r"(acc[31]),"+r"(acc[32]),"+r"(acc[33]),"+r"(acc[34]),
          "+r"(acc[35]),"+r"(acc[36]),"+r"(acc[37]),"+r"(acc[38]),"+r"(acc[39]),
          "+r"(acc[40]),"+r"(acc[41]),"+r"(acc[42]),"+r"(acc[43]),"+r"(acc[44]),
          "+r"(acc[45]),"+r"(acc[46]),"+r"(acc[47]),"+r"(acc[48]),"+r"(acc[49]),
          "+r"(acc[50]),"+r"(acc[51]),"+r"(acc[52]),"+r"(acc[53]),"+r"(acc[54]),
          "+r"(acc[55]),"+r"(acc[56]),"+r"(acc[57]),"+r"(acc[58]),"+r"(acc[59]),
          "+r"(acc[60]),"+r"(acc[61]),"+r"(acc[62]),"+r"(acc[63]),
          "+r"(acc[64]),"+r"(acc[65]),"+r"(acc[66]),"+r"(acc[67]),"+r"(acc[68]),
          "+r"(acc[69]),"+r"(acc[70]),"+r"(acc[71]),"+r"(acc[72]),"+r"(acc[73]),
          "+r"(acc[74]),"+r"(acc[75]),"+r"(acc[76]),"+r"(acc[77]),"+r"(acc[78]),
          "+r"(acc[79]),"+r"(acc[80]),"+r"(acc[81]),"+r"(acc[82]),"+r"(acc[83]),
          "+r"(acc[84]),"+r"(acc[85]),"+r"(acc[86]),"+r"(acc[87]),"+r"(acc[88]),
          "+r"(acc[89]),"+r"(acc[90]),"+r"(acc[91]),"+r"(acc[92]),"+r"(acc[93]),
          "+r"(acc[94]),"+r"(acc[95]),"+r"(acc[96]),"+r"(acc[97]),"+r"(acc[98]),
          "+r"(acc[99]),"+r"(acc[100]),"+r"(acc[101]),"+r"(acc[102]),"+r"(acc[103]),
          "+r"(acc[104]),"+r"(acc[105]),"+r"(acc[106]),"+r"(acc[107]),"+r"(acc[108]),
          "+r"(acc[109]),"+r"(acc[110]),"+r"(acc[111]),"+r"(acc[112]),"+r"(acc[113]),
          "+r"(acc[114]),"+r"(acc[115]),"+r"(acc[116]),"+r"(acc[117]),"+r"(acc[118]),
          "+r"(acc[119]),"+r"(acc[120]),"+r"(acc[121]),"+r"(acc[122]),"+r"(acc[123]),
          "+r"(acc[124]),"+r"(acc[125]),"+r"(acc[126]),"+r"(acc[127])
          :"l"(md_input),"l"(md_weight1):"memory");
        md_input += 64 * 32 >> 4;
        md_weight1 += 8 * 32 >> 4;
      }
      asm volatile ("wgmma.commit_group.sync.aligned;\n":::"memory");
    };
    for (unsigned k128id = 0; k128id < infeaturesdiv128; ++k128id) {
      /* barrier synchronization code omitted */
      commit_k128();
      buffer_id++;
      if (buffer_id == 4) buffer_id = 0;
    }
    const size_t ldc = gridDim.x * 512;
    output += ldc * blockIdx.y * 512 + blockIdx.x * 512; //cluster_base
    output += ldc * blockIdx.z * 64 + threadIdx.y * 256; // down to wg_base
    output += ldc * (((threadIdx.x & 28) >> 2) + ((threadIdx.x & 96) >> 1)) +
      (threadIdx.x & 3) * 2; //down to thread start
    int *output2 = output + 8 * ldc;
    asm volatile ("wgmma.wait_group.sync.aligned 0;\n"
      :"+r"(acc[0]),"+r"(acc[1]),"+r"(acc[2]),"+r"(acc[3]),"+r"(acc[4]),
      "+r"(acc[5]),"+r"(acc[6]),"+r"(acc[7]),"+r"(acc[8]),"+r"(acc[9]),
      "+r"(acc[10]),"+r"(acc[11]),"+r"(acc[12]),"+r"(acc[13]),"+r"(acc[14]),
      "+r"(acc[15]),"+r"(acc[16]),"+r"(acc[17]),"+r"(acc[18]),"+r"(acc[19]),
      "+r"(acc[20]),"+r"(acc[21]),"+r"(acc[22]),"+r"(acc[23]),"+r"(acc[24]),
      "+r"(acc[25]),"+r"(acc[26]),"+r"(acc[27]),"+r"(acc[28]),"+r"(acc[29]),
      "+r"(acc[30]),"+r"(acc[31]),"+r"(acc[32]),"+r"(acc[33]),"+r"(acc[34]),
      "+r"(acc[35]),"+r"(acc[36]),"+r"(acc[37]),"+r"(acc[38]),"+r"(acc[39]),
      "+r"(acc[40]),"+r"(acc[41]),"+r"(acc[42]),"+r"(acc[43]),"+r"(acc[44]),
      "+r"(acc[45]),"+r"(acc[46]),"+r"(acc[47]),"+r"(acc[48]),"+r"(acc[49]),
      "+r"(acc[50]),"+r"(acc[51]),"+r"(acc[52]),"+r"(acc[53]),"+r"(acc[54]),
      "+r"(acc[55]),"+r"(acc[56]),"+r"(acc[57]),"+r"(acc[58]),"+r"(acc[59]),
      "+r"(acc[60]),"+r"(acc[61]),"+r"(acc[62]),"+r"(acc[63]),
      "+r"(acc[64]),"+r"(acc[65]),"+r"(acc[66]),"+r"(acc[67]),"+r"(acc[68]),
      "+r"(acc[69]),"+r"(acc[70]),"+r"(acc[71]),"+r"(acc[72]),"+r"(acc[73]),
      "+r"(acc[74]),"+r"(acc[75]),"+r"(acc[76]),"+r"(acc[77]),"+r"(acc[78]),
      "+r"(acc[79]),"+r"(acc[80]),"+r"(acc[81]),"+r"(acc[82]),"+r"(acc[83]),
      "+r"(acc[84]),"+r"(acc[85]),"+r"(acc[86]),"+r"(acc[87]),"+r"(acc[88]),
      "+r"(acc[89]),"+r"(acc[90]),"+r"(acc[91]),"+r"(acc[92]),"+r"(acc[93]),
      "+r"(acc[94]),"+r"(acc[95]),"+r"(acc[96]),"+r"(acc[97]),"+r"(acc[98]),
      "+r"(acc[99]),"+r"(acc[100]),"+r"(acc[101]),"+r"(acc[102]),"+r"(acc[103]),
      "+r"(acc[104]),"+r"(acc[105]),"+r"(acc[106]),"+r"(acc[107]),"+r"(acc[108]),
      "+r"(acc[109]),"+r"(acc[110]),"+r"(acc[111]),"+r"(acc[112]),"+r"(acc[113]),
      "+r"(acc[114]),"+r"(acc[115]),"+r"(acc[116]),"+r"(acc[117]),"+r"(acc[118]),
      "+r"(acc[119]),"+r"(acc[120]),"+r"(acc[121]),"+r"(acc[122]),"+r"(acc[123]),
      "+r"(acc[124]),"+r"(acc[125]),"+r"(acc[126]),"+r"(acc[127])::"memory");
    for (int nd8 = 0; nd8 < 32; ++nd8) {
      *(int2*)output = make_int2(acc[nd8*4], acc[nd8*4+1]);
      *(int2*)output2 = make_int2(acc[nd8*4+2], acc[nd8*4+3]);
      output += 8;
      output2 += 8;
    }
  } else {
    /* producer code omitted */
    asm volatile ("setmaxnreg.dec.sync.aligned.u32 40;\n");
  }
}

This problem exists for CUDA-12.8 and CUDA-13.2 (other versions not tested).

2

Here is the error message from PTXAS:

$ nvcc --compile -gencode arch=compute_90a,code=sm_90a -Xptxas -v -keep test.cu
ptxas fatal   : (C7602) Insufficient registers (128) to compile instruction at line 332 in function _Z6kernelPKcS0_Pij
. Try to compile with register target of 154 or higher.
ptxas fatal   : Ptx assembly aborted due to errors

Interestingly, after I split 1 m64n256k32 wgmma.mma_async instruction into 2 m64n128k32 wgmma.mma_async ones (each use 64 accumulate registers individually), this error disappears and the program get compiled normally! However such split may harm performance because of increased workloads of shared memory.

3

BTW, the documentation of PTX 9.2 about wgmma.mma_async is imperfect: it does not include m64n256k32 for available shapes of int8 inputs, while actually m64n256k32 is supported.

4

The cause of this problem may be one of the following:

  1. The validation process of individual wgmma instructions is buggy, which takes initial regcount regardless of the effects of setmaxnreg.
  2. The register requirements for m64n256k32 shape may be very special, which does not accept dynamically allocated registers.

I hope some specialists from NVIDIA could help checking the real reason.

5

I haven’t studied the initial posting here carefully, but I note that if I remove the launch bounds, the code will compile, as is, albeit with some ptxas info generated. Amongst other reported info, this results in the ptxas info “setmaxnreg ignored”.

If instead of removing the launch bounds, I change it to (32,1) (or, for that matter, (256,1), see below), then the code compiles without any message about the setmaxnreg directive.

It appears to me that your launch bounds is simply incompatible with other register requirements, whether they arise from the setmaxnreg directive, or the code, or both. 200x512 = 100,000 registers (approximately) which exceeds the 90a SM capacity (65536 registers max).

(128 is the register requirement, as a maximum, that would arise from your launch bounds directive)

If you’ve already understood all this, and are asking, for example, “why can’t I use setmaxnreg.inc to increase the available registers beyond the SM limit?” or another question, then please just ignore my post. I don’t have any comments about the wgmma instruction itself, splitting it, or its register requirements.

6

Thanks for your reply. Maybe I failed to make myself clear. Here’ s what happened in the test code:

  1. There are 4 warpgroups.
  2. Each thread has 128 registers initially.
  3. Warpgroup 0 and 1 are consumers. They request rising regcount to 200 before issueing m64n256k32 wgmma instructions.
  4. Warpgroup 2 and 3 are producers. They request decreasing regcount to 40 by setmaxnreg.dec, so as to let consumers use more registers. They won’t issue wgmma instructions.

So the test code can run on a SM with only 65536 registers.

I’ve already fired a bug report here, which has been noticed by the maintenance team.

7

Have you tried using (512, 1, 1) as threadblock shape? It may be a matter of the compiler being unable to guarantee that threadIdx.y < 2 would be a warpgroup-aligned condition, whereas it can do so for threadIdx.x < 256.

8

FYI. 6089012 is received and internally tracking this. We will bring back conclusion here when it is finished.