Method Compute

oneflow/user/kernels/dim_gather_kernels.cpp:41–61

 private:
  void Compute(KernelComputeContext* ctx) const override {
    const Tensor* input_tensor = ctx->Tensor4ArgNameAndIndex("input", 0);
    if (input_tensor->shape_view().elem_cnt() == 0) { return; }
    const Tensor* index_tensor = ctx->Tensor4ArgNameAndIndex("index", 0);
    Tensor* out_tensor = ctx->Tensor4ArgNameAndIndex("output", 0);
    const int32_t dim = ctx->Attr<int32_t>("dim");

    const IN_T* input = input_tensor->dptr<IN_T>();
    const IDX_T* index = index_tensor->dptr<IDX_T>();
    IN_T* output = out_tensor->mut_dptr<IN_T>();

    const Shape in_shape = ExpandDimIf0D(input_tensor->shape_view());
    const auto ndim = in_shape.NumAxes();
    const auto dim_length = in_shape.At(dim);

    DimOpIndexNdHelper<IDX_T> input_nd_helper(in_shape.data(), ndim);
    DimOpIndexNdHelper<IDX_T> index_nd_helper(index_tensor->shape_view().data(), ndim);
    DimGatherFunctor<device_type, IN_T, IDX_T>()(ctx->stream(), input_nd_helper, index_nd_helper,
                                                 ndim, index_tensor->shape_view().elem_cnt(),
                                                 dim_length, dim, index, input, output);
  }
  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
};