add transpose block read variant

Author: bashbaug

samples/99_matrixexperimentsi8/main.cpp

@@ -142,6 +142,23 @@ static void compute_reference(
     }
 }
 
+template <typename DstT, typename SrcT>
+static void compute_reference_TN(
+    std::vector<DstT>& C,
+    const std::vector<SrcT>& A, const std::vector<SrcT>& B,
+    size_t M, size_t N, size_t K)
+{
+    for (size_t m = 0; m < M; m++) {
+        for (size_t n = 0; n < N; n++) {
+            DstT sum = 0;
+            for (size_t k = 0; k < K; k++) {
+                sum = A[k * K + m] * B[k * N + n] + sum;
+            }
+            C[m * N + n] = sum;
+        }
+    }
+}
+
 template <typename T>
 void check_results(
     size_t M,
@@ -660,6 +677,107 @@ static void i8_dpas_blockread_vnni_tiled(
     }
 }
 
+static void i8_naive_TN(
+    cl::Context& context, cl::Program& program, cl::CommandQueue& queue,
+    cl::Buffer& C, cl::Buffer& A, cl::Buffer& B,
+    size_t M, size_t N, size_t K,
+    const std::vector<int>& C_ref)
+{
+    printf("%80s: ", makeTestName(__FUNCTION__, M, N, K).c_str()); fflush(stdout);
+
+    cl::Kernel kernel{program, "i8_naive_TN"};
+    if (kernel() == nullptr) {
+        printf("unsupported.\n");
+    } else {
+        kernel.setArg(0, C);
+        kernel.setArg(1, A);
+        kernel.setArg(2, B);
+        kernel.setArg(3, static_cast<cl_int>(K));
+
+        if (!skipinit) {
+            queue.enqueueFillBuffer(C, 0, 0, C_ref.size() * sizeof(C_ref[0]));
+        }
+
+        float best = 999.0f;
+        for (int test = 0; test < testIterations; test++) {
+            cl::Event event;
+            auto start = test_clock::now();
+            queue.enqueueNDRangeKernel(kernel, cl::NullRange,
+                cl::NDRange{N, M}, cl::NullRange, nullptr, &event);
+            queue.finish();
+            auto end = test_clock::now();
+            std::chrono::duration<float> sw_time = end - start;
+            auto elapsed = wallclock ? sw_time.count() : hw_time(event);
+            best = std::min(best, elapsed);
+        }
+        auto gops = 2.0 * M * N * K / best / 1e9;
+        printf("Best in %f seconds (%f gops)\n", best, gops);
+
+        if (validate) {
+            printf("Checking results... "); fflush(stdout);
+            std::vector<int> C_check(C_ref.size());
+            queue.enqueueReadBuffer(C, CL_TRUE, 0, C_check.size() * sizeof(C_check[0]), C_check.data());
+            check_results(M, N, C_check, C_ref);
+            printf(" done!\n");
+        }
+    }
+}
+
+template<int tM, int tN>
+static void i8_dpas_blockread_rowmajor_TN(
+    cl::Context& context, cl::Program& program, cl::CommandQueue& queue,
+    cl::Buffer& C, cl::Buffer& A, cl::Buffer& B,
+    size_t M, size_t N, size_t K,
+    const std::vector<int>& C_ref)
+{
+    printf("%80s: ", makeTestName(__FUNCTION__, tM, tN, M, N, K).c_str()); fflush(stdout);
+
+    std::string kernelName = "i8_dpas_blockread_rowmajor_TN";
+    kernelName += "_m" + std::to_string(tM);
+    kernelName += "_n" + std::to_string(tN);
+    cl::Kernel kernel{program, kernelName.c_str()};
+    if (kernel() == nullptr) {
+        printf("unsupported.\n");
+    } else if (K < 64 || N < 64/4) {
+        printf("matrix pitch for block reads must be >= 64 bytes.\n");
+    } else {
+        kernel.setArg(0, C);
+        kernel.setArg(1, A);
+        kernel.setArg(2, B);
+        kernel.setArg(3, static_cast<cl_int>(K));
+        if (roundRobin) {
+            setRoundRobin(kernel);
+        }
+
+        if (!skipinit) {
+            queue.enqueueFillBuffer(C, 0, 0, C_ref.size() * sizeof(C_ref[0]));
+        }
+
+        float best = 999.0f;
+        for (int test = 0; test < testIterations; test++) {
+            cl::Event event;
+            auto start = test_clock::now();
+            queue.enqueueNDRangeKernel(kernel, cl::NullRange,
+                cl::NDRange{N, M/tM}, cl::NullRange, nullptr, &event);
+            queue.finish();
+            auto end = test_clock::now();
+            std::chrono::duration<float> sw_time = end - start;
+            auto elapsed = wallclock ? sw_time.count() : hw_time(event);
+            best = std::min(best, elapsed);
+        }
+        auto gops = 2.0 * M * N * K / best / 1e9;
+        printf("Best in %f seconds (%f gops)\n", best, gops);
+
+        if (validate) {
+            printf("Checking results... "); fflush(stdout);
+            std::vector<int> C_check(C_ref.size());
+            queue.enqueueReadBuffer(C, CL_TRUE, 0, C_check.size() * sizeof(C_check[0]), C_check.data());
+            check_results(M, N, C_check, C_ref);
+            printf(" done!\n");
+        }
+    }
+}
+
 int main(int argc, char** argv)
 {
     int platformIndex = 0;
@@ -784,6 +902,7 @@ int main(int argc, char** argv)
     std::vector<int8_t> Bvnni_vec(K * N);
 
     std::vector<int> C_ref(M * N);
+    std::vector<int> C_TN_ref(M * N);
 
     printf("Initializing source matrices...\n");
     fill_matrix(A_vec, M, K);
@@ -794,6 +913,8 @@ int main(int argc, char** argv)
     if (validate) {
         printf("Computing reference...\n");
         compute_reference(C_ref, A_vec, B_vec, M, N, K);
+        printf("Computing transposed reference...\n");
+        compute_reference_TN(C_TN_ref, A_vec, B_vec, M, N, K);
     }
 
     printf("Creating source buffers...\n");
@@ -910,6 +1031,11 @@ int main(int argc, char** argv)
         i8_dpas_blockread_vnni_tiled<8, 16, 4, 4>(context, program, queue, C, A, Bvnni, M, N, K, C_ref);
     }
 
