diff --git a/cmake/onnxruntime_providers.cmake b/cmake/onnxruntime_providers.cmake index 6e387d1710..a6b991af21 100644 --- a/cmake/onnxruntime_providers.cmake +++ b/cmake/onnxruntime_providers.cmake @@ -464,7 +464,7 @@ if (onnxruntime_USE_CUDA) endif() add_dependencies(onnxruntime_providers_cuda onnxruntime_providers_shared ${onnxruntime_EXTERNAL_DEPENDENCIES}) - target_link_libraries(onnxruntime_providers_cuda PRIVATE cublas cudnn curand cufft ${ABSEIL_LIBS} ${ONNXRUNTIME_PROVIDERS_SHARED}) + target_link_libraries(onnxruntime_providers_cuda PRIVATE cublasLt cublas cudnn curand cufft ${ABSEIL_LIBS} ${ONNXRUNTIME_PROVIDERS_SHARED}) if(onnxruntime_CUDNN_HOME) target_include_directories(onnxruntime_providers_cuda PRIVATE ${onnxruntime_CUDNN_HOME}/include) endif() diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index e35bc53033..e1a8d23a0a 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -56,6 +56,9 @@ Do not modify directly.* * com.microsoft.QLinearReduceMean * com.microsoft.QLinearSigmoid * com.microsoft.QLinearSoftmax + * com.microsoft.QOrderedGelu + * com.microsoft.QOrderedLayerNormalization + * com.microsoft.QOrderedMatMul * com.microsoft.QuantizeLinear * com.microsoft.Range * com.microsoft.ReduceSumInteger @@ -2862,6 +2865,163 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.QOrderedGelu** + + Ordered Quantize Gelu. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
order_X : int
+
cublasLt order of input X. Default is ROW MAJOR.
+
order_Y : int
+
cublasLt order of matrix Y, must be same as order_X. Default is ROW MAJOR.
+
+ +#### Inputs + +
+
X : Q
+
N-dimensional input A
+
scale_X : S
+
scale of the input A
+
scale_Y : S
+
scale of the output Y
+
+ +#### Outputs + +
+
Y : Q
+
Output of the Gelu
+
+ +#### Type Constraints + +
+
Q : tensor(int8)
+
Constrain input and output types to int8 tensors.
+
S : tensor(float)
+
Constrain scales to float32
+
+ + +### **com.microsoft.QOrderedLayerNormalization** + + QOrderedLayerNormalization + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
axis : int
+
The first normalization dimension: normalization will be performed along dimensions axis : rank(inputs).
+
epsilon : float
+
The epsilon value to use to avoid division by zero.
+
order_X : int
+
cublasLt order of input X. Default is ROW MAJOR.
+
order_Y : int
+
cublasLt order of matrix Y, must be same as order_X. Default is ROW MAJOR.
+
+ +#### Inputs + +
+
X : Q
+
Input data tensor from the previous layer.
+
scale_X : S
+
scale of the quantized X
+
scale : F
+
Scale tensor, i.e., gamma vector.
+
B (optional) : F
+
Bias tensor.
+
scale_Y : S
+
scale of the quantized X
+
+ +#### Outputs + +
+
Y : Q
+
Output data tensor.
+
+ +#### Type Constraints + +
+
F : tensor(float16), tensor(float)
+
Constrain input gamma and bias could be float16/float tensors. float may get better precision, float16 runs faster.
+
S : tensor(float)
+
quantization scale must be float tensors.
+
Q : tensor(int8)
+
quantization tensor must be int8 tensors.
+
+ + +### **com.microsoft.QOrderedMatMul** + + TODO + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
order_A : int
+
cublasLt order of matrix A. Default is ROW MAJOR.
+
order_B : int
+
cublasLt order of matrix B. Default is ROW MAJOR.
+
order_Y : int
+
cublasLt order of matrix Y and optional matrix C. Default is ROW MAJOR.
+
+ +#### Inputs (5 - 8) + +
+
A : Q
+
3-dimensional matrix A
+
scale_A : S
+
scale of the input A
+
B : Q
+
2-dimensional matrix B
+
scale_B : S
+
scale of the input B
+
scale_Y : S
+
scale of the output Y
+
bias (optional) : S
+
1d bias
+
C (optional) : Q
+
3d or 2d matrix C. if 2d expand to 3d first. Shape[0] should be 1 or same as A.shape[0]
+
scale_C (optional) : S
+
scale of the input A
+
+ +#### Outputs + +
+
Y : Q
+
Matrix multiply results from A * B
+
+ +#### Type Constraints + +
+
Q : tensor(int8)
+
Constrain input and output types to int8 tensors.
+
S : tensor(float)
+
Constrain bias and scales to float32
+
+ + ### **com.microsoft.QuantizeLinear** The linear quantization operator. It consumes a full precision data, a scale, a zero point to compute the low precision / quantized tensor. diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 0b5ae058a3..370b52ab38 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -780,6 +780,9 @@ Do not modify directly.* |LongformerAttention|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* mask:**T**
*in* global_weight:**T**
*in* global_bias:**T**
*in* global:**G**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| |NGramRepeatBlock|*in* input_ids:**Tid**
*in* scores:**T**
*out* scores_out:**T**|1+|**T** = tensor(float)
**Tid** = tensor(int64)| |QAttention|*in* input:**T1**
*in* weight:**T2**
*in* bias:**T3**
*in* input_scale:**T3**
*in* weight_scale:**T3**
*in* mask_index:**T4**
*in* input_zero_point:**T1**
*in* weight_zero_point:**T2**
*in* past:**T3**
*out* output:**T3**
*out* present:**T3**|1+|**T1** = tensor(int8)
**T2** = tensor(int8)
**T3** = tensor(float), tensor(float16)
**T4** = tensor(int32)| +|QOrderedGelu|*in* X:**Q**
*in* scale_X:**S**
*in* scale_Y:**S**
*out* Y:**Q**|1+|**Q** = tensor(int8)
**S** = tensor(float)| +|QOrderedLayerNormalization|*in* X:**Q**
*in* scale_X:**S**
*in* scale:**F**
*in* B:**F**
*in* scale_Y:**S**
*out* Y:**Q**|1+|**F** = tensor(float), tensor(float16)
**Q** = tensor(int8)
**S** = tensor(float)| +|QOrderedMatMul|*in* A:**Q**
*in* scale_A:**S**
*in* B:**Q**
*in* scale_B:**S**
*in* scale_Y:**S**
*in* bias:**S**
*in* C:**Q**
*in* scale_C:**S**
*out* Y:**Q**|1+|**Q** = tensor(int8)
**S** = tensor(float)| |QuantizeLinear|*in* x:**T1**
*in* y_scale:**T1**
*in* y_zero_point:**T2**
*out* y:**T2**|1+|**T1** = tensor(float16)
**T2** = tensor(int8), tensor(uint8)| |Rfft|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| |SkipLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* beta:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**|1+|**T** = tensor(float), tensor(float16)| diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 7062398dda..d9e9e6dc55 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -97,6 +97,9 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, TransposeMatMul); // backward compatibility class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, FusedMatMul); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, BFloat16_float_BFloat16, LayerNormalization); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, QOrderedMatMul); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, QOrderedLayerNormalization); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, QOrderedGelu); #ifdef ENABLE_ATEN class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kPytorchAtenDomain, 1, ATen); @@ -201,6 +204,9 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, #ifdef ENABLE_ATEN BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_common.cuh b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_common.cuh new file mode 100644 index 0000000000..ad3647e2a5 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_common.cuh @@ -0,0 +1,29 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/cuda/cu_inc/common.cuh" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +__device__ inline float ToFloat(const __half h) { return __half2float(h); } + +__device__ inline float ToFloat(const float f) { return f; } + +template +__inline__ __device__ T +WarpReduceSum(T val) { + val += __shfl_xor_sync(0xFFFFFFFF, val, 1); + val += __shfl_xor_sync(0xFFFFFFFF, val, 2); + val += __shfl_xor_sync(0xFFFFFFFF, val, 4); + val += __shfl_xor_sync(0xFFFFFFFF, val, 8); + val += __shfl_xor_sync(0xFFFFFFFF, val, 16); + return val; +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_layer_norm.cc b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_layer_norm.cc new file mode 100644 index 0000000000..aba3d4759a --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_layer_norm.cc @@ -0,0 +1,100 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "qordered_layer_norm.h" +#include "qordered_layer_norm_impl.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +ONNX_OPERATOR_KERNEL_EX( + QOrderedLayerNormalization, + kMSDomain, + 1, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("Q", DataTypeImpl::GetTensorType()) + .TypeConstraint("S", DataTypeImpl::GetTensorType()) + .TypeConstraint("F", BuildKernelDefConstraints()) + .InputMemoryType(OrtMemTypeCPUInput, 1) // scale_X + .InputMemoryType(OrtMemTypeCPUInput, 4), // scale_Y + QOrderedLayerNormalization); + +QOrderedLayerNormalization::QOrderedLayerNormalization(const OpKernelInfo& op_kernel_info) + : CudaKernel(op_kernel_info) { + ORT_ENFORCE(op_kernel_info.GetAttr("axis", &axis_).IsOK()); + + float tmp_epsilon; + ORT_ENFORCE(op_kernel_info.GetAttr("epsilon", &tmp_epsilon).IsOK()); + epsilon_ = tmp_epsilon; + + ORT_ENFORCE(op_kernel_info.GetAttr("order_X", &order_X_).IsOK()); + + ORT_ENFORCE(op_kernel_info.GetAttr("order_Y", &order_Y_).IsOK()); + + ORT_ENFORCE(order_X_ == 1, "QOrderedLayerNormlalization: Only Row major data ordering is currently supported"); + + ORT_ENFORCE(order_X_ == order_Y_, "QOrderedLayerNormlalization: Input ordering should match the output ordering"); +} + +Status QOrderedLayerNormalization::ComputeInternal(OpKernelContext* ctx) const { + typedef typename ToCudaType::MappedType CudaQ; + + // Inputs + const Tensor* X = ctx->Input(0); + const auto* X_data = reinterpret_cast(X->Data()); + const TensorShape& x_shape = X->Shape(); + ORT_ENFORCE(x_shape.GetDims().size() == 3, + "QOrderedLayerNormlalization: Input shape must be {batch, rows, cols}"); + + const Tensor* scale = ctx->Input(2); + const void* scale_data = scale->DataRaw(); + + const Tensor* bias = ctx->Input(3); + const void* bias_data = (nullptr == bias) ? nullptr : bias->DataRaw(); + + const int64_t axis = HandleNegativeAxis(axis_, x_shape.NumDimensions()); + ORT_ENFORCE(axis == 2, + "QOrderedLayerNormlalization: Implementation only " + "supports on normalizing on innermost axis"); + + unsigned int batch = gsl::narrow(x_shape.GetDims()[0]); + unsigned int rows = gsl::narrow(x_shape.GetDims()[1]); + unsigned int cols = gsl::narrow(x_shape.GetDims()[2]); + + if (cols & 0x03) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "QOrderedLayerNormlalization: Cols MUST be a multiple of 4"); + } + + // Outputs + Tensor* Y = ctx->Output(0, x_shape); + + if (x_shape.Size() == 0) { + return Status::OK(); + } + + auto* Y_data = reinterpret_cast(Y->MutableData()); + + const float* scale_x = ctx->Input(1)->Data(); + const float* scale_y = ctx->Input(4)->Data(); + + if (scale->IsDataType()) { + QOrderedLayerNorm(Stream(), GetDeviceProp(), static_cast(order_X_), + X_data, *scale_x, Y_data, *scale_y, static_cast(scale_data), + static_cast(bias_data), + static_cast(epsilon_), batch, rows, cols); + } else { + QOrderedLayerNorm(Stream(), GetDeviceProp(), static_cast(order_X_), + X_data, *scale_x, Y_data, *scale_y, static_cast(scale_data), + static_cast(bias_data), + static_cast(epsilon_), batch, rows, cols); + } + + return Status::OK(); +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_layer_norm.h b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_layer_norm.h new file mode 100644 index 0000000000..f66a0597ca --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_layer_norm.h @@ -0,0 +1,29 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/cuda/cuda_kernel.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +class QOrderedLayerNormalization final : public CudaKernel { + public: + explicit QOrderedLayerNormalization(const OpKernelInfo& op_kernel_info); + + Status ComputeInternal(OpKernelContext* ctx) const override; + + private: + double epsilon_; + int64_t axis_; + int64_t order_X_; + int64_t order_Y_; +}; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_layer_norm_impl.cu new file mode 100644 index 0000000000..8ba13961f5 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_layer_norm_impl.cu @@ -0,0 +1,120 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/cuda/cu_inc/common.cuh" +#include "qordered_common.cuh" +#include "qordered_layer_norm_impl.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +constexpr static unsigned QORDER_LAYERNORM_ROWS_PER_BLOCK = 8; + +__device__ inline int8_t QuantizeFloatS8(const float val, const float inverse_scale) { + float dqval = fmaxf(fminf(127.0f, val * inverse_scale), -128.0f); + return static_cast(__float2int_rn(dqval)); +} + +// TODO: Const Ref ? +__device__ inline char4 QuantizeFloat4Char4(const float4 val4, const float rscale) { + return char4{QuantizeFloatS8(val4.x, rscale), QuantizeFloatS8(val4.y, rscale), + QuantizeFloatS8(val4.z, rscale), QuantizeFloatS8(val4.w, rscale)}; +} + +// TODO: Const Ref ? +__device__ inline int32_t Dp4a_Defined(const char4 input_1, const char4 input_2) { + return static_cast(input_1.x) * static_cast(input_2.x) + + static_cast(input_1.y) * static_cast(input_2.y) + + static_cast(input_1.z) * static_cast(input_2.z) + + static_cast(input_1.w) * static_cast(input_2.w); +} + +template +__global__ void QOrderedLayerNormRowKernel(const int8_t* __restrict__ src, const float src_scale, + int8_t* __restrict__ dst, const float dst_scale, + const T* __restrict__ gamma, const T* __restrict__ beta, const float epsilon, + const unsigned rows, const unsigned cols) { + int32_t sum = 0; + int32_t square_sum = 0; + + unsigned r = blockIdx.x * QORDER_LAYERNORM_ROWS_PER_BLOCK + threadIdx.y; + + if (rows <= r) { + return; + } + + const size_t batch_row_index = static_cast(blockIdx.y) * (rows * cols) + r * cols; + src += batch_row_index; + dst += batch_row_index; + for (unsigned c = threadIdx.x << 2; c < cols; c += 128) { + char4 ch4 = __ldg(reinterpret_cast(src + c)); + sum += (static_cast(ch4.x) + static_cast(ch4.y) + + static_cast(ch4.z) + static_cast(ch4.w)); +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 610 + square_sum = __dp4a(ch4, ch4, square_sum); +#else + square_sum = Dp4a_Defined(ch4, ch4); +#endif + } + + sum = WarpReduceSum(sum); + square_sum = WarpReduceSum(square_sum); + + const float mean = __double2float_rn(src_scale * (double)sum / cols); + + const float rvar = rsqrtf(src_scale * src_scale * __double2float_rn(static_cast(square_sum) - (static_cast(sum) * static_cast(sum) / static_cast(cols))) / cols + epsilon); + + const float dst_rscale = 1.0f / dst_scale; + + float4 f4; + for (unsigned c = threadIdx.x << 2; c < cols; c += 128) { + char4 ch4 = __ldg(reinterpret_cast(src + c)); + + f4.x = (src_scale * ch4.x - mean) * rvar * ToFloat(gamma[c]); + f4.y = (src_scale * ch4.y - mean) * rvar * ToFloat(gamma[c + 1]); + f4.z = (src_scale * ch4.z - mean) * rvar * ToFloat(gamma[c + 2]); + f4.w = (src_scale * ch4.w - mean) * rvar * ToFloat(gamma[c + 3]); + + if (beta) { + f4.x += ToFloat(beta[c]); + f4.y += ToFloat(beta[c + 1]); + f4.z += ToFloat(beta[c + 2]); + f4.w += ToFloat(beta[c + 3]); + } + + *reinterpret_cast(dst + c) = QuantizeFloat4Char4(f4, dst_rscale); + } +} + +template +void QOrderedLayerNorm(cudaStream_t stream, const cudaDeviceProp& /*device_prop*/, cublasLtOrder_t order, + const int8_t* src, const float src_scale, int8_t* dst, const float dst_scale, + const T* gamma, const T* beta, const float epsilon, + const unsigned batch, const unsigned rows, const unsigned cols) { + // The implementation only supports Row major tensor data ordering for now + ORT_ENFORCE(order == CUBLASLT_ORDER_ROW); + + dim3 threads(32, QORDER_LAYERNORM_ROWS_PER_BLOCK, 1); + + dim3 blocks(static_cast(rows + QORDER_LAYERNORM_ROWS_PER_BLOCK - 1) / QORDER_LAYERNORM_ROWS_PER_BLOCK, + static_cast(batch), 1); + + QOrderedLayerNormRowKernel<<>>(src, src_scale, + dst, dst_scale, gamma, beta, + epsilon, rows, cols); +} + +template void QOrderedLayerNorm(cudaStream_t stream, const cudaDeviceProp& /*device_prop*/, cublasLtOrder_t order, + const int8_t* src, const float src_scale, int8_t* dst, const float dst_scale, + const float* gamma, const float* beta, const float epsilon, + const unsigned batch, const unsigned rows, const unsigned cols); + +template void QOrderedLayerNorm<__half>(cudaStream_t stream, const cudaDeviceProp& /*device_prop*/, cublasLtOrder_t order, + const int8_t* src, const float src_scale, int8_t* dst, const float dst_scale, + const __half* gamma, const __half* beta, const float epsilon, + const unsigned batch, const unsigned rows, const unsigned cols); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_layer_norm_impl.h b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_layer_norm_impl.h new file mode 100644 index 0000000000..1ad8b46a68 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_layer_norm_impl.h @@ -0,0 +1,21 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/cuda/cuda_common.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +template +void QOrderedLayerNorm( + cudaStream_t stream, const cudaDeviceProp& device_prop, cublasLtOrder_t order, + const int8_t* src, const float src_scale, int8_t* dst, const float dst_scale, + const T* gamma, const T* beta, const float epsilon, + unsigned batch, unsigned rows, unsigned cols); + +} +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_matmul.cc b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_matmul.cc new file mode 100644 index 0000000000..34198b7b9f --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_matmul.cc @@ -0,0 +1,131 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "qordered_matmul.h" +#include "qordered_matmul_utils.h" + +#include + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +ONNX_OPERATOR_KERNEL_EX( + QOrderedMatMul, + kMSDomain, + 1, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("Q", DataTypeImpl::GetTensorType()) + .TypeConstraint("S", DataTypeImpl::GetTensorType()) + .InputMemoryType(OrtMemTypeCPUInput, 1) // scale_A + .InputMemoryType(OrtMemTypeCPUInput, 3) // scale_B + .InputMemoryType(OrtMemTypeCPUInput, 4) // scale_Y + .InputMemoryType(OrtMemTypeCPUInput, 7), // scale_C + QOrderedMatMul); + +static Status ParseRowMajorTensorMetadata(const Tensor& input_tensor, int64_t& rows, + int64_t& cols, int64_t& batch_count, int64_t& element_count) { + const auto& dims = input_tensor.Shape().GetDims(); + + cols = dims.back(); + rows = (dims.size() <= 1 ? 1LL : dims[dims.size() - 2]); + batch_count = (dims.size() <= 2 + ? 1LL + : std::accumulate(dims.begin(), dims.begin() + (dims.size() - 2), + 1LL, std::multiplies())); + + element_count = cols * rows * batch_count; + + return Status::OK(); +} + +QOrderedMatMul::QOrderedMatMul(const OpKernelInfo& info) : CudaKernel(info) { + ORT_ENFORCE(info.GetAttr("order_A", &order_A_).IsOK()); + + ORT_ENFORCE(info.GetAttr("order_B", &order_B_).IsOK()); + + ORT_ENFORCE(info.GetAttr("order_Y", &order_Y_).IsOK()); + + ORT_ENFORCE(order_B_ == CUBLASLT_ORDER_COL && + order_A_ == CUBLASLT_ORDER_ROW && + order_Y_ == CUBLASLT_ORDER_ROW, + "QOrderedMatMul: Input 1 and output's data ordering should be ROW_MAJOR and " + "Input 2's data ordering should be COL_MAJOR"); +} + +Status QOrderedMatMul::QOrderedMatMul::ComputeInternal(OpKernelContext* context) const { + int64_t rows_A = 0, cols_A = 0, batch_A = 1, elements_A = 0; + int64_t rows_B = 0, cols_B = 0, batch_B = 1, elements_B = 0; + int64_t rows_C = 0, cols_C = 0, batch_C = 1, elements_C = 0; + + const Tensor& tensor_A = *context->Input(0); + const Tensor& tensor_B = *context->Input(2); + + // Support General case only. No broadcasting, is handled now. + ORT_ENFORCE(tensor_A.Shape().NumDimensions() == 2 || tensor_A.Shape().NumDimensions() == 3); + ORT_ENFORCE(tensor_B.Shape().NumDimensions() == 2 || tensor_B.Shape().NumDimensions() == 3); + + ORT_RETURN_IF_ERROR(ParseRowMajorTensorMetadata(tensor_A, rows_A, cols_A, batch_A, elements_A)); + ORT_RETURN_IF_ERROR(ParseRowMajorTensorMetadata(tensor_B, rows_B, cols_B, batch_B, elements_B)); + + const float* scale_A = context->Input(1)->Data(); + const float* scale_B = context->Input(3)->Data(); + const float* scale_Y = context->Input(4)->Data(); + ORT_ENFORCE(*scale_Y > 0.0f && *scale_A > 0.0f && *scale_B > 0.0f); + + const Tensor* tensor_bias = context->Input(5); + ORT_ENFORCE(tensor_bias == nullptr || + (tensor_bias->Shape().NumDimensions() == 1 && tensor_bias->Shape()[0] == cols_B)); + + const float* bias = (tensor_bias == nullptr) ? nullptr : tensor_bias->Data(); + + ORT_ENFORCE(batch_A == batch_B || batch_B == 1, "Batch count for matrix A and matrix B does not match"); + ORT_ENFORCE(cols_A == rows_B, "MatMul shape mis-match"); + + TensorShape output_shape(tensor_A.Shape()); + output_shape[output_shape.NumDimensions() - 1] = cols_B; + + const float zero = 0.0f; + const float* scale_C = &zero; + + const int8_t* C = nullptr; + const Tensor* tensor_C = context->Input(6); + + if (tensor_C != nullptr) { + ORT_ENFORCE(tensor_C->Shape().NumDimensions() == 2 || tensor_C->Shape().NumDimensions() == 3); + ORT_RETURN_IF_ERROR(ParseRowMajorTensorMetadata(*tensor_C, rows_C, cols_C, batch_C, elements_C)); + + ORT_ENFORCE(batch_C == batch_A || batch_C == 1); + ORT_ENFORCE(rows_C == rows_A && cols_C == cols_B); + + const Tensor* tensor_scale_C = context->Input(7); + ORT_ENFORCE(tensor_scale_C != nullptr); + scale_C = tensor_scale_C->Data(); + + C = tensor_C->Data(); + } + + Tensor* tensor_Y = context->Output(0, output_shape); + cublasLtHandle_t cublasLt = CublasLtHandle(); + cudaStream_t stream = Stream(); + auto& device_prop = GetDeviceProp(); + + const float alpha = *scale_A * *scale_B / *scale_Y; + const float beta = *scale_C / *scale_Y; + + ORT_RETURN_IF_ERROR(QOrdered_MatMul(cublasLt, stream, device_prop, + static_cast(batch_A), rows_A, cols_B, cols_A, + &alpha, tensor_A.Data(), tensor_B.Data(), + static_cast(batch_B), bias, + &beta, C, static_cast(batch_C), + tensor_Y->MutableData(), static_cast(order_B_))); + + return Status::OK(); +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_matmul.h b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_matmul.h new file mode 100644 index 0000000000..829cd06625 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_matmul.h @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/cuda/cuda_kernel.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +class QOrderedMatMul final : public CudaKernel { + public: + explicit QOrderedMatMul(const OpKernelInfo& info); + + Status ComputeInternal(OpKernelContext* context) const override; + + private: + int64_t order_A_; + int64_t order_B_; + int64_t order_Y_; +}; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_matmul_utils.cc b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_matmul_utils.cc new file mode 100644 index 0000000000..12021c99e6 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_matmul_utils.cc @@ -0,0 +1,209 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "qordered_matmul_utils.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +int64_t CalcLeadingDimensionLt(int64_t rows, int64_t cols, cublasLtOrder_t order) { + switch (order) { + case CUBLASLT_ORDER_ROW: + return cols; + case CUBLASLT_ORDER_COL: + return rows; + // TODO: Support other CUBLASLT ordering + default: + return 0; + } +} + +static Status cublasLtMatMulInt8SetupAlgo(cublasLtHandle_t cublasLt_handle, cublasLtMatmulAlgo_t& algo, + int algo_id, int swizzle, + int custom_option, int tile, int splitk_val, + int reduction_scheme, int stages) { + CUBLAS_RETURN_IF_ERROR(cublasLtMatmulAlgoInit(cublasLt_handle, CUBLAS_COMPUTE_32I, CUDA_R_32F, + CUDA_R_8I, CUDA_R_8I, CUDA_R_8I, CUDA_R_8I, algo_id, &algo)); + + CUBLAS_RETURN_IF_ERROR(cublasLtMatmulAlgoConfigSetAttribute(&algo, + CUBLASLT_ALGO_CONFIG_CUSTOM_OPTION, + &(custom_option), sizeof(custom_option))); + + CUBLAS_RETURN_IF_ERROR(cublasLtMatmulAlgoConfigSetAttribute(&algo, + CUBLASLT_ALGO_CONFIG_TILE_ID, + &(tile), sizeof(tile))); + + CUBLAS_RETURN_IF_ERROR(cublasLtMatmulAlgoConfigSetAttribute(&algo, CUBLASLT_ALGO_CONFIG_SPLITK_NUM, + &(splitk_val), sizeof(splitk_val))); + + CUBLAS_RETURN_IF_ERROR(cublasLtMatmulAlgoConfigSetAttribute(&algo, CUBLASLT_ALGO_CONFIG_CTA_SWIZZLING, + &(swizzle), sizeof(swizzle))); + + CUBLAS_RETURN_IF_ERROR(cublasLtMatmulAlgoConfigSetAttribute(&algo, CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME, + &(reduction_scheme), sizeof(int))); + + CUBLAS_RETURN_IF_ERROR(cublasLtMatmulAlgoConfigSetAttribute(&algo, CUBLASLT_ALGO_CONFIG_STAGES_ID, + &(stages), sizeof(stages))); + + return Status::OK(); +} + +static inline std::string AlgoKey(const cudaDeviceProp& /*device_prop*/, + int batch_count, int m, int n, int k, + cublasLtOrder_t weight_order, cublasLtOrder_t input_output_order) { + std::stringstream ss; + ss << batch_count << "-" << m << "_" << n << "_" << k << "-" + << static_cast(weight_order) << "-" << static_cast(input_output_order); + return ss.str(); +} + +CublasLtMMAlgoMap& CublasLtMMAlgoMap::Instance() { + static CublasLtMMAlgoMap instance; + return instance; +} + +void CublasLtMMAlgoMap::GetAlgo(cublasLtHandle_t cublasLt_handle, cublasLtMatmulAlgo_t& algo, + const cudaDeviceProp& device_prop, + int batch_count, int m, int n, int k, + cublasLtOrder_t weight_order, + cublasLtOrder_t input_output_order) const { + ORT_ENFORCE(input_output_order == CUBLASLT_ORDER_ROW, "Input/Output should be ORDER_ROW"); + ORT_ENFORCE(weight_order == CUBLASLT_ORDER_COL, "Weight should be ORDER_COL"); + + const std::string& key = AlgoKey(device_prop, batch_count, m, n, k, weight_order, input_output_order); + + // TODO: Pre-load the best_algos_ from a config file which will be generated offline (or) + // Find the best algo dynamically in the warm-up run and cache it in `best_algos_` + auto algo_it = best_algos_.find(key); + if (algo_it != best_algos_.end() && algo_it->second.workspace_size == 0) { + const auto& algo_info = algo_it->second; + ORT_THROW_IF_ERROR(cublasLtMatMulInt8SetupAlgo(cublasLt_handle, algo, algo_info.algo_id, + algo_info.swizzle, algo_info.custom_option, + algo_info.tile, algo_info.splitk_val, + algo_info.reduction_scheme, algo_info.stages)); + } else { + // Default algo + int algo_id = 21; + int stages = 0; + ORT_THROW_IF_ERROR(cublasLtMatMulInt8SetupAlgo(cublasLt_handle, algo, algo_id, 0, 0, 20, 0, 0, stages)); + } +} + +static Status CreateLtMatrixLayout(cublasLtMatrixLayout_t& layout_desc, + int const batch_count, int64_t const rows_after_op, int64_t const cols_after_op, + cudaDataType_t const data_type, cublasLtOrder_t const mat_order, + cublasOperation_t const mat_trans) { + if (mat_trans == CUBLAS_OP_T) { + CUBLAS_RETURN_IF_ERROR(cublasLtMatrixLayoutCreate(&layout_desc, data_type, cols_after_op, rows_after_op, + CalcLeadingDimensionLt(cols_after_op, rows_after_op, mat_order))); + } else { + CUBLAS_RETURN_IF_ERROR(cublasLtMatrixLayoutCreate(&layout_desc, data_type, rows_after_op, cols_after_op, + CalcLeadingDimensionLt(rows_after_op, cols_after_op, mat_order))); + } + + CUBLAS_RETURN_IF_ERROR(cublasLtMatrixLayoutSetAttribute(layout_desc, + CUBLASLT_MATRIX_LAYOUT_ORDER, + &mat_order, sizeof(mat_order))); + + CUBLAS_RETURN_IF_ERROR(cublasLtMatrixLayoutSetAttribute(layout_desc, + CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, + &batch_count, sizeof(batch_count))); + + if (batch_count > 1) { + int64_t batch_stride = rows_after_op * cols_after_op; + CUBLAS_RETURN_IF_ERROR(cublasLtMatrixLayoutSetAttribute(layout_desc, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, + &batch_stride, sizeof(batch_stride))); + } + + return Status::OK(); +} + +Status QOrdered_MatMul(cublasLtHandle_t cublasLt_handle, cudaStream_t stream, + const cudaDeviceProp& device_prop, + int32_t batch_count, int64_t m, int64_t n, int64_t k, + const float* alpha, const int8_t* A, const int8_t* B, int32_t batch_B, + const float* bias, const float* beta, + const int8_t* C, int32_t batch_C, + int8_t* D, cublasLtOrder_t weight_order) { + ORT_ENFORCE(weight_order == CUBLASLT_ORDER_COL, "Weight should be ORDER_COL"); + + const cublasOperation_t transpose_op = CUBLAS_OP_T; + + cublasLtMatmulDesc_t matmul_desc = nullptr; + auto clean_matmul_desc = gsl::finally([&matmul_desc]() {if (matmul_desc) cublasLtMatmulDescDestroy(matmul_desc); }); + + cublasLtMatrixLayout_t desc_A = nullptr; + auto clean_desc_A = gsl::finally([&desc_A]() {if (desc_A) cublasLtMatrixLayoutDestroy(desc_A); }); + + cublasLtMatrixLayout_t desc_B = nullptr; + auto clean_desc_B = gsl::finally([&desc_B]() {if (desc_B) cublasLtMatrixLayoutDestroy(desc_B); }); + + cublasLtMatrixLayout_t desc_C = nullptr; + auto clean_desc_C = gsl::finally([&desc_C]() {if (desc_C) cublasLtMatrixLayoutDestroy(desc_C); }); + + cublasLtMatrixLayout_t desc_D = nullptr; + auto clean_desc_D = gsl::finally([&desc_D]() {if (desc_D) cublasLtMatrixLayoutDestroy(desc_D); }); + + const float beta_zero = 0.0f; + beta = (C == nullptr ? &beta_zero : beta); + + CUBLAS_RETURN_IF_ERROR(cublasLtMatmulDescCreate(&matmul_desc, CUBLAS_COMPUTE_32I, CUDA_R_32F)); + + CUBLAS_RETURN_IF_ERROR(cublasLtMatmulDescSetAttribute(matmul_desc, + CUBLASLT_MATMUL_DESC_TRANSA, + &transpose_op, sizeof(transpose_op))); + + if (bias != nullptr) { + cublasLtEpilogue_t epilogue_bias = CUBLASLT_EPILOGUE_BIAS; + CUBLAS_RETURN_IF_ERROR(cublasLtMatmulDescSetAttribute(matmul_desc, + CUBLASLT_MATMUL_DESC_EPILOGUE, + &epilogue_bias, sizeof(epilogue_bias))); + + CUBLAS_RETURN_IF_ERROR(cublasLtMatmulDescSetAttribute(matmul_desc, + CUBLASLT_MATMUL_DESC_BIAS_POINTER, + &bias, sizeof(bias))); + } + + ORT_RETURN_IF_ERROR(CreateLtMatrixLayout(desc_A, batch_count, k, m, CUDA_R_8I, + CUBLASLT_ORDER_COL, CUBLAS_OP_N)); // for A' + + ORT_RETURN_IF_ERROR(CreateLtMatrixLayout(desc_B, batch_B, k, n, CUDA_R_8I, + CUBLASLT_ORDER_COL, CUBLAS_OP_N)); // For B + + CUBLAS_RETURN_IF_ERROR(cublasLtMatrixLayoutSetAttribute(desc_B, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, + &batch_count, sizeof(batch_count))); + + ORT_RETURN_IF_ERROR(CreateLtMatrixLayout(desc_D, batch_count, n, m, + CUDA_R_8I, CUBLASLT_ORDER_COL, + CUBLAS_OP_N)); // For D' + + if (C != nullptr) { + ORT_RETURN_IF_ERROR(CreateLtMatrixLayout(desc_C, batch_C, n, m, CUDA_R_8I, + CUBLASLT_ORDER_COL, CUBLAS_OP_N)); // For C' + + CUBLAS_RETURN_IF_ERROR(cublasLtMatrixLayoutSetAttribute(desc_C, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, + &batch_count, sizeof(batch_count))); + } + + cublasLtMatmulAlgo_t algo; + CublasLtMMAlgoMap::Instance().GetAlgo(cublasLt_handle, algo, device_prop, batch_count, + static_cast(m), static_cast(n), static_cast(k), + weight_order, CUBLASLT_ORDER_ROW); + + CUBLAS_RETURN_IF_ERROR(cublasLtMatmul(cublasLt_handle, matmul_desc, alpha, + B, desc_B, + A, desc_A, + beta, + C == nullptr ? D : C, C == nullptr ? desc_D : desc_C, + D, desc_D, + &algo, nullptr, 0, // algo, workspace, workspace_size + stream)); + return Status::OK(); +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_matmul_utils.h b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_matmul_utils.h new file mode 100644 index 0000000000..143cdb19dd --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_matmul_utils.h @@ -0,0 +1,60 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/cuda/cuda_common.h" + +#include +#include + +namespace onnxruntime { + +namespace contrib { +namespace cuda { + +class CublasLtMMAlgoMap { + public: + static CublasLtMMAlgoMap& Instance(); + + void GetAlgo(cublasLtHandle_t cublasLt_handle, cublasLtMatmulAlgo_t& algo, + const cudaDeviceProp& device_prop, + int batch_count, int m, int n, int k, + cublasLtOrder_t weight_order, + cublasLtOrder_t input_output_order = CUBLASLT_ORDER_ROW) const; + + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(CublasLtMMAlgoMap); + + private: + CublasLtMMAlgoMap() = default; + + ~CublasLtMMAlgoMap() = default; + + struct CublasLtMatmulAlgoInfo { + int algo_id, custom_option, tile, splitk_val, + swizzle, reduction_scheme, workspace_size, stages; + + float exec_time; + }; + + std::unordered_map best_algos_; +}; + +Status QOrdered_MatMul(cublasLtHandle_t cublasLt_handle, cudaStream_t stream, + const cudaDeviceProp& device_prop, + int32_t batch_count, int64_t m, int64_t n, int64_t k, + const float* alpha, const int8_t* A, const int8_t* B, + int32_t batch_B, + const float* bias, const float* beta, + const int8_t* C, int32_t batch_C, + int8_t* D, cublasLtOrder_t weight_order); + +Status QOrdered_MatMul(cublasLtHandle_t cublasLt_handle, cudaStream_t stream, + const cudaDeviceProp& device_prop, + int32_t batch_count, int64_t m, int64_t n, int64_t k, + const float* alpha, const int8_t* A, const int8_t* B, + const float* bias, int8_t* D, cublasLtOrder_t order_weight); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_unary_ops.cc b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_unary_ops.cc new file mode 100644 index 0000000000..8cb710d0cb --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_unary_ops.cc @@ -0,0 +1,43 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "qordered_unary_ops.h" +#include "qordered_unary_ops_impl.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +ONNX_OPERATOR_KERNEL_EX( + QOrderedGelu, + kMSDomain, + 1, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("Q", DataTypeImpl::GetTensorType()) + .TypeConstraint("S", DataTypeImpl::GetTensorType()) + .InputMemoryType(OrtMemTypeCPUInput, 1) // scale_input + .InputMemoryType(OrtMemTypeCPUInput, 2), // scale_output + QOrderedGelu); + +Status QOrderedGelu::ComputeInternal(OpKernelContext* context) const { + const Tensor& input = *context->Input(0); + + const float* scale_input = context->Input(1)->Data(); + const float* scale_output = context->Input(2)->Data(); + + const auto& shape = input.Shape(); + + Tensor* output = context->Output(0, shape); + + QOrderedUnarySharedMemory_Gelu(Stream(), input.Data(), scale_input, output->MutableData(), + scale_output, static_cast(shape.Size())); + + return Status::OK(); +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_unary_ops.h b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_unary_ops.h new file mode 100644 index 0000000000..e32e9eb38d --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_unary_ops.h @@ -0,0 +1,23 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/cuda/cuda_kernel.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +class QOrderedGelu final : public CudaKernel { + public: + explicit QOrderedGelu(const OpKernelInfo& info) : CudaKernel(info) {} + + Status ComputeInternal(OpKernelContext* context) const override; +}; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_unary_ops_impl.cu b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_unary_ops_impl.cu new file mode 100644 index 0000000000..95dfc26036 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_unary_ops_impl.cu @@ -0,0 +1,79 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "qordered_unary_ops_impl.h" +#include "core/providers/cuda/cu_inc/common.cuh" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +constexpr int kNumLinePerThread = 4; +constexpr int kNumThreadsPerBlock = 256; +constexpr int kNumElementsPerBlockLine = sizeof(char4) * kNumThreadsPerBlock; +constexpr int kNumElementsPerBlock = sizeof(char4) * kNumLinePerThread * kNumThreadsPerBlock; + +template +__global__ void QOrderedUnaryElementWiseSharedMemoryKernel( + const int8_t* input_data, float input_scale, int8_t* output_data, + float inverse_output_scale, const FuncT& functor, CUDA_LONG N) { + __shared__ char table[256]; + + const int calc_id = static_cast(threadIdx.x) - 128; + float gelu_value = inverse_output_scale * functor(input_scale * calc_id); + gelu_value = fmaxf(-128.0f, fmin(127.0f, gelu_value)); + table[threadIdx.x] = static_cast(__float2int_rn(gelu_value)); + __syncthreads(); + + CUDA_LONG id = kNumElementsPerBlock * blockIdx.x + threadIdx.x * (CUDA_LONG)sizeof(char4); + +#pragma unroll + for (int line = 0; line < kNumLinePerThread; line++) { + if (id < N) { + char4 i4 = *(const char4*)(input_data + id); + i4.x = table[128 + i4.x]; + i4.y = table[128 + i4.y]; + i4.z = table[128 + i4.z]; + i4.w = table[128 + i4.w]; + *reinterpret_cast(output_data + id) = i4; + id += kNumElementsPerBlockLine; + } + } +} + +template +void QOrderedUnaryElementWiseSharedMemoryImpl( + cudaStream_t stream, + const int8_t* input_data, + const float* input_scale, + int8_t* output_data, + const float* output_scale, + const FuncT& func, + size_t count) { + if (count > 0) { + float inverse_output_scale = 1.0f / *output_scale; + int blocksPerGrid = static_cast(CeilDiv(count, kNumElementsPerBlock)); + + QOrderedUnaryElementWiseSharedMemoryKernel<<>>( + input_data, *input_scale, output_data, inverse_output_scale, func, static_cast(count)); + } +} + +struct QOrderedUnaryOpGelu { + static constexpr float sqrt2 = 1.4142135623730950488016887242097f; + + __device__ __inline__ float operator()(const float& x) const { + return x * 0.5f * (1.0f + erff(x / sqrt2)); + } +}; + +QORDERED_UNARY_OP_SHARED_MEMORY_DECLARATION(Gelu) { + QOrderedUnaryElementWiseSharedMemoryImpl( + stream, input_data, input_scale, output_data, output_scale, QOrderedUnaryOpGelu(), count); +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_unary_ops_impl.h b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_unary_ops_impl.h new file mode 100644 index 0000000000..8c1ad09021 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/qordered_ops/qordered_unary_ops_impl.h @@ -0,0 +1,25 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/cuda/cuda_common.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +#define QORDERED_UNARY_OP_SHARED_MEMORY_DECLARATION(name) \ + void QOrderedUnarySharedMemory_##name( \ + cudaStream_t stream, \ + const int8_t* input_data, \ + const float* input_scale, \ + int8_t* output_data, \ + const float* output_scale, \ + size_t count) + +QORDERED_UNARY_OP_SHARED_MEMORY_DECLARATION(Gelu); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index c6850eb8e3..23a54d325e 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -32,6 +32,9 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QLinearSigmoid); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QLinearSoftmax); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QuantizeLinear); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ReduceSumInteger); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QOrderedMatMul); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QOrderedGelu); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QOrderedLayerNormalization); //Others class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Attention); @@ -102,6 +105,9 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); + fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/core/graph/contrib_ops/quantization_defs.cc b/onnxruntime/core/graph/contrib_ops/quantization_defs.cc index 34fb5a16ad..ccc3dd3b93 100644 --- a/onnxruntime/core/graph/contrib_ops/quantization_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/quantization_defs.cc @@ -1131,5 +1131,89 @@ If mask is provided, mask index (that is position of first 0 in mask, or number .TypeConstraint("T2", {"tensor(int8)", "tensor(uint8)"}, "Constrain input and output types to int8 tensors.") .TypeConstraint("T", {"tensor(float)"}, "Constrain input and output types to float32 tensors.") .TypeAndShapeInferenceFunction(EmbedLayerNormalizationShapeInference)); + + constexpr const char* QOrderedMatMul_ver1_doc = R"DOC(TODO)DOC"; + + ONNX_MS_OPERATOR_SET_SCHEMA(QOrderedMatMul, 1, + OpSchema() + .SetDoc(QOrderedMatMul_ver1_doc) + .Attr("order_A", "cublasLt order of matrix A. Default is ROW MAJOR.", + AttributeProto::INT, static_cast(1)) + .Attr("order_B", "cublasLt order of matrix B. Default is ROW MAJOR.", + AttributeProto::INT, static_cast(1)) + .Attr("order_Y", + "cublasLt order of matrix Y and optional matrix C. Default is ROW MAJOR.", + AttributeProto::INT, static_cast(1)) + .Input(0, "A", "3-dimensional matrix A", "Q") + .Input(1, "scale_A", "scale of the input A", "S") + .Input(2, "B", "2-dimensional matrix B", "Q") + .Input(3, "scale_B", "scale of the input B", "S") + .Input(4, "scale_Y", "scale of the output Y", "S") + .Input(5, "bias", "1d bias", "S", OpSchema::Optional) + .Input(6, "C", + "3d or 2d matrix C. if 2d expand to 3d first. " + "Shape[0] should be 1 or same as A.shape[0] ", + "Q", OpSchema::Optional) + .Input(7, "scale_C", "scale of the input A", "S", OpSchema::Optional) + .Output(0, "Y", "Matrix multiply results from A * B", "Q") + .TypeConstraint("Q", {"tensor(int8)"}, + "Constrain input and output types to int8 tensors.") + .TypeConstraint("S", {"tensor(float)"}, "Constrain bias and scales to float32") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + ONNX_NAMESPACE::matmulShapeInference(ctx, 0, 2); + })); + + ONNX_MS_OPERATOR_SET_SCHEMA(QOrderedLayerNormalization, 1, + OpSchema() + .SetDoc("QOrderedLayerNormalization") + .Attr("axis", + "The first normalization dimension: normalization " + "will be performed along dimensions axis " + ": rank(inputs).", + AttributeProto::INT, + static_cast(-1)) + .Attr("epsilon", "The epsilon value to use to avoid division by zero.", + AttributeProto::FLOAT, 1e-5f) + .Attr("order_X", "cublasLt order of input X. Default is ROW MAJOR.", + AttributeProto::INT, static_cast(1)) + .Attr("order_Y", + "cublasLt order of matrix Y, must be same as order_X. " + "Default is ROW MAJOR.", + AttributeProto::INT, static_cast(1)) + .AllowUncheckedAttributes() + .Input(0, "X", "Input data tensor from the previous layer.", "Q") + .Input(1, "scale_X", "scale of the quantized X", "S") + .Input(2, "scale", "Scale tensor, i.e., gamma vector.", "F") + .Input(3, "B", "Bias tensor.", "F", OpSchema::Optional) + .Input(4, "scale_Y", "scale of the quantized X", "S") + .Output(0, "Y", "Output data tensor.", "Q") + .TypeConstraint("F", {"tensor(float16)", "tensor(float)"}, + "Constrain input gamma and bias could be float16/float tensors. " + "float may get better precision, float16 runs faster.") + .TypeConstraint("S", {"tensor(float)"}, "quantization scale must be float tensors.") + .TypeConstraint("Q", {"tensor(int8)"}, "quantization tensor must be int8 tensors.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateShapeAndTypeFromFirstInput(ctx); + propagateElemTypeFromInputToOutput(ctx, 0, 0); + })); + + ONNX_MS_OPERATOR_SET_SCHEMA(QOrderedGelu, 1, + OpSchema() + .SetDoc(R"DOC(Ordered Quantize Gelu.)DOC") + .Attr("order_X", "cublasLt order of input X. Default is ROW MAJOR.", + AttributeProto::INT, static_cast(1)) + .Attr("order_Y", + "cublasLt order of matrix Y, must be same as order_X. " + "Default is ROW MAJOR.", + AttributeProto::INT, static_cast(1)) + .Input(0, "X", "N-dimensional input A", "Q") + .Input(1, "scale_X", "scale of the input A", "S") + .Input(2, "scale_Y", "scale of the output Y", "S") + .Output(0, "Y", "Output of the Gelu", "Q") + .TypeConstraint("Q", {"tensor(int8)"}, + "Constrain input and output types to int8 tensors.") + .TypeConstraint("S", {"tensor(float)"}, "Constrain scales to float32") + .TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput)); } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index f18865498e..6a0c3283b0 100755 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -149,6 +149,7 @@ CUDAExecutionProvider::PerThreadContext::PerThreadContext(OrtDevice::DeviceId de stream_ = stream; CUBLAS_CALL_THROW(cublasCreate(&cublas_handle_)); + CUBLAS_CALL_THROW(cublasLtCreate(&cublas_lt_handle_)); CUBLAS_CALL_THROW(cublasSetStream(cublas_handle_, stream)); CUDNN_CALL_THROW(cudnnCreate(&cudnn_handle_)); @@ -164,6 +165,7 @@ CUDAExecutionProvider::PerThreadContext::PerThreadContext(OrtDevice::DeviceId de CUDAExecutionProvider::PerThreadContext::~PerThreadContext() { ORT_IGNORE_RETURN_VALUE(CUBLAS_CALL(cublasDestroy(cublas_handle_))); + ORT_IGNORE_RETURN_VALUE(CUBLAS_CALL(cublasLtDestroy(cublas_lt_handle_))); ORT_IGNORE_RETURN_VALUE(CUDNN_CALL(cudnnDestroy(cudnn_handle_))); } diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.h b/onnxruntime/core/providers/cuda/cuda_execution_provider.h index cb126f203b..806be894f6 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.h +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.h @@ -45,6 +45,10 @@ class CUDAExecutionProvider : public IExecutionProvider { return GetPerThreadContext().CublasHandle(); } + cublasLtHandle_t PerThreadCublasLtHandle() { + return GetPerThreadContext().CublasLtHandle(); + } + cudnnHandle_t PerThreadCudnnHandle() { return GetPerThreadContext().CudnnHandle(); } @@ -130,6 +134,10 @@ class CUDAExecutionProvider : public IExecutionProvider { return cudnn_handle_; } + cublasLtHandle_t CublasLtHandle() const { + return cublas_lt_handle_; + } + cudaEvent_t& GetCurrentDeferredReleaseEvent() { return current_deferred_release_event_; } @@ -166,18 +174,19 @@ class CUDAExecutionProvider : public IExecutionProvider { } #if defined(CUDA_VERSION) && CUDA_VERSION >= 10000 - bool IsGraphCaptureAllowed() const; - void CaptureBegin(); - void CaptureEnd(); - bool IsGraphCaptured() const; - Status ReplayGraph(); - void IncrementRegularRunCountBeforeGraphCapture(); + bool IsGraphCaptureAllowed() const; + void CaptureBegin(); + void CaptureEnd(); + bool IsGraphCaptured() const; + Status ReplayGraph(); + void IncrementRegularRunCountBeforeGraphCapture(); #endif private: cudaStream_t stream_ = nullptr; cublasHandle_t cublas_handle_ = nullptr; cudnnHandle_t cudnn_handle_ = nullptr; + cublasLtHandle_t cublas_lt_handle_ = nullptr; // deferred release for temporary CPU pinned memory used in cudaMemcpyAsync // note that cudaEvent will be assigned at OnRunEnd() when PerThreadContext destory @@ -197,10 +206,9 @@ class CUDAExecutionProvider : public IExecutionProvider { CUDAGraph cuda_graph_; bool is_graph_captured_ = false; int regular_run_count_before_graph_capture_ = 0; - const int min_num_runs_before_cuda_graph_capture_ = 1; // required min regular runs before graph capture for the necessary memory allocations. + const int min_num_runs_before_cuda_graph_capture_ = 1; // required min regular runs before graph capture for the necessary memory allocations. #endif - }; using PerThreadContextMap = std::unordered_map>; diff --git a/onnxruntime/core/providers/cuda/cuda_kernel.h b/onnxruntime/core/providers/cuda/cuda_kernel.h index 7ea2a07ae0..2ce72d7112 100644 --- a/onnxruntime/core/providers/cuda/cuda_kernel.h +++ b/onnxruntime/core/providers/cuda/cuda_kernel.h @@ -135,6 +135,10 @@ class CudaKernel : public OpKernel { return provider_->PerThreadCublasHandle(); } + inline cublasLtHandle_t CublasLtHandle() const { + return provider_->PerThreadCublasLtHandle(); + } + inline cudnnHandle_t CudnnHandle() const { return provider_->PerThreadCudnnHandle(); } diff --git a/onnxruntime/core/providers/cuda/cuda_pch.h b/onnxruntime/core/providers/cuda/cuda_pch.h index a671b91fb1..bdb0e78d77 100644 --- a/onnxruntime/core/providers/cuda/cuda_pch.h +++ b/onnxruntime/core/providers/cuda/cuda_pch.h @@ -15,6 +15,7 @@ #include #include #include +#include #ifdef ORT_USE_NCCL #include diff --git a/onnxruntime/test/contrib_ops/qordered_gelu_op_test.cc b/onnxruntime/test/contrib_ops/qordered_gelu_op_test.cc new file mode 100644 index 0000000000..15e36233e0 --- /dev/null +++ b/onnxruntime/test/contrib_ops/qordered_gelu_op_test.cc @@ -0,0 +1,47 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "qordered_test_utils.h" + +namespace onnxruntime { +namespace test { + +// Only the CUDA EP supports ordered quantized ops for now +#ifdef USE_CUDA + +static void +RunQOrdered_Gelu_Test(std::vector const& shape, float scale_x, + float scale_y, OrderCublasLt order) { + static const float sqrt_of_2 = std::sqrt(2.0f); + + int64_t N = std::accumulate(shape.begin(), shape.end(), int64_t{1LL}, std::multiplies()); + std::vector vec_x = GenData(shape, 1.0f); + std::vector vec_y(N); + for (int64_t i = 0; i < N; i++) { + float x = scale_x * static_cast(vec_x[i]); + float r = (x * (0.5f * (1.0f + std::erff(x / sqrt_of_2)))) / scale_y; + int8_t q = static_cast(std::nearbyintf(std::min(127.0f, std::max(-128.0f, r)))); + vec_y[i] = q; + } + + OpTester test_qorder("QOrderedGelu", 1, onnxruntime::kMSDomain); + test_qorder.AddAttribute("order_X", static_cast(order)); + test_qorder.AddAttribute("order_Y", static_cast(order)); + test_qorder.AddInput("X", shape, vec_x); + test_qorder.AddInput("scale_X", {}, {scale_x}); + test_qorder.AddInput("scale_Y", {}, {scale_y}); + test_qorder.AddOutput("Y", shape, vec_y, false, 0.0f, 0.0f /* abs error */); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test_qorder.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST(QOrderedTest, Gelu_3x11x12) { + RunQOrdered_Gelu_Test({3, 11, 12}, 1.0f / 32.0f, 1.0f / 128.0f, ORDER_ROW); +} + +#endif + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/qordered_layer_norm_op_test.cc b/onnxruntime/test/contrib_ops/qordered_layer_norm_op_test.cc new file mode 100644 index 0000000000..824f4324dd --- /dev/null +++ b/onnxruntime/test/contrib_ops/qordered_layer_norm_op_test.cc @@ -0,0 +1,112 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "qordered_test_utils.h" + +namespace onnxruntime { +namespace test { + +// Only the CUDA EP supports ordered quantized ops for now +#ifdef USE_CUDA + +template // MLFloat16 or float +static void RunQOrdered_LayerNorm_RowMajor(std::vector const& shape, int axis, float epsilon, + const std::vector& vec_x, float scale_x, + const std::vector& gamma, const std::vector* beta, + float scale_y, const std::vector& vec_y) { + std::vector bias_shape = {shape.back()}; + OpTester test_qorder("QOrderedLayerNormalization", 1, onnxruntime::kMSDomain); + test_qorder.AddAttribute("axis", (int64_t)axis); + test_qorder.AddAttribute("epsilon", epsilon); + test_qorder.AddInput("X", shape, vec_x); + test_qorder.AddInput("scale_X", {}, {scale_x}); + test_qorder.AddInput("scale", bias_shape, gamma); + if (beta) { + test_qorder.AddInput("B", bias_shape, *beta); + } else { + test_qorder.AddOptionalInputEdge(); + } + test_qorder.AddInput("scale_Y", {}, {scale_y}); + test_qorder.AddOutput("Y", shape, vec_y, false, 0.0f, 0.0f /* abs error */); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test_qorder.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST(QOrderedTest, LayerNormalization_RowMajor) { + float scale_x = 1.0; + + int64_t batch = 2; + int64_t sequence = 2; + int64_t hidden = 32; + + std::vector vec_x = { + -103, -24, -11, 125, 103, -117, 44, -109, 27, 122, 113, + 0, -7, 26, 110, -34, 28, -61, -107, -35, -54, 69, + -3, -58, -119, 61, -19, -115, 35, 76, 16, 54, + + -103, -24, -11, 125, 103, -117, 44, -109, 27, 122, 113, + 0, -7, 26, 110, -34, 28, -61, -107, -35, -54, 69, + -3, -58, -119, 61, -19, -115, 35, 76, 16, 54, + + -103, -24, -11, 125, 103, -117, 44, -109, 27, 122, 113, + 0, -7, 26, 110, -34, 28, -61, -107, -35, -54, 69, + -3, -58, -119, 61, -19, -115, 35, 76, 16, 54, + + -103, -24, -11, 125, 103, -117, 44, -109, 27, 122, 113, + 0, -7, 26, 110, -34, 28, -61, -107, -35, -54, 69, + -3, -58, -119, 61, -19, -115, 35, 76, 16, 54}; + + std::vector gamma_fp32 = { + 0.9536133f, 0.9160156f, 1.3310547f, 1.0068359f, 0.8095703f, + 1.0126953f, 0.8876953f, 1.4316406f, 1.0947266f, 1.0498047f, + 0.7373047f, 1.0615234f, 1.015625f, 1.0751953f, 1.0068359f, + 1.0908203f, 1.2011719f, 1.1962891f, 0.91796875f, 1.0947266f, + 1.3183594f, 1.0185547f, 1.0791016f, 1.0273438f, 0.8364258f, + 0.94873047f, 1.0292969f, 1.09375f, 1.0371094f, 1.1240234f, + 1.4384766f, 1.0068359f}; + + auto gamma_fp16 = ToFloat16(gamma_fp32); + + std::vector beta_fp32 = { + 0.01411438f, 0.18273926f, -0.12414551f, 0.09887695f, -0.1114502f, + -0.08227539f, -0.04598999f, -0.11322021f, 0.12731934f, -0.06591797f, + -0.00662994f, 0.04962158f, -0.04281616f, 0.07476807f, 0.23010254f, + 0.1036377f, 0.10852051f, 0.10919189f, -0.02905273f, -0.0512085f, + -0.1194458f, 0.02661133f, 0.05789185f, -0.05239868f, 0.17907715f, + -0.01765442f, -0.12255859f, -0.09729004f, 0.06591797f, 0.02258301f, + -0.01844788f, -0.11999512f}; + + auto beta_fp16 = ToFloat16(beta_fp32); + + std::vector vec_y = { + -84, -8, -22, 114, 64, -108, 30, -128, 33, 105, 71, + 2, -10, 28, 109, -26, 35, -57, -87, -37, -70, 61, + 0, -56, -75, 48, -26, -115, 35, 74, 17, 38, + + -84, -8, -22, 114, 64, -108, 30, -128, 33, 105, 71, + 2, -10, 28, 109, -26, 35, -57, -87, -37, -70, 61, + 0, -56, -75, 48, -26, -115, 35, 74, 17, 38, + + -84, -8, -22, 114, 64, -108, 30, -128, 33, 105, 71, + 2, -10, 28, 109, -26, 35, -57, -87, -37, -70, 61, + 0, -56, -75, 48, -26, -115, 35, 74, 17, 38, + + -84, -8, -22, 114, 64, -108, 30, -128, 33, 105, 71, + 2, -10, 28, 109, -26, 35, -57, -87, -37, -70, 61, + 0, -56, -75, 48, -26, -115, 35, 74, 17, 38}; + + float scale_y = 1.0 / 64.0f; + + RunQOrdered_LayerNorm_RowMajor({batch, sequence, hidden}, -1, 0.00001f, vec_x, scale_x, + gamma_fp16, &beta_fp16, scale_y, vec_y); + + RunQOrdered_LayerNorm_RowMajor({batch, sequence, hidden}, -1, 0.00001f, vec_x, scale_x, + gamma_fp32, &beta_fp32, scale_y, vec_y); +} + +#endif + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/qordered_matmul_op_test.cc b/onnxruntime/test/contrib_ops/qordered_matmul_op_test.cc new file mode 100644 index 0000000000..a6fa1c7e53 --- /dev/null +++ b/onnxruntime/test/contrib_ops/qordered_matmul_op_test.cc @@ -0,0 +1,209 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "qordered_test_utils.h" + +namespace onnxruntime { +namespace test { + +// Only the CUDA EP supports ordered quantized ops for now +#ifdef USE_CUDA + +static void RunQOrdered_MatMul_Test( + std::vector const& shape_A, + std::vector const& shape_B, + std::vector const& shape_Y, + OrderCublasLt weight_order, + float scale_A, float scale_B, float scale_C, float scale_Y, + bool add_bias = false, bool broadcast_c_batch = false) { + // Needs Turing or higher architecture + if (NeedSkipIfCudaArchLowerThan(750)) { + return; + } + + scale_A = MLFloat16(scale_A).ToFloat(); + scale_B = MLFloat16(scale_B).ToFloat(); + scale_C = MLFloat16(scale_C).ToFloat(); + scale_Y = MLFloat16(scale_Y).ToFloat(); + int64_t num_elements_Y = std::accumulate(shape_Y.begin(), shape_Y.end(), int64_t{1LL}, std::multiplies()); + std::vector vec_A = GenData(shape_A, 1.0f); + std::vector vec_B = GenData(shape_B, 1.0f); + std::vector vec_Y(num_elements_Y); + + int64_t cols_A = 0, rows_A = 0, batch_A = 0; + BatchRowColFromShape(shape_A, batch_A, rows_A, cols_A); + OrderedIndex indexerA(ORDER_ROW, rows_A, cols_A); + + int64_t cols_B = 0, rows_B = 0, batch_B = 0; + BatchRowColFromShape(shape_B, batch_B, rows_B, cols_B); + OrderedIndex indexerB(ORDER_ROW, cols_B, rows_B); // B need Transpose + + int64_t cols_Y = 0, rows_Y = 0, batch_Y = 0; + BatchRowColFromShape(shape_Y, batch_Y, rows_Y, cols_Y); + OrderedIndex indexerY(ORDER_ROW, rows_Y, cols_Y); + + std::vector bias_shape = {cols_Y}; + std::vector bias; + if (add_bias) { // make bias not too big + float bias_scale = MLFloat16(scale_Y / 27.0f).ToFloat(); + bias = GenData(bias_shape, bias_scale); + } + std::vector vec_C; + std::vector shape_C = {broadcast_c_batch ? 1 : batch_Y, rows_Y, cols_Y}; + if (scale_C != 0.0f) { + vec_C = GenData(shape_C, 1.0f); + } + + // TODO: broadcasting higher dims + float alpha = scale_A * scale_B / scale_Y; + float beta = scale_C / scale_Y; + int8_t const* A = vec_A.data(); + int8_t const* B = vec_B.data(); + int8_t* Y = vec_Y.data(); + int8_t* C = vec_C.data(); + for (int64_t b = 0; b < batch_Y; b++) { + for (int64_t m = 0; m < rows_A; m++) { + for (int64_t n = 0; n < cols_B; n++) { + int32_t isum = 0; + for (int64_t k = 0; k < cols_A; k++) { + auto posA = indexerA(m, k); + auto posB = indexerB(n, k); // Transpose B + isum += (A[posA] * B[posB]); + } + float sum = alpha * isum; + if (add_bias) { + sum += bias[n]; + } + auto posY = indexerY(m, n); + if (scale_C != 0.0f) { + sum += beta * C[posY]; + } + Y[posY] = static_cast(std::nearbyintf(std::min(127.0f, std::max(-128.0f, sum)))); + } + } + A += (batch_A <= 1 ? int64_t{0} : (rows_A * cols_A)); + B += (batch_B <= 1 ? int64_t{0} : (rows_B * cols_B)); + Y += (batch_Y <= 1 ? int64_t{0} : (rows_Y * cols_Y)); + C += (shape_C[0] == 1 ? int64_t{0} : (rows_Y * cols_Y)); + } + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + + OpTester test_qorder("QOrderedMatMul", 1, onnxruntime::kMSDomain); + test_qorder.AddAttribute("order_B", static_cast(weight_order)); + test_qorder.AddAttribute("order_A", static_cast(ORDER_ROW)); + test_qorder.AddAttribute("order_Y", static_cast(ORDER_ROW)); + test_qorder.AddInput("A", shape_A, vec_A); + test_qorder.AddInput("scale_A", {}, {scale_A}); + test_qorder.AddInput("B", shape_B, vec_B); + test_qorder.AddInput("scale_B", {}, {scale_B}); + test_qorder.AddInput("scale_Y", {}, {scale_Y}); + if (add_bias) { + test_qorder.AddInput("bias", bias_shape, bias); + } else { + test_qorder.AddOptionalInputEdge(); + } + + if (scale_C != 0.0f) { + test_qorder.AddInput("C", shape_C, vec_C); + test_qorder.AddInput("scale_C", {}, {scale_C}); + } + test_qorder.AddOutput("Y", shape_Y, vec_Y, false, 0.0f, 0.0f /* abs error */); + test_qorder.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST(QOrderedTest, MatMul_COL_16x64x32) { + std::vector shape_A = {16, 32}; + std::vector shape_B = {32, 64}; + std::vector shape_Y = {16, 64}; + RunQOrdered_MatMul_Test(shape_A, shape_B, shape_Y, ORDER_COL, + 1.0f / 32.0f, 1.0f / 32.0f, 0.0f /*scaleC*/, 2.0f, + false /* add bias */, false /* broadcast batch c */); +} + +TEST(QOrderedTest, MatMul_bias_COL_16x64x32) { + std::vector shape_A = {16, 32}; + std::vector shape_B = {32, 64}; + std::vector shape_Y = {16, 64}; + RunQOrdered_MatMul_Test(shape_A, shape_B, shape_Y, ORDER_COL, + 1.0f / 32.0f, 1.0f / 32.0f, 0.0f /*scaleC*/, 2.0f, + true /* add bias */, false /* broadcast batch c */); +} + +TEST(QOrderedTest, MatMul_addC_COL_16x64x32) { + std::vector shape_A = {16, 32}; + std::vector shape_B = {32, 64}; + std::vector shape_Y = {16, 64}; + RunQOrdered_MatMul_Test(shape_A, shape_B, shape_Y, ORDER_COL, + 1.0f / 32.0f, 1.0f / 32.0f, 4.0f /*scaleC*/, 2.0f, + false /* add bias */, false /* broadcast batch c */); +} + +TEST(QOrderedTest, MatMul_bias_addC_COL_16x64x32) { + std::vector shape_A = {16, 32}; + std::vector shape_B = {32, 64}; + std::vector shape_Y = {16, 64}; + RunQOrdered_MatMul_Test(shape_A, shape_B, shape_Y, ORDER_COL, + 1.0f / 32.0f, 1.0f / 32.0f, 4.0f /*scaleC*/, 2.0f, + true /* add bias */, true /* broadcast batch c */); +} + +TEST(QOrderedTest, MatMul_COL_16x64x32_b3_1) { + std::vector shape_A = {3, 16, 32}; + std::vector shape_B = {1, 32, 64}; + std::vector shape_Y = {3, 16, 64}; + RunQOrdered_MatMul_Test(shape_A, shape_B, shape_Y, ORDER_COL, + 1.0f / 32.0f, 1.0f / 32.0f, 0.0f /*scaleC*/, 2.0f, + false /* add bias */, false /* broadcast batch c */); +} + +TEST(QOrderedTest, MatMul_bias_COL_16x64x32_b2_1) { + std::vector shape_A = {2, 16, 32}; + std::vector shape_B = {1, 32, 64}; + std::vector shape_Y = {2, 16, 64}; + RunQOrdered_MatMul_Test(shape_A, shape_B, shape_Y, ORDER_COL, + 1.0f / 32.0f, 1.0f / 32.0f, 0.0f /*scaleC*/, 2.0f, + true /* add bias */, false /* broadcast batch c */); +} + +TEST(QOrderedTest, MatMul_addC_COL_16x64x32_b2_1) { + std::vector shape_A = {2, 16, 32}; + std::vector shape_B = {1, 32, 64}; + std::vector shape_Y = {2, 16, 64}; + RunQOrdered_MatMul_Test(shape_A, shape_B, shape_Y, ORDER_COL, + 1.0f / 32.0f, 1.0f / 32.0f, 4.0f /*scaleC*/, 2.0f, + false /* add bias */, false /* broadcast batch c */); +} + +TEST(QOrderedTest, MatMul_addC_broadcastC_COL_16x64x32_b2_1) { + std::vector shape_A = {2, 16, 32}; + std::vector shape_B = {1, 32, 64}; + std::vector shape_Y = {2, 16, 64}; + RunQOrdered_MatMul_Test(shape_A, shape_B, shape_Y, ORDER_COL, + 1.0f / 32.0f, 1.0f / 32.0f, 0.0f /*scaleC*/, 2.0f, + false /* add bias */, true /* broadcast batch c */); +} + +TEST(QOrderedTest, MatMul_addC_bias_COL_16x64x32_b2_1) { + std::vector shape_A = {2, 16, 32}; + std::vector shape_B = {1, 32, 64}; + std::vector shape_Y = {2, 16, 64}; + RunQOrdered_MatMul_Test(shape_A, shape_B, shape_Y, ORDER_COL, + 1.0f / 32.0f, 1.0f / 32.0f, 0.0f /*scaleC*/, 2.0f, + true /* add bias */, false /* broadcast batch c */); +} + +TEST(QOrderedTest, MatMul_bias_addC_broadcastC_COL_16x64x32_b2_1) { + std::vector shape_A = {2, 16, 32}; + std::vector shape_B = {1, 32, 64}; + std::vector shape_Y = {2, 16, 64}; + RunQOrdered_MatMul_Test(shape_A, shape_B, shape_Y, ORDER_COL, + 1.0f / 32.0f, 1.0f / 32.0f, 0.0f /*scaleC*/, 2.0f, + true /* add bias */, true /* broadcast batch c */); +} + +#endif + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/qordered_test_utils.h b/onnxruntime/test/contrib_ops/qordered_test_utils.h new file mode 100644 index 0000000000..f60a48611f --- /dev/null +++ b/onnxruntime/test/contrib_ops/qordered_test_utils.h @@ -0,0 +1,65 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "test/common/tensor_op_test_utils.h" +#include "test/common/cuda_op_test_utils.h" +#include "test/providers/provider_test_utils.h" + +#include +#include + +namespace onnxruntime { +namespace test { + +enum OrderCublasLt { + ORDER_COL = 0, + ORDER_ROW = 1, +}; + +class OrderedIndex { + OrderCublasLt order_; + int64_t rows_; + int64_t cols_; + + public: + OrderedIndex(OrderCublasLt order, int64_t rows, int64_t cols) + : order_(order), rows_(rows), cols_(cols) {} + + int64_t operator()(int64_t r, int64_t c) { + switch (order_) { + case ORDER_ROW: + return r * cols_ + c; + case ORDER_COL: + return c * rows_ + r; + default: + return 0; + } + } +}; + +template +inline std::vector GenData(std::vector const& shape, float scale) { + int64_t n = std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies()); + + scale = std::is_same::value ? 1.0f : scale; // using scale = 1.0f to generate int8_t data, + std::vector r(n); + RandomValueGenerator random{}; + std::vector tmp = random.Uniform(shape, -128, 127); + for (int64_t i = 0; i < n; i++) { + r[i] = static_cast(tmp[i] * scale); + } + return r; +} + +inline void BatchRowColFromShape(std::vector const& shape, int64_t& batch, int64_t& rows, int64_t& cols) { + cols = shape.back(); + rows = (shape.size() > 1 ? shape[shape.size() - 2] : 1LL); + batch = (shape.size() <= 2) + ? 1LL + : std::accumulate(shape.data(), shape.data() + (shape.size() - 2), 1LL, std::multiplies()); +} + +} // namespace test +} // namespace onnxruntime diff --git a/tools/ci_build/amd_hipify.py b/tools/ci_build/amd_hipify.py index 1310e15629..3302ee1693 100644 --- a/tools/ci_build/amd_hipify.py +++ b/tools/ci_build/amd_hipify.py @@ -47,6 +47,19 @@ contrib_ops_excluded_files = [ "quantization/attention_quantization_impl.cu", "quantization/attention_quantization_impl.cuh", "quantization/quantize_dequantize_linear.cc", + "quantization/qordered_ops/qordered_common.cuh", + "quantization/qordered_ops/qordered_layer_norm.h", + "quantization/qordered_ops/qordered_layer_norm.cc", + "quantization/qordered_ops/qordered_layer_norm_impl.h", + "quantization/qordered_ops/qordered_layer_norm_impl.cu", + "quantization/qordered_ops/qordered_matmul.h", + "quantization/qordered_ops/qordered_matmul.cc", + "quantization/qordered_ops/qordered_matmul_utils.h", + "quantization/qordered_ops/qordered_matmul_utils.cc", + "quantization/qordered_ops/qordered_unary_ops.h", + "quantization/qordered_ops/qordered_unary_ops.cc", + "quantization/qordered_ops/qordered_unary_ops_impl.h", + "quantization/qordered_ops/qordered_unary_ops_impl.cu", "tensor/crop.cc", "tensor/crop.h", "tensor/crop_impl.cu",