Method Compute

oneflow/user/kernels/lerp_kernel.cpp:79–105

 private:
  void Compute(user_op::KernelComputeContext* ctx) const override {
    const user_op::Tensor* start = ctx->Tensor4ArgNameAndIndex("start", 0);
    const user_op::Tensor* end = ctx->Tensor4ArgNameAndIndex("end", 0);
    const user_op::Tensor* weight = ctx->Tensor4ArgNameAndIndex("weight", 0);
    const user_op::Tensor* out_diff = ctx->Tensor4ArgNameAndIndex("out_diff", 0);
    user_op::Tensor* start_diff = ctx->Tensor4ArgNameAndIndex("start_diff", 0);
    user_op::Tensor* end_diff = ctx->Tensor4ArgNameAndIndex("end_diff", 0);
    user_op::Tensor* weight_diff = ctx->Tensor4ArgNameAndIndex("weight_diff", 0);

    const ShapeView& start_shape = start->shape_view();
    const ShapeView& end_shape = end->shape_view();
    const ShapeView& weight_shape = weight->shape_view();
    CHECK_EQ(start_shape, end_shape);
    CHECK_EQ(start_shape, weight_shape);

    const T* start_ptr = start->dptr<T>();
    const T* end_ptr = end->dptr<T>();
    const T* weight_ptr = weight->dptr<T>();
    const T* out_diff_ptr = out_diff->dptr<T>();
    T* start_diff_ptr = start_diff->mut_dptr<T>();
    T* end_diff_ptr = end_diff->mut_dptr<T>();
    T* weight_diff_ptr = weight_diff->mut_dptr<T>();

    LerpKernelUtil<device_type, T>::Backward(ctx->stream(), start_shape.elem_cnt(), start_ptr,
                                             weight_ptr, end_ptr, out_diff_ptr, start_diff_ptr,
                                             weight_diff_ptr, end_diff_ptr);
  }
  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
};