cudnn_utils.h

/*************************************************************************
 * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * See LICENSE for license information.
 ************************************************************************/

#ifndef TRANSFORMER_ENGINE_CUDNN_UTILS_H_
#define TRANSFORMER_ENGINE_CUDNN_UTILS_H_

#ifndef __HIP_PLATFORM_AMD__
#include <cudnn.h>
#include <cudnn_frontend.h>
#include <cudnn_frontend_utils.h>
#include <cudnn_graph.h>
#endif

#include "transformer_engine/transformer_engine.h"
#include "util/handle_manager.h"

namespace transformer_engine {
#ifndef __HIP_PLATFORM_AMD__
namespace detail {

void CreateCuDNNHandle(cudnnHandle_t* handle);

}  // namespace detail

cudnnDataType_t get_cudnn_dtype(const transformer_engine::DType t);

cudnn_frontend::DataType_t get_cudnn_fe_dtype(const transformer_engine::DType t);

using cudnnExecutionPlanManager = detail::HandleManager<cudnnHandle_t, detail::CreateCuDNNHandle>;
#endif
}  // namespace transformer_engine

#endif  //  TRANSFORMER_ENGINE_CUDNN_UTILS_H_