Method Compute

tensorflow/core/kernels/mkl_cwise_ops_common.cc:209–346

      : BinaryOp<Device, Functor>(context) {}

  void Compute(OpKernelContext* context) override {
    const int kNumInputs = 2;
    std::vector<Tensor> inputs(kNumInputs);
    std::vector<MklDnnShape> mkl_input_shapes(kNumInputs);
    std::vector<TensorShape> tf_input_shapes(kNumInputs);
    std::vector<T*> inputs_data(kNumInputs);

    for (int i = 0; i < kNumInputs; i++) {
      inputs[i] = context->input(i);
      inputs_data[i] = inputs[i].flat<T>().data();
      GetMklShape(context, i, &(mkl_input_shapes[i]));

      tf_input_shapes[i] = mkl_input_shapes[i].IsMklTensor()
                               ? mkl_input_shapes[i].GetTfShape()
                               : inputs[i].shape();
    }

    VLOG(1) << "MklBinaryOp: Inputs shapes " << tf_input_shapes[0].DebugString()
            << " _and_ " << tf_input_shapes[1].DebugString();

    if (!ShouldFallback(tf_input_shapes[0], tf_input_shapes[1])) {
      VLOG(1) << "MklBinaryOp: Run oneDNN primitive";

      std::vector<memory::dims> srcs_dims(kNumInputs);
      for (int i = 0; i < kNumInputs; i++) {
        auto& mkl_shape = mkl_input_shapes[i];
        if (mkl_shape.IsMklTensor()) {
          srcs_dims[i] = mkl_shape.GetSizesAsMklDnnDims();
        } else {
          srcs_dims[i] = TFShapeToMklDnnDims(tf_input_shapes[i]);
        }
      }

      // oneDNN only supports inputs[1] bcast to inputs[0]. So if inputs[1]
      // has more elements than inputs[0], swap the 2 inputs.
      // Use an index to indicate the swapped result.
      const int kFirst = (tf_input_shapes[1].num_elements() >
                          tf_input_shapes[0].num_elements()) ||
                         (tf_input_shapes[1].num_elements() ==
                              tf_input_shapes[0].num_elements() &&
                          srcs_dims[1].size() > srcs_dims[0].size());
      const int kSecond = 1 - kFirst;

      // oneDNN only supports inputs with same rank size, so expand dimension
      // if they are not consistent.
      // E.g. 8x4 * 4 --> 8x4 * 1x4.
      if (srcs_dims[0].size() != srcs_dims[1].size()) {
        const int kSmall = srcs_dims[0].size() > srcs_dims[1].size();
        ExpandDim(srcs_dims[kSmall], srcs_dims[1 - kSmall].size());
      }

      std::vector<memory::desc> srcs_md(kNumInputs);
      for (int i = 0; i < kNumInputs; i++) {
        auto& mkl_shape = mkl_input_shapes[i];
        if (mkl_shape.IsMklTensor()) {
          srcs_md[i] = mkl_shape.GetMklLayout();
        } else {
          auto src_strides = CalculateTFStrides(srcs_dims[i]);