Method Compute

oneflow/user/kernels/lerp_kernel.cpp:187–216

 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* 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);

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

    const T* start_ptr = start->dptr<T>();
    const T* end_ptr = end->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>();

    Scalar scalar_operand;
    if (ctx->Attr<bool>("has_int_operand")) {
      scalar_operand = ctx->Attr<int64_t>("int_operand");
    } else if (ctx->Attr<bool>("has_float_operand")) {
      scalar_operand = ctx->Attr<double>("float_operand");
    } else {
      UNIMPLEMENTED();
    }

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