Method Compute

oneflow/user/kernels/fold_kernel.cpp:63–81

 private:
  void Compute(KernelComputeContext* ctx) const override {
    const Tensor* input = ctx->Tensor4ArgNameAndIndex("x", 0);
    Tensor* output = ctx->Tensor4ArgNameAndIndex("y", 0);

    const std::vector<int32_t> output_size = ctx->Attr<std::vector<int32_t>>("output_size");
    const std::vector<int32_t> kernel_size = ctx->Attr<std::vector<int32_t>>("kernel_size");
    const std::vector<int32_t> dilation = ctx->Attr<std::vector<int32_t>>("dilation_rate");
    const std::vector<int32_t> padding = ctx->Attr<std::vector<int32_t>>("padding");
    const std::vector<int32_t> stride = ctx->Attr<std::vector<int32_t>>("strides");

    const auto& state_ptr = CreateFoldOpKernelState<INDEX_T, NDIM, SDIM>(
        input->shape_view(), output_size, kernel_size, padding, stride, dilation);
    const FoldParams<INDEX_T, NDIM, SDIM> params = state_ptr->params();
    size_t out_bytes_size =
        output->shape_view().elem_cnt() * GetSizeOfDataType(output->data_type());
    Memset<device_type>(ctx->stream(), output->mut_dptr<T>(), 0, out_bytes_size);
    FoldKernelUtil<device_type, T, INDEX_T, NDIM, SDIM>::Forward(
        ctx->stream(), &params, input->dptr<T>(), output->mut_dptr<T>());
  }

  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
};