[QST] Where MMA operation is performed?

[IzanCatalan]

What is your question?
Hi, I am trying to understand how it is performed the MMA operation inside the convolution kernel, specifically in this part of the code:

cutlass/include/cutlass/conv/kernel/implicit_gemm_convolution.h

Line 350 in 62750a2

mma(params.gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators, params.gemm_k_iterations_per_channel);

What work is done by each thread inside this code? Which channel filters are computed for each thread using iteratorB? Where is the result, and is it in shared memory accessed by all threads?

A little review of the workflow from this line of code to the end would greatly help me understand how the convolution is implemented in GPU.

Thank you.

[hwu36]

It calls here https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/conv/threadblock/implicit_gemm_multistage.h#L266

It is basically the same as the gemm operation. You could check our ampere gtc talk to understand. the iterators are reimplmented for implicit gemm. fitler iterator is here: https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_analytic.h

[hwu36]

implicit gemm doc is here: https://github.com/NVIDIA/cutlass/blob/main/media/docs/implicit_gemm_convolution.md

[IzanCatalan]

@hwu36 Thank you for the information; however, I cannot understand where the product is performed in the code . I mean this part, according to the Gemm doc:

 for (int gemm_k = 0; gemm_k < GEMM_K; ++gemm_k) {

      int k = gemm_j;

      int c = gemm_k / (RS);
      int crs_residual = gemm_k % (RS);

      int r = crs_residual / S;
      int s = crs_residual % S;

      int h = f(p, r);
      int w = g(q, s);

      ElementA a = tensor_A.at({n, h, w, c});
      ElementB b = tensor_B.at({k, r, s, c});

      accum += a * b;
    }

My doubt is if I can perform something similar to this piece of code that I have implemented in

cutlass/include/cutlass/conv/kernel/implicit_gemm_convolution.h

Line 281 in 62750a2

void operator()(Params const &params, SharedStorage &shared_storage) {

long int K = params.problem_size.K, C = params.problem_size.C, 
              R = params.problem_size.R, S = params.problem_size.S, 
              HW = R * S, total_channels = K * C;

      for (long int channel_idx = thread_idx; channel_idx < total_channels; channel_idx += total_threads) {

             long int base_index = (channel_idx / C) * (C * HW) + (channel_idx % C);
              uint32_t protected_mask0 = 0, protected_mask1 = 0;
              memcpy(&protected_mask0, &params.ptr_B_mask0s[channel_idx], sizeof(ElementB));
              memcpy(&protected_mask1, &params.ptr_B_mask1s[channel_idx], sizeof(ElementB));
                for (long int i = 0; i < HW; i++) {
                  uint32_t original_bits = 0, val_check = 0;
                  memcpy(&original_bits, &params.ptr_B[base_index + i * C], sizeof(ElementB));
                  val_check = (original_bits & protected_mask1) | protected_mask0;
                  params.ptr_B[base_index + i * C] = *reinterpret_cast<ElementB *>(&val_check);
          }
    }

It performs an & and OR operation over the filters of the convolution (params.ptr_B) using some masks saved in another pointer. My idea is to optimise this process, perform this two-bitwise operation inside the MMA code and, just before the multiplication, something similar to this:

 for (int gemm_k = 0; gemm_k < GEMM_K; ++gemm_k) {

      int k = gemm_j;

      int c = gemm_k / (RS);
      int crs_residual = gemm_k % (RS);

      int r = crs_residual / S;
      int s = crs_residual % S;

      int h = f(p, r);
      int w = g(q, s);

      ElementA a = tensor_A.at({n, h, w, c});
      ElementB b = tensor_B.at({k, r, s, c});

      ElementB mask0s = tensor_B_mask0.at({k,r,s,c})
      ElementB mask1s = tensor_B_mask1.at({k,r,s,c})
      uint32_t protected_mask0 = 0, protected_mask1 = 0;
      memcpy(&protected_mask0, &mask0s, sizeof(ElementB));
       memcpy(&protected_mask1, &mask1s, sizeof(ElementB)); <
      resB = ( b & protected_mask1) | protected_mask0; 

      accum += a * resB;
    }

Something like this is possible? And if it is, where exactly is the part of the code in which you think I should do it?

[hwu36]

warp level mma happens in https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/conv/threadblock/implicit_gemm_multistage.h#L463. I think you can implement your own warp_mma.transform to hack B operand. one example is here: https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/gemm/warp/mma_tensor_op.h#L357

[thakkarV]

closing as answered. please reopen if you still have questions.