Method Compute

oneflow/user/kernels/gather_kernel.cpp:86–108

 private:
  void Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState*,
               const user_op::OpKernelCache* cache) const override {
    const user_op::Tensor* in = ctx->Tensor4ArgNameAndIndex("in", 0);
    const user_op::Tensor* indices = ctx->Tensor4ArgNameAndIndex("indices", 0);
    const int64_t axis = ctx->Attr<int64_t>("axis");
    const int64_t num_indices = indices->shape_view().elem_cnt();
    user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0);
    if (out->shape_view().elem_cnt() == 0) { return; }

    const Shape in_shape = ExpandDimIf0D(in->shape_view());

    int64_t offset = 0;
    if (cache != nullptr) {
      auto* gather_cache = dynamic_cast<const GatherOpKernelCache*>(cache);
      CHECK_NOTNULL(gather_cache);
      CHECK_EQ(in_shape.At(axis), gather_cache->upper() - gather_cache->lower());
      offset = gather_cache->lower();
    }

    GatherKernelUtilImpl<device_type, T, K>::Forward(ctx->stream(), indices->dptr<K>(), num_indices,
                                                     in->dptr<T>(), GetFlatShape(in_shape, axis),
                                                     out->mut_dptr<T>(), offset);
  }
  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
};