Stable Diffusion CUDA Optimizations

cmake/onnxruntime_rocm_hipify.cmake

@@ -27,6 +27,10 @@ set(contrib_ops_excluded_files
   "bert/tensorrt_fused_multihead_attention/*"
   "bert/transformer_common.h"
   "bert/transformer_common.cc"
+  "diffusion/group_norm.h"
+  "diffusion/group_norm.cc"
+  "diffusion/group_norm_impl.cu"
+  "diffusion/group_norm_impl.h"
   "math/complex_mul.cc"
   "math/complex_mul.h"
   "math/complex_mul_impl.cu"

onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc

@@ -71,6 +71,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain,
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, EmbedLayerNormalization);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, EmbedLayerNormalization);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, GreedySearch);
+class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, GroupNorm);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, float, ImageScaler);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, double, ImageScaler);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16, ImageScaler);
@@ -192,6 +193,7 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, EmbedLayerNormalization)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, EmbedLayerNormalization)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, GreedySearch)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, GroupNorm)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, float, ImageScaler)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, double, ImageScaler)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16, ImageScaler)>,

onnxruntime/contrib_ops/cuda/diffusion/group_norm.cc

@@ -0,0 +1,106 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/cuda/cuda_common.h"
+#include "contrib_ops/cuda/diffusion/group_norm.h"
+#include "contrib_ops/cuda/diffusion/group_norm_impl.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+#define REGISTER_KERNEL_TYPED(T)                                  \
+  ONNX_OPERATOR_TYPED_KERNEL_EX(                                  \
+      GroupNorm,                                                  \
+      kMSDomain,                                                  \
+      1,                                                          \
+      T,                                                          \
+      kCudaExecutionProvider,                                     \
+      (*KernelDefBuilder::Create())                               \
+          .TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
+      GroupNorm<T>);
+
+REGISTER_KERNEL_TYPED(MLFloat16);
+
+using namespace ONNX_NAMESPACE;
+
+template <typename T>
+GroupNorm<T>::GroupNorm(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) {
+  ORT_ENFORCE(op_kernel_info.GetAttr<float>("epsilon", &epsilon_).IsOK());
+  ORT_ENFORCE(epsilon_ >= 0);
+
+  int64_t num_groups;
+  ORT_ENFORCE(op_kernel_info.GetAttr<int64_t>("groups", &num_groups).IsOK());
+  ORT_ENFORCE(num_groups >= 0);
+  num_groups_ = static_cast<int>(num_groups);
+
+
+  ORT_ENFORCE(op_kernel_info.GetAttr<bool>("swish", &swish_).IsOK());
+}
+
+template <typename T>
+Status GroupNorm<T>::ComputeInternal(OpKernelContext* context) const {
+  const Tensor* input = context->Input<Tensor>(0);
+  const Tensor* gamma = context->Input<Tensor>(1);
+  const Tensor* beta = context->Input<Tensor>(2);
+  Tensor* output = context->Output(0, input->Shape());
+
+  const auto& input_dims = input->Shape().GetDims();
+  if (input_dims.size() != 4) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "input is expected to have 4 dimensions, got ", input_dims.size());
+  }
+
+  const auto& gamma_dims = gamma->Shape().GetDims();
+  if (gamma_dims.size() != 1) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "gamma is expected to have 1 dimension, got ", gamma_dims.size());
+  }
+  if (gamma_dims[0] != input_dims[2]) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "Last dimension of gamma and input does not match");
+  }
+
+  const auto& beta_dims = beta->Shape().GetDims();
+  if (beta_dims.size() != 1) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                            "beta is expected to have 1 dimension, got ", beta_dims.size());
+  }
+  if (beta_dims[0] != input_dims[2]) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                            "Last dimension of beta and input does not match");
+  }
+
+  int batch_size = static_cast<int>(input_dims[0]);
+  int num_channels = static_cast<int>(input_dims[1]);
+  int height = static_cast<int>(input_dims[2]);
+  int width = static_cast<int>(input_dims[3]);
+
+  if (num_channels % num_groups_ != 0) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                            "number of channels should be divisiable by num_groups");
+  }
+
+  auto workspace = GetScratchBuffer<void>(GetGroupNormWorkspaceSizeInBytes(), context->GetComputeStream());
+
+  typedef typename ToCudaType<T>::MappedType CudaT;
+
+  return LaunchGroupNormKernel<CudaT>(
+      Stream(context),
+      reinterpret_cast<CudaT*>(output->MutableData<T>()),
+      reinterpret_cast<const CudaT*>(input->Data<T>()),
+      reinterpret_cast<const float*>(gamma->Data<T>()),
+      reinterpret_cast<const float*>(beta->Data<T>()),
+      reinterpret_cast<float*>(workspace.get()),
+      epsilon_,
+      batch_size,
+      num_channels,
+      height,
+      width,
+      num_groups_,
+      swish_);
+}
+
+}  // namespace cuda
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cuda/diffusion/group_norm.h

@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+#include "core/common/common.h"
+#include "core/providers/cuda/cuda_kernel.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+using namespace onnxruntime::cuda;
+
+template <typename T>
+class GroupNorm final : public CudaKernel {
+ public:
+  GroupNorm(const OpKernelInfo& op_kernel_info);
+  Status ComputeInternal(OpKernelContext* context) const override;
+
+ private:
+  bool swish_;
+  float epsilon_;
+  int num_groups_;
+};
+
+}  // namespace cuda
+}  // namespace contrib
+}  // namespace onnxruntime