 
 template <typename T>
 __global__ void transpose_front_kernel(size_t M, size_t N, size_t K, const T* x, T* y) {
-    auto m  = threadIdx.x + blockIdx.x * blockDim.x;
+    const auto mk  = threadIdx.x + blockIdx.x * blockDim.x;
 
-    if (m < M) {
+    if (mk < M * K) {
+        // Note: Ideally, we would use 2D Indexing. But I can't get it to be as fast as the 1D index
+        const size_t m = mk / K;
+        const size_t k = mk % K;
+
         for (size_t n = 0; n < N; ++n) {
-            // y(n, m) = x(m, n) 
+            // y(n, m) = x(m, n)
             // x[M, N, K]
             // y[N, M, K]
 
-            for (size_t k = 0; k < K; ++k) {
-                y[n * (M * K) + m * K + k] = x[m * (N * K) + n * K + k];
-            }
+            y[n * (M * K) + m * K + k] = x[m * (N * K) + n * K + k];
         }
     }
 }


    cudaOccupancyMaxPotentialBlockSize(&minGridSize, &blockSize, &transpose_front_kernel<double>, 0, 0);

    int gridSize = (m + blockSize - 1) / blockSize;

    transpose_front_kernel<<<gridSize, blockSize>>>(m, n, k, x, y);
}


    cudaOccupancyMaxPotentialBlockSize(&minGridSize, &blockSize, &transpose_front_kernel<float>, 0, 0);

    int gridSize = (m * k + blockSize - 1) / blockSize;

    transpose_front_kernel<<<gridSize, blockSize>>>(m, n, k, x, y);
}
