Add Conv and ConvTrans to ROCm EP (#9338)

Added support for Conv and ConvTrans operators
in the ROCm execution provider. Doubles not currently
supported.

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<int64_t>& filter_dims, miopenDataType_t data_type) {
+  if (!desc_)
+    MIOPEN_RETURN_IF_ERROR(miopenCreateTensorDescriptor(&desc_));
+
+  int rank = gsl::narrow_cast<int>(filter_dims.size());
+  std::vector<int> w_dims(rank);
+  for (int i = 0; i < rank; i++) {
+    w_dims[i] = gsl::narrow_cast<int>(filter_dims[i]);
+  }
+
+  MIOPEN_RETURN_IF_ERROR(miopenSetTensorDescriptor(desc_,
+                                                   data_type,
+                                                   rank,
+                                                   w_dims.data(),
+						   nullptr));
+  return Status::OK();
+}
+
 template <typename ElemType>
 miopenDataType_t MiopenTensor::GetDataType() {
   ORT_THROW("miopen engine currently supports only single/half/int32/int8 precision data types.");

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<int64_t>& filter_dims, miopenDataType_t data_typ);
+
+  operator miopenTensorDescriptor_t() const { return desc_; }
+
+ private:
+  miopenTensorDescriptor_t desc_;
+};
+
 template <typename ElemType>
 struct Consts {
   static const ElemType Zero;

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<T>()), \
+      Conv<T>);                                                                            \
+  ONNX_OPERATOR_TYPED_KERNEL_EX(                                                           \
+      Conv,                                                                                \
+      kOnnxDomain,                                                                         \
+      11,                                                                                  \
+      T,                                                                                   \
+      kRocmExecutionProvider,                                                              \
+      (*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
+      Conv<T>);
+
+REGISTER_KERNEL_TYPED(float)
+// not yet supported in MIOpen
+//REGISTER_KERNEL_TYPED(double)
+REGISTER_KERNEL_TYPED(MLFloat16)
+
+template <typename T>
+const miopenConvFwdAlgorithm_t Conv<T>::kAllAlgos[] = {
+    miopenConvolutionFwdAlgoGEMM,
+    miopenConvolutionFwdAlgoDirect,
+    miopenConvolutionFwdAlgoFFT,
+    miopenConvolutionFwdAlgoWinograd,
+    miopenConvolutionFwdAlgoImplicitGEMM
+};
+  miopenStatus_t GetWorkspaceSize(const MiopenConvState<miopenConvAlgoPerf_t>& 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<miopenConvAlgoPerf_t>& 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<size_t>(static_cast<double>(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<int64_t>& input_dims,
+                                     void* output_data,
+                                     const std::vector<int64_t>& output_dims,
+                                     std::vector<int64_t> starts,
+                                     const std::vector<int64_t>& ends,
+                                     const std::vector<int64_t>& 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 <typename T>
+Status Conv<T>::UpdateState(OpKernelContext* context, bool bias_expected) const {
+  //set X
+  const Tensor* X = context->Input<Tensor>(0);
+  const TensorShape& x_shape = X->Shape();
+  const auto& x_dims = x_shape.GetDims();
+  s_.x_data = reinterpret_cast<const HipT*>(X->template Data<T>());
+  s_.element_size = X->DataType()->Size();
+  //set W
+  const Tensor* W = context->Input<Tensor>(1);
+  const TensorShape& w_shape = W->Shape();
+  std::vector<int64_t> w_dims = w_shape.GetDims();
+  s_.w_data = reinterpret_cast<const HipT*>(W->template Data<T>());
+  //set B
+  if (context->InputCount() >= 3) {
+    const Tensor* B = context->Input<Tensor>(2);
+    s_.b_data = reinterpret_cast<const HipT*>(B->template Data<T>());
+  } else {
+    s_.b_data = nullptr;
+  }
+  //set Z
+  if (context->InputCount() >= 4) {
+    const Tensor* Z = context->Input<Tensor>(3);
+    ORT_RETURN_IF_ERROR(s_.z_tensor.Set(Z->Shape().GetDims(), MiopenTensor::GetDataType<HipT>()));
+    s_.z_data = reinterpret_cast<const HipT*>(Z->template Data<T>());
+  } 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<int64_t> kernel_shape;
+    ORT_RETURN_IF_ERROR(conv_attrs_.ComputeKernelShape(W->Shape(), kernel_shape));
+    auto rank = kernel_shape.size();
+    std::vector<int64_t> pads(conv_attrs_.pads);
+    if (pads.empty()) {
+      pads.resize(rank * 2, 0);
+    }
+    std::vector<int64_t> dilations(conv_attrs_.dilations);
+    if (dilations.empty()) {
+      dilations.resize(rank, 1);
+    }
+    std::vector<int64_t> strides(conv_attrs_.strides);
+    if (strides.empty()) {
+      strides.resize(rank, 1);
+    }
+
+    std::vector<int64_t> 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<int64_t> 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<int64_t> slice_starts;
+    slice_starts.reserve(rank);
+
+    std::vector<int64_t> slice_ends;
+    slice_ends.reserve(rank);
+
+    std::vector<int64_t> 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<void>(TensorShape(y_dims_with_adjusted_pads).Size() * s_.element_size);
+      s_.y_data = reinterpret_cast<HipT*>(s_.memory_for_miopen_conv_results.get());
+    } else {
+      // No post slicing needed. Fill the output tensor's buffer directly.
+      s_.y_data = reinterpret_cast<HipT*>(s_.Y->template MutableData<T>());
+    }
+
+    std::vector<int64_t> x_dims_miopen = x_dims;
+    std::vector<int64_t> 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<HipT>()));
+    }
+    ORT_RETURN_IF_ERROR(s_.x_tensor.Set(x_dims_miopen, MiopenTensor::GetDataType<HipT>()));
+    ORT_RETURN_IF_ERROR(s_.y_tensor.Set(y_dims_miopen, MiopenTensor::GetDataType<HipT>()));
+    ORT_RETURN_IF_ERROR(s_.conv_desc.Set(kernel_shape.size(), pads, strides, dilations,
+                                         gsl::narrow_cast<int>(conv_attrs_.group),
+                                         miopenConvolution, MiopenTensor::GetDataType<HipT>()));
+
+    if (context->InputCount() >= 3) {
+      const Tensor* B = context->Input<Tensor>(2);
+      const auto& b_shape = B->Shape();
+      ORT_RETURN_IF_NOT(b_shape.NumDimensions() == 1, "bias should be 1D");
+      std::vector<int64_t> 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<HipT>()));
+    } else if (bias_expected) {
+      std::vector<int64_t> 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<HipT>()));
+      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<const ROCMExecutionProvider*>(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<void> algo_search_workspace = GetTransientScratchBuffer<void>(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<void>(TensorShape(s_.y_dims_with_adjusted_pads).Size() * s_.element_size);
+      s_.y_data = reinterpret_cast<HipT*>(s_.memory_for_miopen_conv_results.get());
+    } else {
+      s_.y_data = reinterpret_cast<HipT*>(s_.Y->template MutableData<T>());
+    }
+  }
+  return Status::OK();
+}
+
+template <typename T>
+Status Conv<T>::ComputeInternal(OpKernelContext* context) const {
+  std::lock_guard<OrtMutex> lock(s_.mutex);
+  ORT_RETURN_IF_ERROR(UpdateState(context));
+  if (s_.Y->Shape().Size() == 0) {
+    return Status::OK();
+  }
+  const auto alpha = Consts<HipT>::One;
+  const auto beta = Consts<HipT>::Zero;
+  IAllocatorUniquePtr<void> 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<int64_t>& pads,
+    const std::vector<int64_t>& strides,
+    const std::vector<int64_t>& dilations,
+    int groups,
+    miopenConvolutionMode_t mode,
+    miopenDataType_t data_type) {
+  if (!desc_)
+    MIOPEN_RETURN_IF_ERROR(miopenCreateConvolutionDescriptor(&desc_));
+
+  std::vector<int> pad_dims(rank);
+  std::vector<int> stride_dims(rank);
+  std::vector<int> dilation_dims(rank);
+  for (size_t i = 0; i < rank; i++) {
+    pad_dims[i] = gsl::narrow_cast<int>(pads[i]);
+    stride_dims[i] = gsl::narrow_cast<int>(strides[i]);
+    dilation_dims[i] = gsl::narrow_cast<int>(dilations[i]);
+  }
+
+ MIOPEN_RETURN_IF_ERROR(miopenInitConvolutionNdDescriptor(
+      desc_,
+      gsl::narrow_cast<int>(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

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 <list>
+
+namespace onnxruntime {
+namespace rocm {
+
+class MiopenConvolutionDescriptor final {
+ public:
+  MiopenConvolutionDescriptor();
+  ~MiopenConvolutionDescriptor();
+
+  Status Set(size_t rank,
+             const std::vector<int64_t>& pads,
+             const std::vector<int64_t>& strides,
+             const std::vector<int64_t>& dilations,
+             int groups,
+             miopenConvolutionMode_t mode,
+             miopenDataType_t data_type);
+
+  operator miopenConvolutionDescriptor_t() const { return desc_; }
+
+ private:
+  miopenConvolutionDescriptor_t desc_;
+};
+
+template <typename T>
+struct vector_hash {
+  std::size_t operator()(const std::vector<T>& values) const {
+    std::size_t seed = values.size();
+    for (auto& val : values)
+      seed ^= std::hash<T>()(val) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
+    return seed;
+  }
+};
+
+template <typename Key, typename T,
+          typename Hash = std::hash<Key>,
+          typename KeyEqual = std::equal_to<Key>,
+          typename ListAllocator = std::allocator<Key>>
+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<Key, ListAllocator>;
+  using iterator_type = typename list_type::iterator;
+  struct value_type {
+    T value;
+    iterator_type lru_iterator;
+  };
+  using MapAllocator = std::allocator<std::pair<const Key, value_type>>;
+
+  void move_to_front(iterator_type it) {
+    lru_list_.splice(lru_list_.begin(), lru_list_, it);
+  }
+
+  size_t max_size_;
+  std::unordered_map<Key, value_type, Hash, KeyEqual, MapAllocator> items_;
+  list_type lru_list_;
+};
+
+// cached miopen descriptors
+constexpr size_t MAX_CACHED_ALGO_PERF_RESULTS = 10000;
+
+template <typename AlgoPerfType>
+struct MiopenConvState {
+  miopenHandle_t handle;
+
+  // if x/w dims changed, update algo and miopenTensors
+  std::vector<int64_t> last_x_dims;
+  std::vector<int64_t> last_w_dims;
+
+  // these would be recomputed if x/w dims change
+  std::vector<int64_t> y_dims;
+  std::vector<int64_t> 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<std::vector<int64_t>, PerfFwdResultParams, vector_hash<int64_t>> cached_benchmark_fwd_results{MAX_CACHED_ALGO_PERF_RESULTS};
+  lru_unordered_map<std::vector<int64_t>, PerfBwdResultParams, vector_hash<int64_t>> cached_benchmark_bwd_results{MAX_CACHED_ALGO_PERF_RESULTS};
+
+  // Some properties needed to support asymmetric padded Conv nodes
+  bool post_slicing_required;
+  std::vector<int64_t> slice_starts;
+  std::vector<int64_t> slice_ends;
+  std::vector<int64_t> slice_axes;
+
+  // note that conv objects are shared between execution frames, and a lock is needed to avoid multi-thread racing
+  OrtMutex mutex;
+  IAllocatorUniquePtr<void> 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 <typename T>
+class Conv : public RocmKernel {
+ public:
+  using HipT = typename ToHipType<T>::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<void> GetWorkSpace() const {
+    return GetScratchBuffer<void>(s_.workspace_bytes);
+  }
+
+  Status UpdateState(OpKernelContext* context, bool bias_expected = false) const;
+  ConvAttributes conv_attrs_;
+  mutable MiopenConvState<miopenConvAlgoPerf_t> s_;
+  constexpr static auto kDefaultConvAlgo = miopenConvolutionFwdAlgoGEMM;
+  static const miopenConvFwdAlgorithm_t kAllAlgos[];
+};
+
+Status SliceOutUnwantedOutputSection(hipStream_t stream,
+                                     const void* input_data,
+                                     const std::vector<int64_t>& input_dims,
+                                     void* output_data,
+                                     const std::vector<int64_t>& output_dims,
+                                     std::vector<int64_t> starts,
+                                     const std::vector<int64_t>& ends,
+                                     const std::vector<int64_t>& axes,
+                                     size_t element_size);
+}  // namespace rocm
+}  // namespace onnxruntime

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<T>()), \
+      ConvTranspose<T>);                                                                   \
+  ONNX_OPERATOR_TYPED_KERNEL_EX(                                                           \
+      ConvTranspose,                                                                       \
+      kOnnxDomain,                                                                         \
+      11,                                                                                  \
+      T,                                                                                   \
+      kRocmExecutionProvider,                                                              \
+      (*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
+      ConvTranspose<T>);
+
+REGISTER_KERNEL_TYPED(float)
+// not yet supported in MIOpen
+//REGISTER_KERNEL_TYPED(double)
+REGISTER_KERNEL_TYPED(MLFloat16)
+
+template <typename T>
+Status ConvTranspose<T>::ComputeInternal(OpKernelContext* context) const {
+  return DoConvTranspose(context, false);
+}
+
+template <typename T>
+Status ConvTranspose<T>::DoConvTranspose(OpKernelContext* context, bool dynamic_padding) const {
+  typedef typename ToHipType<T>::MappedType HipT;
+
+  const Tensor* X = context->Input<Tensor>(0);
+  const TensorShape& x_shape = X->Shape();
+  auto x_dims = x_shape.GetDims();
+  auto x_data = reinterpret_cast<const HipT*>(X->template Data<T>());
+
+  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<Tensor>(1);
+  const TensorShape& w_shape = W->Shape();
+  std::vector<int64_t> w_dims = w_shape.GetDims();
+  auto w_data = reinterpret_cast<const HipT*>(W->template Data<T>());
+
+  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<OrtMutex> 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<HipT>()));
+	}
+
+      // 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<HipT>()));
+      ORT_RETURN_IF_ERROR(s_.y_tensor.Set(y_dims, MiopenTensor::GetDataType<HipT>()));
+
+      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<int>(conv_transpose_attrs_.group),
+                                           mode, MiopenTensor::GetDataType<HipT>()));
+
+      if (has_bias) {
+        const auto& b_shape = p.B->Shape();
+        ORT_RETURN_IF_NOT(b_shape.NumDimensions() == 1, "bias should be 1D");
+        std::vector<int64_t> 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<HipT>()));
+      }
+
+      y_data = reinterpret_cast<HipT*>(p.Y->template MutableData<T>());
+
+      if (!s_.cached_benchmark_bwd_results.contains(x_dims)) {
+        IAllocatorUniquePtr<void> algo_search_workspace = GetScratchBuffer<void>(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<HipT*>(Y->template MutableData<T>());
+
+      // Bail out early if one of the output dimensions is zero.
+      if (Y->Shape().Size() == 0) {
+        return Status::OK();
+      }
+    }
+
+    const auto alpha = Consts<HipT>::One;
+    const auto beta = Consts<HipT>::Zero;
+
+    IAllocatorUniquePtr<void> workspace = GetScratchBuffer<void>(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<Tensor>(3) : context->Input<Tensor>(2);
+      auto b_data = reinterpret_cast<const HipT*>(B->template Data<T>());
+      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

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 <typename T>
+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<miopenConvAlgoPerf_t> s_;
+};
+
+}  // namespace rocm
+}  // namespace onnxruntime

