summaryrefslogtreecommitdiff
path: root/llama.cpp/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp
diff options
context:
space:
mode:
Diffstat (limited to 'llama.cpp/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp')
-rw-r--r--llama.cpp/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp124
1 files changed, 124 insertions, 0 deletions
diff --git a/llama.cpp/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp b/llama.cpp/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp
new file mode 100644
index 0000000..beea529
--- /dev/null
+++ b/llama.cpp/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_nc.comp
@@ -0,0 +1,124 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : enable
+#extension GL_EXT_shader_16bit_storage : require
+
+#define BLOCK_SIZE 32
+#define FLOAT_TYPE float
+
+layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+#include "mul_mat_vec_iface.glsl"
+
+layout (push_constant) uniform parameter
+{
+ uint ncols_x;
+ uint nrows_x;
+ uint row_stride_x;
+ uint channel_stride_x;
+ uint channel_stride_y;
+ uint channel_x_divisor;
+ uint ne12;
+ uint b_offset;
+ uint d_offset;
+ uint nb03;
+ uint nb13;
+ uint nb23;
+ uint fusion_flags;
+} p;
+
+shared FLOAT_TYPE tmp[BLOCK_SIZE];
+
+void main() {
+ const uint tid = gl_LocalInvocationID.x;
+ const uint row_x = gl_GlobalInvocationID.y;
+ const uint channel = gl_GlobalInvocationID.z;
+ const uint i3 = gl_WorkGroupID.x;
+ const uint channel_x = channel / p.channel_x_divisor;
+ const uint channel_y = channel % p.ne12;
+
+ const uint nrows_y = p.ncols_x;
+ const uint nrows_dst = p.nrows_x;
+ const uint row_dst = row_x;
+
+ const uint idst = i3*p.nb23 + channel*nrows_dst + row_dst;
+
+ FLOAT_TYPE temp = 0.0f;
+
+ // Detect alignment for vector loads
+ bool is_aligned = (p.ncols_x % 4) == 0 && (p.row_stride_x % 4) == 0 && (p.channel_stride_x % 4) == 0;
+
+ for (uint col_x0 = 0; col_x0 < p.ncols_x;) {
+
+ // Unroll 2x and do vec4 loads if aligned
+ const uint unroll_count = 2;
+ if (col_x0 + unroll_count * 4 * BLOCK_SIZE <= p.ncols_x && is_aligned) {
+ [[unroll]] for (uint i = 0; i < unroll_count; ++i) {
+ const uint col_x = col_x0 + 4*tid;
+
+ const uint row_y = col_x;
+
+ const uint ix = i3*p.nb03 + channel_x*p.channel_stride_x + row_x*p.row_stride_x + col_x;
+ const uint iy = i3*p.nb13 + channel_y*p.channel_stride_y + row_y;
+
+ const vec4 av4 = vec4(data_a_v4[ix / 4]);
+ const vec4 bv4 = vec4(data_b_v4[iy / 4]);
+
+ temp += dot(av4, bv4);
+
+ col_x0 += 4*BLOCK_SIZE;
+ }
+ // do vec4 loads if aligned
+ } else if (col_x0 + 4*BLOCK_SIZE <= p.ncols_x && is_aligned) {
+ const uint col_x = col_x0 + 4*tid;
+
+ const uint row_y = col_x;
+
+ const uint ix = i3*p.nb03 + channel_x*p.channel_stride_x + row_x*p.row_stride_x + col_x;
+ const uint iy = i3*p.nb13 + channel_y*p.channel_stride_y + row_y;
+
+ const vec4 av4 = vec4(data_a_v4[ix / 4]);
+ const vec4 bv4 = vec4(data_b_v4[iy / 4]);
+
+ temp += dot(av4, bv4);
+
+ col_x0 += 4*BLOCK_SIZE;
+ } else {
+ const uint col_x = col_x0 + tid;
+ if (col_x >= p.ncols_x) {
+ break;
+ }
+
+ const uint row_y = col_x;
+
+ const uint ix = i3*p.nb03 + channel_x*p.channel_stride_x + row_x*p.row_stride_x + col_x;
+ const uint iy = i3*p.nb13 + channel_y*p.channel_stride_y + row_y;
+
+ const FLOAT_TYPE xi = FLOAT_TYPE(data_a[ix]);
+
+ temp = fma(xi, FLOAT_TYPE(data_b[iy]), temp);
+ col_x0 += BLOCK_SIZE;
+ }
+ }
+
+ tmp[tid] = temp;
+
+ // sum up partial sums and write back result
+ barrier();
+ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+ if (tid < s) {
+ tmp[tid] += tmp[tid + s];
+ }
+ barrier();
+ }
+
+ if (tid == 0) {
+ if ((p.fusion_flags & MAT_VEC_FUSION_FLAGS_BIAS0) != 0) {
+ tmp[0] += FLOAT_TYPE(data_fuse0[idst]);
+ }
+ if ((p.fusion_flags & MAT_VEC_FUSION_FLAGS_BIAS1) != 0) {
+ tmp[0] += FLOAT_TYPE(data_fuse1[idst]);
+ }
+ data_d[idst] = tmp[0];
+ }
+}