+    if (mask & 0x2000) {
+        //i8_naive_TN(context, program, queue, C, A, B, M, N, K, C_TN_ref);
+        i8_dpas_blockread_rowmajor_TN<4, 16>(context, program, queue, C, A, B, M, N, K, C_TN_ref);
+    }
+
     printf("Done.\n");
 
     return 0;

samples/99_matrixexperimentsi8/matrix_helpers_i8.cl

@@ -742,4 +742,20 @@ void intel_subgroup_block_write_u32_m8k16(__global void* base_address, int width
     __builtin_IB_subgroup_block_write_flat_u32_m8k16v1(as_long(base_address), width - 1, height - 1, pitch - 1, coord, data);
 }
 
+uint    __builtin_IB_subgroup_block_read_flat_transpose_u32_k1(long baseoffset, int width_minus_one, int height_minus_one, int pitch_minus_one, int2 coord);
+uint2   __builtin_IB_subgroup_block_read_flat_transpose_u32_m32k1(long baseoffset, int width_minus_one, int height_minus_one, int pitch_minus_one, int2 coord);
+
+void intel_sub_group_2d_block_read_transpose_32b_16r1x1c(global void* base_address, int width, int height, int pitch, int2 coord, private uint* destination)
+{
+    uint temp = __builtin_IB_subgroup_block_read_flat_transpose_u32_k1(as_long(base_address), width - 1, height - 1, pitch - 1, coord);
+    destination[0] = temp;
+}
+
+void intel_sub_group_2d_block_read_transpose_32b_32r1x1c(global void* base_address, int width, int height, int pitch, int2 coord, private uint* destination)
+{
+    uint2 temp = __builtin_IB_subgroup_block_read_flat_transpose_u32_m32k1(as_long(base_address), width - 1, height - 1, pitch - 1, coord);
+    destination[0] = temp.s0;
+    destination[1] = temp.s1;
+}
+
 #endif // cl_intel_subgroup_extended_block_read

