         for (n = 0; n < BLK_N / DIM_YB; n++) {
             #pragma unroll
             for (m = 0; m < BLK_K / DIM_XB; m++) {
-                rb[n][m] = offs_dB[n * DIM_YB * K + m * DIM_XB];
+                rb[n][m] = fetch(offs_dB, K, m * DIM_XB, n * DIM_YB, boundB);
             }
         }
 

    for (n = 0; n < BLK_K; n += DIM_YA) {
        #pragma unroll
        for (m = 0; m < BLK_M; m += DIM_XA) {
            sA[n + idyA][m + idxA] = offs_dA[n * M + m];
        }
    }
    // blockwise transpose to transpose load

    for (n = 0; n < BLK_N; n += DIM_YB) {
        #pragma unroll
        for (m = 0; m < BLK_K; m += DIM_XB) {
            sB[n + idyB][m + idxB] = fetch(offs_dB, K, m, n, boundB);
        }
    }
    __syncthreads();
