@@ -6546,29 +6546,29 @@ void ggml_compute_forward_im2col_back_f32(
65466546 }
65476547}
65486548
6549- static void ggml_call_mul_mat (
6550- const ggml_compute_params * params,
6551- int64_t m, int64_t n, int64_t k,
6552- void * a, void * b, void * c) {
6553-
6549+ static void ggml_call_mul_mat (ggml_type T, const ggml_compute_params * params, int64_t m, int64_t n, int64_t k,
6550+ void * a, void * b, void * c) {
6551+ const ggml_type_traits * traits = ggml_get_type_traits (T);
65546552 struct ggml_tensor src1 = {};
6553+ src1.type = T;
65556554 src1.ne [0 ] = k;
65566555 src1.ne [1 ] = m;
65576556 src1.ne [2 ] = 1 ;
65586557 src1.ne [3 ] = 1 ;
6559- src1.nb [0 ] = sizeof ( float ) ;
6560- src1.nb [1 ] = k * sizeof ( float ) ;
6558+ src1.nb [0 ] = traits-> type_size ;
6559+ src1.nb [1 ] = k * traits-> type_size ;
65616560 src1.nb [2 ] = src1.nb [1 ];
65626561 src1.nb [3 ] = src1.nb [2 ];
65636562 src1.data = a;
65646563
65656564 struct ggml_tensor src0 = {};
6565+ src0.type = T;
65666566 src0.ne [0 ] = k;
65676567 src0.ne [1 ] = n;
65686568 src0.ne [2 ] = 1 ;
65696569 src0.ne [3 ] = 1 ;
6570- src0.nb [0 ] = sizeof ( float ) ;
6571- src0.nb [1 ] = k * sizeof ( float ) ;
6570+ src0.nb [0 ] = traits-> type_size ;
6571+ src0.nb [1 ] = k * traits-> type_size ;
65726572 src0.nb [2 ] = src0.nb [1 ];
65736573 src0.nb [3 ] = src0.nb [2 ];
65746574 src0.data = b;
@@ -6589,17 +6589,18 @@ static void ggml_call_mul_mat(
65896589 ggml_compute_forward_mul_mat (params, &dst);
65906590}
65916591
6592-
65936592// ggml_compute_forward_conv_2d
65946593
6595- static void ggml_compute_forward_conv_2d_f32 (
6596- const ggml_compute_params * params,
6597- const ggml_tensor * kernel , // [KW, KH, IC, OC] - fp32
6598- const ggml_tensor * src , // [W, H, C , N]
6599- ggml_tensor * dst) { // [OW, OH, OC, N]
6594+ static void ggml_compute_forward_conv_2d_impl ( const ggml_compute_params * params,
6595+ const ggml_tensor * kernel, // [KW, KH, IC, OC]
6596+ const ggml_tensor * src , // [W, H, C, N]
6597+ ggml_tensor * dst , // [OW, OH, OC , N]
6598+ ggml_type kernel_type) {
66006599
66016600 GGML_ASSERT (ggml_is_contiguous (kernel));
6602- GGML_ASSERT (kernel->type == GGML_TYPE_F32);
6601+ GGML_ASSERT (kernel->type == kernel_type);
6602+
6603+ const ggml_type_traits * traits = ggml_get_type_traits (kernel_type);
66036604
66046605 const int32_t stride_x = dst->op_params [0 ];
66056606 const int32_t stride_y = dst->op_params [1 ];
@@ -6620,20 +6621,20 @@ static void ggml_compute_forward_conv_2d_f32(
66206621 const int64_t dst_h = dst->ne [1 ];
66216622
66226623 float * src_data = (float *) src->data ;
6623- float * knl_data = ( float *) kernel->data ;
6624+ void * knl_data = kernel->data ;
66246625 float * dst_data = (float *) dst->data ;
66256626
66266627 const int64_t knl_n = knl_w * knl_h * c_in;
66276628 const int64_t patch_total = dst->ne [3 ] * dst_w * dst_h;
66286629
6629- const int64_t space_per_patch = knl_n * sizeof ( float ) + c_out * sizeof (float );
6630+ const int64_t space_per_patch = knl_n * traits-> type_size + c_out * sizeof (float );
66306631 const int64_t batch_size = params->wsize / space_per_patch;
66316632 const int64_t patches_per_batch = batch_size > 8 ? (batch_size / 8 ) * 8 : batch_size;
66326633 const int64_t batch_n = (patch_total + patches_per_batch - 1 ) / patches_per_batch;
66336634
66346635 GGML_ASSERT (patches_per_batch > 0 && batch_size >= 1 );
66356636
6636- float * tmp = ( float *) params->wdata ;
6637+ void * tmp = params->wdata ;
66376638
66386639 for (int64_t batch_i = 0 ; batch_i < batch_n; ++batch_i) {
66396640
@@ -6653,7 +6654,7 @@ static void ggml_compute_forward_conv_2d_f32(
66536654 const int64_t src_y = p % dst_w;
66546655
66556656 float * src_base = (float *)((char *)src_data + batch_n * src->nb [3 ]);
6656- float * dst_row = tmp + (p % patches_per_batch) * knl_n;
6657+ char * dst_row = ( char *) tmp + (p % patches_per_batch) * knl_n * traits-> type_size ;
66576658
66586659 for (int64_t ic = 0 ; ic < c_in; ++ic) {
66596660 for (int64_t ky = 0 ; ky < knl_h; ++ky) {
@@ -6663,11 +6664,19 @@ static void ggml_compute_forward_conv_2d_f32(
66636664
66646665 int64_t dst_idx = ic * (knl_h * knl_w) + ky * knl_w + kx;
66656666
6667+ float src_val;
66666668 if (sy < 0 || sy >= src_h || sx < 0 || sx >= src_w) {
6667- dst_row[dst_idx] = 0 .0f ;
6669+ src_val = 0 .0f ;
66686670 } else {
66696671 float * src_ptr = (float *)((char *)src_base + sx * src->nb [0 ] + sy * src->nb [1 ] + ic * src->nb [2 ]);
6670- dst_row[dst_idx] = *src_ptr;
6672+ src_val = *src_ptr;
6673+ }
6674+
6675+ char * element_ptr = dst_row + dst_idx * traits->type_size ;
6676+ if (kernel_type == GGML_TYPE_F32) {
6677+ *(float *) element_ptr = src_val;
6678+ } else if (kernel_type == GGML_TYPE_F16) {
6679+ *(ggml_fp16_t *) element_ptr = GGML_FP32_TO_FP16 (src_val);
66716680 }
66726681 }
66736682 }
@@ -6676,11 +6685,10 @@ static void ggml_compute_forward_conv_2d_f32(
66766685
66776686 ggml_barrier (params->threadpool );
66786687
6679- float * gemm_output = tmp + patches_per_batch * knl_n;
6688+ float * gemm_output = ( float *) (( char *) tmp + patches_per_batch * knl_n * traits-> type_size ) ;
66806689
66816690 // GEMM: patches[patch_n, knl_n] × kernel[knl_n, c_out] = output[patch_n, c_out]
6682- ggml_call_mul_mat (params, patch_n, c_out, knl_n,
6683- tmp, knl_data, gemm_output);
6691+ ggml_call_mul_mat (kernel_type, params, patch_n, c_out, knl_n, tmp, knl_data, gemm_output);
66846692
66856693 ggml_barrier (params->threadpool );
66866694
@@ -6698,7 +6706,6 @@ static void ggml_compute_forward_conv_2d_f32(
66986706
66996707 for (int64_t oc = 0 ; oc < c_out; ++oc) {
67006708 const float value = gemm_output[i * c_out + oc];
6701- // Write to WHCN layout: dst[w, h, c, n]
67026709 float * dst_ptr = (float *)((char *)dst_data + dst_x * dst->nb [0 ] + dst_y * dst->nb [1 ] + oc * dst->nb [2 ] + batch_n * dst->nb [3 ]);
67036710 *dst_ptr = value;
67046711 }
@@ -6713,11 +6720,7 @@ void ggml_compute_forward_conv_2d(
67136720 const ggml_tensor * src0 = dst->src [0 ];
67146721 const ggml_tensor * src1 = dst->src [1 ];
67156722
6716- if (src0->type == GGML_TYPE_F16) {
6717- GGML_ASSERT (false && " F16 not supported yet" );
6718- } else {
6719- ggml_compute_forward_conv_2d_f32 (params, src0, src1, dst);
6720- }
6723+ ggml_compute_forward_conv_2d_impl (params, src0, src1, dst, src0->type );
67216724}
67226725
67236726// ggml_compute_forward_conv_transpose_2d
0 commit comments