samples/99_matrixexperimentsi8/matrix_kernels_i8.cl

@@ -27,6 +27,24 @@ kernel void i8_naive(global int* C, global char* A, global char* B, int K)
     C[m * N + n] = sum;
 }
 
+kernel void i8_naive_TN(global int* C, global char* A, global char* B, int K)
+{
+    const int N = get_global_size(0);
+    const int m = get_global_id(1);
+    const int n = get_global_id(0);
+
+    int sum = 0;
+    for (int k = 0; k < K; k++) {
+        sum = A[k * K + m] * B[k * N + n] + sum;
+        if (get_global_id(0) == 1 && get_global_id(1) == 0) {
+            printf("after iteration %d: sum is %d\n", k, sum);
+        }
+    }
+
+    sum = activation(sum);
+    C[m * N + n] = sum;
+}
+
 // For all i8 kernels tK == 32:
 #define tK 32
 
@@ -574,6 +592,50 @@ kernel void i8_dpas_blockread_vnni_m8_n16(global int* C, global char* A, global
     intel_sub_group_2d_block_write_32b_8r16x1c(C, N * sizeof(float), M, N * sizeof(float), (int2)(n, m), (uint*)&sum);
 }
 
+__attribute__((intel_reqd_sub_group_size(16))) __attribute__((reqd_work_group_size(16, 1, 1)))
+kernel void i8_dpas_blockread_rowmajor_TN_m4_n16(global int* C, global char* A, global char* B, int K)
+{
+    __builtin_assume(K > 0);    // Always at least one K iteration.
+    const int tM = 4;
+    const int tN = 16;
+    const int M = get_global_size(1) * tM;
+    const int N = get_global_size(0);
+    const int m = get_group_id(1) * tM;
+    const int n = get_group_id(0) * tN;
+    const int sglid = get_sub_group_local_id();
+
+    int4 sum = 0;
+    for (int k = 0; k < K; k += tK) {
+        int2    readData;
+        intel_sub_group_2d_block_read_transpose_32b_32r1x1c(A, M * sizeof(char), K, M * sizeof(char), (int2)(m / 4, k), (uint*)&readData);
+
+        // Note: after the transpose block read:
+        //      readData.s0 contains row  0-15
+        //      readData.s1 contains row 16-31
+        // So, WI0 has rows 0 and 16, WI1 has rows 1 and 17, etc.
+        // We want WI0 to have rows 0 and 1, WI1 to have rows 2 and 3, etc.
+        int shuffledData0 = (sglid < 8) ?
+            sub_group_shuffle(readData.s0, (sglid * 2)) :
+            sub_group_shuffle(readData.s1, (sglid * 2) % 16);
+        int shuffledData1 = (sglid < 8) ?
+            sub_group_shuffle(readData.s0, (sglid * 2) + 1) :
+            sub_group_shuffle(readData.s1, (sglid * 2) % 16 + 1);
+
+        short4  aData;
+        aData.s0 = as_short((char2)(as_char4(shuffledData0).s0, as_char4(shuffledData1).s0));
+        aData.s1 = as_short((char2)(as_char4(shuffledData0).s1, as_char4(shuffledData1).s1));
+        aData.s2 = as_short((char2)(as_char4(shuffledData0).s2, as_char4(shuffledData1).s2));
+        aData.s3 = as_short((char2)(as_char4(shuffledData0).s3, as_char4(shuffledData1).s3));
+
+        int8    bData;
+        intel_sub_group_2d_block_read_transform_8b_32r16x1c(B, N * sizeof(char), K, N * sizeof(char), (int2)(n, k), (uint*)&bData);
+        sum = mat_mul_sg16(aData, bData, sum);
+    }
+
+    sum = activation(sum);
+    intel_sub_group_2d_block_write_32b_4r16x1c(C, N * sizeof(float), M, N * sizeof(float), (int2)(n, m), (uint*)&sum);
+}
+
 #endif // cl_intel_subgroup_2d_block_io
 
 #if 0 // disable the tiled cases for now