CudaGraphPrimitive (#6264)

Co-authored-by: Noneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

oneflow/core/primitive/cuda/cast.cu

@@ -15,6 +15,7 @@ limitations under the License.
 */
 #include "oneflow/core/primitive/cast.h"
 #include "oneflow/core/primitive/cuda/type_seq.h"
+#include "oneflow/core/primitive/cuda/cuda_graph_support.h"
 #include "oneflow/core/cuda/elementwise.cuh"
 #include "oneflow/core/stream/cuda_stream_context.h"
 
@@ -73,7 +74,7 @@ void LaunchCast(cudaStream_t stream, const void* from, void* to, size_t count) {
 using LaunchFn = std::function<void(cudaStream_t /*stream*/, const void* /*from*/, void* /*to*/,
                                     size_t /*count*/)>;
 
-class CastImpl : public Cast {
+class CastImpl : public Cast, public CudaGraphSupport {
  public:
   OF_DISALLOW_COPY_AND_MOVE(CastImpl);
   explicit CastImpl(LaunchFn launch_fn) : launch_fn_(std::move(launch_fn)) {}

oneflow/core/primitive/cuda/cuda_graph_support.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/primitive/cuda/cuda_graph_support.h"
+
+namespace oneflow {
+
+namespace primitive {
+
+bool IsCudaGraphPrimitive(const Primitive* primitive) {
+  auto* cuda_graph_support = dynamic_cast<const CudaGraphSupport*>(primitive);
+  if (cuda_graph_support == nullptr) { return false; }
+  return cuda_graph_support->IsCudaGraphSupported();
+}
+
+}  // namespace primitive
+
+}  // namespace oneflow

oneflow/core/primitive/cuda/cuda_graph_support.h

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#ifndef ONEFLOW_CORE_PRIMITIVE_CUDA_CUDA_GRAPH_SUPPORT_H_
+#define ONEFLOW_CORE_PRIMITIVE_CUDA_CUDA_GRAPH_SUPPORT_H_
+
+#include "oneflow/core/primitive/primitive.h"
+
+namespace oneflow {
+
+namespace primitive {
+
+class CudaGraphSupport {
+ public:
+  OF_DISALLOW_COPY_AND_MOVE(CudaGraphSupport);
+  CudaGraphSupport() = default;
+  virtual ~CudaGraphSupport() = default;
+
+  virtual bool IsCudaGraphSupported() const { return true; }
+};
+
+bool IsCudaGraphPrimitive(const Primitive* primitive);
+
+}  // namespace primitive
+
+}  // namespace oneflow
+
+#endif  // ONEFLOW_CORE_PRIMITIVE_CUDA_CUDA_GRAPH_SUPPORT_H_

oneflow/core/primitive/cuda/memcpy.cpp

@@ -17,6 +17,7 @@ limitations under the License.
 
 #include "oneflow/core/primitive/memcpy.h"
 #include "oneflow/core/stream/cuda_stream_context.h"
+#include "oneflow/core/primitive/cuda/cuda_graph_support.h"
 #include <cuda_runtime.h>
 
 namespace oneflow {
@@ -25,7 +26,7 @@ namespace primitive {
 
 namespace {
 
-class MemcpyImpl : public Memcpy {
+class MemcpyImpl : public Memcpy, public CudaGraphSupport {
  public:
   OF_DISALLOW_COPY_AND_MOVE(MemcpyImpl);
   MemcpyImpl() = default;

oneflow/core/primitive/cuda/memset.cpp

@@ -17,6 +17,7 @@ limitations under the License.
 
 #include "oneflow/core/primitive/memset.h"
 #include "oneflow/core/stream/cuda_stream_context.h"
+#include "oneflow/core/primitive/cuda/cuda_graph_support.h"
 #include <cuda_runtime.h>
 
 namespace oneflow {
@@ -25,7 +26,7 @@ namespace primitive {
 
 namespace {
 
-class MemsetImpl : public Memset {
+class MemsetImpl : public Memset, public CudaGraphSupport {
  public:
   OF_DISALLOW_COPY_AND_MOVE(MemsetImpl);
   MemsetImpl() = default;