Method ForwardShape

oneflow/core/kernel/user_kernel.cpp:786–817

}

void UserKernel::ForwardShape(KernelContext* ctx) const {
  const auto BnInOp2Blob = [ctx](const std::string& bn) { return ctx->BnInOp2Blob(bn); };
  infer_ctx_->UpdateArg2Tensor(BnInOp2Blob);
  infer_cache_->UpdateCacheKey(infer_ctx_.get());
  if (!infer_cache_->IsCacheHit()) {
    auto* op_infer_ctx = dynamic_cast<UserKernelOpInferContext*>(infer_ctx_->MutOpInferContext());
    CHECK_NOTNULL(op_infer_ctx);
    op_infer_ctx->UpdateArg2TensorDesc(BnInOp2Blob);
    kernel_->InferShape(infer_ctx_.get());
    for (const auto& out_arg_pair : infer_ctx_->outputs()) {
      const Shape& static_shape =
          infer_ctx_->TensorDesc4ArgNameAndIndex(out_arg_pair.first, out_arg_pair.second)->shape();
      const ShapeView& shape_view =
          infer_ctx_->ShapeView4ArgNameAndIndex(out_arg_pair.first, out_arg_pair.second);
      CHECK_LE(shape_view.elem_cnt(), static_shape.elem_cnt())
          << "InferShape of OpKernel (op_type_name: " << op_conf().user_conf().op_type_name()
          << ", op_name: " << op_conf().name()
          << ") raise error, output arg's (name: " << out_arg_pair.first
          << ", index: " << out_arg_pair.second << ") runtime shape " << shape_view.ToString()
          << " surpass the limit of static shape " << static_shape.ToString();
    }
    infer_cache_->UpdateCacheValue(infer_ctx_.get());
  } else {
    std::shared_ptr<const OpInferCacheValue> cache_value_ptr = infer_cache_->GetCacheValue();
    FOR_RANGE(int, i, 0, infer_ctx_->outputs().size()) {
      const auto& out_arg_pair = infer_ctx_->outputs().at(i);
      MutShapeView mut_shape_view =
          infer_ctx_->MutShapeView4ArgNameAndIndex(out_arg_pair.first, out_arg_pair.second);
      mut_shape_view.set_shape(*cache_value_ptr->obn_idx2shape_sym.at(i));
    }
  }
}

bool UserKernel::IsStateless() const { return !kernel_->AlwaysComputeWhenAllOutputsEmpty(); }
NEW_REGISTER_KERNEL(OperatorConf::kUserConf, UserKernel).SetIsMatchedPred([](const KernelConf&) {