Method HandleConvolution

tensorflow/compiler/xla/service/hlo_cost_analysis.cc:540–663

}

Status HloCostAnalysis::HandleConvolution(const HloInstruction* convolution) {
  auto lhs = convolution->operand(0);
  auto rhs = convolution->operand(1);
  Window window = convolution->window();
  const auto& result_shape = convolution->shape();
  const Shape& lhs_shape = lhs->shape();
  const Shape& rhs_shape = rhs->shape();

  const auto& dnums = convolution->convolution_dimension_numbers();

  const int64 input_batch_dim = dnums.input_batch_dimension();
  const int64 input_feature_dim = dnums.input_feature_dimension();
  const int64 output_feature_dim = dnums.output_feature_dimension();
  const int64 input_feature =
      ShapeUtil::GetDimension(lhs_shape, input_feature_dim);
  const int64 output_feature =
      ShapeUtil::GetDimension(result_shape, output_feature_dim);
  const int64 batch = ShapeUtil::GetDimension(lhs_shape, input_batch_dim);

  DimensionVector kernel_limits;
  DimensionVector output_limits;
  DimensionVector input_limits;
  if (window.dimensions().empty()) {
    window = window_util::MakeWindow({1});
    kernel_limits.push_back(1);
    output_limits.push_back(1);
    input_limits.push_back(1);
  } else {
    for (int64 spatial_dimension = 0;
         spatial_dimension < window.dimensions_size(); ++spatial_dimension) {
      // Spatial dimension number for kernel (rhs).
      const int64 kernel_spatial_dim =
          dnums.kernel_spatial_dimensions(spatial_dimension);
      const int64 kernel_limit = rhs_shape.dimensions(kernel_spatial_dim);
      kernel_limits.push_back(kernel_limit);

      // Spatial dimension number for output.
      const int64 output_spatial_dim =
          dnums.output_spatial_dimensions(spatial_dimension);
      const int64 output_limit = result_shape.dimensions(output_spatial_dim);
      output_limits.push_back(output_limit);

      // Spatial dimension number for input (lhs).
      const int64 input_spatial_dim =
          dnums.input_spatial_dimensions(spatial_dimension);
      const int64 input_limit = lhs_shape.dimensions(input_spatial_dim);
      input_limits.push_back(input_limit);
    }
  }

  DimensionVector valid_position_counts;

  // Loop over each spatial dimension.
  for (int64 spatial_dimension = 0;
       spatial_dimension < window.dimensions_size(); ++spatial_dimension) {
    const auto& window_dim = window.dimensions(spatial_dimension);
    // These two conditions will create an N^2 iteration pattern with only N
    // valid elements. This is a performance optimization and produces the same