Method Compute

oneflow/user/kernels/same_padding_kernel.cpp:59–98

 private:
  void Compute(user_op::KernelComputeContext* ctx) const override {
    const user_op::Tensor* x = ctx->Tensor4ArgNameAndIndex("x", 0);
    user_op::Tensor* y = ctx->Tensor4ArgNameAndIndex("y", 0);
    const int64_t num_axes = x->shape_view().NumAxes();
    const std::string& padding = ctx->Attr<std::string>("padding");
    const std::string& data_format = ctx->Attr<std::string>("data_format");
    const std::vector<int32_t> kernel_size = ctx->Attr<std::vector<int32_t>>("kernel_size");
    const std::vector<int32_t> strides = ctx->Attr<std::vector<int32_t>>("strides");
    const std::vector<int32_t> dilation_rate = ctx->Attr<std::vector<int32_t>>("dilation_rate");
    std::vector<int64_t> padding_before(num_axes, 0);
    const size_t idx_offset = IdxOffset(data_format);
    const int32_t num_spatial_dims = x->shape_view().NumAxes() - 2;
    for (int32_t i = 0; i < num_spatial_dims; ++i) {
      int32_t padding_small = 0;
      int32_t padding_large = 0;
      CHECK_JUST(CalcSamePadding(x->shape_view().At(idx_offset + i), kernel_size.at(i),  // NOLINT
                                 dilation_rate.at(i), strides.at(i), &padding_small,     // NOLINT
                                 &padding_large));                                       // NOLINT
      if (padding == "same_lower") {
        padding_before[idx_offset + i] = padding_large;
      } else if (padding == "same_upper") {
        padding_before[idx_offset + i] = padding_small;
      } else {
        UNIMPLEMENTED();
      }
      CHECK_EQ(y->shape_view().At(idx_offset + i),
               x->shape_view().At(idx_offset + i) + padding_small + padding_large);
    }
    CHECK_EQ(padding_before.size(), num_axes);
    std::unique_ptr<ep::primitive::Fill> fill_primitive = NewFillPrimitive(ctx);
    CHECK(fill_primitive);
    fill_primitive->Launch(ctx->stream(), y->mut_dptr(), Scalar(0), y->shape_view().elem_cnt());
    DimVector src_pos_vec(num_axes, 0);
    DimVector dst_pos_vec(padding_before.cbegin(), padding_before.cend());
    std::unique_ptr<ep::primitive::CopyNd> copy_nd_primitive = NewCopyNdPrimitive(ctx);
    CHECK(copy_nd_primitive);
    copy_nd_primitive->Launch(ctx->stream(), x->data_type(), num_axes, y->mut_dptr(),
                              y->shape_view().ptr(), dst_pos_vec.data(), x->dptr(),
                              x->shape_view().ptr(), src_pos_vec.data(), x->shape_view().ptr());
  }
  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
};