add tiled block read kernels for PVC

bashbaug · Jan 15, 2024 · c7edcd6 · c7edcd6
1 parent 4e89026
commit c7edcd6
Showing 1 changed file with 86 additions and 0 deletions.
diff --git a/samples/99_matrixexperiments/matrix_kernel_tiled.cl b/samples/99_matrixexperiments/matrix_kernel_tiled.cl
@@ -180,3 +180,89 @@ kernel void MM_KERNEL_NAME(bfloat16_dpas_vnni_tiled, 8, 16, MM, NN)(global float
         }
     }
 }
+
+#ifdef cl_intel_subgroup_extended_block_read
+
+__attribute__((intel_reqd_sub_group_size(16))) __attribute__((reqd_work_group_size(16, 1, 1)))
+kernel void MM_KERNEL_NAME(bfloat16_dpas_rowmajor_tiled, 8, 16, MM, NN)(global float* C, global ushort* A, global ushort* B, int K)
+{
+    const int tM = 8;
+    const int tN = 16;
+    const int N = get_global_size(0) * NN;
+    const int m = get_group_id(1) * tM * MM;
+    const int n = get_group_id(0) * tN * NN;
+
+    float8 sum[MM][NN];
+    for (int mm = 0; mm < MM; mm++) {
+        for (int nn = 0; nn < NN; nn++) {
+            sum[mm][nn] = 0;
+        }
+    }
+
+    for (int k = 0; k < K; k += tK) {
+        short8  aData[MM];
+        for (int mm = 0; mm < MM; mm++) {
+            aData[mm] = as_short8(intel_subgroup_block_read_u16_m8k16(A, K * sizeof(ushort), M, K * sizeof(ushort), (int2)(k, m + mm * tM)));
+        }
+
+        int8    bData[NN];
+        for (int nn = 0; nn < NN; nn++) {
+            bData[nn] = load_b_rowmajor_d16_k16_nx(B, k, n + nn * tN, N);
+        }
+
+        for (int mm = 0; mm < MM; mm++) {
+            for (int nn = 0; nn < NN; nn++) {
+                sum[mm][nn] = mat_mul_sg16(aData[mm], bData[nn], sum[mm][nn]);
+            }
+        }
+    }
+
+    for (int mm = 0; mm < MM; mm++) {
+        for (int nn = 0; nn < NN; nn++) {
+            intel_subgroup_block_write_u32_m8k16v1(C, N * sizeof(float), M, N * sizeof(float), (int2)(n + nn * tN, m + mm * tM), as_uint8(sum[mm][nn]));
+        }
+    }
+}
+
+__attribute__((intel_reqd_sub_group_size(16))) __attribute__((reqd_work_group_size(16, 1, 1)))
+kernel void MM_KERNEL_NAME(bfloat16_dpas_vnni_tiled, 8, 16, MM, NN)(global float* C, global ushort* A, global ushort* B, int K)
+{
+    const int tM = 8;
+    const int tN = 16;
+    const int N = get_global_size(0) * NN;
+    const int m = get_group_id(1) * tM * MM;
+    const int n = get_group_id(0) * tN * NN;
+
+    float8 sum[MM][NN];
+    for (int mm = 0; mm < MM; mm++) {
+        for (int nn = 0; nn < NN; nn++) {
+            sum[mm][nn] = 0;
+        }
+    }
+
+    for (int k = 0; k < K; k += tK) {
+        short8  aData[MM];
+        for (int mm = 0; mm < MM; mm++) {
+            aData[mm] = as_short8(intel_subgroup_block_read_u16_m8k16(A, K * sizeof(ushort), M, K * sizeof(ushort), (int2)(k, m + mm * tM)));
+        }
+
+        int8    bData[NN];
+        for (int nn = 0; nn < NN; nn++) {
+            bData[nn] = as_int8(intel_subgroup_block_read_transform_u16_k16(B, N * sizeof(ushort), K, N * sizeof(ushort), (int2)(n + nn * tN, k)));
+        }
+
+        for (int mm = 0; mm < MM; mm++) {
+            for (int nn = 0; nn < NN; nn++) {
+                sum[mm][nn] = mat_mul_sg16(aData[mm], bData[nn], sum[mm][nn]);
+            }
+        }
+    }
+
+    for (int mm = 0; mm < MM; mm++) {
+        for (int nn = 0; nn < NN; nn++) {
+            intel_subgroup_block_write_u32_m8k16v1(C, N * sizeof(float), M, N * sizeof(float), (int2)(n + nn * tN, m + mm * tM), as_uint8(sum[mm][nn]));
+        }
+    }
+}
+
+#endif // defined(cl_intel_subgroups) && defined(cl_intel_subgroups_short) && defined(cl_intel_required_subgroup_size)