av1/encoder/x86/ml_avx2.c - aom - Git at Google

 /*
  * Copyright (c) 2023, Alliance for Open Media. All rights reserved
  *
  * This source code is subject to the terms of the BSD 2 Clause License and
  * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
  * was not distributed with this source code in the LICENSE file, you can
  * obtain it at www.aomedia.org/license/software. If the Alliance for Open
  * Media Patent License 1.0 was not distributed with this source code in the
  * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
  */

 #include <stdbool.h>
 #include <assert.h>
 #include <immintrin.h>

 #include "config/av1_rtcd.h"
 #include "av1/encoder/ml.h"
 #include "av1/encoder/x86/ml_sse3.h"

 #define CALC_OUTPUT_FOR_2ROWS                                               \
   const int index = weight_idx + (2 * i * tot_num_inputs);                  \
   const __m256 weight0 = _mm256_loadu_ps(&weights[index]);                  \
   const __m256 weight1 = _mm256_loadu_ps(&weights[index + tot_num_inputs]); \
   const __m256 mul0 = _mm256_mul_ps(inputs256, weight0);                    \
   const __m256 mul1 = _mm256_mul_ps(inputs256, weight1);                    \
   hadd[i] = _mm256_hadd_ps(mul0, mul1);

 static INLINE void nn_propagate_8to1(
     const float *const inputs, const float *const weights,
     const float *const bias, int num_inputs_to_process, int tot_num_inputs,
     int num_outputs, float *const output_nodes, int is_clip_required) {
   // Process one output row at a time.
   for (int out = 0; out < num_outputs; out++) {
     __m256 in_result = _mm256_setzero_ps();
     float bias_val = bias[out];
     for (int in = 0; in < num_inputs_to_process; in += 8) {
       const __m256 inputs256 = _mm256_loadu_ps(&inputs[in]);
       const int weight_idx = in + (out * tot_num_inputs);
       const __m256 weight0 = _mm256_loadu_ps(&weights[weight_idx]);
       const __m256 mul0 = _mm256_mul_ps(inputs256, weight0);
       in_result = _mm256_add_ps(in_result, mul0);
     }
     const __m128 low_128 = _mm256_castps256_ps128(in_result);
     const __m128 high_128 = _mm256_extractf128_ps(in_result, 1);
     const __m128 sum_par_0 = _mm_add_ps(low_128, high_128);
     const __m128 sum_par_1 = _mm_hadd_ps(sum_par_0, sum_par_0);
     const __m128 sum_tot =
         _mm_add_ps(_mm_shuffle_ps(sum_par_1, sum_par_1, 0x99), sum_par_1);

     bias_val += _mm_cvtss_f32(sum_tot);
     if (is_clip_required) bias_val = AOMMAX(bias_val, 0);
     output_nodes[out] = bias_val;
   }
 }

 static INLINE void nn_propagate_8to4(
     const float *const inputs, const float *const weights,
     const float *const bias, int num_inputs_to_process, int tot_num_inputs,
     int num_outputs, float *const output_nodes, int is_clip_required) {
   __m256 hadd[2];
   for (int out = 0; out < num_outputs; out += 4) {
     __m128 bias_reg = _mm_loadu_ps(&bias[out]);
     __m128 in_result = _mm_setzero_ps();
     for (int in = 0; in < num_inputs_to_process; in += 8) {
       const __m256 inputs256 = _mm256_loadu_ps(&inputs[in]);
       const int weight_idx = in + (out * tot_num_inputs);
       // Process two output row at a time.
       for (int i = 0; i < 2; i++) {
         CALC_OUTPUT_FOR_2ROWS
       }

       const __m256 sum_par = _mm256_hadd_ps(hadd[0], hadd[1]);
       const __m128 low_128 = _mm256_castps256_ps128(sum_par);
       const __m128 high_128 = _mm256_extractf128_ps(sum_par, 1);
       const __m128 result = _mm_add_ps(low_128, high_128);

       in_result = _mm_add_ps(in_result, result);
     }

     in_result = _mm_add_ps(in_result, bias_reg);
     if (is_clip_required) in_result = _mm_max_ps(in_result, _mm_setzero_ps());
     _mm_storeu_ps(&output_nodes[out], in_result);
   }
 }

 static INLINE void nn_propagate_8to8(
     const float *const inputs, const float *const weights,
     const float *const bias, int num_inputs_to_process, int tot_num_inputs,
     int num_outputs, float *const output_nodes, int is_clip_required) {
   __m256 hadd[4];
   for (int out = 0; out < num_outputs; out += 8) {
     __m256 bias_reg = _mm256_loadu_ps(&bias[out]);
     __m256 in_result = _mm256_setzero_ps();
     for (int in = 0; in < num_inputs_to_process; in += 8) {
       const __m256 inputs256 = _mm256_loadu_ps(&inputs[in]);
       const int weight_idx = in + (out * tot_num_inputs);
       // Process two output rows at a time.
       for (int i = 0; i < 4; i++) {
         CALC_OUTPUT_FOR_2ROWS
       }
       const __m256 hh0 = _mm256_hadd_ps(hadd[0], hadd[1]);
       const __m256 hh1 = _mm256_hadd_ps(hadd[2], hadd[3]);

       __m256 ht_0 = _mm256_permute2f128_ps(hh0, hh1, 0x20);
       __m256 ht_1 = _mm256_permute2f128_ps(hh0, hh1, 0x31);

       __m256 result = _mm256_add_ps(ht_0, ht_1);
       in_result = _mm256_add_ps(in_result, result);
     }
     in_result = _mm256_add_ps(in_result, bias_reg);
     if (is_clip_required)
       in_result = _mm256_max_ps(in_result, _mm256_setzero_ps());
     _mm256_storeu_ps(&output_nodes[out], in_result);
   }
 }

 static INLINE void nn_propagate_input_multiple_of_8(
     const float *const inputs, const float *const weights,
     const float *const bias, int num_inputs_to_process, int tot_num_inputs,
     bool is_output_layer, int num_outputs, float *const output_nodes) {
   // The saturation of output is considered for hidden layer which is not equal
   // to final hidden layer.
   const int is_clip_required =
       !is_output_layer && num_inputs_to_process == tot_num_inputs;
   if (num_outputs % 8 == 0) {
     nn_propagate_8to8(inputs, weights, bias, num_inputs_to_process,
                       tot_num_inputs, num_outputs, output_nodes,
                       is_clip_required);
   } else if (num_outputs % 4 == 0) {
     nn_propagate_8to4(inputs, weights, bias, num_inputs_to_process,
                       tot_num_inputs, num_outputs, output_nodes,
                       is_clip_required);
   } else {
     nn_propagate_8to1(inputs, weights, bias, num_inputs_to_process,
                       tot_num_inputs, num_outputs, output_nodes,
                       is_clip_required);
   }
 }

 void av1_nn_predict_avx2(const float *input_nodes,
                          const NN_CONFIG *const nn_config, int reduce_prec,
                          float *const output) {
   float buf[2][NN_MAX_NODES_PER_LAYER];
   int buf_index = 0;
   int num_inputs = nn_config->num_inputs;
   assert(num_inputs > 0 && num_inputs <= NN_MAX_NODES_PER_LAYER);

   for (int layer = 0; layer <= nn_config->num_hidden_layers; layer++) {
     const float *layer_weights = nn_config->weights[layer];
     const float *layer_bias = nn_config->bias[layer];
     bool is_output_layer = layer == nn_config->num_hidden_layers;
     float *const output_nodes = is_output_layer ? output : &buf[buf_index][0];
     const int num_outputs = is_output_layer
                                 ? nn_config->num_outputs
                                 : nn_config->num_hidden_nodes[layer];
     assert(num_outputs > 0 && num_outputs <= NN_MAX_NODES_PER_LAYER);

     // Process input multiple of 8 using AVX2 intrinsic.
     if (num_inputs % 8 == 0) {
       nn_propagate_input_multiple_of_8(input_nodes, layer_weights, layer_bias,
                                        num_inputs, num_inputs, is_output_layer,
                                        num_outputs, output_nodes);
     } else {
       // When number of inputs is not multiple of 8, use hybrid approach of AVX2
       // and SSE3 based on the need.
       const int in_mul_8 = num_inputs / 8;
       const int num_inputs_to_process = in_mul_8 * 8;
       int bias_is_considered = 0;
       if (in_mul_8) {
         nn_propagate_input_multiple_of_8(
             input_nodes, layer_weights, layer_bias, num_inputs_to_process,
             num_inputs, is_output_layer, num_outputs, output_nodes);
         bias_is_considered = 1;
       }

       const float *out_temp = bias_is_considered ? output_nodes : layer_bias;
       const int input_remaining = num_inputs % 8;
       if (input_remaining % 4 == 0 && num_outputs % 8 == 0) {
         for (int out = 0; out < num_outputs; out += 8) {
           __m128 out_h = _mm_loadu_ps(&out_temp[out + 4]);
           __m128 out_l = _mm_loadu_ps(&out_temp[out]);
           for (int in = in_mul_8 * 8; in < num_inputs; in += 4) {
             av1_nn_propagate_4to8_sse3(&input_nodes[in],
                                        &layer_weights[out * num_inputs + in],
                                        &out_h, &out_l, num_inputs);
           }
           if (!is_output_layer) {
             const __m128 zero = _mm_setzero_ps();
             out_h = _mm_max_ps(out_h, zero);
             out_l = _mm_max_ps(out_l, zero);
           }
           _mm_storeu_ps(&output_nodes[out + 4], out_h);
           _mm_storeu_ps(&output_nodes[out], out_l);
         }
       } else if (input_remaining % 4 == 0 && num_outputs % 4 == 0) {
         for (int out = 0; out < num_outputs; out += 4) {
           __m128 outputs = _mm_loadu_ps(&out_temp[out]);
           for (int in = in_mul_8 * 8; in < num_inputs; in += 4) {
             av1_nn_propagate_4to4_sse3(&input_nodes[in],
                                        &layer_weights[out * num_inputs + in],
                                        &outputs, num_inputs);
           }
           if (!is_output_layer) outputs = _mm_max_ps(outputs, _mm_setzero_ps());
           _mm_storeu_ps(&output_nodes[out], outputs);
         }
       } else if (input_remaining % 4 == 0) {
         for (int out = 0; out < num_outputs; out++) {
           __m128 outputs = _mm_load1_ps(&out_temp[out]);
           for (int in = in_mul_8 * 8; in < num_inputs; in += 4) {
             av1_nn_propagate_4to1_sse3(&input_nodes[in],
                                        &layer_weights[out * num_inputs + in],
                                        &outputs);
           }
           if (!is_output_layer) outputs = _mm_max_ps(outputs, _mm_setzero_ps());
           output_nodes[out] = _mm_cvtss_f32(outputs);
         }
       } else {
         // Use SSE instructions for scalar operations to avoid the latency
         // of swapping between SIMD and FPU modes.
         for (int out = 0; out < num_outputs; out++) {
           __m128 outputs = _mm_load1_ps(&out_temp[out]);
           for (int in_node = in_mul_8 * 8; in_node < num_inputs; in_node++) {
             __m128 input = _mm_load1_ps(&input_nodes[in_node]);
             __m128 weight =
                 _mm_load1_ps(&layer_weights[num_inputs * out + in_node]);
             outputs = _mm_add_ps(outputs, _mm_mul_ps(input, weight));
           }
           if (!is_output_layer) outputs = _mm_max_ps(outputs, _mm_setzero_ps());
           output_nodes[out] = _mm_cvtss_f32(outputs);
         }
       }
     }
     // Before processing the next layer, treat the output of current layer as
     // input to next layer.
     input_nodes = output_nodes;
     num_inputs = num_outputs;
     buf_index = 1 - buf_index;
   }
   if (reduce_prec) av1_nn_output_prec_reduce(output, nn_config->num_outputs);
 }
	/*
	* Copyright (c) 2023, Alliance for Open Media. All rights reserved
	*
	* This source code is subject to the terms of the BSD 2 Clause License and
	* the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
	* was not distributed with this source code in the LICENSE file, you can
	* obtain it at www.aomedia.org/license/software. If the Alliance for Open
	* Media Patent License 1.0 was not distributed with this source code in the
	* PATENTS file, you can obtain it at www.aomedia.org/license/patent.
	*/

	#include <stdbool.h>
	#include <assert.h>
	#include <immintrin.h>

	#include "config/av1_rtcd.h"
	#include "av1/encoder/ml.h"
	#include "av1/encoder/x86/ml_sse3.h"

	#define CALC_OUTPUT_FOR_2ROWS \
	const int index = weight_idx + (2 * i * tot_num_inputs); \
	const __m256 weight0 = _mm256_loadu_ps(&weights[index]); \
	const __m256 weight1 = _mm256_loadu_ps(&weights[index + tot_num_inputs]); \
	const __m256 mul0 = _mm256_mul_ps(inputs256, weight0); \
	const __m256 mul1 = _mm256_mul_ps(inputs256, weight1); \
	hadd[i] = _mm256_hadd_ps(mul0, mul1);

	static INLINE void nn_propagate_8to1(
	const float const inputs, const float const weights,
	const float *const bias, int num_inputs_to_process, int tot_num_inputs,
	int num_outputs, float *const output_nodes, int is_clip_required) {
	// Process one output row at a time.
	for (int out = 0; out < num_outputs; out++) {
	__m256 in_result = _mm256_setzero_ps();
	float bias_val = bias[out];
	for (int in = 0; in < num_inputs_to_process; in += 8) {
	const __m256 inputs256 = _mm256_loadu_ps(&inputs[in]);
	const int weight_idx = in + (out * tot_num_inputs);
	const __m256 weight0 = _mm256_loadu_ps(&weights[weight_idx]);
	const __m256 mul0 = _mm256_mul_ps(inputs256, weight0);
	in_result = _mm256_add_ps(in_result, mul0);
	}
	const __m128 low_128 = _mm256_castps256_ps128(in_result);
	const __m128 high_128 = _mm256_extractf128_ps(in_result, 1);
	const __m128 sum_par_0 = _mm_add_ps(low_128, high_128);
	const __m128 sum_par_1 = _mm_hadd_ps(sum_par_0, sum_par_0);
	const __m128 sum_tot =
	_mm_add_ps(_mm_shuffle_ps(sum_par_1, sum_par_1, 0x99), sum_par_1);

	bias_val += _mm_cvtss_f32(sum_tot);
	if (is_clip_required) bias_val = AOMMAX(bias_val, 0);
	output_nodes[out] = bias_val;
	}
	}

	static INLINE void nn_propagate_8to4(
	const float const inputs, const float const weights,
	const float *const bias, int num_inputs_to_process, int tot_num_inputs,
	int num_outputs, float *const output_nodes, int is_clip_required) {
	__m256 hadd[2];
	for (int out = 0; out < num_outputs; out += 4) {
	__m128 bias_reg = _mm_loadu_ps(&bias[out]);
	__m128 in_result = _mm_setzero_ps();
	for (int in = 0; in < num_inputs_to_process; in += 8) {
	const __m256 inputs256 = _mm256_loadu_ps(&inputs[in]);
	const int weight_idx = in + (out * tot_num_inputs);
	// Process two output row at a time.
	for (int i = 0; i < 2; i++) {
	CALC_OUTPUT_FOR_2ROWS
	}

	const __m256 sum_par = _mm256_hadd_ps(hadd[0], hadd[1]);
	const __m128 low_128 = _mm256_castps256_ps128(sum_par);
	const __m128 high_128 = _mm256_extractf128_ps(sum_par, 1);
	const __m128 result = _mm_add_ps(low_128, high_128);

	in_result = _mm_add_ps(in_result, result);
	}

	in_result = _mm_add_ps(in_result, bias_reg);
	if (is_clip_required) in_result = _mm_max_ps(in_result, _mm_setzero_ps());
	_mm_storeu_ps(&output_nodes[out], in_result);
	}
	}

	static INLINE void nn_propagate_8to8(
	const float const inputs, const float const weights,
	const float *const bias, int num_inputs_to_process, int tot_num_inputs,
	int num_outputs, float *const output_nodes, int is_clip_required) {
	__m256 hadd[4];
	for (int out = 0; out < num_outputs; out += 8) {
	__m256 bias_reg = _mm256_loadu_ps(&bias[out]);
	__m256 in_result = _mm256_setzero_ps();
	for (int in = 0; in < num_inputs_to_process; in += 8) {
	const __m256 inputs256 = _mm256_loadu_ps(&inputs[in]);
	const int weight_idx = in + (out * tot_num_inputs);
	// Process two output rows at a time.
	for (int i = 0; i < 4; i++) {
	CALC_OUTPUT_FOR_2ROWS
	}
	const __m256 hh0 = _mm256_hadd_ps(hadd[0], hadd[1]);
	const __m256 hh1 = _mm256_hadd_ps(hadd[2], hadd[3]);

	__m256 ht_0 = _mm256_permute2f128_ps(hh0, hh1, 0x20);
	__m256 ht_1 = _mm256_permute2f128_ps(hh0, hh1, 0x31);

	__m256 result = _mm256_add_ps(ht_0, ht_1);
	in_result = _mm256_add_ps(in_result, result);
	}
	in_result = _mm256_add_ps(in_result, bias_reg);
	if (is_clip_required)
	in_result = _mm256_max_ps(in_result, _mm256_setzero_ps());
	_mm256_storeu_ps(&output_nodes[out], in_result);
	}
	}

	static INLINE void nn_propagate_input_multiple_of_8(
	const float const inputs, const float const weights,
	const float *const bias, int num_inputs_to_process, int tot_num_inputs,
	bool is_output_layer, int num_outputs, float *const output_nodes) {
	// The saturation of output is considered for hidden layer which is not equal
	// to final hidden layer.
	const int is_clip_required =
	!is_output_layer && num_inputs_to_process == tot_num_inputs;
	if (num_outputs % 8 == 0) {
	nn_propagate_8to8(inputs, weights, bias, num_inputs_to_process,
	tot_num_inputs, num_outputs, output_nodes,
	is_clip_required);
	} else if (num_outputs % 4 == 0) {
	nn_propagate_8to4(inputs, weights, bias, num_inputs_to_process,
	tot_num_inputs, num_outputs, output_nodes,
	is_clip_required);
	} else {
	nn_propagate_8to1(inputs, weights, bias, num_inputs_to_process,
	tot_num_inputs, num_outputs, output_nodes,
	is_clip_required);
	}
	}

	void av1_nn_predict_avx2(const float *input_nodes,
	const NN_CONFIG *const nn_config, int reduce_prec,
	float *const output) {
	float buf[2][NN_MAX_NODES_PER_LAYER];
	int buf_index = 0;
	int num_inputs = nn_config->num_inputs;
	assert(num_inputs > 0 && num_inputs <= NN_MAX_NODES_PER_LAYER);

	for (int layer = 0; layer <= nn_config->num_hidden_layers; layer++) {
	const float *layer_weights = nn_config->weights[layer];
	const float *layer_bias = nn_config->bias[layer];
	bool is_output_layer = layer == nn_config->num_hidden_layers;
	float *const output_nodes = is_output_layer ? output : &buf[buf_index][0];
	const int num_outputs = is_output_layer
	? nn_config->num_outputs
	: nn_config->num_hidden_nodes[layer];
	assert(num_outputs > 0 && num_outputs <= NN_MAX_NODES_PER_LAYER);

	// Process input multiple of 8 using AVX2 intrinsic.
	if (num_inputs % 8 == 0) {
	nn_propagate_input_multiple_of_8(input_nodes, layer_weights, layer_bias,
	num_inputs, num_inputs, is_output_layer,
	num_outputs, output_nodes);
	} else {
	// When number of inputs is not multiple of 8, use hybrid approach of AVX2
	// and SSE3 based on the need.
	const int in_mul_8 = num_inputs / 8;
	const int num_inputs_to_process = in_mul_8 * 8;
	int bias_is_considered = 0;
	if (in_mul_8) {
	nn_propagate_input_multiple_of_8(
	input_nodes, layer_weights, layer_bias, num_inputs_to_process,
	num_inputs, is_output_layer, num_outputs, output_nodes);
	bias_is_considered = 1;
	}

	const float *out_temp = bias_is_considered ? output_nodes : layer_bias;
	const int input_remaining = num_inputs % 8;
	if (input_remaining % 4 == 0 && num_outputs % 8 == 0) {
	for (int out = 0; out < num_outputs; out += 8) {
	__m128 out_h = _mm_loadu_ps(&out_temp[out + 4]);
	__m128 out_l = _mm_loadu_ps(&out_temp[out]);
	for (int in = in_mul_8 * 8; in < num_inputs; in += 4) {
	av1_nn_propagate_4to8_sse3(&input_nodes[in],
	&layer_weights[out * num_inputs + in],
	&out_h, &out_l, num_inputs);
	}
	if (!is_output_layer) {
	const __m128 zero = _mm_setzero_ps();
	out_h = _mm_max_ps(out_h, zero);
	out_l = _mm_max_ps(out_l, zero);
	}
	_mm_storeu_ps(&output_nodes[out + 4], out_h);
	_mm_storeu_ps(&output_nodes[out], out_l);
	}
	} else if (input_remaining % 4 == 0 && num_outputs % 4 == 0) {
	for (int out = 0; out < num_outputs; out += 4) {
	__m128 outputs = _mm_loadu_ps(&out_temp[out]);
	for (int in = in_mul_8 * 8; in < num_inputs; in += 4) {
	av1_nn_propagate_4to4_sse3(&input_nodes[in],
	&layer_weights[out * num_inputs + in],
	&outputs, num_inputs);
	}
	if (!is_output_layer) outputs = _mm_max_ps(outputs, _mm_setzero_ps());
	_mm_storeu_ps(&output_nodes[out], outputs);
	}
	} else if (input_remaining % 4 == 0) {
	for (int out = 0; out < num_outputs; out++) {
	__m128 outputs = _mm_load1_ps(&out_temp[out]);
	for (int in = in_mul_8 * 8; in < num_inputs; in += 4) {
	av1_nn_propagate_4to1_sse3(&input_nodes[in],
	&layer_weights[out * num_inputs + in],
	&outputs);
	}
	if (!is_output_layer) outputs = _mm_max_ps(outputs, _mm_setzero_ps());
	output_nodes[out] = _mm_cvtss_f32(outputs);
	}
	} else {
	// Use SSE instructions for scalar operations to avoid the latency
	// of swapping between SIMD and FPU modes.
	for (int out = 0; out < num_outputs; out++) {
	__m128 outputs = _mm_load1_ps(&out_temp[out]);
	for (int in_node = in_mul_8 * 8; in_node < num_inputs; in_node++) {
	__m128 input = _mm_load1_ps(&input_nodes[in_node]);
	__m128 weight =
	_mm_load1_ps(&layer_weights[num_inputs * out + in_node]);
	outputs = _mm_add_ps(outputs, _mm_mul_ps(input, weight));
	}
	if (!is_output_layer) outputs = _mm_max_ps(outputs, _mm_setzero_ps());
	output_nodes[out] = _mm_cvtss_f32(outputs);
	}
	}
	}
	// Before processing the next layer, treat the output of current layer as
	// input to next layer.
	input_nodes = output_nodes;
	num_inputs = num_outputs;
	buf_index = 1 - buf_index;
	}
	if (reduce_prec) av1_nn_output_prec_reduce(output, nn_config->num_outputs);
	}