Method Compute

tensorflow/core/kernels/mkl_dequantize_op.cc:52–198

  }

  void Compute(OpKernelContext* ctx) override {
    try {
      // Using CPU device
      auto cpu_engine = engine(ENGINE_CPU, 0);

      // Get the inputs
      const Tensor& src_tensor = ctx->input(kSrcIndex);
      const float min_range = ctx->input(1).flat<float>()(0);
      const float max_range = ctx->input(2).flat<float>()(0);

      // Get MklShape
      MklDnnShape src_mkl_shape;
      GetMklShape(ctx, kSrcIndex, &src_mkl_shape);
      auto src_tf_shape = src_mkl_shape.IsMklTensor()
                              ? src_mkl_shape.GetTfShape()
                              : src_tensor.shape();

      // src_dims is the dimension of src_tensor
      // output_dims are same as src_dims
      auto src_dims = src_mkl_shape.IsMklTensor()
                          ? src_mkl_shape.GetSizesAsMklDnnDims()
                          : TFShapeToMklDnnDims(src_tensor.shape());
      auto output_dims = src_dims;

      // Create reorder memory for src and dst
      MklDnnData<T> src(&cpu_engine);
      MklDnnData<float> dst(&cpu_engine);

      std::shared_ptr<stream> reorder_stream;
      MklDnnThreadPool eigen_tp(ctx);
      reorder_stream.reset(CreateStream(&eigen_tp, cpu_engine));

      memory::format_tag dst_layout_type;
      switch (src_tf_shape.dims()) {
        case 1:
          dst_layout_type = memory::format_tag::x;
          break;
        case 2:
          dst_layout_type = memory::format_tag::nc;
          break;
        case 3:
          dst_layout_type = memory::format_tag::tnc;
          break;
        case 4:
          dst_layout_type = memory::format_tag::nhwc;
          break;
        case 5:
          dst_layout_type = memory::format_tag::ndhwc;
          break;
        default:
          OP_REQUIRES_OK(ctx,
                         errors::Aborted("Input dims must be <= 5 and >= 1"));
          return;
      }

      // If input is in OneDNN layout, then simply grab input layout; otherwise,
      // construct input TF layout. For TF layout, although input shape
      // (src_dims) required is in OneDNN order, the layout is Tensorflow's