text: #include <stdbool.h>
  #include <stdint.h>
  #include <stddef.h>
  
  #include <nnpack.h>
  #include <nnpack/macros.h>
  #include <nnpack/utils.h>
  #include <nnpack/system.h>
  
  #include <nnpack/hwinfo.h>
  #include <nnpack/validation.h>
  
  
  struct NNP_CACHE_ALIGN input_transform_context {
  	size_t tuple_elements;
  	size_t input_elements;
  	size_t batch_block_size;
  	size_t input_channels;
  	size_t input_stride;
  	uint32_t row_offset;
  	uint32_t column_offset;
  	uint32_t row_count;
  	uint32_t column_count;
  	const float* input;
  	float* input_transform;
  	nnp_transform_2d_with_offset transform_function;
  };
  
  static void compute_input_transform(
  	const struct input_transform_context context[restrict static 1],
  	size_t batch_block_offset,       size_t input_channels_subblock_start,
  	size_t batch_block_offset_range, size_t input_channels_subblock_size)
  {
  	const size_t tuple_elements                  = context->tuple_elements;
  	const size_t input_elements                  = context->input_elements;
  	const size_t batch_block_size                = context->batch_block_size;
  	const size_t input_channels                  = context->input_channels;
  	const size_t input_stride                    = context->input_stride;
  	const uint32_t row_count                     = context->row_count;
  	const uint32_t column_count                  = context->column_count;
  	const uint32_t row_offset                    = context->row_offset;
  	const uint32_t column_offset                 = context->column_offset;
  	const float* input                           = context->input;
  	float* input_transform                       = context->input_transform;
  	const nnp_transform_2d_with_offset transform = context->transform_function;
  
  	for (size_t input_channels_subblock_offset = 0; input_channels_subblock_offset < input_channels_subblock_size; input_channels_subblock_offset += 1) {
  		const size_t input_channel = input_channels_subblock_start + input_channels_subblock_offset;
  		transform(
  			input +
  				(batch_block_offset * input_channels + input_channel) * input_elements,
  			input_transform +
  				(input_channels_subblock_start * batch_block_size + batch_block_offset * input_channels_subblock_size + input_channels_subblock_offset) * tuple_elements,
  			input_stride,
  			batch_block_size * input_channels * tuple_elements * sizeof(float),
  			row_count, column_count, row_offset, column_offset);
  	}
  }
  
  struct NNP_CACHE_ALIGN grad_output_transform_context {
  	size_t tuple_elements;
  	size_t output_elements;
  	size_t batch_block_size;
  	size_t output_channels;
  	size_t grad_output_stride;
  	uint32_t row_count;
  	uint32_t column_count;
  	const float* grad_output;
  	float* grad_output_transform;
  	nnp_transform_2d_with_offset transform_function;
  };
  
  static void compute_grad_output_transform(
  	const struct grad_output_transform_context context[restrict static 1],
  	size_t batch_block_offset,       size_t output_channels_subblock_start,
  	size_t batch_block_offset_range, size_t output_channels_subblock_size)
  {
  	const size_t tuple_elements                  = context->tuple_elements;
  	const size_t output_elements                 = context->output_elements;
  	const size_t batch_block_size                = context->batch_block_size;
  	const size_t output_channels                 = context->output_channels;
  	const size_t grad_output_stride              = context->grad_output_stride;
  	const uint32_t row_count                     = context->row_count;
  	const uint32_t column_count                  = context->column_count;
  	const float* grad_output                     = context->grad_output;
  	float* grad_output_transform                 = context->grad_output_transform;
  	const nnp_transform_2d_with_offset transform = context->transform_function;
  
  	for (size_t output_channels_subblock_offset = 0; output_channels_subblock_offset < output_channels_subblock_size; output_channels_subblock_offset += 1) {
  		const size_t output_channel = output_channels_subblock_start + output_channels_subblock_offset;
  		transform(
  			grad_output +
  				(batch_block_offset * output_channels + output_channel) * output_elements,
  			grad_output_transform +
  				(output_channels_subblock_start * batch_block_size + batch_block_offset * output_channels_subblock_size + output_channels_subblock_offset) * tuple_elements,
  			grad_output_stride,
  			batch_block_size * output_channels * tuple_elements * sizeof(float),
  			row_count, column_count, 0, 0);
  	}
  }
  
  struct NNP_CACHE_ALIGN grad_kernel_transform_context {
  	size_t tuple_elements;
  	size_t input_channels;
  	size_t output_channels;
  	size_t output_channels_block_max;
  	struct nnp_size kernel_size;
  	const float* grad_kernel_transform;
  	float* grad_kernel;
  	nnp_transform_2d_with_offset transform_function;
  };
  
  static void compute_grad_kernel_transform(
  	const struct grad_kernel_transform_context context[restrict static 1],
  	size_t output_channel,       size_t input_channels_subblock_start,
  	size_t output_channel_range, size_t input_channels_subblock_size)
  {
  	const size_t tuple_elements                  = context->tuple_elements;
  	const size_t input_channels                  = context->input_channels;
  	const size_t output_channels                 = context->output_channels;
  	const size_t output_channels_block_max       = context->output_channels_block_max;
  	const struct nnp_size kernel_size            = context->kernel_size;
  	const float* grad_kernel_transform           = context->grad_kernel_transform;
  	float* grad_kernel                           = context->grad_kernel;
  	const nnp_transform_2d_with_offset transform = context->transform_function;
  
  	const size_t output_channels_block_start  = round_down(output_channel, output_channels_block_max);
  	const size_t output_channels_block_size   = min(output_channels - output_channels_block_start, output_channels_block_max);
  	const size_t output_channels_block_offset = output_channel - output_channels_block_start;
  	const size_t kernel_elements = kernel_size.height * kernel_size.width;
  
  	for (size_t input_channels_subblock_offset = 0; input_channels_subblock_offset < input_channels_subblock_size; input_channels_subblock_offset += 1) {
  		const size_t input_channel = input_channels_subblock_start + input_channels_subblock_offset;
  		transform(
  			grad_kernel_transform +
  				(output_channels_block_start * input_channels + input_channels_subblock_start * output_channels_block_size + output_channels_block_offset * input_channels_subblock_size + input_channels_subblock_offset) * tuple_elements,
  			grad_kernel +
  				(output_channel * input_channels + input_channel) * kernel_elements,
  			output_channels * input_channels * tuple_elements * sizeof(float),
  			kernel_size.width,
  			kernel_size.height, kernel_size.width, 0, 0);
  	}
  }
  
  struct NNP_CACHE_ALIGN matrix_multiplication_context {
  	size_t tuple_elements;
  	size_t batch_size;
  	size_t batch_block_size;
  	size_t batch_block_update;
  	size_t input_channels;
  	size_t input_channels_block_start;
  	size_t input_channels_block_size;
  	size_t input_channels_subblock_max;
  	size_t output_channels;
  	size_t output_channels_subblock_max;
  	const float* grad_output_transform;
  	const float* input_transform;
  	float* grad_kernel_transform;
  
  	nnp_fast_tuple_gemm_function fast_gemm;
  	nnp_full_tuple_gemm_function full_gemm;
  };
  
  static void compute_matrix_multiplication(
  	const struct matrix_multiplication_context context[restrict static 1],
  	size_t output_channels_block_start, size_t input_channels_subblock_start,
  	size_t output_channels_block_size,  size_t input_channels_subblock_size)
  {
  	const size_t tuple_elements               = context->tuple_elements;
  	const size_t batch_size                   = context->batch_size;
  	const size_t batch_block_size             = context->batch_block_size;
  	const size_t batch_block_update           = context->batch_block_update;
  	const size_t input_channels               = context->input_channels;
  	const size_t input_channels_block_start   = context->input_channels_block_start;
  	const size_t input_channels_block_size    = context->input_channels_block_size;
  	const size_t input_channels_subblock_max  = context->input_channels_subblock_max;
  	const size_t output_channels              = context->output_channels;
  	const size_t output_channels_subblock_max = context->output_channels_subblock_max;
  
  	const float* grad_output_transform = context->grad_output_transform +
  		output_channels_block_start * batch_block_size * tuple_elements;
  	const float* input_transform       = context->input_transform +
  		(input_channels_block_start + input_channels_subblock_start) * batch_block_size * tuple_elements;
  	float* grad_kernel_transform       = context->grad_kernel_transform +
  		(output_channels_block_start * input_channels + (input_channels_block_start + input_channels_subblock_start) * output_channels_block_size) * tuple_elements;
  
  	if (input_channels_subblock_size == input_channels_subblock_max) {
  		const nnp_fast_tuple_gemm_function fast_gemm = context->fast_gemm;
  		while (output_channels_block_size >= output_channels_subblock_max) {
  			output_channels_block_size -= output_channels_subblock_max;
  
  			fast_gemm(
  				batch_block_size, batch_block_update,
  				input_transform,
  				grad_output_transform,
  				grad_kernel_transform,
  				input_channels_subblock_size * tuple_elements);
  
  			grad_output_transform += output_channels_subblock_max * batch_block_size * tuple_elements;
  			grad_kernel_transform += output_channels_subblock_max * input_channels_subblock_size * tuple_elements;
  		}
  	}
  
  	const nnp_full_tuple_gemm_function full_gemm = context->full_gemm;
  	while (output_channels_block_size != 0) {
  		const size_t output_channels_subblock_size = min(output_channels_block_size, output_channels_subblock_max);
  		output_channels_block_size -= output_channels_subblock_size;
  
  		full_gemm(
  			input_channels_subblock_size, output_channels_subblock_size,
  			batch_block_size, batch_block_update,
  			input_transform,
  			grad_output_transform,
  			grad_kernel_transform,
  			input_channels_subblock_size * tuple_elements);
  
  		grad_output_transform += output_channels_subblock_max * batch_block_size * tuple_elements;
  		grad_kernel_transform += output_channels_subblock_max * input_channels_subblock_size * tuple_elements;
  	}
  }
  
  static enum nnp_status compute_fast_convolution_kernel_gradient(
  	size_t batch_size,
  	size_t input_channels,
  	size_t output_channels,
  	struct nnp_size tile_size,
  	struct nnp_size input_size,
  	struct nnp_padding input_padding,
  	struct nnp_size kernel_size,
  	struct nnp_size output_size,
  	const float* input,
  	const float* grad_output,
  	float* grad_kernel,
  	void* workspace_buffer,
  	size_t* workspace_size,
  	nnp_transform_2d_with_offset input_transform_function,
  	nnp_transform_2d_with_offset grad_output_transform_function,
  	nnp_transform_2d_with_offset grad_kernel_transform_function,
  	pthreadpool_t threadpool,
  	struct nnp_profile* profile)
  {
  	void* memory_block = NULL;
  	const size_t simd_width = nnp_hwinfo.simd_width;
  	const size_t tuple_elements = simd_width * 2;
  	const size_t tile_elements = tile_size.height * tile_size.width;
  	const size_t tuple_count = tile_elements / tuple_elements;
  
  	const struct nnp_size output_tile = {
  		.height = tile_size.height - kernel_size.height + 1,
  		.width = tile_size.width - kernel_size.width + 1
  	};
  
  	/* Calculate cache blocking parameters */
  	const size_t cache_elements_l1 = nnp_hwinfo.blocking.l1 / (tuple_elements * sizeof(float));
  	const size_t cache_elements_l2 = nnp_hwinfo.blocking.l2 / (tuple_elements * sizeof(float));
  	const size_t cache_elements_l3 = nnp_hwinfo.blocking.l3 / (tuple_elements * sizeof(float));
  
  	const size_t input_channels_subblock_max = nnp_hwinfo.cxgemm.mr;
  	const size_t output_channels_subblock_max = nnp_hwinfo.cxgemm.nr;
  
  	const size_t batch_block_max =
  		round_down(cache_elements_l1 / (input_channels_subblock_max + output_channels_subblock_max), 2);
  	const size_t input_channels_block_max =
  		round_down(cache_elements_l3 / batch_block_max, input_channels_subblock_max);
  	const size_t output_channels_block_max =
  		round_down(cache_elements_l2 / batch_block_max, output_channels_subblock_max);
  
  	/* Calculate memory footprint and allocate memory */
  	const size_t input_transform_size = min(batch_size, batch_block_max) * input_channels * tile_elements * sizeof(float);
  	const size_t grad_output_transform_size = min(batch_size, batch_block_max) * output_channels * tile_elements * sizeof(float);
  	const size_t grad_kernel_transform_size = output_channels * input_channels * tile_elements * sizeof(float);
  	const size_t memory_size = input_transform_size + grad_output_transform_size + grad_kernel_transform_size;
  
  	if (workspace_buffer == NULL) {
  		if (workspace_size == NULL) {
  			memory_block = allocate_memory(memory_size);
  			if (memory_block == NULL) {
  				return nnp_status_out_of_memory;
  			}
  		} else {
  			*workspace_size = memory_size;
  			return nnp_status_success;
  		}
  	} else {
  		if (*workspace_size < memory_size) {
  			return nnp_status_insufficient_buffer;
  		}
  		memory_block = workspace_buffer;
  	}
  
  	float* input_transform = memory_block;
  	float* grad_output_transform = memory_block + input_transform_size;
  	float* grad_kernel_transform = memory_block + input_transform_size + grad_output_transform_size;
  
  	for (size_t y = 0; y < output_size.height; y += output_tile.height) {
  		const size_t input_y = min(doz(y, input_padding.top), input_size.height);
  		for (size_t x = 0; x < output_size.width; x += output_tile.width) {
  			const size_t input_x = min(doz(x, input_padding.left), input_size.width);
  
  			for (size_t batch_block_start = 0; batch_block_start < batch_size; batch_block_start += batch_block_max) {
  				const size_t batch_block_size = min(batch_size - batch_block_start, batch_block_max);
  
  				/* Input transform */
  				NNP_INPUT_TRANSFORM_START(profile)
  				struct input_transform_context input_transform_context = {
  					.tuple_elements = tuple_elements,
  					.input_elements = input_size.height * input_size.width,
  					.batch_block_size = batch_block_size,
  					.input_channels = input_channels,
  					.input_stride = input_size.width,
  					.row_offset = doz(input_padding.top, y),
  					.column_offset = doz(input_padding.left, x),
  					.row_count = min(input_size.height - input_y,
  						tile_size.height - input_transform_context.row_offset),
  					.column_count = min(input_size.width - input_x,
  						tile_size.width - input_transform_context.column_offset),
  					.input = input + (batch_block_start * input_channels * input_size.height + input_y) * input_size.width + input_x,
  					.input_transform = input_transform,
  					.transform_function = input_transform_function,
  				};
  				pthreadpool_compute_2d_tiled(threadpool,
  					(pthreadpool_function_2d_tiled_t) compute_input_transform,
  					&input_transform_context,
  					batch_block_size, input_channels,
  					1,                input_channels_subblock_max);
  				NNP_INPUT_TRANSFORM_END(profile)
  
  				/* Grad output transform */
  				NNP_OUTPUT_TRANSFORM_START(profile)
  				struct grad_output_transform_context grad_output_transform_context = {
  					.tuple_elements = tuple_elements,
  					.output_elements = output_size.height * output_size.width,
  					.batch_block_size = batch_block_size,
  					.output_channels = output_channels,
  					.grad_output_stride = output_size.width,
  					.row_count = min(output_tile.height, output_size.height - y),
  					.column_count = min(output_tile.width, output_size.width - x),
  					.grad_output = grad_output + (batch_block_start * output_channels * output_size.height + y) * output_size.width + x,
  					.grad_output_transform = grad_output_transform,
  					.transform_function = grad_output_transform_function,
  				};
  				pthreadpool_compute_2d_tiled(threadpool,
  					(pthreadpool_function_2d_tiled_t) compute_grad_output_transform,
  					&grad_output_transform_context,
  					batch_block_size, output_channels,
  					1,                output_channels_subblock_max);
  				NNP_OUTPUT_TRANSFORM_END(profile)
  
  				NNP_BLOCK_MULTIPLICATION_START(profile)
  				for (size_t tuple_index = 0; tuple_index < tuple_count; tuple_index += 1) {
  					for (size_t input_channels_block_start = 0; input_channels_block_start < input_channels; input_channels_block_start += input_channels_block_max) {
  						const size_t input_channels_block_size = min(input_channels - input_channels_block_start, input_channels_block_max);
  
  						struct matrix_multiplication_context matrix_multiplication_context = {
  							.tuple_elements = tuple_elements,
  							.batch_size = batch_size,
  							.batch_block_size = batch_block_size,
  							.batch_block_update = batch_block_start | x | y,
  							.input_channels = input_channels,
  							.input_channels_block_start = input_channels_block_start,
  							.input_channels_block_size = input_channels_block_size,
  							.input_channels_subblock_max = input_channels_subblock_max,
  							.output_channels = output_channels,
  							.output_channels_subblock_max = output_channels_subblock_max,
  							.grad_output_transform = grad_output_transform +
  								tuple_index * tuple_elements * batch_block_size * output_channels,
  							.input_transform = input_transform +
  								tuple_index * tuple_elements * batch_block_size * input_channels,
  							.grad_kernel_transform = grad_kernel_transform +
  								tuple_index * tuple_elements * output_channels * input_channels,
  						};
  						if (tuple_index < NNP_COMPLEX_TUPLE_INDEX) {
  							matrix_multiplication_context.fast_gemm = nnp_hwinfo.cxgemm.s4cX_conjb_transc_only_mr_x_nr;
  							matrix_multiplication_context.full_gemm = nnp_hwinfo.cxgemm.s4cX_conjb_transc_upto_mr_x_nr;
  						} else {
  							matrix_multiplication_context.fast_gemm = nnp_hwinfo.cxgemm.cX_conjb_transc_only_mr_x_nr;
  							matrix_multiplication_context.full_gemm = nnp_hwinfo.cxgemm.cX_conjb_transc_upto_mr_x_nr;
  						}
  						pthreadpool_compute_2d_tiled(threadpool,
  							(pthreadpool_function_2d_tiled_t) compute_matrix_multiplication,
  							&matrix_multiplication_context,
  							output_channels,           input_channels_block_size,
  							output_channels_block_max, input_channels_subblock_max);
  					}
  				}
  				NNP_BLOCK_MULTIPLICATION_END(profile)
  
  
  			}
  		}
  	}
  	/* Grad kernel transform */
  	NNP_KERNEL_TRANSFORM_START(profile)
  	struct grad_kernel_transform_context grad_kernel_transform_context = {
  		.tuple_elements = tuple_elements,
  		.input_channels = input_channels,
  		.output_channels = output_channels,
  		.output_channels_block_max = output_channels_block_max,
  		.kernel_size = kernel_size,
  		.grad_kernel = grad_kernel,
  		.grad_kernel_transform = grad_kernel_transform,
  		.transform_function = grad_kernel_transform_function,
  	};
  	pthreadpool_compute_2d_tiled(threadpool,
  		(pthreadpool_function_2d_tiled_t) compute_grad_kernel_transform,
  		&grad_kernel_transform_context,
  		output_channels, input_channels,
  		1,               input_channels_subblock_max);
  	NNP_KERNEL_TRANSFORM_END(profile)
  
  	if (memory_block != workspace_buffer) {
  		release_memory(memory_block, memory_size);
  	}
  	return nnp_status_success;
  }
  
  enum nnp_status nnp_convolution_kernel_gradient(
  	enum nnp_convolution_algorithm algorithm,
  	size_t batch_size,
  	size_t input_channels,
  	size_t output_channels,
  	struct nnp_size input_size,
  	struct nnp_padding input_padding,
  	struct nnp_size kernel_size,
  	const float* input,
  	const float* grad_output,
  	float* grad_kernel,
  	void* workspace_buffer,
  	size_t* workspace_size,
  	enum nnp_activation activation,
  	const void* activation_parameters,
  	pthreadpool_t threadpool,
  	struct nnp_profile* profile)
  {
  	NNP_TOTAL_START(profile)
  
  	/* Basic validation of parameters. This check detects invalid, but not unsupported parameters. */
  	enum nnp_status status = validate_convolution_arguments(
  		batch_size, input_channels, output_channels,
  		input_size, input_padding, kernel_size, (struct nnp_size) { 1, 1 },
  		activation, activation_parameters);
  	if (status != nnp_status_success) {
  		goto cleanup;
  	}
  
  	if (activation != nnp_activation_identity) {
  		status = nnp_status_unsupported_activation;
  		goto cleanup;
  	}
  
  	if (activation_parameters != NULL) {
  		status = nnp_status_unsupported_activation_parameters;
  		goto cleanup;
  	}
  
  	const struct nnp_size output_size = {
  		.width = input_padding.left + input_size.width + input_padding.right - kernel_size.width + 1,
  		.height = input_padding.top + input_size.height + input_padding.bottom - kernel_size.height + 1
  	};
  
  	/* If requested, choose optimal convolution algorithm */
  	if (algorithm == nnp_convolution_algorithm_auto) {
  		if (max(kernel_size.width, kernel_size.height) > 8) {
  			algorithm = nnp_convolution_algorithm_ft16x16;
  		} else {
  			const size_t tile_count_8x8 =
  				divide_round_up(output_size.height, 8 - kernel_size.height + 1) *
  				divide_round_up(output_size.width, 8 - kernel_size.width + 1);
  			const size_t tile_count_16x16 =
  				divide_round_up(output_size.height, 16 - kernel_size.height + 1) *
  				divide_round_up(output_size.width, 16 - kernel_size.width + 1);
  			if (tile_count_8x8 <= 4 * tile_count_16x16) {
  				/* 8x8 tiles are more efficient */
  				algorithm = nnp_convolution_algorithm_ft8x8;
  			} else {
  				algorithm = nnp_convolution_algorithm_ft16x16;
  			}
  		}
  	}
  
  	/* Choose tiling parameters and transform functions depending on convolution algorithm */
  	struct nnp_size tile_size;
  	nnp_transform_2d_with_offset input_transform_function;
  	nnp_transform_2d_with_offset grad_output_transform_function;
  	nnp_transform_2d_with_offset grad_kernel_transform_function;
  	switch (algorithm) {
  		case nnp_convolution_algorithm_ft8x8:
  			input_transform_function = nnp_hwinfo.transforms.fft8x8_with_offset_and_stream;
  			grad_output_transform_function = nnp_hwinfo.transforms.fft8x8_with_offset_and_stream;
  			grad_kernel_transform_function = nnp_hwinfo.transforms.ifft8x8_with_offset;
  			tile_size = (struct nnp_size) { .height = 8, .width = 8 };
  			break;
  		case nnp_convolution_algorithm_ft16x16:
  			input_transform_function = nnp_hwinfo.transforms.fft16x16_with_offset_and_stream;
  			grad_output_transform_function = nnp_hwinfo.transforms.fft16x16_with_offset_and_stream;
  			grad_kernel_transform_function = nnp_hwinfo.transforms.ifft16x16_with_offset;
  			tile_size = (struct nnp_size) { .height = 16, .width = 16 };
  			break;
  		case nnp_convolution_algorithm_wt8x8:
  			/*
  			 * Winograd transform is not supported for this operation:
  			 * it needs F(5x5, 4x4) transform and presently we implement only F(3x3, 6x6)
  			 */
  			status = nnp_status_unsupported_algorithm;
  			goto cleanup;
  		case nnp_convolution_algorithm_auto:
  			NNP_UNREACHABLE;
  		default:
  			status = nnp_status_invalid_algorithm;
  			goto cleanup;
  	}
  
  	switch (algorithm) {
  		case nnp_convolution_algorithm_wt8x8:
  		case nnp_convolution_algorithm_ft8x8:
  		case nnp_convolution_algorithm_ft16x16:
  			if (kernel_size.height > tile_size.height || kernel_size.width > tile_size.width) {
  				status = nnp_status_unsupported_algorithm;
  				goto cleanup;
  			}
  			status = compute_fast_convolution_kernel_gradient(
  				batch_size, input_channels, output_channels,
  				tile_size, input_size, input_padding, kernel_size, output_size,
  				input, grad_output, grad_kernel, workspace_buffer, workspace_size,
  				input_transform_function, grad_output_transform_function, grad_kernel_transform_function,
  				threadpool, profile);
  			break;
  		case nnp_convolution_algorithm_implicit_gemm:
  		case nnp_convolution_algorithm_direct:
  		case nnp_convolution_algorithm_auto:
  			NNP_UNREACHABLE;
  	}
  
  cleanup:
  	NNP_TOTAL_END(profile)
  	return status;
  }
  