| 752 | } |
| 753 | |
| 754 | void UserKernel::VirtualKernelInit(KernelContext* ctx) { |
| 755 | InitUserKernel(ctx->stream()); |
| 756 | CHECK(opkernel_state_.get() == nullptr); |
| 757 | opkernel_state_ = CreateOpKernelState(ctx); |
| 758 | kernel_->InitOpKernelCacheWithFlags(cache_ctx_.get(), user_op::OpKernelCache::kAllMayChanged, |
| 759 | &opkernel_cache_); |
| 760 | #ifdef WITH_CUDA_GRAPHS |
| 761 | if (ParseBooleanFromEnv("ONEFLOW_KERNEL_ENABLE_CUDA_GRAPH", false) |
| 762 | && (!ParseBooleanFromEnv("ONEFLOW_GRAPH_ENABLE_STREAM_ORDERED_MEMORY_ALLOCATION", false))) { |
| 763 | UserKernelInitContext init_ctx(ctx->stream(), kernel_conf()); |
| 764 | auto* cuda_stream = dynamic_cast<ep::CudaStream*>(ctx->stream()); |
| 765 | const auto* cuda_graph_support = dynamic_cast<const user_op::CudaGraphSupport*>(kernel_.get()); |
| 766 | if (cuda_stream != nullptr) { |
| 767 | if (cuda_graph_support != nullptr |
| 768 | && cuda_graph_support->IsCudaGraphSupported(&init_ctx, opkernel_state_.get())) { |
| 769 | cuda_graph_exec_.reset(new ep::CudaGraphExecutable()); |
| 770 | VLOG(3) << "CUDA Graphs Kernel: " << op_conf().name() << " (" |
| 771 | << op_conf().user_conf().op_type_name() << ")"; |
| 772 | } else { |
| 773 | VLOG(3) << "CUDA Graphs not supported: " << op_conf().name() << " (" |
| 774 | << op_conf().user_conf().op_type_name() << ")"; |
| 775 | } |
| 776 | } |
| 777 | } |
| 778 | #endif // WITH_CUDA_GRAPHS |
| 779 | } |
| 780 | |
| 781 | void UserKernel::ForwardDataContent(KernelContext* ctx) const { |
| 782 | const auto BnInOp2Blob = [ctx](const std::string& bn) { return ctx->BnInOp2Blob(bn); }; |
nothing calls this directly
no test coverage detected