Method Compute

tensorflow/core/kernels/mkl_aggregate_ops.cc:109–252

  }

  void Compute(OpKernelContext* ctx) override {
    // Each input tensor in OneDNN layout has additional meta-tensor carrying
    // layout information. So the number of actual tensors is half the total
    // number of inputs.
    const int num_inputs = ctx->num_inputs() / 2;

    MklDnnShape mkl_shape;
    const size_t kSrc0Idx = 0;
    const size_t kOutputIdx = 0;

    if (num_inputs == 1) {
      GetMklShape(ctx, kSrc0Idx, &mkl_shape);
      bool input_in_mkl_format = mkl_shape.IsMklTensor();

      if (input_in_mkl_format) {
        ForwardMklTensorInToOut(ctx, kSrc0Idx, kOutputIdx);
      } else {
        ForwardTfTensorInToOut(ctx, kSrc0Idx, kOutputIdx);
      }
      return;
    }

    // Check if the input shape is same
    if (!CheckInputShape(ctx)) return;

    try {
      TensorShape output_tf_shape;
      MklDnnShape output_mkl_shape;
      const Tensor& src_tensor = MklGetInput(ctx, kSrc0Idx);

      Tensor* dst_tensor = nullptr;

      // Nothing to compute, return.
      if (src_tensor.shape().num_elements() == 0) {
        output_mkl_shape.SetMklTensor(false);
        output_tf_shape = src_tensor.shape();
        AllocateOutputSetMklShape(ctx, kOutputIdx, &dst_tensor, output_tf_shape,
                                  output_mkl_shape);
        return;
      }

      if (src_tensor.dims() == 0) {
        ComputeScalar(ctx);
        return;
      }

      auto cpu_engine = engine(ENGINE_CPU, 0);
      std::vector<float> coeff(num_inputs, 1.0);
      std::vector<MEMORY_PRIMITIVE_DESC> srcs_pd;
      std::vector<memory> inputs;

      MklDnnData<T> dst(&cpu_engine);
      MklDnnData<T> src(&cpu_engine);
      bool has_mkl_input = false;
      int mkl_input_index = FindMKLInputIndex(ctx);
      MKL_TENSOR_FORMAT mkl_data_format;
      TensorFormat tf_data_format;
      MEMORY_FORMAT dnn_fmt = MEMORY_FORMAT::any;