Method Compute

oneflow/user/kernels/dim_scatter_kernels.cpp:30–73

 private:
  void Compute(KernelComputeContext* ctx) const override {
    const Tensor* input_tensor = ctx->Tensor4ArgNameAndIndex("input", 0);
    const Tensor* index_tensor = ctx->Tensor4ArgNameAndIndex("index", 0);
    Tensor* out_tensor = ctx->Tensor4ArgNameAndIndex("output", 0);
    const Tensor* src_tensor = ctx->Tensor4ArgNameAndIndex("src", 0);
    const int32_t dim = ctx->Attr<int32_t>("dim");

    const IDX_T* index = index_tensor->dptr<IDX_T>();
    IN_T* output = out_tensor->mut_dptr<IN_T>();
    size_t out_bytes_size =
        out_tensor->shape_view().elem_cnt() * GetSizeOfDataType(out_tensor->data_type());

    Tensor* like_tensor = ctx->Tensor4ArgNameAndIndex("like", 0);
    const IN_T* src = src_tensor->dptr<IN_T>();

    if (input_tensor) {
      Memcpy<device_type>(ctx->stream(), output, input_tensor->dptr<IN_T>(), out_bytes_size);
    } else if (like_tensor) {
      Memset<device_type>(ctx->stream(), output, 0, out_bytes_size);
    } else {
      UNIMPLEMENTED() << "Input tensor and like tensor cannot be empty simultaneously.";
    }

    const Shape src_shape = ExpandDimIf0D(src_tensor->shape_view());
    const Shape index_shape = ExpandDimIf0D(index_tensor->shape_view());
    const int ndim = src_shape.NumAxes();
    DimOpIndexNdHelper<IDX_T> src_nd_helper(src_shape.data(), ndim);
    DimOpIndexNdHelper<IDX_T> idx_nd_helper(index_shape.data(), ndim);
    DimOpIndexNdHelper<IDX_T> output_nd_helper(out_tensor->shape_view().data(), ndim);

    const int64_t upper_bound = [&]() {
      if (input_tensor) {
        const Shape input_shape = ExpandDimIf0D(input_tensor->shape_view());
        return input_shape.At(dim);
      } else {
        const Shape like_shape = ExpandDimIf0D(like_tensor->shape_view());
        return like_shape.At(dim);
      }
    }();

    DimScatterFunctor<device_type, IN_T, IDX_T, Opt>()(
        ctx->stream(), src_nd_helper, idx_nd_helper, output_nd_helper, ndim, index_shape.elem_cnt(),
        dim, upper_bound, index, src, output);
  }
  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
};