From 905fe365998c4b1288d0814f286e2fe7ee05d2fe Mon Sep 17 00:00:00 2001 From: groenenboomj Date: Wed, 13 Oct 2021 16:18:08 -0500 Subject: [PATCH] Add Conv and ConvTrans to ROCm EP (#9338) Added support for Conv and ConvTrans operators in the ROCm execution provider. Doubles not currently supported. --- .../core/providers/rocm/miopen_common.cc | 29 ++ .../core/providers/rocm/miopen_common.h | 16 + onnxruntime/core/providers/rocm/nn/conv.cc | 357 ++++++++++++++++++ onnxruntime/core/providers/rocm/nn/conv.h | 211 +++++++++++ .../core/providers/rocm/nn/conv_transpose.cc | 207 ++++++++++ .../core/providers/rocm/nn/conv_transpose.h | 29 ++ .../providers/rocm/rocm_execution_provider.cc | 19 +- .../providers/rocm/rocm_execution_provider.h | 12 +- .../rocm/rocm_execution_provider_info.cc | 4 +- .../rocm/rocm_execution_provider_info.h | 1 + onnxruntime/core/providers/rocm/rocm_kernel.h | 10 + 11 files changed, 883 insertions(+), 12 deletions(-) create mode 100644 onnxruntime/core/providers/rocm/nn/conv.cc create mode 100644 onnxruntime/core/providers/rocm/nn/conv.h create mode 100644 onnxruntime/core/providers/rocm/nn/conv_transpose.cc create mode 100644 onnxruntime/core/providers/rocm/nn/conv_transpose.h diff --git a/onnxruntime/core/providers/rocm/miopen_common.cc b/onnxruntime/core/providers/rocm/miopen_common.cc index 06a48e458a..542eb409be 100644 --- a/onnxruntime/core/providers/rocm/miopen_common.cc +++ b/onnxruntime/core/providers/rocm/miopen_common.cc @@ -47,6 +47,35 @@ Status MiopenTensor::Set(const MiopenTensor& x_desc, miopenBatchNormMode_t mode) return Status::OK(); } +MiopenTensorDescriptor::MiopenTensorDescriptor() : desc_(nullptr) { + miopenCreateTensorDescriptor(&desc_); +} + +MiopenTensorDescriptor::~MiopenTensorDescriptor() { + if (desc_ != nullptr) { + miopenCreateTensorDescriptor(&desc_); + desc_ = nullptr; + } +} + +Status MiopenTensorDescriptor::Set(const std::vector& filter_dims, miopenDataType_t data_type) { + if (!desc_) + MIOPEN_RETURN_IF_ERROR(miopenCreateTensorDescriptor(&desc_)); + + int rank = gsl::narrow_cast(filter_dims.size()); + std::vector w_dims(rank); + for (int i = 0; i < rank; i++) { + w_dims[i] = gsl::narrow_cast(filter_dims[i]); + } + + MIOPEN_RETURN_IF_ERROR(miopenSetTensorDescriptor(desc_, + data_type, + rank, + w_dims.data(), + nullptr)); + return Status::OK(); +} + template miopenDataType_t MiopenTensor::GetDataType() { ORT_THROW("miopen engine currently supports only single/half/int32/int8 precision data types."); diff --git a/onnxruntime/core/providers/rocm/miopen_common.h b/onnxruntime/core/providers/rocm/miopen_common.h index 9a1382d6eb..edb90ba453 100644 --- a/onnxruntime/core/providers/rocm/miopen_common.h +++ b/onnxruntime/core/providers/rocm/miopen_common.h @@ -14,6 +14,8 @@ const double MIOPEN_BN_MIN_EPSILON = 1e-5; namespace onnxruntime { namespace rocm { +#define MIOPEN_CONVOLUTION_FWD_ALGO_COUNT 6 + class MiopenTensor final { public: MiopenTensor(); @@ -34,6 +36,20 @@ class MiopenTensor final { miopenTensorDescriptor_t tensor_; }; +class MiopenTensorDescriptor final { + public: + MiopenTensorDescriptor(); + ~MiopenTensorDescriptor(); + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(MiopenTensorDescriptor); + + Status Set(const std::vector& filter_dims, miopenDataType_t data_typ); + + operator miopenTensorDescriptor_t() const { return desc_; } + + private: + miopenTensorDescriptor_t desc_; +}; + template struct Consts { static const ElemType Zero; diff --git a/onnxruntime/core/providers/rocm/nn/conv.cc b/onnxruntime/core/providers/rocm/nn/conv.cc new file mode 100644 index 0000000000..22567916a0 --- /dev/null +++ b/onnxruntime/core/providers/rocm/nn/conv.cc @@ -0,0 +1,357 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/rocm/rocm_common.h" +#include "core/providers/rocm/nn/conv.h" +#include "core/providers/rocm/shared_inc/fpgeneric.h" +#include "core/providers/rocm/tensor/slice.h" + +namespace onnxruntime { +namespace rocm { + +// Op Set 11 for Conv only update document to clearify default dilations and strides value. +// which are already convered by op set 11 cpu versoin, so simply add declaration. +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \ + Conv, \ + kOnnxDomain, \ + 1, 10, \ + T, \ + kRocmExecutionProvider, \ + (*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + Conv); \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + Conv, \ + kOnnxDomain, \ + 11, \ + T, \ + kRocmExecutionProvider, \ + (*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + Conv); + +REGISTER_KERNEL_TYPED(float) +// not yet supported in MIOpen +//REGISTER_KERNEL_TYPED(double) +REGISTER_KERNEL_TYPED(MLFloat16) + +template +const miopenConvFwdAlgorithm_t Conv::kAllAlgos[] = { + miopenConvolutionFwdAlgoGEMM, + miopenConvolutionFwdAlgoDirect, + miopenConvolutionFwdAlgoFFT, + miopenConvolutionFwdAlgoWinograd, + miopenConvolutionFwdAlgoImplicitGEMM +}; + miopenStatus_t GetWorkspaceSize(const MiopenConvState& s, miopenConvFwdAlgorithm_t algo, + size_t* sz) { + + return miopenConvolutionForwardGetWorkSpaceSize(s.handle, s.w_desc, s.x_tensor, s.conv_desc, s.y_tensor, sz); +} + +size_t GetMaxWorkspaceSize(const MiopenConvState& s, + const miopenConvFwdAlgorithm_t* algo, int n_algo) { + // TODO: get maximum available size from memory areana + size_t free, total; + HIP_CALL_THROW(hipMemGetInfo(&free, &total)); + // Assuming 10% of fragmentation + free = static_cast(static_cast(free) * 0.9); + size_t max_ws_size = 0; + for (int i = 0; i < n_algo; i++) { + miopenStatus_t err; + size_t sz; + err = GetWorkspaceSize(s, algo[i], &sz); + if (miopenStatusSuccess != err || sz == 0 || sz < max_ws_size || sz > free) continue; + max_ws_size = sz; + } + return max_ws_size; +} + +Status SliceOutUnwantedOutputSection(hipStream_t stream, + const void* input_data, const std::vector& input_dims, + void* output_data, + const std::vector& output_dims, + std::vector starts, + const std::vector& ends, + const std::vector& axes, + size_t element_size) { + SliceOp::PrepareForComputeMetadata compute_metadata(input_dims); + + SliceBase::PrepareForCompute(starts, ends, axes, compute_metadata); + + // As a sanity check, ensure that the slice operator's output shape matches with the expected output shape + ORT_ENFORCE(compute_metadata.output_dims_ == output_dims); + + return SliceRocm::Impl(stream, input_data, input_dims, output_data, compute_metadata, element_size); +} + +template +Status Conv::UpdateState(OpKernelContext* context, bool bias_expected) const { + //set X + const Tensor* X = context->Input(0); + const TensorShape& x_shape = X->Shape(); + const auto& x_dims = x_shape.GetDims(); + s_.x_data = reinterpret_cast(X->template Data()); + s_.element_size = X->DataType()->Size(); + //set W + const Tensor* W = context->Input(1); + const TensorShape& w_shape = W->Shape(); + std::vector w_dims = w_shape.GetDims(); + s_.w_data = reinterpret_cast(W->template Data()); + //set B + if (context->InputCount() >= 3) { + const Tensor* B = context->Input(2); + s_.b_data = reinterpret_cast(B->template Data()); + } else { + s_.b_data = nullptr; + } + //set Z + if (context->InputCount() >= 4) { + const Tensor* Z = context->Input(3); + ORT_RETURN_IF_ERROR(s_.z_tensor.Set(Z->Shape().GetDims(), MiopenTensor::GetDataType())); + s_.z_data = reinterpret_cast(Z->template Data()); + } else { + s_.z_data = nullptr; + } + bool input_dims_changed = (s_.last_x_dims != x_dims); + bool w_dims_changed = (s_.last_w_dims != w_dims); + if (input_dims_changed || w_dims_changed) { + if (input_dims_changed) + s_.last_x_dims = x_dims; + + if (w_dims_changed) { + s_.last_w_dims = w_dims; + s_.cached_benchmark_fwd_results.clear(); + } + + const int64_t N = X->Shape()[0]; + const int64_t M = W->Shape()[0]; + + ORT_RETURN_IF_ERROR(conv_attrs_.ValidateInputShape(X, W)); + + std::vector kernel_shape; + ORT_RETURN_IF_ERROR(conv_attrs_.ComputeKernelShape(W->Shape(), kernel_shape)); + auto rank = kernel_shape.size(); + std::vector pads(conv_attrs_.pads); + if (pads.empty()) { + pads.resize(rank * 2, 0); + } + std::vector dilations(conv_attrs_.dilations); + if (dilations.empty()) { + dilations.resize(rank, 1); + } + std::vector strides(conv_attrs_.strides); + if (strides.empty()) { + strides.resize(rank, 1); + } + + std::vector y_dims; + y_dims.reserve(2 + rank); // rank indicates number of feature dimensions - so add 2 to account for 'N' and 'C' + y_dims.insert(y_dims.begin(), {N, M}); + + std::vector y_dims_with_adjusted_pads; + y_dims_with_adjusted_pads.reserve(2 + rank); // rank indicates number of feature dimensions - so add 2 to account for 'N' and 'C' + y_dims_with_adjusted_pads.insert(y_dims_with_adjusted_pads.begin(), {N, M}); + + bool post_slicing_required = false; + std::vector slice_starts; + slice_starts.reserve(rank); + + std::vector slice_ends; + slice_ends.reserve(rank); + + std::vector slice_axes; + slice_axes.reserve(rank); + + ORT_RETURN_IF_ERROR(conv_attrs_.InferOutputShapeWithAdjustedPads(x_shape.Slice(2), kernel_shape, + strides, dilations, pads, y_dims, y_dims_with_adjusted_pads, + post_slicing_required, slice_starts, slice_ends, slice_axes)); + ORT_ENFORCE(y_dims.size() == y_dims_with_adjusted_pads.size()); + s_.y_dims = y_dims; + s_.y_dims_with_adjusted_pads = y_dims_with_adjusted_pads; + s_.post_slicing_required = post_slicing_required; + s_.slice_starts = slice_starts; + s_.slice_ends = slice_ends; + s_.slice_axes = slice_axes; + + s_.Y = context->Output(0, TensorShape(s_.y_dims)); + if (s_.Y->Shape().Size() == 0) { + return Status::OK(); + } + if (post_slicing_required) { + // Post slicing needed. Create and fill in the Conv results in an intermediate buffer. + s_.memory_for_miopen_conv_results = GetScratchBuffer(TensorShape(y_dims_with_adjusted_pads).Size() * s_.element_size); + s_.y_data = reinterpret_cast(s_.memory_for_miopen_conv_results.get()); + } else { + // No post slicing needed. Fill the output tensor's buffer directly. + s_.y_data = reinterpret_cast(s_.Y->template MutableData()); + } + + std::vector x_dims_miopen = x_dims; + std::vector y_dims_miopen = !post_slicing_required ? y_dims : y_dims_with_adjusted_pads; + if (rank < 2) { + // TODO: Remove asym padding correction. + x_dims_miopen.push_back(1); + y_dims_miopen.push_back(1); + w_dims.push_back(1); + pads.insert(pads.begin() + rank, 0); + pads.insert(pads.end(), 0); + kernel_shape.push_back(1); + strides.push_back(1); + dilations.push_back(1); + } + + if (w_dims_changed) { + ORT_RETURN_IF_ERROR(s_.w_desc.Set(w_dims, MiopenTensor::GetDataType())); + } + ORT_RETURN_IF_ERROR(s_.x_tensor.Set(x_dims_miopen, MiopenTensor::GetDataType())); + ORT_RETURN_IF_ERROR(s_.y_tensor.Set(y_dims_miopen, MiopenTensor::GetDataType())); + ORT_RETURN_IF_ERROR(s_.conv_desc.Set(kernel_shape.size(), pads, strides, dilations, + gsl::narrow_cast(conv_attrs_.group), + miopenConvolution, MiopenTensor::GetDataType())); + + if (context->InputCount() >= 3) { + const Tensor* B = context->Input(2); + const auto& b_shape = B->Shape(); + ORT_RETURN_IF_NOT(b_shape.NumDimensions() == 1, "bias should be 1D"); + std::vector b_dims(2 + kernel_shape.size(), 1); + b_dims[1] = b_shape[0]; + ORT_RETURN_IF_ERROR(s_.b_tensor.Set(b_dims, MiopenTensor::GetDataType())); + } else if (bias_expected) { + std::vector b_dims(2 + kernel_shape.size(), 1); + b_dims[1] = w_dims[0]; + auto malloc_size = b_dims[1] * sizeof(HipT); + ORT_RETURN_IF_ERROR(s_.b_tensor.Set(b_dims, MiopenTensor::GetDataType())); + if (s_.b_zero) { + HIP_CALL_THROW(hipFree(s_.b_zero)); + s_.b_zero = nullptr; + } + HIP_CALL_THROW(hipMalloc(&s_.b_zero, malloc_size)); + HIP_CALL_THROW(hipMemsetAsync(s_.b_zero, 0, malloc_size, Stream())); + } + + if (!s_.cached_benchmark_fwd_results.contains(x_dims_miopen)) { + + miopenConvAlgoPerf_t perf; + int algo_count = 1; + const ROCMExecutionProvider* rocm_ep = static_cast(this->Info().GetExecutionProvider()); + static constexpr int num_algos = MIOPEN_CONVOLUTION_FWD_ALGO_COUNT; + size_t max_ws_size = rocm_ep->GetMiopenConvUseMaxWorkspace() ? GetMaxWorkspaceSize(s_, kAllAlgos, num_algos) + : AlgoSearchWorkspaceSize; + IAllocatorUniquePtr algo_search_workspace = GetTransientScratchBuffer(max_ws_size); + MIOPEN_RETURN_IF_ERROR(miopenFindConvolutionForwardAlgorithm( + s_.handle, + s_.x_tensor, + s_.x_data, + s_.w_desc, + s_.w_data, + s_.conv_desc, + s_.y_tensor, + s_.y_data, + 1, // requestedAlgoCount + &algo_count, // returnedAlgoCount + &perf, + algo_search_workspace.get(), + max_ws_size, + false)); // Do not do exhaustive algo search. + s_.cached_benchmark_fwd_results.insert(x_dims_miopen, {perf.fwd_algo, perf.memory}); + } + const auto& perf = s_.cached_benchmark_fwd_results.at(x_dims_miopen); + s_.fwd_algo = perf.fwd_algo; + s_.workspace_bytes = perf.memory; + } else { + //set Y + s_.Y = context->Output(0, TensorShape(s_.y_dims)); + if (s_.Y->Shape().Size() == 0) { + return Status::OK(); + } + if (s_.post_slicing_required) { + s_.memory_for_miopen_conv_results = GetScratchBuffer(TensorShape(s_.y_dims_with_adjusted_pads).Size() * s_.element_size); + s_.y_data = reinterpret_cast(s_.memory_for_miopen_conv_results.get()); + } else { + s_.y_data = reinterpret_cast(s_.Y->template MutableData()); + } + } + return Status::OK(); +} + +template +Status Conv::ComputeInternal(OpKernelContext* context) const { + std::lock_guard lock(s_.mutex); + ORT_RETURN_IF_ERROR(UpdateState(context)); + if (s_.Y->Shape().Size() == 0) { + return Status::OK(); + } + const auto alpha = Consts::One; + const auto beta = Consts::Zero; + IAllocatorUniquePtr workspace = GetWorkSpace(); + MIOPEN_RETURN_IF_ERROR(miopenConvolutionForward(s_.handle, + &alpha, + s_.x_tensor, + s_.x_data, + s_.w_desc, + s_.w_data, + s_.conv_desc, + s_.fwd_algo, + &beta, + s_.y_tensor, + s_.y_data, + workspace.get(), + s_.workspace_bytes)); + if (nullptr != s_.b_data) { + MIOPEN_RETURN_IF_ERROR(miopenConvolutionForwardBias(s_.handle, &alpha, s_.b_tensor, s_.b_data, + &beta, s_.y_tensor, s_.y_data)); + } + // To deal with asymmetric padding, we may have over-padded on one or both sides of the spatial dimensions + // This may have lead to extra results that are unnecessary and hence we slice that off here + if (s_.post_slicing_required) { + SliceOutUnwantedOutputSection(Stream(), s_.y_data, s_.y_dims_with_adjusted_pads, s_.Y->MutableDataRaw(), + s_.y_dims, s_.slice_starts, s_.slice_ends, s_.slice_axes, s_.element_size); + } + return Status::OK(); +} + +MiopenConvolutionDescriptor::MiopenConvolutionDescriptor() : desc_(nullptr) { +} + +MiopenConvolutionDescriptor::~MiopenConvolutionDescriptor() { + if (desc_ != nullptr) { + miopenDestroyConvolutionDescriptor(desc_); + desc_ = nullptr; + } +} + +Status MiopenConvolutionDescriptor::Set( + size_t rank, + const std::vector& pads, + const std::vector& strides, + const std::vector& dilations, + int groups, + miopenConvolutionMode_t mode, + miopenDataType_t data_type) { + if (!desc_) + MIOPEN_RETURN_IF_ERROR(miopenCreateConvolutionDescriptor(&desc_)); + + std::vector pad_dims(rank); + std::vector stride_dims(rank); + std::vector dilation_dims(rank); + for (size_t i = 0; i < rank; i++) { + pad_dims[i] = gsl::narrow_cast(pads[i]); + stride_dims[i] = gsl::narrow_cast(strides[i]); + dilation_dims[i] = gsl::narrow_cast(dilations[i]); + } + + MIOPEN_RETURN_IF_ERROR(miopenInitConvolutionNdDescriptor( + desc_, + gsl::narrow_cast(rank), + pad_dims.data(), + stride_dims.data(), + dilation_dims.data(), + mode)); + + MIOPEN_RETURN_IF_ERROR(miopenSetConvolutionGroupCount(desc_, groups)); + + return Status::OK(); +} + +} // namespace rocm +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/rocm/nn/conv.h b/onnxruntime/core/providers/rocm/nn/conv.h new file mode 100644 index 0000000000..1087cd05d3 --- /dev/null +++ b/onnxruntime/core/providers/rocm/nn/conv.h @@ -0,0 +1,211 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/platform/ort_mutex.h" +#include "core/providers/rocm/rocm_kernel.h" +#include "core/providers/rocm/miopen_common.h" +#include "core/providers/cpu/nn/conv_attributes.h" +#include + +namespace onnxruntime { +namespace rocm { + +class MiopenConvolutionDescriptor final { + public: + MiopenConvolutionDescriptor(); + ~MiopenConvolutionDescriptor(); + + Status Set(size_t rank, + const std::vector& pads, + const std::vector& strides, + const std::vector& dilations, + int groups, + miopenConvolutionMode_t mode, + miopenDataType_t data_type); + + operator miopenConvolutionDescriptor_t() const { return desc_; } + + private: + miopenConvolutionDescriptor_t desc_; +}; + +template +struct vector_hash { + std::size_t operator()(const std::vector& values) const { + std::size_t seed = values.size(); + for (auto& val : values) + seed ^= std::hash()(val) + 0x9e3779b9 + (seed << 6) + (seed >> 2); + return seed; + } +}; + +template , + typename KeyEqual = std::equal_to, + typename ListAllocator = std::allocator> +class lru_unordered_map { + public: + lru_unordered_map(size_t max_size) : max_size_(max_size) {} + + void insert(const Key& key, const T& value) { + auto it = items_.find(key); + if (it != items_.end()) { + it->second.value = value; + move_to_front(it->second.lru_iterator); + return; + } + + while (size() + 1 > max_size_) { + items_.erase(lru_list_.back()); + lru_list_.pop_back(); + } + + lru_list_.emplace_front(key); + items_.emplace(key, value_type{value, lru_list_.begin()}); + } + + T& at(const Key& key) { + auto it = items_.find(key); + if (it == items_.end()) { + throw std::out_of_range("There is no such key in cache"); + } + move_to_front(it->second.lru_iterator); + return it->second.value; + } + + bool contains(const Key& key) const { + return items_.find(key) != items_.end(); + } + + size_t size() const { + return items_.size(); + } + + void clear() { + items_.clear(); + lru_list_.clear(); + } + + private: + using list_type = std::list; + using iterator_type = typename list_type::iterator; + struct value_type { + T value; + iterator_type lru_iterator; + }; + using MapAllocator = std::allocator>; + + void move_to_front(iterator_type it) { + lru_list_.splice(lru_list_.begin(), lru_list_, it); + } + + size_t max_size_; + std::unordered_map items_; + list_type lru_list_; +}; + +// cached miopen descriptors +constexpr size_t MAX_CACHED_ALGO_PERF_RESULTS = 10000; + +template +struct MiopenConvState { + miopenHandle_t handle; + + // if x/w dims changed, update algo and miopenTensors + std::vector last_x_dims; + std::vector last_w_dims; + + // these would be recomputed if x/w dims change + std::vector y_dims; + std::vector y_dims_with_adjusted_pads; + size_t workspace_bytes; + decltype(AlgoPerfType().bwd_data_algo) bwd_data_algo; + decltype(AlgoPerfType().fwd_algo) fwd_algo; + MiopenTensor x_tensor; + const void* x_data = nullptr; + size_t element_size = 0; + MiopenTensorDescriptor w_desc; + const void* w_data = nullptr; + MiopenTensor b_tensor; + const void* b_data = nullptr; + void* b_zero = nullptr; + MiopenTensor y_tensor; + Tensor* Y = nullptr; + void* y_data = nullptr; + MiopenTensor z_tensor; + const void* z_data = nullptr; + MiopenConvolutionDescriptor conv_desc; + + struct PerfFwdResultParams { + decltype(AlgoPerfType().fwd_algo) fwd_algo; + decltype(AlgoPerfType().memory) memory; + }; + + struct PerfBwdResultParams { + decltype(AlgoPerfType().bwd_data_algo) bwd_data_algo; + decltype(AlgoPerfType().memory) memory; + }; + + lru_unordered_map, PerfFwdResultParams, vector_hash> cached_benchmark_fwd_results{MAX_CACHED_ALGO_PERF_RESULTS}; + lru_unordered_map, PerfBwdResultParams, vector_hash> cached_benchmark_bwd_results{MAX_CACHED_ALGO_PERF_RESULTS}; + + // Some properties needed to support asymmetric padded Conv nodes + bool post_slicing_required; + std::vector slice_starts; + std::vector slice_ends; + std::vector slice_axes; + + // note that conv objects are shared between execution frames, and a lock is needed to avoid multi-thread racing + OrtMutex mutex; + IAllocatorUniquePtr memory_for_miopen_conv_results; + + ~MiopenConvState() { + if (b_zero) { + HIP_CALL_THROW(hipFree(b_zero)); + b_zero = nullptr; + } + } +}; + +enum : size_t { + AlgoSearchWorkspaceSize = 32 * 1024 * 1024, +}; + +template +class Conv : public RocmKernel { + public: + using HipT = typename ToHipType::MappedType; + + Conv(const OpKernelInfo& info) : RocmKernel(info), conv_attrs_(info) { + auto pads_size = conv_attrs_.pads.size(); + ORT_ENFORCE(pads_size % 2 == 0); + s_.handle = MiopenHandle(); + } + + Status ComputeInternal(OpKernelContext* context) const override; + + protected: + inline IAllocatorUniquePtr GetWorkSpace() const { + return GetScratchBuffer(s_.workspace_bytes); + } + + Status UpdateState(OpKernelContext* context, bool bias_expected = false) const; + ConvAttributes conv_attrs_; + mutable MiopenConvState s_; + constexpr static auto kDefaultConvAlgo = miopenConvolutionFwdAlgoGEMM; + static const miopenConvFwdAlgorithm_t kAllAlgos[]; +}; + +Status SliceOutUnwantedOutputSection(hipStream_t stream, + const void* input_data, + const std::vector& input_dims, + void* output_data, + const std::vector& output_dims, + std::vector starts, + const std::vector& ends, + const std::vector& axes, + size_t element_size); +} // namespace rocm +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/rocm/nn/conv_transpose.cc b/onnxruntime/core/providers/rocm/nn/conv_transpose.cc new file mode 100644 index 0000000000..9b9b0f189b --- /dev/null +++ b/onnxruntime/core/providers/rocm/nn/conv_transpose.cc @@ -0,0 +1,207 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "conv_transpose.h" + +namespace onnxruntime { +namespace rocm { + +// Op Set 11 for ConvTranspose only update document to clearify default dilations and strides value. +// which are already covered by op set 11 cpu version, so simply add declaration. +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \ + ConvTranspose, \ + kOnnxDomain, \ + 1, 10, \ + T, \ + kRocmExecutionProvider, \ + (*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + ConvTranspose); \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + ConvTranspose, \ + kOnnxDomain, \ + 11, \ + T, \ + kRocmExecutionProvider, \ + (*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + ConvTranspose); + +REGISTER_KERNEL_TYPED(float) +// not yet supported in MIOpen +//REGISTER_KERNEL_TYPED(double) +REGISTER_KERNEL_TYPED(MLFloat16) + +template +Status ConvTranspose::ComputeInternal(OpKernelContext* context) const { + return DoConvTranspose(context, false); +} + +template +Status ConvTranspose::DoConvTranspose(OpKernelContext* context, bool dynamic_padding) const { + typedef typename ToHipType::MappedType HipT; + + const Tensor* X = context->Input(0); + const TensorShape& x_shape = X->Shape(); + auto x_dims = x_shape.GetDims(); + auto x_data = reinterpret_cast(X->template Data()); + + auto x_dimensions = X->Shape().NumDimensions(); + if (x_dimensions < 3 || x_dimensions > 5) { + // TODO: the error message should tell which operator raises it. + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input X must be 3-, 4- or 5-dimensional.", + " X: ", X->Shape().ToString().c_str()); + } + const Tensor* W = context->Input(1); + const TensorShape& w_shape = W->Shape(); + std::vector w_dims = w_shape.GetDims(); + auto w_data = reinterpret_cast(W->template Data()); + + size_t num_inputs = OpKernel::Node().InputDefs().size(); + bool has_bias = dynamic_padding ? num_inputs == 4 : num_inputs == 3; + + HipT* y_data = nullptr; + if (x_dimensions == 3) { + x_dims.insert(x_dims.begin() + 2, 1); + w_dims.insert(w_dims.begin() + 2, 1); + } + + { + std::lock_guard lock(s_.mutex); + // TODO: add a global cache if need to handle cases for multiple frames running simultaneously with different batch_size + bool input_dims_changed = (s_.last_x_dims != x_dims); + bool w_dims_changed = (s_.last_w_dims != w_dims); + if (input_dims_changed || w_dims_changed) { + if (input_dims_changed) + s_.last_x_dims = x_dims; + + if (w_dims_changed) { + s_.last_w_dims = w_dims; + s_.cached_benchmark_bwd_results.clear(); + } + + ConvTransposeAttributes::Prepare p; + ORT_RETURN_IF_ERROR(conv_transpose_attrs_.PrepareForCompute(context, has_bias, p, dynamic_padding)); + + auto y_dims = p.Y->Shape().GetDims(); + if (x_dimensions == 3) { + y_dims.insert(y_dims.begin() + 2, 1); + p.kernel_shape.insert(p.kernel_shape.begin(), 1); + p.pads.insert(p.pads.begin(), 0); + p.pads.insert(p.pads.begin() + 2, 0); + p.strides.insert(p.strides.begin(), 1); + p.dilations.insert(p.dilations.begin(), 1); + } + s_.y_dims = y_dims; + + if (w_dims_changed) + { + ORT_RETURN_IF_ERROR(s_.w_desc.Set(w_dims, MiopenTensor::GetDataType())); + } + + // Special case when there is a dim value of 0 in the shape. + // Return only after we have cached the following for subsequent runs : + // 1) `w_dims` in the `w_desc` + // 2) `y_dims` in s_.y_dims + if (p.Y->Shape().Size() == 0) { + return Status::OK(); + } + + ORT_RETURN_IF_ERROR(s_.x_tensor.Set(x_dims, MiopenTensor::GetDataType())); + ORT_RETURN_IF_ERROR(s_.y_tensor.Set(y_dims, MiopenTensor::GetDataType())); + + miopenConvolutionMode_t mode = miopenConvolution; + ORT_RETURN_IF_ERROR(s_.conv_desc.Set(p.kernel_shape.size(), p.pads, p.strides, p.dilations, + gsl::narrow_cast(conv_transpose_attrs_.group), + mode, MiopenTensor::GetDataType())); + + if (has_bias) { + const auto& b_shape = p.B->Shape(); + ORT_RETURN_IF_NOT(b_shape.NumDimensions() == 1, "bias should be 1D"); + std::vector b_dims(2 + p.kernel_shape.size()); + b_dims[0] = 1; // N + b_dims[1] = b_shape[0]; // C + for (size_t i = 0; i < p.kernel_shape.size(); i++) + b_dims[2 + i] = 1; + + ORT_RETURN_IF_ERROR(s_.b_tensor.Set(b_dims, MiopenTensor::GetDataType())); + } + + y_data = reinterpret_cast(p.Y->template MutableData()); + + if (!s_.cached_benchmark_bwd_results.contains(x_dims)) { + IAllocatorUniquePtr algo_search_workspace = GetScratchBuffer(AlgoSearchWorkspaceSize); + + miopenConvAlgoPerf_t perf; + int algo_count = 1; + MIOPEN_RETURN_IF_ERROR(miopenFindConvolutionBackwardDataAlgorithm( + MiopenHandle(), + s_.x_tensor, + x_data, + s_.w_desc, + w_data, + s_.conv_desc, + s_.y_tensor, + y_data, + 1, + &algo_count, + &perf, + algo_search_workspace.get(), + AlgoSearchWorkspaceSize, + false)); + s_.cached_benchmark_bwd_results.insert(x_dims, {perf.bwd_data_algo, perf.memory}); + } + + const auto& perf = s_.cached_benchmark_bwd_results.at(x_dims); + s_.bwd_data_algo = perf.bwd_data_algo; + s_.workspace_bytes = perf.memory; + } + + // The following block will be executed in case there has been no change in the shapes of the + // input and the filter compared to the previous run + if (!y_data) { + auto y_dims = s_.y_dims; + if (x_dimensions == 3) { + y_dims.erase(y_dims.begin() + 2); + } + Tensor* Y = context->Output(0, TensorShape(y_dims)); + y_data = reinterpret_cast(Y->template MutableData()); + + // Bail out early if one of the output dimensions is zero. + if (Y->Shape().Size() == 0) { + return Status::OK(); + } + } + + const auto alpha = Consts::One; + const auto beta = Consts::Zero; + + IAllocatorUniquePtr workspace = GetScratchBuffer(s_.workspace_bytes); + + MIOPEN_RETURN_IF_ERROR( + miopenConvolutionBackwardData( + MiopenHandle(), + &alpha, + s_.x_tensor, + x_data, + s_.w_desc, + w_data, + s_.conv_desc, + s_.bwd_data_algo, + &beta, + s_.y_tensor, + y_data, + workspace.get(), + s_.workspace_bytes)); + + if (has_bias) { + const Tensor* B = dynamic_padding ? context->Input(3) : context->Input(2); + auto b_data = reinterpret_cast(B->template Data()); + MIOPEN_RETURN_IF_ERROR((miopenConvolutionForwardBias(MiopenHandle(), &alpha, s_.b_tensor, b_data, &beta, s_.y_tensor, y_data))); + } + } + + return Status::OK(); +} + +} // namespace rocm +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/rocm/nn/conv_transpose.h b/onnxruntime/core/providers/rocm/nn/conv_transpose.h new file mode 100644 index 0000000000..b4b80aeec9 --- /dev/null +++ b/onnxruntime/core/providers/rocm/nn/conv_transpose.h @@ -0,0 +1,29 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/rocm/rocm_common.h" +#include "core/providers/rocm/rocm_kernel.h" +#include "core/providers/rocm/miopen_common.h" +#include "core/providers/rocm/nn/conv.h" +#include "core/providers/cpu/nn/conv_transpose_attributes.h" + +namespace onnxruntime { +namespace rocm { + +template +class ConvTranspose : public RocmKernel { + public: + ConvTranspose(const OpKernelInfo& info) : RocmKernel(info), conv_transpose_attrs_(info){}; + Status ComputeInternal(OpKernelContext* context) const override; + Status DoConvTranspose(OpKernelContext* context, bool dynamic_padding) const; + + private: + ConvTransposeAttributes conv_transpose_attrs_; + + mutable MiopenConvState s_; +}; + +} // namespace rocm +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/rocm/rocm_execution_provider.cc b/onnxruntime/core/providers/rocm/rocm_execution_provider.cc index 11c0a70832..93a0eb7383 100644 --- a/onnxruntime/core/providers/rocm/rocm_execution_provider.cc +++ b/onnxruntime/core/providers/rocm/rocm_execution_provider.cc @@ -1229,12 +1229,12 @@ static Status RegisterRocmKernels(KernelRegistry& kernel_registry) { // BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, // BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo, @@ -1502,12 +1502,12 @@ static Status RegisterRocmKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, // BuildKernelCreateInfo, BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, // BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo, @@ -1869,8 +1869,7 @@ ROCMExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph, bool force_inside = false; // for some compute heavy ops, we'll force it to run inside ROCM if ("LSTM" == node.OpType() || "RNN" == node.OpType() || - "GRU" == node.OpType() || - "ConvTranspose" == node.OpType()) { + "GRU" == node.OpType()) { not_supported = true; force_inside = !not_supported; } else if ("Cast" == node.OpType()) { diff --git a/onnxruntime/core/providers/rocm/rocm_execution_provider.h b/onnxruntime/core/providers/rocm/rocm_execution_provider.h index 40e606dcc8..901a51427a 100644 --- a/onnxruntime/core/providers/rocm/rocm_execution_provider.h +++ b/onnxruntime/core/providers/rocm/rocm_execution_provider.h @@ -85,10 +85,20 @@ class ROCMExecutionProvider : public IExecutionProvider { int GetDeviceId() const override { return info_.device_id; } const hipDeviceProp_t& GetDeviceProp() const { return device_prop_; }; + bool GetMiopenConvUseMaxWorkspace() const { return info_.miopen_conv_use_max_workspace; } + ProviderOptions GetProviderOptions() const override { return ROCMExecutionProviderInfo::ToProviderOptions(info_); } - + + template + IAllocatorUniquePtr GetTransientScratchBuffer(size_t count_or_bytes) const { + if (count_or_bytes == 0) + return nullptr; + + return IAllocator::MakeUniquePtr(GetAllocator(info_.device_id, OrtMemTypeDefault), count_or_bytes, true); + } + void RegisterAllocator(std::shared_ptr allocator_manager) override; static AllocatorPtr CreateRocmAllocator(OrtDevice::DeviceId device_id, size_t rocm_mem_limit, ArenaExtendStrategy arena_extend_strategy, ROCMExecutionProviderExternalAllocatorInfo external_alloc_info); diff --git a/onnxruntime/core/providers/rocm/rocm_execution_provider_info.cc b/onnxruntime/core/providers/rocm/rocm_execution_provider_info.cc index 705daf2356..96c4eac9c4 100644 --- a/onnxruntime/core/providers/rocm/rocm_execution_provider_info.cc +++ b/onnxruntime/core/providers/rocm/rocm_execution_provider_info.cc @@ -17,6 +17,7 @@ constexpr const char* kConvExhaustiveSearch = "conv_exhaustive_search"; constexpr const char* kGpuExternalAlloc = "gpu_external_alloc"; constexpr const char* kGpuExternalFree = "gpu_external_free"; constexpr const char* kGpuExternalEmptyCache = "gpu_external_empty_cache"; +constexpr const char* kMiopenConvUseMaxWorkspace = "miopen_conv_use_max_workspace"; } // namespace provider_option_names } // namespace rocm @@ -71,9 +72,10 @@ ROCMExecutionProviderInfo ROCMExecutionProviderInfo::FromProviderOptions(const P "Invalid device ID: ", info.device_id, ", must be between 0 (inclusive) and ", num_devices, " (exclusive)."); return Status::OK(); - }) + }) .AddAssignmentToReference(rocm::provider_option_names::kMemLimit, info.gpu_mem_limit) .AddAssignmentToReference(rocm::provider_option_names::kConvExhaustiveSearch, info.miopen_conv_exhaustive_search) + .AddAssignmentToReference(rocm::provider_option_names::kMiopenConvUseMaxWorkspace, info.miopen_conv_use_max_workspace) .AddAssignmentToEnumReference( rocm::provider_option_names::kArenaExtendStrategy, arena_extend_strategy_mapping, info.arena_extend_strategy) diff --git a/onnxruntime/core/providers/rocm/rocm_execution_provider_info.h b/onnxruntime/core/providers/rocm/rocm_execution_provider_info.h index be9d233eaf..e88a8b2473 100644 --- a/onnxruntime/core/providers/rocm/rocm_execution_provider_info.h +++ b/onnxruntime/core/providers/rocm/rocm_execution_provider_info.h @@ -42,6 +42,7 @@ struct ROCMExecutionProviderInfo { bool has_user_compute_stream{false}; void* user_compute_stream{nullptr}; ROCMExecutionProviderExternalAllocatorInfo external_allocator_info{}; + bool miopen_conv_use_max_workspace{false}; static ROCMExecutionProviderInfo FromProviderOptions(const ProviderOptions& options); static ProviderOptions ToProviderOptions(const ROCMExecutionProviderInfo& info); diff --git a/onnxruntime/core/providers/rocm/rocm_kernel.h b/onnxruntime/core/providers/rocm/rocm_kernel.h index 6c63ded1dd..412db4175f 100644 --- a/onnxruntime/core/providers/rocm/rocm_kernel.h +++ b/onnxruntime/core/providers/rocm/rocm_kernel.h @@ -52,6 +52,16 @@ class RocmKernel : public OpKernel { return provider_->GetScratchBuffer(count_or_bytes); } + // Different from GetScratchBuffer which use IAllocator::Alloc() to allocate memory, + // this GetTransientScratchBuffer will call IAllocator::Reserve() to allocate memory. + // IAllocator::Reserve() optionally implement some allocation logic that by-passes any arena-based + // logic (or similar for different allocator) that may be housed in the Alloc() implementation. + template + inline IAllocatorUniquePtr GetTransientScratchBuffer(size_t count_or_bytes) const { + return provider_->GetTransientScratchBuffer(count_or_bytes); + } + + inline void AddDeferredReleaseCPUPtr(void* p) const { provider_->AddDeferredReleaseCPUPtr(p); }