| /* |
| * Copyright 2024 Google LLC |
| * |
| * Licensed under the Apache License, Version 2.0 (the "License"); |
| * you may not use this file except in compliance with the License. |
| * You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| |
| // Convolution based on Kelvin ops |
| // Data types: input: s8, filter: s8, bias: s32 |
| // Special case for 48x3x1x48 filter |
| |
| #include "tflm/opt/conv_util.h" |
| |
| namespace kelvin::opt { |
| |
| void ConvS8K3x1D48( |
| const tflite::ConvParams& params, const int32_t* output_multiplier, |
| const int32_t* output_shift, const tflite::RuntimeShape& input_shape, |
| const int8_t* input_data, const tflite::RuntimeShape& filter_shape, |
| const int8_t* filter_data, const tflite::RuntimeShape& bias_shape, |
| const int32_t* bias_data, const tflite::RuntimeShape& output_shape, |
| int8_t* output_data) { |
| const auto batches = MatchingDim(input_shape, 0, output_shape, 0); |
| const int stride_width = params.stride_width; |
| const int stride_height = params.stride_height; |
| const int dilation_width_factor = params.dilation_width_factor; |
| const int dilation_height_factor = params.dilation_height_factor; |
| const int pad_width = params.padding_values.width; |
| const int pad_height = params.padding_values.height; |
| const int input_width = input_shape.Dims(2); |
| const int input_depth = input_shape.Dims(3); |
| const int32_t input_offset = params.input_offset; |
| const int filter_height = filter_shape.Dims(1); |
| const int filter_width = filter_shape.Dims(2); |
| const int filter_depth = filter_shape.Dims(3); |
| const int output_height = output_shape.Dims(1); |
| const int output_depth = output_shape.Dims(3); |
| const int32_t output_offset = params.output_offset; |
| const int32_t output_activation_min = params.quantized_activation_min; |
| const int32_t output_activation_max = params.quantized_activation_max; |
| |
| TFLITE_DCHECK(batches == 1); |
| TFLITE_DCHECK(filter_depth == input_depth); |
| TFLITE_DCHECK(filter_height == 3); |
| TFLITE_DCHECK(filter_width == 1); |
| TFLITE_DCHECK(input_width == 1); |
| TFLITE_DCHECK(stride_width == 1); |
| TFLITE_DCHECK(stride_height == 1); |
| TFLITE_DCHECK(dilation_width_factor == 1); |
| TFLITE_DCHECK(dilation_height_factor == 1); |
| TFLITE_DCHECK(pad_width == 0); |
| TFLITE_DCHECK(pad_height == 0); |
| |
| int32_t bias[48 * 4]; |
| int32_t mult[48 * 4]; |
| int32_t shft[48 * 4]; |
| Swizzle<48>(bias_data, bias); |
| Swizzle<48>(output_multiplier, mult); |
| Swizzle<48, true>(output_shift, shft); |
| |
| int8_t juggled_filter_data[48 / 8][3][1][48 / 4][8][4]; |
| Filter_N_H_W_M<48>(filter_data, juggled_filter_data[0][0][0][0][0], 3, 1, 48); |
| union { |
| vconv_u8_t conv; |
| uint32_t raw; |
| } cmds; |
| cmds.conv.mode = 0; |
| cmds.conv.start = 0; |
| cmds.conv.stop = 7; |
| cmds.conv.sbias1 = input_offset; |
| cmds.conv.sdata1 = true; |
| cmds.conv.sbias2 = 0; |
| cmds.conv.sdata2 = true; |
| |
| union { |
| vconv_u8_t conv; |
| uint32_t raw; |
| } cmds16; |
| cmds16.conv.mode = 0; |
| cmds16.conv.start = 0; |
| cmds16.conv.stop = 3; |
| cmds16.conv.sbias1 = input_offset; |
| cmds16.conv.sdata1 = true; |
| cmds16.conv.sbias2 = 0; |
| cmds16.conv.sdata2 = true; |
| |
| for (int zo_hi = 0; zo_hi < output_depth; zo_hi += 8) { |
| // For each pixel, the general flow for this kernel looks like: |
| // 1) Reset accumulator and load activations into [v32, v46] |
| // 2) For each group of 32 scalars in the pixel fan-in, run MAC pipeline |
| // 2a) Load subset of activations from [v32, v46] to [v0, v7] |
| // 2b) Load subset of weights |
| // 2c) Run aconv |
| // 3) Run the output pipeline and store. |
| // |
| // For step 1, we'll alias [v32, v46] to [L0, LE]. For most iterations, |
| // we load all of these registers (10 pixels). For remainder iterations, |
| // we load a subset and pad the rest with 0's. The data will be stored as |
| // follows, where each letter represents 16 bytes of a pixel stored into |
| // a register (capitalization used to help distinguish channels in a pixel): |
| // L0 L1 L2 L3 L4 L5 L6 L7 L8 L9 LA LB LC LD LE |
| // Aa AB bB Cc CD dD Ee EF fF Gg GH hH Ii IJ jJ |
| #define L0 v32 |
| #define L1 v33 |
| #define L2 v34 |
| #define L3 v35 |
| #define L4 v36 |
| #define L5 v37 |
| #define L6 v38 |
| #define L7 v39 |
| #define L8 v40 |
| #define L9 v41 |
| #define LA v42 |
| #define LB v43 |
| #define LC v44 |
| #define LD v45 |
| #define LE v46 |
| |
| // We run 5 iterations of step 2, 4 full iterations and one half iteration. |
| // Because each pixel takes 1.5 registers, we have to interleave vmv_v and |
| // vsliden_w_4_vv instructions to ensure the same output channels are stored |
| // in each register per-pixel. As a refresher, vsliden_w_4_vv takes two |
| // register arguments (X and Y), and returns the concatenation of the last |
| // half of X and the first half of Y. ie: |
| // L1 L2 |
| // AB bB |
| // vsliden_w_4_vv(v1, L1, L2); -> v1 = Bb |
| #define CONV_PER_CHANNEL_B8_3X1_48C_MAC_PIPELINE(p_flt) \ |
| { \ |
| /* 1/5 */ \ |
| /* Ky = 0, IC:[0-31] */ \ |
| vmv_v(v0, L0); /* Aa */ \ |
| vsliden_w_4_vv(v1, L1, L2); /* Bb */ \ |
| vmv_v(v2, L3); /* Cc */ \ |
| vsliden_w_4_vv(v3, L4, L5); /* Dd */ \ |
| vmv_v(v4, L6); /* Ee */ \ |
| vsliden_w_4_vv(v5, L7, L8); /* Ff */ \ |
| vmv_v(v6, L9); /* Gg */ \ |
| vsliden_w_4_vv(v7, LA, LB); /* Hh */ \ |
| vld_b_x_m(v56, p_flt + 128 * 0); \ |
| vld_b_x_m(v60, p_flt + 128 * 1); \ |
| aconv_vxv(v48, v0, cmds, v56); \ |
| \ |
| /* 2/5 */ \ |
| /* Ky = 0, IC:[32-47]; Ky = 1, IC:[0-15] */ \ |
| vmv_v(v0, L1); /* AB */ \ |
| vsliden_w_4_vv(v1, L2, L3); /* BC */ \ |
| vmv_v(v2, L4); /* CD */ \ |
| vsliden_w_4_vv(v3, L5, L6); /* DE */ \ |
| vmv_v(v4, L7); /* EF */ \ |
| vsliden_w_4_vv(v5, L8, L9); /* FG */ \ |
| vmv_v(v6, LA); /* GH */ \ |
| vsliden_w_4_vv(v7, LB, LC); /* HI */ \ |
| vld_b_x_m(v56, p_flt + 128 * 2); \ |
| vld_b_x_m(v60, p_flt + 128 * 3); \ |
| aconv_vxv(v48, v0, cmds, v56); \ |
| \ |
| /* 3/5 */ \ |
| /* Ky = 1, IC:[16-47] */ \ |
| vmv_v(v0, L2); /* bB */ \ |
| vsliden_w_4_vv(v1, L3, L4); /* cC */ \ |
| vmv_v(v2, L5); /* dD */ \ |
| vsliden_w_4_vv(v3, L6, L7); /* eE */ \ |
| vmv_v(v4, L8); /* fF */ \ |
| vsliden_w_4_vv(v5, L9, LA); /* gG */ \ |
| vmv_v(v6, LB); /* hH */ \ |
| vsliden_w_4_vv(v7, LC, LD); /* iI */ \ |
| vld_b_x_m(v56, p_flt + 128 * 4); \ |
| vld_b_x_m(v60, p_flt + 128 * 5); \ |
| aconv_vxv(v48, v0, cmds, v56); \ |
| \ |
| /* 4/5 */ \ |
| /* Ky = 2, IC:[0-31] */ \ |
| vmv_v(v0, L3); /* Cc */ \ |
| vsliden_w_4_vv(v1, L4, L5); /* Dd */ \ |
| vmv_v(v2, L6); /* Ee */ \ |
| vsliden_w_4_vv(v3, L4, L5); /* Ff */ \ |
| vmv_v(v4, L9); /* Gg */ \ |
| vsliden_w_4_vv(v5, LA, LB); /* Hh */ \ |
| vmv_v(v6, LC); /* Ii */ \ |
| vsliden_w_4_vv(v7, LD, LE); /* Jj */ \ |
| vld_b_x_m(v56, p_flt + 128 * 6); \ |
| vld_b_x_m(v60, p_flt + 128 * 7); \ |
| aconv_vxv(v48, v0, cmds, v56); \ |
| \ |
| /* 5/5 */ \ |
| /* Ky = 2, IC:[32-47] half iteration */ \ |
| vmv_v(v0, L4); /* C(D- ignored) */ \ |
| vsliden_w_4_vv(v1, L5, L6); /* D(E- ignored) */ \ |
| vmv_v(v2, L7); /* E(F- ignored) */ \ |
| vsliden_w_4_vv(v3, L8, L9); /* F(G- ignored) */ \ |
| vmv_v(v4, LA); /* G(H- ignored) */ \ |
| vsliden_w_4_vv(v5, LB, LC); /* H(I- ignored) */ \ |
| vmv_v(v6, LD); /* I(J- ignored) */ \ |
| /* Pad last iteration with first pixel. Gets ignored by cmd16 */ \ |
| vsliden_w_4_vv(v7, LE, L0); /* J(A- ignored) */ \ |
| vld_b_x_m(v56, p_flt + 128 * 8); /*Load once half iteration*/ \ |
| /* cmds16 runs subset of outer product */ \ |
| aconv_vxv(v48, v0, cmds16, v56); \ |
| } |
| |
| // Iterate over outputs |
| int out_y = 0; |
| for (; out_y + 8 <= output_height; out_y += 8) { |
| // Reset accumulator |
| vdup_w_x_m(v48, 0); |
| vdup_w_x_m(v52, 0); |
| acset_v(v48, v48); |
| |
| const int8_t* p_flt = juggled_filter_data[zo_hi / 8][0][0][0][0]; |
| const int8_t* p_in = input_data + (out_y * input_width * input_depth); |
| |
| // Load 10*48 activations into 10*48*32 = 15 registers |
| vld_b_x_m(L0, p_in); |
| vld_b_x_m(L4, p_in + 32 * 4); |
| vld_b_x_m(L8, p_in + 32 * 8); |
| vld_b_x(LC, p_in + 32 * 12); |
| vld_b_x(LD, p_in + 32 * 13); |
| vld_b_x(LE, p_in + 32 * 14); |
| |
| // MAC pipeline |
| CONV_PER_CHANNEL_B8_3X1_48C_MAC_PIPELINE(p_flt); |
| |
| // Output pipeline |
| INT32_TO_INT8_OUTPUT_PIPELINE(bias + zo_hi * 4, mult + zo_hi * 4, |
| shft + zo_hi * 4, output_activation_min, |
| output_activation_max, output_offset, v36, |
| v40, v44); |
| int8_t* p_out = |
| output_data + tflite::Offset(output_shape, 0, out_y, 0, zo_hi); |
| vstq_b_sp_xx(v48, p_out, output_depth); |
| vstq_b_sp_xx(v52, p_out, output_depth); |
| } |
| |
| // Left over minibatch |
| int remainder = output_height - out_y; |
| if (remainder != 0) { |
| // Reset accumulator |
| vdup_w_x_m(v48, 0); |
| vdup_w_x_m(v52, 0); |
| acset_v(v48, v48); |
| |
| const int8_t* p_flt = juggled_filter_data[zo_hi / 8][0][0][0][0]; |
| const int8_t* p_in = input_data + (out_y * input_width * input_depth); |
| |
| // Load (remainder + 2) * 48 activations |
| // L0 L1 L2 L3 L4 L5 L6 L7 L8 L9 LA LB LC LD |
| // AA AB BB CC CD DD EE EF FF GG GH HH II I- |
| vld_b_x_m(L0, p_in); |
| vdup_w_x_m(L4, 0); |
| vdup_w_x_m(L8, 0); |
| vdup_w_x_m(LC, 0); |
| switch (remainder) { |
| case 7: |
| vld_b_x(LD, p_in + 32 * 13); |
| vld_b_x(LC, p_in + 32 * 12); |
| case 6: |
| vld_b_x(LB, p_in + 32 * 11); |
| case 5: |
| vld_b_x(LA, p_in + 32 * 10); |
| vld_b_x(L9, p_in + 32 * 9); |
| case 4: |
| vld_b_x(L8, p_in + 32 * 8); |
| case 3: |
| vld_b_x(L7, p_in + 32 * 7); |
| vld_b_x(L6, p_in + 32 * 6); |
| case 2: |
| vld_b_x(L5, p_in + 32 * 5); |
| default: |
| break; |
| } |
| vld_b_x(L4, p_in + 32 * 4); |
| |
| // MAC pipeline |
| CONV_PER_CHANNEL_B8_3X1_48C_MAC_PIPELINE(p_flt); |
| |
| // Output pipeline |
| INT32_TO_INT8_OUTPUT_PIPELINE(bias + zo_hi * 4, mult + zo_hi * 4, |
| shft + zo_hi * 4, output_activation_min, |
| output_activation_max, output_offset, v36, |
| v40, v44); |
| |
| int8_t* p_out = |
| output_data + tflite::Offset(output_shape, 0, out_y, 0, zo_hi); |
| uint8_t local_data[64]; |
| vst_b_x(v0, local_data); |
| vst_b_x(v1, local_data + 32); |
| for (int i = 0; i < remainder; i++) { |
| memcpy(p_out + (i * output_depth), local_data + (i * 8), 8); |
| } |
| } |
| |
| #undef CONV_PER_CHANNEL_B8_3X1_48C_MAC_PIPELINE |
| #undef L0 |
| #undef L1 |
| #undef L2 |
| #undef L3 |
| #undef L4 |
| #undef L5 |
| #undef L6 |
| #undef L7 |
| #undef L8 |
| #undef L9 |
| #undef LA |
| #undef LB |
| #undef LC |
| #undef LD |
| #undef LE |
| } |
| } |
| |
| } // namespace kelvin::opt |