onnxruntime/core/providers/rocm/rocm_execution_provider.cc

@@ -1229,12 +1229,12 @@ static Status RegisterRocmKernels(KernelRegistry& kernel_registry) {
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, float, LRN)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, double, LRN)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, MLFloat16, LRN)>,
-      // BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 10, float, Conv)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 10, float, Conv)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 10, double, Conv)>,
-      // BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 10, MLFloat16, Conv)>,
-      // BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 10, float, ConvTranspose)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 10, MLFloat16, Conv)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 10, float, ConvTranspose)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 10, double, ConvTranspose)>,
-      // BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 10, MLFloat16, ConvTranspose)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 10, MLFloat16, ConvTranspose)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 7, 9, float, AveragePool)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 7, 9, double, AveragePool)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 7, 9, MLFloat16, AveragePool)>,
@@ -1502,12 +1502,12 @@ static Status RegisterRocmKernels(KernelRegistry& kernel_registry) {
       BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, 12, Squeeze)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, TopK)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, 12, Unsqueeze)>,
-      // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, float, Conv)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, float, Conv)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, double, Conv)>,
-      // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, MLFloat16, Conv)>,
-      // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, float, ConvTranspose)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, MLFloat16, Conv)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, float, ConvTranspose)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, double, ConvTranspose)>,
-      // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, MLFloat16, ConvTranspose)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, MLFloat16, ConvTranspose)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, float, AveragePool)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, double, AveragePool)>,
       // BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, MLFloat16, AveragePool)>,
