ggml: replace conv 1D - 2D stage_0 and stage_1 with im2col and mul_mat

include/ggml/ggml.h

@@ -403,13 +403,8 @@ extern "C" {
  GGML_OP_ROPE_BACK,
  GGML_OP_ALIBI,
  GGML_OP_CLAMP,
- GGML_OP_CONV_1D,
- GGML_OP_CONV_1D_STAGE_0, // internal
- GGML_OP_CONV_1D_STAGE_1, // internal
  GGML_OP_CONV_TRANSPOSE_1D,
- GGML_OP_CONV_2D,
- GGML_OP_CONV_2D_STAGE_0, // internal
- GGML_OP_CONV_2D_STAGE_1, // internal
+ GGML_OP_IM2COL,
  GGML_OP_CONV_TRANSPOSE_2D,
  GGML_OP_POOL_1D,
  GGML_OP_POOL_2D,
@@ -1398,6 +1393,18 @@ extern "C" {
  float min,
  float max);
 
+ GGML_API struct ggml_tensor * ggml_im2col(
+ struct ggml_context * ctx,
+ struct ggml_tensor * a,
+ struct ggml_tensor * b,
+ int s0,
+ int s1,
+ int p0,
+ int p1,
+ int d0,
+ int d1,
+ bool is_2D);
+
  GGML_API struct ggml_tensor * ggml_conv_1d(
  struct ggml_context * ctx,
  struct ggml_tensor * a,

src/ggml-cuda.cu

@@ -39,6 +39,7 @@
 #define cudaDeviceCanAccessPeer hipDeviceCanAccessPeer
 #define cudaDeviceDisablePeerAccess hipDeviceDisablePeerAccess
 #define cudaDeviceEnablePeerAccess hipDeviceEnablePeerAccess
+#define cudaDeviceGetMemPool hipDeviceGetMemPool
 #define cudaDeviceProp hipDeviceProp_t
 #define cudaDeviceSynchronize hipDeviceSynchronize
 #define cudaError_t hipError_t
@@ -48,13 +49,15 @@
 #define cudaEvent_t hipEvent_t
 #define cudaEventDestroy hipEventDestroy
 #define cudaFree hipFree
+#define cudaFreeAsync hipFreeAsync
 #define cudaFreeHost hipHostFree
 #define cudaGetDevice hipGetDevice
 #define cudaGetDeviceCount hipGetDeviceCount
 #define cudaGetDeviceProperties hipGetDeviceProperties
 #define cudaGetErrorString hipGetErrorString
 #define cudaGetLastError hipGetLastError
 #define cudaMalloc hipMalloc
+#define cudaMallocFromPoolAsync hipMallocFromPoolAsync
 #define cudaMallocHost(ptr, size) hipHostMalloc(ptr, size, hipHostMallocDefault)
 #define cudaMemcpy hipMemcpy
 #define cudaMemcpy2DAsync hipMemcpy2DAsync
@@ -63,6 +66,9 @@
 #define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost
 #define cudaMemcpyHostToDevice hipMemcpyHostToDevice
 #define cudaMemcpyKind hipMemcpyKind
+#define cudaMemPool_t hipMemPool_t
+#define cudaMemPoolAttrReleaseThreshold hipMemPoolAttrReleaseThreshold
+#define cudaMemPoolSetAttribute hipMemPoolSetAttribute
 #define cudaMemset hipMemset
 #define cudaMemsetAsync hipMemsetAsync
 #define cudaOccupancyMaxPotentialBlockSize hipOccupancyMaxPotentialBlockSize
@@ -4470,6 +4476,13 @@ static __device__ void cpy_1_f32_f16(const char * cxi, char * cdsti) {
  *dsti = __float2half(*xi);
 }
 
+static __device__ void cpy_1_f16_f16(const char * cxi, char * cdsti) {
+ const half * xi = (const half *) cxi;
+ half * dsti = (half *) cdsti;
+
+ *dsti = *xi;
+}
+
 template <cpy_kernel_t cpy_1>
 static __global__ void cpy_f32_f16(const char * cx, char * cdst, const int ne,
  const int ne00, const int ne01, const int nb00, const int nb01, const int nb02,
@@ -4723,6 +4736,25 @@ static __global__ void clamp_f32(const float * x, float * dst, const float min,
  dst[i] = x[i] < min ? min : (x[i] > max ? max : x[i]);
 }
 
+static __global__ void im2col_f32_f16(
+ const float * x, half * dst,
+ int ofs0, int ofs1, int IW, int IH, int CHW,
+ int s0, int s1, int p0, int p1, int d0, int d1) {
+ const int iiw = blockIdx.z * s0 + threadIdx.z * d0 - p0;
+ const int iih = blockIdx.y * s1 + threadIdx.y * d1 - p1;
+
+ const int offset_dst =
+ (threadIdx.x * gridDim.y * gridDim.z + blockIdx.y * gridDim.z + blockIdx.z) * CHW +
+ (blockIdx.x * (blockDim.y * blockDim.z) + threadIdx.y * blockDim.z + threadIdx.z);
+
+ if (iih < 0 || iih >= IH || iiw < 0 || iiw >= IW) {
+ dst[offset_dst] = __float2half(0.0f);
+ } else {
+ const int offset_src = threadIdx.x * ofs0 + blockIdx.x * ofs1;
+ dst[offset_dst] = __float2half(x[offset_src + iih * IW + iiw]);
+ }
+}
+
 template<int qk, int qr, dequantize_kernel_t dq>
 static void get_rows_cuda(const void * x, const int32_t * y, float * dst, const int nrows, const int ncols, cudaStream_t stream) {
  const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1);
@@ -5612,6 +5644,16 @@ static void ggml_cpy_f32_f16_cuda(
  (cx, cdst, ne, ne00, ne01, nb00, nb01, nb02, ne10, ne11, nb10, nb11, nb12);
 }
 
+static void ggml_cpy_f16_f16_cuda(
+ const char * cx, char * cdst, const int ne,
+ const int ne00, const int ne01, const int nb00, const int nb01, const int nb02,
+ const int ne10, const int ne11, const int nb10, const int nb11, const int nb12, cudaStream_t stream) {
+
+ const int num_blocks = (ne + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
+ cpy_f32_f16<cpy_1_f16_f16><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
+ (cx, cdst, ne, ne00, ne01, nb00, nb01, nb02, ne10, ne11, nb10, nb11, nb12);
+}
+
 static void scale_f32_cuda(const float * x, float * dst, const float scale, const int k, cudaStream_t stream) {
  const int num_blocks = (k + CUDA_SCALE_BLOCK_SIZE - 1) / CUDA_SCALE_BLOCK_SIZE;
  scale_f32<<<num_blocks, CUDA_SCALE_BLOCK_SIZE, 0, stream>>>(x, dst, scale, k);
@@ -5695,6 +5737,15 @@ static void soft_max_f32_cuda(const float * x, float * dst, const int ncols_x, c
  soft_max_f32<<<block_nums, block_dims, 0, stream>>>(x, dst, ncols_x);
 }
 
+static void im2col_f32_f16_cuda(const float * x, half * dst,
+ int OH, int IW, int IH, int OW, int IC,
+ int KH, int KW, int N, int ofs0, int ofs1,
+ int s0, int s1, int p0, int p1, int d0, int d1, cudaStream_t stream) {
+ dim3 block_nums(IC, OH, OW);
+ dim3 block_dims(N, KH, KW);
+ im2col_f32_f16<<<block_nums, block_dims, 0, stream>>>(x, dst, ofs0, ofs1, IW, IH, (IC * KH * KW), s0, s1, p0, p1, d0, d1);
+}
+
 // buffer pool for cuda
 #define MAX_CUDA_BUFFERS 256
 
@@ -6477,7 +6528,7 @@ inline void ggml_cuda_op_mul_mat_cublas(
  src1_as_f16 = (half *) ggml_cuda_pool_malloc_async(ne * sizeof(half), &src1_as, id, stream);
  to_fp16_cuda(src1_ddf_i, src1_as_f16, ne, stream);
  }
- const half * src1_ptr = src1->type == GGML_TYPE_F16 ? (const half *) src1_ddq_i : src1_as_f16;
+ const half * src1_ptr = src1->type == GGML_TYPE_F16 ? (const half *) src1_ddf_i : src1_as_f16;
  size_t dst_f16_as = 0;
  half * dst_f16 = (half *) ggml_cuda_pool_malloc_async(row_diff*src1_ncols * sizeof(half), &dst_f16_as, id, stream);
 
@@ -6653,6 +6704,45 @@ inline void ggml_cuda_op_alibi(
  (void) src1_dd;
 }
 
+inline void ggml_cuda_op_im2col(
+ const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst,
+ const float * src0_dd, const float * src1_dd, float * dst_dd, const cudaStream_t & main_stream) {
+
+ GGML_ASSERT(src0->type == GGML_TYPE_F16);
+ GGML_ASSERT(src1->type == GGML_TYPE_F32);
+ GGML_ASSERT( dst->type == GGML_TYPE_F16);
+
+ const int32_t s0 = ((const int32_t*)(dst->op_params))[0];
+ const int32_t s1 = ((const int32_t*)(dst->op_params))[1];
+ const int32_t p0 = ((const int32_t*)(dst->op_params))[2];
+ const int32_t p1 = ((const int32_t*)(dst->op_params))[3];
+ const int32_t d0 = ((const int32_t*)(dst->op_params))[4];
+ const int32_t d1 = ((const int32_t*)(dst->op_params))[5];
+
+ const bool is_2D = ((const int32_t*)(dst->op_params))[6] == 1;
+
+ const int64_t N = src1->ne[is_2D ? 3 : 2];
+ const int64_t IC = src1->ne[is_2D ? 2 : 1];
+ const int64_t IH = is_2D ? src1->ne[1] : 1;
+ const int64_t IW = src1->ne[0];
+
+ const int64_t KH = is_2D ? src0->ne[1] : 1;
+ const int64_t KW = src0->ne[0];
+
+ const int64_t OH = is_2D ? dst->ne[2] : 1;
+ const int64_t OW = dst->ne[1];
+
+ const size_t ofs0 = src1->nb[is_2D ? 3 : 2] / 4; // nb is byte offset, src is type float32
+ const size_t ofs1 = src1->nb[is_2D ? 2 : 1] / 4; // nb is byte offset, src is type float32
+
+ im2col_f32_f16_cuda(src1_dd, (half*) dst_dd,
+ OH, IW, IH, OW, IC, KH, KW, N,
+ ofs0, ofs1, s0, s1, p0, p1, d0, d1, main_stream);
+
+ (void) src0;
+ (void) src0_dd;
+}
+
 inline void ggml_cuda_op_diag_mask_inf(
  const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst,
  const float * src0_dd, const float * src1_dd, float * dst_dd, const cudaStream_t & main_stream) {
@@ -7543,6 +7633,9 @@ static void ggml_cuda_cpy(const ggml_tensor * src0, const ggml_tensor * src1, gg
  } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F16) {
  ggml_cpy_f32_f16_cuda(src0_ddc, src1_ddc, ne, ne00, ne01, nb00, nb01, nb02,
  ne10, ne11, nb10, nb11, nb12, main_stream);
+ } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16) {
+ ggml_cpy_f16_f16_cuda(src0_ddc, src1_ddc, ne, ne00, ne01, nb00, nb01, nb02,
+ ne10, ne11, nb10, nb11, nb12, main_stream);
  } else {
  fprintf(stderr, "%s: unsupported type combination (%s to %s)\n", __func__,
  ggml_type_name(src0->type), ggml_type_name(src1->type));
@@ -7574,6 +7667,10 @@ static void ggml_cuda_alibi(const ggml_tensor * src0, const ggml_tensor * src1,
  ggml_cuda_op_flatten(src0, src1, dst, ggml_cuda_op_alibi);
 }
 
+void ggml_cuda_im2col(const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+ ggml_cuda_op_flatten(src0, src1, dst, ggml_cuda_op_im2col);
+}
+
 static void ggml_cuda_nop(const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
  (void) src0;
  (void) src1;
@@ -7937,6 +8034,9 @@ bool ggml_cuda_compute_forward(struct ggml_compute_params * params, struct ggml_
  case GGML_OP_ALIBI:
  func = ggml_cuda_alibi;
  break;
+ case GGML_OP_IM2COL:
+ func = ggml_cuda_im2col;
+ break;
  default:
  return false;
  }

src/ggml-metal.m

@@ -86,6 +86,7 @@
  GGML_METAL_DECL_KERNEL(rms_norm);
  GGML_METAL_DECL_KERNEL(norm);
  GGML_METAL_DECL_KERNEL(mul_mv_f32_f32);
+ GGML_METAL_DECL_KERNEL(mul_mv_f16_f16);
  GGML_METAL_DECL_KERNEL(mul_mv_f16_f32);
  GGML_METAL_DECL_KERNEL(mul_mv_f16_f32_1row);
  GGML_METAL_DECL_KERNEL(mul_mv_f16_f32_l4);
@@ -114,6 +115,7 @@
  GGML_METAL_DECL_KERNEL(rope_f32);
  GGML_METAL_DECL_KERNEL(rope_f16);
  GGML_METAL_DECL_KERNEL(alibi_f32);
+ GGML_METAL_DECL_KERNEL(im2col_f16);
  GGML_METAL_DECL_KERNEL(cpy_f32_f16);
  GGML_METAL_DECL_KERNEL(cpy_f32_f32);
  GGML_METAL_DECL_KERNEL(cpy_f16_f16);
@@ -287,6 +289,7 @@ static void ggml_metal_log(enum ggml_log_level level, const char* format, ...){
  GGML_METAL_ADD_KERNEL(rms_norm);
  GGML_METAL_ADD_KERNEL(norm);
  GGML_METAL_ADD_KERNEL(mul_mv_f32_f32);
+ GGML_METAL_ADD_KERNEL(mul_mv_f16_f16);
  GGML_METAL_ADD_KERNEL(mul_mv_f16_f32);
  GGML_METAL_ADD_KERNEL(mul_mv_f16_f32_1row);
  GGML_METAL_ADD_KERNEL(mul_mv_f16_f32_l4);
@@ -317,6 +320,7 @@ static void ggml_metal_log(enum ggml_log_level level, const char* format, ...){
  GGML_METAL_ADD_KERNEL(rope_f32);
  GGML_METAL_ADD_KERNEL(rope_f16);
  GGML_METAL_ADD_KERNEL(alibi_f32);
+ GGML_METAL_ADD_KERNEL(im2col_f16);
  GGML_METAL_ADD_KERNEL(cpy_f32_f16);
  GGML_METAL_ADD_KERNEL(cpy_f32_f32);
  GGML_METAL_ADD_KERNEL(cpy_f16_f16);
@@ -386,6 +390,7 @@ void ggml_metal_free(struct ggml_metal_context * ctx) {
  GGML_METAL_DEL_KERNEL(rms_norm);
  GGML_METAL_DEL_KERNEL(norm);
  GGML_METAL_DEL_KERNEL(mul_mv_f32_f32);
+ GGML_METAL_DEL_KERNEL(mul_mv_f16_f16);
  GGML_METAL_DEL_KERNEL(mul_mv_f16_f32);
  GGML_METAL_DEL_KERNEL(mul_mv_f16_f32_1row);
  GGML_METAL_DEL_KERNEL(mul_mv_f16_f32_l4);
@@ -416,6 +421,7 @@ void ggml_metal_free(struct ggml_metal_context * ctx) {
  GGML_METAL_DEL_KERNEL(rope_f32);
  GGML_METAL_DEL_KERNEL(rope_f16);
  GGML_METAL_DEL_KERNEL(alibi_f32);
+ GGML_METAL_DEL_KERNEL(im2col_f16);
  GGML_METAL_DEL_KERNEL(cpy_f32_f16);
  GGML_METAL_DEL_KERNEL(cpy_f32_f32);
  GGML_METAL_DEL_KERNEL(cpy_f16_f16);
@@ -1030,7 +1036,7 @@ void ggml_metal_graph_compute(
  [encoder setBytes:&ne00 length:sizeof(ne00) atIndex:2];
  [encoder setBytes:&ne01 length:sizeof(ne01) atIndex:3];
  [encoder setBytes:&ne02 length:sizeof(ne02) atIndex:4];
- [encoder setThreadgroupMemoryLength:nth/32*sizeof(float) atIndex:0];
+ [encoder setThreadgroupMemoryLength:GGML_PAD(nth/32*sizeof(float), 16) atIndex:0];
 
  [encoder dispatchThreadgroups:MTLSizeMake(ne01*ne02*ne03, 1, 1) threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
  } break;
@@ -1139,20 +1145,26 @@ void ggml_metal_graph_compute(
  switch (src0t) {
  case GGML_TYPE_F32:
  {
+ GGML_ASSERT(src1t == GGML_TYPE_F32);
  [encoder setComputePipelineState:ctx->pipeline_mul_mv_f32_f32];
  nrows = 4;
  } break;
  case GGML_TYPE_F16:
  {
  nth0 = 32;
  nth1 = 1;
- if (ne11 * ne12 < 4) {
- [encoder setComputePipelineState:ctx->pipeline_mul_mv_f16_f32_1row];
- } else if (ne00 >= 128 && ne01 >= 8 && ne00%4 == 0) {
- [encoder setComputePipelineState:ctx->pipeline_mul_mv_f16_f32_l4];
- nrows = ne11;
+ if (src1t == GGML_TYPE_F32) {
+ if (ne11 * ne12 < 4) {
+ [encoder setComputePipelineState:ctx->pipeline_mul_mv_f16_f32_1row];
+ } else if (ne00 >= 128 && ne01 >= 8 && ne00%4 == 0) {
+ [encoder setComputePipelineState:ctx->pipeline_mul_mv_f16_f32_l4];
+ nrows = ne11;
+ } else {
+ [encoder setComputePipelineState:ctx->pipeline_mul_mv_f16_f32];
+ nrows = 4;
+ }
  } else {
- [encoder setComputePipelineState:ctx->pipeline_mul_mv_f16_f32];
+ [encoder setComputePipelineState:ctx->pipeline_mul_mv_f16_f16];
  nrows = 4;
  }
  } break;
@@ -1342,7 +1354,7 @@ void ggml_metal_graph_compute(
  [encoder setBytes:&ne00 length:sizeof( int64_t) atIndex:2];
  [encoder setBytes:&nb01 length:sizeof(uint64_t) atIndex:3];
  [encoder setBytes:&eps length:sizeof( float) atIndex:4];
- [encoder setThreadgroupMemoryLength:nth/32*sizeof(float) atIndex:0];
+ [encoder setThreadgroupMemoryLength:GGML_PAD(nth/32*sizeof(float), 16) atIndex:0];
 
  const int64_t nrows = ggml_nrows(src0);
 
@@ -1361,7 +1373,7 @@ void ggml_metal_graph_compute(
  [encoder setBytes:&ne00 length:sizeof( int64_t) atIndex:2];
  [encoder setBytes:&nb01 length:sizeof(uint64_t) atIndex:3];
  [encoder setBytes:&eps length:sizeof( float) atIndex:4];
- [encoder setThreadgroupMemoryLength:nth*sizeof(float) atIndex:0];
+ [encoder setThreadgroupMemoryLength:GGML_PAD(nth*sizeof(float), 16) atIndex:0];
 
  const int64_t nrows = ggml_nrows(src0);
 
@@ -1464,6 +1476,58 @@ void ggml_metal_graph_compute(
 
  [encoder dispatchThreadgroups:MTLSizeMake(ne01, ne02, ne03) threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
  } break;
+ case GGML_OP_IM2COL:
+ {
+ GGML_ASSERT(src0->type == GGML_TYPE_F16);
+ GGML_ASSERT(src1->type == GGML_TYPE_F32);
+ GGML_ASSERT( dst->type == GGML_TYPE_F16);
+
+ const int32_t s0 = ((const int32_t *)(dst->op_params))[0];
+ const int32_t s1 = ((const int32_t *)(dst->op_params))[1];
+ const int32_t p0 = ((const int32_t *)(dst->op_params))[2];
+ const int32_t p1 = ((const int32_t *)(dst->op_params))[3];
+ const int32_t d0 = ((const int32_t *)(dst->op_params))[4];
+ const int32_t d1 = ((const int32_t *)(dst->op_params))[5];
+ const bool is_2D = ((const int32_t *)(dst->op_params))[6] == 1;
+
+ const int32_t N = src1->ne[is_2D ? 3 : 2];
+ const int32_t IC = src1->ne[is_2D ? 2 : 1];
+ const int32_t IH = is_2D ? src1->ne[1] : 1;
+ const int32_t IW = src1->ne[0];
+
+ const int32_t KH = is_2D ? src0->ne[1] : 1;
+ const int32_t KW = src0->ne[0];
+
+ const int32_t OH = is_2D ? dst->ne[2] : 1;
+ const int32_t OW = dst->ne[1];
+
+ const int32_t CHW = IC * KH * KW;
+
+ const int32_t ofs0 = src1->nb[is_2D ? 3 : 2] / 4;
+ const int32_t ofs1 = src1->nb[is_2D ? 2 : 1] / 4;
+
+ switch (src0->type) {
+ case GGML_TYPE_F32: GGML_ASSERT(false && "not implemented"); break;
+ case GGML_TYPE_F16: [encoder setComputePipelineState:ctx->pipeline_im2col_f16]; break;
+ default: GGML_ASSERT(false);
+ };
+
+ [encoder setBuffer:id_src1 offset:offs_src1 atIndex:0];
+ [encoder setBuffer:id_dst offset:offs_dst atIndex:1];
+ [encoder setBytes:&ofs0 length:sizeof( int32_t) atIndex:2];
+ [encoder setBytes:&ofs1 length:sizeof( int32_t) atIndex:3];
+ [encoder setBytes:&IW length:sizeof( int32_t) atIndex:4];
+ [encoder setBytes:&IH length:sizeof( int32_t) atIndex:5];
+ [encoder setBytes:&CHW length:sizeof( int32_t) atIndex:6];
+ [encoder setBytes:&s0 length:sizeof( int32_t) atIndex:7];
+ [encoder setBytes:&s1 length:sizeof( int32_t) atIndex:8];
+ [encoder setBytes:&p0 length:sizeof( int32_t) atIndex:9];
+ [encoder setBytes:&p1 length:sizeof( int32_t) atIndex:10];
+ [encoder setBytes:&d0 length:sizeof( int32_t) atIndex:11];
+ [encoder setBytes:&d1 length:sizeof( int32_t) atIndex:12];
+
+ [encoder dispatchThreadgroups:MTLSizeMake(IC, OH, OW) threadsPerThreadgroup:MTLSizeMake(N, KH, KW)];
+ } break;
  case GGML_OP_DUP:
  case GGML_OP_CPY:
  case GGML_OP_CONT: