Method InitOpKernelCache

oneflow/user/kernels/slice_kernel.cpp:290–325

  ~SliceKernel() = default;

  std::shared_ptr<user_op::OpKernelCache> InitOpKernelCache(
      user_op::KernelCacheContext* ctx) const override {
    SliceContext slice_ctx;
    if (ctx->parallel_ctx().parallel_num() == 1) {
      // split_axis == SPLIT_AXIS_FOR_NON_SPLIT means the sbp attribute is not 'split'
      CHECK_JUST(slice_ctx.PushSplitInfo(SPLIT_AXIS_FOR_NON_SPLIT, 0, 0, 0));
    } else {
      const Shape& parallel_hierarchy = *ctx->parallel_desc().hierarchy();
      NdSbp in_nd_sbp = ctx->NdSbp4ArgNameAndIndex("x", 0);
      {
        const NdSbp& y_nd_sbp = ctx->NdSbp4ArgNameAndIndex("y", 0);
        // If x and y both split in the same axis(must be full slice),
        // we can consider the physical tensor is broadcast in this axis.
        FOR_RANGE(int32_t, i, 0, parallel_hierarchy.NumAxes()) {
          const SbpParallel& x_sbp = in_nd_sbp.sbp_parallel(i);
          const SbpParallel& y_sbp = y_nd_sbp.sbp_parallel(i);
          if (x_sbp.has_split_parallel() && y_sbp.has_split_parallel()) {
            CHECK_EQ(x_sbp.split_parallel().axis(), y_sbp.split_parallel().axis());
            in_nd_sbp.mutable_sbp_parallel(i)->clear_split_parallel();
            in_nd_sbp.mutable_sbp_parallel(i)->mutable_broadcast_parallel();
          }
        }
      }
      const Shape& logical_shape = ctx->LogicalTensorDesc4ArgNameAndIndex("x", 0)->shape();
      const int64_t parallel_id = ctx->parallel_ctx().parallel_id();
      const TensorSliceView& slice_view =
          GetTensorSliceView4ParallelId(parallel_hierarchy, in_nd_sbp, logical_shape, parallel_id);
      for (int i = 0; i < logical_shape.NumAxes(); ++i) {
        const Range& range = slice_view.At(i);
        if (range.begin() != 0 || range.end() != logical_shape.At(i)) {
          CHECK_JUST(slice_ctx.PushSplitInfo(i, range.begin(), range.end(), logical_shape.At(i)));
        }
      }
    }
    return std::make_shared<OpKernelCacheWrapper<SliceContext>>(slice_ctx);
  }

 private:
  void Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState*,