@@ -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()) {

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 <typename T>
+  IAllocatorUniquePtr<T> GetTransientScratchBuffer(size_t count_or_bytes) const {
+    if (count_or_bytes == 0)
+      return nullptr;
+
+    return IAllocator::MakeUniquePtr<T>(GetAllocator(info_.device_id, OrtMemTypeDefault), count_or_bytes, true);
+  }
+
   void RegisterAllocator(std::shared_ptr<AllocatorManager> allocator_manager) override;
   static AllocatorPtr CreateRocmAllocator(OrtDevice::DeviceId device_id, size_t rocm_mem_limit, ArenaExtendStrategy arena_extend_strategy,
                                           ROCMExecutionProviderExternalAllocatorInfo external_alloc_info);

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)

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);

onnxruntime/core/providers/rocm/rocm_kernel.h

@@ -52,6 +52,16 @@ class RocmKernel : public OpKernel {
     return provider_->GetScratchBuffer<T>(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 <typename T>
+  inline IAllocatorUniquePtr<T> GetTransientScratchBuffer(size_t count_or_bytes) const {
+    return provider_->GetTransientScratchBuffer<T>(count_or_bytes);
+  }
+
+
   inline void AddDeferredReleaseCPUPtr(void* p) const {
     provider_->AddDeferredReleaseCPUPtr(p);
   }