CUDA performance highly sensitive to memory layout (x/y axis direction).

[BenkangPeng]

The main_kernel interprets the horizontal axis as Y and vertical axis as X, while my_main_kernel interprets horizontal as X and vertical as Y.

The example is below:
Performance diff:

Average time cost: 28.4 ms
4839.4 GFLOPS
Average time cost: 221.7 ms
619.932 GFLOPS
Success

#include <chrono>
#include <cmath>
#include <cstdio>
#include <cuda_runtime.h>
#include <iostream>
#include <random>

#define CUDA_CHECK(call)                                                       \
  {                                                                            \
    cudaError_t err = call;                                                    \
    if (err != cudaSuccess) {                                                  \
      fprintf(stderr, "CUDA error %d: %s\n", err, cudaGetErrorString(err));    \
      exit(err);                                                               \
    }                                                                          \
  }

/// generated by TVM
extern "C" __global__ void __launch_bounds__(1024)
    main_kernel(float *__restrict__ A, float *__restrict__ B,
                float *__restrict__ C) {
  float C_local[1];
  for (int k = 0; k < 4096; ++k) {
    if (k == 0) {
      C_local[0] = 0.000000e+00f;
    }
    C_local[0] =
        (C_local[0] +
         (A[(((((int)blockIdx.x) * 131072) + (((int)threadIdx.x) * 4096)) +
             k)] *
          B[(((k * 4096) + (((int)blockIdx.y) * 32)) + ((int)threadIdx.y))]));
  }
  C[((((((int)blockIdx.x) * 131072) + (((int)threadIdx.x) * 4096)) +
      (((int)blockIdx.y) * 32)) +
     ((int)threadIdx.y))] = C_local[0];
}

/// my index calculation way
__global__ void my_main_kernel(float *__restrict__ A, float *__restrict__ B,
                               float *__restrict__ C) {
  float C_local[1];
  for (int k = 0; k < 4096; ++k) {
    if (k == 0) {
      C_local[0] = 0.000000e+00f;
    }
    C_local[0] =
        (C_local[0] +
         (A[(((((int)blockIdx.y) * 131072) + (((int)threadIdx.y) * 4096)) +
             k)] *
          B[(((k * 4096) + (((int)blockIdx.x) * 32)) + ((int)threadIdx.x))]));
  }
  C[((((((int)blockIdx.y) * 131072) + (((int)threadIdx.y) * 4096)) +
      (((int)blockIdx.x) * 32)) +
     ((int)threadIdx.x))] = C_local[0];
}

std::mt19937 rng(2025);
std::uniform_real_distribution<float> dist(-100.0, 100.0);

void init_matrix(float *A, float *B, int M, int N, int K) {
  for (int i = 0; i < M; i++) {
    for (int j = 0; j < K; j++) {
      A[i * K + j] = dist(rng);
    }
  }
  for (int i = 0; i < K; i++) {
    for (int j = 0; j < N; j++) {
      B[i * N + j] = dist(rng);
    }
  }
}

int main() {
  const int M = 4096;
  const int N = 4096;
  const int K = 4096;

  float *A = (float *)malloc(M * K * sizeof(float));
  float *B = (float *)malloc(K * N * sizeof(float));
  float *C1 = (float *)malloc(M * N * sizeof(float));
  float *C2 = (float *)malloc(M * N * sizeof(float));
  init_matrix(A, B, M, N, K);

  float *A_d, *B_d, *C_d1, *C_d2;
  CUDA_CHECK(cudaMalloc((void **)&A_d, M * K * sizeof(float)));
  CUDA_CHECK(cudaMalloc((void **)&B_d, K * N * sizeof(float)));
  CUDA_CHECK(cudaMalloc((void **)&C_d1, M * N * sizeof(float)));
  CUDA_CHECK(cudaMalloc((void **)&C_d2, M * N * sizeof(float)));

  CUDA_CHECK(cudaMemcpy(A_d, A, M * K * sizeof(float), cudaMemcpyHostToDevice));
  CUDA_CHECK(cudaMemcpy(B_d, B, K * N * sizeof(float), cudaMemcpyHostToDevice));

  dim3 grid(128, 128);
  dim3 block(32, 32);

  int run_times = 10;
  // warmup
  my_main_kernel<<<grid, block>>>(A_d, B_d, C_d1);

  auto start = std::chrono::high_resolution_clock::now();
  for (int i = 0; i < run_times; i++) {
    my_main_kernel<<<grid, block>>>(A_d, B_d, C_d1);
    cudaDeviceSynchronize();
  }
  auto end = std::chrono::high_resolution_clock::now();
  auto duration =
      std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
  std::cout << "Average time cost: "
            << duration.count() / static_cast<double>(run_times) << " ms"
            << std::endl;
  std::cout << static_cast<double>(2) * M * K * N * run_times /
                   duration.count() / 1e6
            << " GFLOPS" << std::endl;

  /// warmup
  main_kernel<<<grid, block>>>(A_d, B_d, C_d2);

  start = std::chrono::high_resolution_clock::now();
  for (int i = 0; i < run_times; i++) {
    main_kernel<<<grid, block>>>(A_d, B_d, C_d2);
    cudaDeviceSynchronize();
  }
  end = std::chrono::high_resolution_clock::now();
  duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
  std::cout << "Average time cost: "
            << duration.count() / static_cast<double>(run_times) << " ms"
            << std::endl;
  std::cout << static_cast<double>(2) * M * K * N * run_times /
                   duration.count() / 1e6
            << " GFLOPS" << std::endl;

  CUDA_CHECK(
      cudaMemcpy(C1, C_d1, M * N * sizeof(float), cudaMemcpyDeviceToHost));
  CUDA_CHECK(
      cudaMemcpy(C2, C_d2, M * N * sizeof(float), cudaMemcpyDeviceToHost));

  for (int i = 0; i < M; i++) {
    for (int j = 0; j < N; j++) {
      if (fabs(C1[i * N + j] - C2[i * N + j]) > 1e-5) {
        printf("Error at (%d, %d): %f != %f\n", i, j, C1[i * N + j],
               C2[i * N + j]);
        return 1;
      }
    }
  }

  printf("Success\n");

  free(A);
  free(B);
  free(C1);
  free(C2);
  CUDA_CHECK(cudaFree(A_d));
  CUDA_CHECK(cudaFree(B_d));
  CUDA_CHECK(cudaFree(C_d1));
  CUDA_CHECK(cudaFree(C_d2));
  return 0;
}

[BenkangPeng]

The kernel performance gap comes from warp memory coalescing access.
my_main_kernel fully exploits coalesced memory access, while main_kernel suffers from non-coalesced memory access.

A warp achieves coalesced memory access, when adjacent threads(e.g., threadIdx.x, threadIdx.x + 1) access consecutive memory addresses(e.g. base, base + 1).

main_kernel: Failed Coalescing
Matrix A Access:

A[(blockIdx.x * 131072) + (threadIdx.x * 4096) + k]

Issue: threadIdx.x → threadIdx.x+1 offsets address by 4096 (not 1).
Consequence: 1024x 32-bit threads require 1024 separate 4-byte loads → zero coalescing.
Matrix B: Irrelevant (all threads access the same address → broadcast).

my_main_kernel: Optimal Coalescing
Matrix A: All threads in warp access identical address (broadcast, 1 transaction).
Matrix B:

B[(k * 4096) + (blockIdx.x * 32) + threadIdx.x]

Key: threadIdx.x → threadIdx.x+1 offsets addresses by 1 (consecutive addresses).
Result: Single merged memory access (32 consecutive 4-byte words → 1 128-byte load).