Update cudnn convolution kernel

cmake/third_party.cmake

@@ -145,6 +145,7 @@ if(BUILD_CUDA)
  include(nccl)
  include(cutlass)
  include(trt_flash_attention)
+ include(cudnn-frontend)
 
  list(APPEND oneflow_third_party_libs ${NCCL_LIBRARIES})
  list(APPEND oneflow_third_party_libs ${CUDNN_LIBRARIES})
@@ -164,6 +165,8 @@ if(BUILD_CUDA)
  list(APPEND oneflow_third_party_dependencies trt_flash_attention)
  list(APPEND oneflow_third_party_libs ${TRT_FLASH_ATTENTION_LIBRARIES})
  list(APPEND ONEFLOW_THIRD_PARTY_INCLUDE_DIRS ${TRT_FLASH_ATTENTION_INCLUDE_DIR})
+ list(APPEND oneflow_third_party_dependencies cudnn_frontend_copy_headers_to_destination)
+ list(APPEND ONEFLOW_THIRD_PARTY_INCLUDE_DIRS ${CUDNN_FRONTEND_INCLUDE_DIR})
 endif()
 
 if(BUILD_RDMA)

cmake/third_party/cudnn-frontend.cmake

@@ -0,0 +1,29 @@
+include(ExternalProject)
+
+set(CUDNN_FRONTEND_URL https://github.com/NVIDIA/cudnn-frontend/archive/refs/tags/v1.1.2.zip)
+set(CUDNN_FRONTEND_MD5 7e16cc2dcaddefa7fd0f3d82b9cf5d73)
+use_mirror(VARIABLE CUDNN_FRONTEND_URL URL ${CUDNN_FRONTEND_URL})
+
+set(CUDNN_FRONTEND_INCLUDE_DIR ${THIRD_PARTY_DIR}/cudnn-frontend/include)
+set(CUDNN_FRONTEND_BASE_DIR ${CMAKE_CURRENT_BINARY_DIR}/cudnn-frontend/src/cudnn-frontend)
+
+if(THIRD_PARTY)
+ ExternalProject_Add(
+ cudnn-frontend
+ PREFIX cudnn-frontend
+ URL ${CUDNN_FRONTEND_URL}
+ URL_MD5 ${CUDNN_FRONTEND_MD5}
+ CONFIGURE_COMMAND ""
+ BUILD_COMMAND ""
+ INSTALL_COMMAND "")
+
+ add_copy_headers_target(
+ NAME
+ cudnn_frontend
+ SRC
+ ${CUDNN_FRONTEND_BASE_DIR}/include/
+ DST
+ ${CUDNN_FRONTEND_INCLUDE_DIR}
+ DEPS
+ cudnn-frontend)
+endif(THIRD_PARTY)

oneflow/core/device/cudnn_conv_util.cpp

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 #ifdef WITH_CUDA
+#include "oneflow/core/framework/infer_util.h"
 #include "oneflow/core/device/cudnn_conv_util.h"
 #include "oneflow/core/device/cuda_util.h"
 #include "oneflow/core/common/cached_caller.h"
@@ -22,6 +23,7 @@ limitations under the License.
 #include "oneflow/core/job/global_for.h"
 #include "oneflow/core/job/global_for.h"
 #include "oneflow/core/framework/op_kernel.h"
+#include "oneflow/core/job/lazy_mode.h"
 
 namespace oneflow {
 
@@ -82,6 +84,7 @@ perf_t GetBestAlgorithm(const CudnnConvArgs& args, CudnnConvResource* res,
  FOR_RANGE(size_t, i, 0, perf_vec.size()) {
  // Note: Shouldn't all returned results be successful?
  CHECK_EQ(perf_vec[i].status, CUDNN_STATUS_SUCCESS);
+ // TODO workspace size limit will lead to dismatch result with pytorch for large tensor
  if (perf_vec[i].memory > args.params.max_ws_size) { continue; }
  if (args.deterministic && perf_vec[i].determinism == CUDNN_NON_DETERMINISTIC) { continue; }
  found_algo_idx = i;
@@ -332,6 +335,22 @@ CudnnConvArgs::CudnnConvArgs(const user_op::KernelComputeContext& ctx, DataType
  params.max_ws_size = max_workspace_size;
 }
 
+CudnnConvArgsV8::CudnnConvArgsV8(const user_op::InferContext& ctx, const user_op::TensorDesc& x,
+ const user_op::TensorDesc& y, const user_op::TensorDesc& w)
+ : xdesc(GetTensorDescriptor(x, 'x')),
+ ydesc(GetTensorDescriptor(y, 'y')),
+ wdesc(GetTensorDescriptor(w, 'w')),
+ cdesc(GetConvDescriptor(ctx, GetCudnnDataType(y.data_type()))),
+ beta(0.0f) {}
+
+CudnnConvArgsV8::CudnnConvArgsV8(const user_op::KernelComputeContext& ctx, const user_op::Tensor* x,
+ const user_op::Tensor* y, const user_op::Tensor* w)
+ : xdesc(GetTensorDescriptor(x, 'x')),
+ ydesc(GetTensorDescriptor(y, 'y')),
+ wdesc(GetTensorDescriptor(w, 'w')),
+ cdesc(GetConvDescriptor(ctx, GetCudnnDataType(y->data_type()))),
+ beta(0.0f) {}
+
 ManagedCudnnConvResource::ManagedCudnnConvResource(const CudnnConvArgs& args)
  : handle_(nullptr), x_dptr_(nullptr), w_dptr_(nullptr), y_dptr_(nullptr), ws_dptr_(nullptr) {
  x_byte_size_ = ByteSize4Tensor(args.params.x_dims, args.params.x_ndim, args.params.x_data_type);
@@ -424,6 +443,245 @@ cudnnStatus_t GetCudnnConvWorkspaceSize(const CudnnConvArgs& args, CudnnConvReso
  args.wdesc.Get(), algo, sz);
 }
 
+void RunSingleConv(const cudnnHandle_t handle, const cudnnBackendDescriptorType_t desc,
+ user_op::Tensor* x, user_op::Tensor* y, user_op::Tensor* w, user_op::Tensor* b,
+ const CudnnConvArgsV8& args) {
+ std::string tag;
+ auto configs =
+ GetConfigs(handle, desc, args.xdesc, args.ydesc, args.wdesc, args.cdesc, args.beta, tag);
+ TryConfigs(handle, x, y, w, b, configs, tag);
+}
+
+cudnn_frontend::EngineConfigList GetConfigs(const cudnnHandle_t handle,
+ const cudnnBackendDescriptorType_t desc,
+ const cudnn_frontend::Tensor& xdesc,
+ const cudnn_frontend::Tensor& ydesc,
+ const cudnn_frontend::Tensor& wdesc,
+ const cudnn_frontend::ConvDesc& cdesc, float beta,
+ std::string& tag) {
+ auto op_graph = BuildConvOpGraph(handle, desc, xdesc, ydesc, wdesc, cdesc, beta);
+ tag = op_graph.getTag();
+ auto sources = GetGeneratorSources(desc);
+ cudnn_frontend::EngineConfigGenerator generator(sources.size(), sources.data());
+ auto configs = generator.generate_engine_config(op_graph);
+ return configs;
+}
+
+cudnn_frontend::OperationGraph BuildConvOpGraph(const cudnnHandle_t handle,
+ const cudnnBackendDescriptorType_t desc,
+ const cudnn_frontend::Tensor& xdesc,
+ const cudnn_frontend::Tensor& ydesc,
+ const cudnn_frontend::Tensor& wdesc,
+ const cudnn_frontend::ConvDesc& cdesc, float beta) {
+ auto conv_op = cudnn_frontend::OperationBuilder(desc)
+ .setxDesc(xdesc)
+ .setyDesc(ydesc)
+ .setwDesc(wdesc)
+ .setcDesc(cdesc)
+ .setBeta(beta)
+ .build();
+ std::array<cudnn_frontend::Operation const*, 1> ops = {&conv_op};
+ auto op_graph = cudnn_frontend::OperationGraphBuilder()
+ .setHandle(handle)
+ .setOperationGraph(ops.size(), ops.data())
+ .build();
+ return op_graph;
+}
+
+cudnn_frontend::Tensor GetTensorDescriptor(const user_op::Tensor* t, const int64_t id) {
+ auto dim = t->shape_view();
+ auto stride = t->stride();
+ return cudnn_frontend::TensorBuilder()
+ .setDim(dim.size(), dim.data())
+ .setStride(stride.size(), stride.data())
+ .setId(id)
+ .setAlignment(32)
+ .setDataType(GetCudnnDataType(t->data_type()))
+ .build();
+}
+
+cudnn_frontend::Tensor GetTensorDescriptor(const user_op::TensorDesc& t, const int64_t id) {
+ auto dim = t.shape();
+ auto stride = t.stride();
+ return cudnn_frontend::TensorBuilder()
+ .setDim(dim.size(), dim.data())
+ .setStride(stride.size(), stride.data())
+ .setId(id)
+ .setAlignment(32)
+ .setDataType(GetCudnnDataType(t.data_type()))
+ .build();
+}
+
+cudnn_frontend::ConvDesc GetConvDescriptor(const user_op::InferContext& ctx,
+ cudnnDataType_t data_type) {
+ if (data_type == CUDNN_DATA_HALF || data_type == CUDNN_DATA_BFLOAT16) {
+ data_type = CUDNN_DATA_FLOAT;
+ }
+
+ std::vector<int64_t> padding;
+ const auto& padding_before = ctx.Attr<std::vector<int32_t>>("padding_before");
+ copy(padding_before.begin(), padding_before.end(), back_inserter(padding));
+
+ std::vector<int64_t> stride;
+ const auto& strides = ctx.Attr<std::vector<int32_t>>("strides");
+ copy(strides.begin(), strides.end(), back_inserter(stride));
+
+ std::vector<int64_t> dilation;
+ const auto& dilation_rate = ctx.Attr<std::vector<int32_t>>("dilation_rate");
+ copy(dilation_rate.begin(), dilation_rate.end(), back_inserter(dilation));
+
+ uint64_t ndim = stride.size();
+ return cudnn_frontend::ConvDescBuilder()
+ .setDataType(data_type)
+ .setMathMode(CUDNN_CROSS_CORRELATION)
+ .setNDims(ndim)
+ .setStrides(ndim, stride.data())
+ .setPrePadding(ndim, padding.data())
+ .setPostPadding(ndim, padding.data())
+ .setDilation(ndim, dilation.data())
+ .build();
+}
+
+cudnn_frontend::ConvDesc GetConvDescriptor(const user_op::KernelComputeContext& ctx,
+ cudnnDataType_t data_type) {
+ if (data_type == CUDNN_DATA_HALF || data_type == CUDNN_DATA_BFLOAT16) {
+ data_type = CUDNN_DATA_FLOAT;
+ }
+
+ std::vector<int64_t> padding;
+ const auto& padding_before = ctx.Attr<std::vector<int32_t>>("padding_before");
+ copy(padding_before.begin(), padding_before.end(), back_inserter(padding));
+
+ std::vector<int64_t> stride;
+ const auto& strides = ctx.Attr<std::vector<int32_t>>("strides");
+ copy(strides.begin(), strides.end(), back_inserter(stride));
+
+ std::vector<int64_t> dilation;
+ const auto& dilation_rate = ctx.Attr<std::vector<int32_t>>("dilation_rate");
+ copy(dilation_rate.begin(), dilation_rate.end(), back_inserter(dilation));
+
+ uint64_t ndim = stride.size();
+ return cudnn_frontend::ConvDescBuilder()
+ .setDataType(data_type)
+ .setMathMode(CUDNN_CROSS_CORRELATION)
+ .setNDims(ndim)
+ .setStrides(ndim, stride.data())
+ .setPrePadding(ndim, padding.data())
+ .setPostPadding(ndim, padding.data())
+ .setDilation(ndim, dilation.data())
+ .build();
+}
+
+std::vector<cudnn_frontend::GeneratorSource> GetGeneratorSources(
+ const cudnnBackendDescriptorType_t desc) {
+ bool deterministic = Singleton<ResourceDesc, ForSession>::Get()
+ ->resource()
+ .cudnn_conf()
+ .cudnn_conv_use_deterministic_algo_only();
+ bool heuristic = ParseBooleanFromEnv("ONEFLOW_CUDNN_USE_HEURISTIC_MODE_B", false);
+ auto heur_mode = heuristic ? CUDNN_HEUR_MODE_B : CUDNN_HEUR_MODE_A;
+ // Method for engine config generator based on heuristics
+ const auto heurgen_method =
+ [deterministic,
+ heur_mode](cudnn_frontend::OperationGraph& opGraph) -> cudnn_frontend::EngineConfigList {
+ auto heuristics = cudnn_frontend::EngineHeuristicsBuilder()
+ .setOperationGraph(opGraph)
+ .setHeurMode(heur_mode)
+ .build();
+ auto& engine_configs = heuristics.getEngineConfig(heuristics.getEngineConfigCount());
+ cudnn_frontend::EngineConfigList filtered_configs;
+ FilterEngineConfigs(engine_configs, filtered_configs, deterministic);
+ return filtered_configs;
+ };
+ // Method for engine config generator based on fallback list
+ const auto fallback_method =
+ [desc,
+ deterministic](cudnn_frontend::OperationGraph& opGraph) -> cudnn_frontend::EngineConfigList {
+ auto fallback = cudnn_frontend::EngineFallbackListBuilder()
+ .setOperationGraph(opGraph)
+ .setOperation(desc)
+ .build();
+ auto& fallback_list = fallback.getFallbackList();
+ cudnn_frontend::EngineConfigList filtered_configs;
+ FilterEngineConfigs(fallback_list, filtered_configs, deterministic);
+ return filtered_configs;
+ };
+ std::vector<cudnn_frontend::GeneratorSource> sources = {heurgen_method, fallback_method};
+ return sources;
+}
+
+void FilterEngineConfigs(cudnn_frontend::EngineConfigList& from,
+ cudnn_frontend::EngineConfigList& to, bool deterministic) {
+ auto filter = [=](cudnnBackendDescriptor_t c) {
+ if (deterministic) {
+ if (cudnn_frontend::hasNumericalNote<CUDNN_NUMERICAL_NOTE_NONDETERMINISTIC>(c)) {
+ return true;
+ }
+ }
+ if (cudnn_frontend::hasNumericalNote<CUDNN_NUMERICAL_NOTE_DOWN_CONVERT_INPUTS>(c)) {
+ return true;
+ }
+ return false;
+ };
+ cudnn_frontend::filter(from, to, filter);
+}
+
+void TryConfigs(const cudnnHandle_t handle, user_op::Tensor* x, user_op::Tensor* y,
+ user_op::Tensor* w, user_op::Tensor* buf, cudnn_frontend::EngineConfigList& configs,
+ const std::string& tag) {
+ for (auto& config : configs) {
+ try {
+ auto plan = cudnn_frontend::ExecutionPlanBuilder()
+ .setHandle(handle)
+ .setEngineConfig(config, tag)
+ .build();
+ if (PlanErrataException(handle, plan.getTag())) { continue; }
+ RunConvPlan(handle, x, y, w, buf, plan);
+ return;
+ } catch (cudnn_frontend::cudnnException& e) {}
+ }
+}
+
+size_t GetCudnnConvWorkspaceSizeV8(const cudnnHandle_t handle,
+ cudnn_frontend::EngineConfigList& configs,
+ const std::string& tag) {
+ for (auto& config : configs) {
+ try {
+ auto plan = cudnn_frontend::ExecutionPlanBuilder()
+ .setHandle(handle)
+ .setEngineConfig(config, tag)
+ .build();
+ if (PlanErrataException(handle, plan.getTag())) { continue; }
+ if (plan.getWorkspaceSize() > 0L) { return plan.getWorkspaceSize(); }
+ } catch (cudnn_frontend::cudnnException& e) {}
+ }
+ return 1L;
+}
+
+bool PlanErrataException(const cudnnHandle_t handle, const std::string& executionPlanTag) {
+ static nlohmann::json errata_json_handle;
+ static bool has_json = cudnn_frontend::load_from_config(errata_json_handle, "");
+ if (!has_json) {
+ return false;
+ } else {
+ return cudnn_frontend::check_errata(errata_json_handle, executionPlanTag, handle,
+ []() { return true; });
+ }
+}
+
+void RunConvPlan(const cudnnHandle_t handle, user_op::Tensor* x, user_op::Tensor* y,
+ user_op::Tensor* w, user_op::Tensor* buf,
+ const cudnn_frontend::ExecutionPlan& plan) {
+ void* data[] = {x->mut_dptr(), y->mut_dptr(), w->mut_dptr()};
+ int64_t ids[] = {'x', 'y', 'w'};
+ auto variantPack = cudnn_frontend::VariantPackBuilder()
+ .setWorkspacePointer(buf->mut_dptr())
+ .setDataPointers(3, data)
+ .setUids(3, ids)
+ .build();
+ OF_CUDNN_CHECK(cudnnBackendExecute(handle, plan.get_raw_desc(), variantPack.get_raw_desc()));
+}
+
 template<>
 struct CudnnConvAlgorithmSearch<cudnnConvolutionFwdAlgoPerf_t> {
  using perf_t = cudnnConvolutionFwdAlgoPerf_t;