mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
Introduce ordered quantization ops for the CUDA EP [1/n] (#12582)
Initial core small set for the ordered quantization ops for cuda EP.
This commit is contained in:
Hariharan Seshadri
GitHub
parent
69f7cc6494
commit
ad69aac491
28 changed files with 1626 additions and 9 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -56,6 +56,9 @@ Do not modify directly.*
|
|||
* <a href="#com.microsoft.QLinearReduceMean">com.microsoft.QLinearReduceMean</a>
|
||||
* <a href="#com.microsoft.QLinearSigmoid">com.microsoft.QLinearSigmoid</a>
|
||||
* <a href="#com.microsoft.QLinearSoftmax">com.microsoft.QLinearSoftmax</a>
|
||||
* <a href="#com.microsoft.QOrderedGelu">com.microsoft.QOrderedGelu</a>
|
||||
* <a href="#com.microsoft.QOrderedLayerNormalization">com.microsoft.QOrderedLayerNormalization</a>
|
||||
* <a href="#com.microsoft.QOrderedMatMul">com.microsoft.QOrderedMatMul</a>
|
||||
* <a href="#com.microsoft.QuantizeLinear">com.microsoft.QuantizeLinear</a>
|
||||
* <a href="#com.microsoft.Range">com.microsoft.Range</a>
|
||||
* <a href="#com.microsoft.ReduceSumInteger">com.microsoft.ReduceSumInteger</a>
|
||||
|
|
@ -2862,6 +2865,163 @@ This version of the operator has been available since version 1 of the 'com.micr
|
|||
</dl>
|
||||
|
||||
|
||||
### <a name="com.microsoft.QOrderedGelu"></a><a name="com.microsoft.qorderedgelu">**com.microsoft.QOrderedGelu**</a>
|
||||
|
||||
Ordered Quantize Gelu.
|
||||
|
||||
#### Version
|
||||
|
||||
This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
|
||||
|
||||
#### Attributes
|
||||
|
||||
<dl>
|
||||
<dt><tt>order_X</tt> : int</dt>
|
||||
<dd>cublasLt order of input X. Default is ROW MAJOR.</dd>
|
||||
<dt><tt>order_Y</tt> : int</dt>
|
||||
<dd>cublasLt order of matrix Y, must be same as order_X. Default is ROW MAJOR.</dd>
|
||||
</dl>
|
||||
|
||||
#### Inputs
|
||||
|
||||
<dl>
|
||||
<dt><tt>X</tt> : Q</dt>
|
||||
<dd>N-dimensional input A</dd>
|
||||
<dt><tt>scale_X</tt> : S</dt>
|
||||
<dd>scale of the input A</dd>
|
||||
<dt><tt>scale_Y</tt> : S</dt>
|
||||
<dd>scale of the output Y</dd>
|
||||
</dl>
|
||||
|
||||
#### Outputs
|
||||
|
||||
<dl>
|
||||
<dt><tt>Y</tt> : Q</dt>
|
||||
<dd>Output of the Gelu</dd>
|
||||
</dl>
|
||||
|
||||
#### Type Constraints
|
||||
|
||||
<dl>
|
||||
<dt><tt>Q</tt> : tensor(int8)</dt>
|
||||
<dd>Constrain input and output types to int8 tensors.</dd>
|
||||
<dt><tt>S</tt> : tensor(float)</dt>
|
||||
<dd>Constrain scales to float32</dd>
|
||||
</dl>
|
||||
|
||||
|
||||
### <a name="com.microsoft.QOrderedLayerNormalization"></a><a name="com.microsoft.qorderedlayernormalization">**com.microsoft.QOrderedLayerNormalization**</a>
|
||||
|
||||
QOrderedLayerNormalization
|
||||
|
||||
#### Version
|
||||
|
||||
This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
|
||||
|
||||
#### Attributes
|
||||
|
||||
<dl>
|
||||
<dt><tt>axis</tt> : int</dt>
|
||||
<dd>The first normalization dimension: normalization will be performed along dimensions axis : rank(inputs).</dd>
|
||||
<dt><tt>epsilon</tt> : float</dt>
|
||||
<dd>The epsilon value to use to avoid division by zero.</dd>
|
||||
<dt><tt>order_X</tt> : int</dt>
|
||||
<dd>cublasLt order of input X. Default is ROW MAJOR.</dd>
|
||||
<dt><tt>order_Y</tt> : int</dt>
|
||||
<dd>cublasLt order of matrix Y, must be same as order_X. Default is ROW MAJOR.</dd>
|
||||
</dl>
|
||||
|
||||
#### Inputs
|
||||
|
||||
<dl>
|
||||
<dt><tt>X</tt> : Q</dt>
|
||||
<dd>Input data tensor from the previous layer.</dd>
|
||||
<dt><tt>scale_X</tt> : S</dt>
|
||||
<dd>scale of the quantized X</dd>
|
||||
<dt><tt>scale</tt> : F</dt>
|
||||
<dd>Scale tensor, i.e., gamma vector.</dd>
|
||||
<dt><tt>B</tt> (optional) : F</dt>
|
||||
<dd>Bias tensor.</dd>
|
||||
<dt><tt>scale_Y</tt> : S</dt>
|
||||
<dd>scale of the quantized X</dd>
|
||||
</dl>
|
||||
|
||||
#### Outputs
|
||||
|
||||
<dl>
|
||||
<dt><tt>Y</tt> : Q</dt>
|
||||
<dd>Output data tensor.</dd>
|
||||
</dl>
|
||||
|
||||
#### Type Constraints
|
||||
|
||||
<dl>
|
||||
<dt><tt>F</tt> : tensor(float16), tensor(float)</dt>
|
||||
<dd>Constrain input gamma and bias could be float16/float tensors. float may get better precision, float16 runs faster.</dd>
|
||||
<dt><tt>S</tt> : tensor(float)</dt>
|
||||
<dd>quantization scale must be float tensors.</dd>
|
||||
<dt><tt>Q</tt> : tensor(int8)</dt>
|
||||
<dd>quantization tensor must be int8 tensors.</dd>
|
||||
</dl>
|
||||
|
||||
|
||||
### <a name="com.microsoft.QOrderedMatMul"></a><a name="com.microsoft.qorderedmatmul">**com.microsoft.QOrderedMatMul**</a>
|
||||
|
||||
TODO
|
||||
|
||||
#### Version
|
||||
|
||||
This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
|
||||
|
||||
#### Attributes
|
||||
|
||||
<dl>
|
||||
<dt><tt>order_A</tt> : int</dt>
|
||||
<dd>cublasLt order of matrix A. Default is ROW MAJOR.</dd>
|
||||
<dt><tt>order_B</tt> : int</dt>
|
||||
<dd>cublasLt order of matrix B. Default is ROW MAJOR.</dd>
|
||||
<dt><tt>order_Y</tt> : int</dt>
|
||||
<dd>cublasLt order of matrix Y and optional matrix C. Default is ROW MAJOR.</dd>
|
||||
</dl>
|
||||
|
||||
#### Inputs (5 - 8)
|
||||
|
||||
<dl>
|
||||
<dt><tt>A</tt> : Q</dt>
|
||||
<dd>3-dimensional matrix A</dd>
|
||||
<dt><tt>scale_A</tt> : S</dt>
|
||||
<dd>scale of the input A</dd>
|
||||
<dt><tt>B</tt> : Q</dt>
|
||||
<dd>2-dimensional matrix B</dd>
|
||||
<dt><tt>scale_B</tt> : S</dt>
|
||||
<dd>scale of the input B</dd>
|
||||
<dt><tt>scale_Y</tt> : S</dt>
|
||||
<dd>scale of the output Y</dd>
|
||||
<dt><tt>bias</tt> (optional) : S</dt>
|
||||
<dd>1d bias</dd>
|
||||
<dt><tt>C</tt> (optional) : Q</dt>
|
||||
<dd>3d or 2d matrix C. if 2d expand to 3d first. Shape[0] should be 1 or same as A.shape[0] </dd>
|
||||
<dt><tt>scale_C</tt> (optional) : S</dt>
|
||||
<dd>scale of the input A</dd>
|
||||
</dl>
|
||||
|
||||
#### Outputs
|
||||
|
||||
<dl>
|
||||
<dt><tt>Y</tt> : Q</dt>
|
||||
<dd>Matrix multiply results from A * B</dd>
|
||||
</dl>
|
||||
|
||||
#### Type Constraints
|
||||
|
||||
<dl>
|
||||
<dt><tt>Q</tt> : tensor(int8)</dt>
|
||||
<dd>Constrain input and output types to int8 tensors.</dd>
|
||||
<dt><tt>S</tt> : tensor(float)</dt>
|
||||
<dd>Constrain bias and scales to float32</dd>
|
||||
</dl>
|
||||
|
||||
|
||||
### <a name="com.microsoft.QuantizeLinear"></a><a name="com.microsoft.quantizelinear">**com.microsoft.QuantizeLinear**</a>
|
||||
|
||||
The linear quantization operator. It consumes a full precision data, a scale, a zero point to compute the low precision / quantized tensor.
|
||||
|
|
|
|||
|
|
@ -780,6 +780,9 @@ Do not modify directly.*
|
|||
|LongformerAttention|*in* input:**T**<br> *in* weight:**T**<br> *in* bias:**T**<br> *in* mask:**T**<br> *in* global_weight:**T**<br> *in* global_bias:**T**<br> *in* global:**G**<br> *out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
|
||||
|NGramRepeatBlock|*in* input_ids:**Tid**<br> *in* scores:**T**<br> *out* scores_out:**T**|1+|**T** = tensor(float)<br/> **Tid** = tensor(int64)|
|
||||
|QAttention|*in* input:**T1**<br> *in* weight:**T2**<br> *in* bias:**T3**<br> *in* input_scale:**T3**<br> *in* weight_scale:**T3**<br> *in* mask_index:**T4**<br> *in* input_zero_point:**T1**<br> *in* weight_zero_point:**T2**<br> *in* past:**T3**<br> *out* output:**T3**<br> *out* present:**T3**|1+|**T1** = tensor(int8)<br/> **T2** = tensor(int8)<br/> **T3** = tensor(float), tensor(float16)<br/> **T4** = tensor(int32)|
|
||||
|QOrderedGelu|*in* X:**Q**<br> *in* scale_X:**S**<br> *in* scale_Y:**S**<br> *out* Y:**Q**|1+|**Q** = tensor(int8)<br/> **S** = tensor(float)|
|
||||
|QOrderedLayerNormalization|*in* X:**Q**<br> *in* scale_X:**S**<br> *in* scale:**F**<br> *in* B:**F**<br> *in* scale_Y:**S**<br> *out* Y:**Q**|1+|**F** = tensor(float), tensor(float16)<br/> **Q** = tensor(int8)<br/> **S** = tensor(float)|
|
||||
|QOrderedMatMul|*in* A:**Q**<br> *in* scale_A:**S**<br> *in* B:**Q**<br> *in* scale_B:**S**<br> *in* scale_Y:**S**<br> *in* bias:**S**<br> *in* C:**Q**<br> *in* scale_C:**S**<br> *out* Y:**Q**|1+|**Q** = tensor(int8)<br/> **S** = tensor(float)|
|
||||
|QuantizeLinear|*in* x:**T1**<br> *in* y_scale:**T1**<br> *in* y_zero_point:**T2**<br> *out* y:**T2**|1+|**T1** = tensor(float16)<br/> **T2** = tensor(int8), tensor(uint8)|
|
||||
|Rfft|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|
||||
|SkipLayerNormalization|*in* input:**T**<br> *in* skip:**T**<br> *in* gamma:**T**<br> *in* beta:**T**<br> *in* bias:**T**<br> *out* output:**T**<br> *out* mean:**U**<br> *out* inv_std_var:**U**|1+|**T** = tensor(float), tensor(float16)|
|
||||
|
|
|
|||
|
|
@ -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<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BFloat16, FusedMatMul)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, BFloat16_float_BFloat16, LayerNormalization)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, FusedConv)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, QOrderedMatMul)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, QOrderedLayerNormalization)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, QOrderedGelu)>,
|
||||
|
||||
#ifdef ENABLE_ATEN
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kPytorchAtenDomain, 1, ATen)>,
|
||||
|
|
|
|||
|
|
@ -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 <typename T>
|
||||
__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
|
||||
|
|
@ -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<int8_t>())
|
||||
.TypeConstraint("S", DataTypeImpl::GetTensorType<float>())
|
||||
.TypeConstraint("F", BuildKernelDefConstraints<float, MLFloat16>())
|
||||
.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<float>("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<int8_t>::MappedType CudaQ;
|
||||
|
||||
// Inputs
|
||||
const Tensor* X = ctx->Input<Tensor>(0);
|
||||
const auto* X_data = reinterpret_cast<const CudaQ*>(X->Data<int8_t>());
|
||||
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<Tensor>(2);
|
||||
const void* scale_data = scale->DataRaw();
|
||||
|
||||
const Tensor* bias = ctx->Input<Tensor>(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<unsigned int>(x_shape.GetDims()[0]);
|
||||
unsigned int rows = gsl::narrow<unsigned int>(x_shape.GetDims()[1]);
|
||||
unsigned int cols = gsl::narrow<unsigned int>(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<CudaQ*>(Y->MutableData<int8_t>());
|
||||
|
||||
const float* scale_x = ctx->Input<Tensor>(1)->Data<float>();
|
||||
const float* scale_y = ctx->Input<Tensor>(4)->Data<float>();
|
||||
|
||||
if (scale->IsDataType<MLFloat16>()) {
|
||||
QOrderedLayerNorm(Stream(), GetDeviceProp(), static_cast<cublasLtOrder_t>(order_X_),
|
||||
X_data, *scale_x, Y_data, *scale_y, static_cast<const __half*>(scale_data),
|
||||
static_cast<const __half*>(bias_data),
|
||||
static_cast<float>(epsilon_), batch, rows, cols);
|
||||
} else {
|
||||
QOrderedLayerNorm(Stream(), GetDeviceProp(), static_cast<cublasLtOrder_t>(order_X_),
|
||||
X_data, *scale_x, Y_data, *scale_y, static_cast<const float*>(scale_data),
|
||||
static_cast<const float*>(bias_data),
|
||||
static_cast<float>(epsilon_), batch, rows, cols);
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace contrib
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -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
|
||||
|
|
@ -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<int8_t>(__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<int8_t>(input_1.x) * static_cast<int8_t>(input_2.x) +
|
||||
static_cast<int8_t>(input_1.y) * static_cast<int8_t>(input_2.y) +
|
||||
static_cast<int8_t>(input_1.z) * static_cast<int8_t>(input_2.z) +
|
||||
static_cast<int8_t>(input_1.w) * static_cast<int8_t>(input_2.w);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__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<size_t>(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<const char4*>(src + c));
|
||||
sum += (static_cast<short>(ch4.x) + static_cast<short>(ch4.y) +
|
||||
static_cast<short>(ch4.z) + static_cast<short>(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<int32_t>(sum);
|
||||
square_sum = WarpReduceSum<int32_t>(square_sum);
|
||||
|
||||
const float mean = __double2float_rn(src_scale * (double)sum / cols);
|
||||
|
||||
const float rvar = rsqrtf(src_scale * src_scale * __double2float_rn(static_cast<double>(square_sum) - (static_cast<double>(sum) * static_cast<double>(sum) / static_cast<double>(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<const char4*>(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<char4*>(dst + c) = QuantizeFloat4Char4(f4, dst_rscale);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
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<unsigned>(rows + QORDER_LAYERNORM_ROWS_PER_BLOCK - 1) / QORDER_LAYERNORM_ROWS_PER_BLOCK,
|
||||
static_cast<unsigned>(batch), 1);
|
||||
|
||||
QOrderedLayerNormRowKernel<T><<<blocks, threads, 0, stream>>>(src, src_scale,
|
||||
dst, dst_scale, gamma, beta,
|
||||
epsilon, rows, cols);
|
||||
}
|
||||
|
||||
template void QOrderedLayerNorm<float>(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
|
||||
|
|
@ -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 <typename T>
|
||||
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
|
||||
|
|
@ -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 <functional>
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace contrib {
|
||||
namespace cuda {
|
||||
|
||||
using namespace onnxruntime::cuda;
|
||||
|
||||
ONNX_OPERATOR_KERNEL_EX(
|
||||
QOrderedMatMul,
|
||||
kMSDomain,
|
||||
1,
|
||||
kCudaExecutionProvider,
|
||||
(*KernelDefBuilder::Create())
|
||||
.TypeConstraint("Q", DataTypeImpl::GetTensorType<int8_t>())
|
||||
.TypeConstraint("S", DataTypeImpl::GetTensorType<float>())
|
||||
.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<int64_t>()));
|
||||
|
||||
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<Tensor>(0);
|
||||
const Tensor& tensor_B = *context->Input<Tensor>(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<Tensor>(1)->Data<float>();
|
||||
const float* scale_B = context->Input<Tensor>(3)->Data<float>();
|
||||
const float* scale_Y = context->Input<Tensor>(4)->Data<float>();
|
||||
ORT_ENFORCE(*scale_Y > 0.0f && *scale_A > 0.0f && *scale_B > 0.0f);
|
||||
|
||||
const Tensor* tensor_bias = context->Input<Tensor>(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<float>();
|
||||
|
||||
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<Tensor>(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<Tensor>(7);
|
||||
ORT_ENFORCE(tensor_scale_C != nullptr);
|
||||
scale_C = tensor_scale_C->Data<float>();
|
||||
|
||||
C = tensor_C->Data<int8_t>();
|
||||
}
|
||||
|
||||
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<int32_t>(batch_A), rows_A, cols_B, cols_A,
|
||||
&alpha, tensor_A.Data<int8_t>(), tensor_B.Data<int8_t>(),
|
||||
static_cast<int32_t>(batch_B), bias,
|
||||
&beta, C, static_cast<int32_t>(batch_C),
|
||||
tensor_Y->MutableData<int8_t>(), static_cast<cublasLtOrder_t>(order_B_)));
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace contrib
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -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
|
||||
|
|
@ -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<int>(weight_order) << "-" << static_cast<int>(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<int>(m), static_cast<int>(n), static_cast<int>(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
|
||||
|
|
@ -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 <string>
|
||||
#include <unordered_set>
|
||||
|
||||
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<std::string, CublasLtMatmulAlgoInfo> 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
|
||||
|
|
@ -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<int8_t>())
|
||||
.TypeConstraint("S", DataTypeImpl::GetTensorType<float>())
|
||||
.InputMemoryType(OrtMemTypeCPUInput, 1) // scale_input
|
||||
.InputMemoryType(OrtMemTypeCPUInput, 2), // scale_output
|
||||
QOrderedGelu);
|
||||
|
||||
Status QOrderedGelu::ComputeInternal(OpKernelContext* context) const {
|
||||
const Tensor& input = *context->Input<Tensor>(0);
|
||||
|
||||
const float* scale_input = context->Input<Tensor>(1)->Data<float>();
|
||||
const float* scale_output = context->Input<Tensor>(2)->Data<float>();
|
||||
|
||||
const auto& shape = input.Shape();
|
||||
|
||||
Tensor* output = context->Output(0, shape);
|
||||
|
||||
QOrderedUnarySharedMemory_Gelu(Stream(), input.Data<int8_t>(), scale_input, output->MutableData<int8_t>(),
|
||||
scale_output, static_cast<size_t>(shape.Size()));
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace contrib
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -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
|
||||
|
|
@ -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 <typename FuncT>
|
||||
__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<int>(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<char>(__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<char4*>(output_data + id) = i4;
|
||||
id += kNumElementsPerBlockLine;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename FuncT>
|
||||
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<int>(CeilDiv(count, kNumElementsPerBlock));
|
||||
|
||||
QOrderedUnaryElementWiseSharedMemoryKernel<FuncT><<<blocksPerGrid, kNumThreadsPerBlock, 0, stream>>>(
|
||||
input_data, *input_scale, output_data, inverse_output_scale, func, static_cast<CUDA_LONG>(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<QOrderedUnaryOpGelu>(
|
||||
stream, input_data, input_scale, output_data, output_scale, QOrderedUnaryOpGelu(), count);
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace contrib
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -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
|
||||
|
|
@ -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<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QLinearSoftmax)>());
|
||||
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QuantizeLinear)>());
|
||||
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ReduceSumInteger)>());
|
||||
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QOrderedMatMul)>());
|
||||
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QOrderedGelu)>());
|
||||
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QOrderedLayerNormalization)>());
|
||||
|
||||
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, Attention)>());
|
||||
fn(GetOpSchema<ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, BeamSearch)>());
|
||||
|
|
|
|||
|
|
@ -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<int64_t>(1))
|
||||
.Attr("order_B", "cublasLt order of matrix B. Default is ROW MAJOR.",
|
||||
AttributeProto::INT, static_cast<int64_t>(1))
|
||||
.Attr("order_Y",
|
||||
"cublasLt order of matrix Y and optional matrix C. Default is ROW MAJOR.",
|
||||
AttributeProto::INT, static_cast<int64_t>(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<int64_t>(-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<int64_t>(1))
|
||||
.Attr("order_Y",
|
||||
"cublasLt order of matrix Y, must be same as order_X. "
|
||||
"Default is ROW MAJOR.",
|
||||
AttributeProto::INT, static_cast<int64_t>(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<int64_t>(1))
|
||||
.Attr("order_Y",
|
||||
"cublasLt order of matrix Y, must be same as order_X. "
|
||||
"Default is ROW MAJOR.",
|
||||
AttributeProto::INT, static_cast<int64_t>(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
|
||||
|
|
|
|||
|
|
@ -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_)));
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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<const CUDAExecutionProvider*, std::weak_ptr<PerThreadContext>>;
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@
|
|||
#include <curand.h>
|
||||
#include <cudnn.h>
|
||||
#include <cufft.h>
|
||||
#include <cublasLt.h>
|
||||
|
||||
#ifdef ORT_USE_NCCL
|
||||
#include <nccl.h>
|
||||
|
|
|
|||
47
onnxruntime/test/contrib_ops/qordered_gelu_op_test.cc
Normal file
47
onnxruntime/test/contrib_ops/qordered_gelu_op_test.cc
Normal file
|
|
@ -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<int64_t> 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<int64_t>());
|
||||
std::vector<int8_t> vec_x = GenData<int8_t>(shape, 1.0f);
|
||||
std::vector<int8_t> vec_y(N);
|
||||
for (int64_t i = 0; i < N; i++) {
|
||||
float x = scale_x * static_cast<float>(vec_x[i]);
|
||||
float r = (x * (0.5f * (1.0f + std::erff(x / sqrt_of_2)))) / scale_y;
|
||||
int8_t q = static_cast<int8_t>(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<int64_t>(order));
|
||||
test_qorder.AddAttribute("order_Y", static_cast<int64_t>(order));
|
||||
test_qorder.AddInput<int8_t>("X", shape, vec_x);
|
||||
test_qorder.AddInput<float>("scale_X", {}, {scale_x});
|
||||
test_qorder.AddInput<float>("scale_Y", {}, {scale_y});
|
||||
test_qorder.AddOutput<int8_t>("Y", shape, vec_y, false, 0.0f, 0.0f /* abs error */);
|
||||
|
||||
std::vector<std::unique_ptr<IExecutionProvider>> 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
|
||||
112
onnxruntime/test/contrib_ops/qordered_layer_norm_op_test.cc
Normal file
112
onnxruntime/test/contrib_ops/qordered_layer_norm_op_test.cc
Normal file
|
|
@ -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 <typename T> // MLFloat16 or float
|
||||
static void RunQOrdered_LayerNorm_RowMajor(std::vector<int64_t> const& shape, int axis, float epsilon,
|
||||
const std::vector<int8_t>& vec_x, float scale_x,
|
||||
const std::vector<T>& gamma, const std::vector<T>* beta,
|
||||
float scale_y, const std::vector<int8_t>& vec_y) {
|
||||
std::vector<int64_t> 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<int8_t>("X", shape, vec_x);
|
||||
test_qorder.AddInput<float>("scale_X", {}, {scale_x});
|
||||
test_qorder.AddInput<T>("scale", bias_shape, gamma);
|
||||
if (beta) {
|
||||
test_qorder.AddInput<T>("B", bias_shape, *beta);
|
||||
} else {
|
||||
test_qorder.AddOptionalInputEdge<T>();
|
||||
}
|
||||
test_qorder.AddInput<float>("scale_Y", {}, {scale_y});
|
||||
test_qorder.AddOutput<int8_t>("Y", shape, vec_y, false, 0.0f, 0.0f /* abs error */);
|
||||
|
||||
std::vector<std::unique_ptr<IExecutionProvider>> 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<int8_t> 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<float> 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<float> 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<int8_t> 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
|
||||
209
onnxruntime/test/contrib_ops/qordered_matmul_op_test.cc
Normal file
209
onnxruntime/test/contrib_ops/qordered_matmul_op_test.cc
Normal file
|
|
@ -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<int64_t> const& shape_A,
|
||||
std::vector<int64_t> const& shape_B,
|
||||
std::vector<int64_t> 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<int64_t>());
|
||||
std::vector<int8_t> vec_A = GenData<int8_t>(shape_A, 1.0f);
|
||||
std::vector<int8_t> vec_B = GenData<int8_t>(shape_B, 1.0f);
|
||||
std::vector<int8_t> 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<int64_t> bias_shape = {cols_Y};
|
||||
std::vector<float> bias;
|
||||
if (add_bias) { // make bias not too big
|
||||
float bias_scale = MLFloat16(scale_Y / 27.0f).ToFloat();
|
||||
bias = GenData<float>(bias_shape, bias_scale);
|
||||
}
|
||||
std::vector<int8_t> vec_C;
|
||||
std::vector<int64_t> shape_C = {broadcast_c_batch ? 1 : batch_Y, rows_Y, cols_Y};
|
||||
if (scale_C != 0.0f) {
|
||||
vec_C = GenData<int8_t>(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<int8_t>(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<std::unique_ptr<IExecutionProvider>> execution_providers;
|
||||
execution_providers.push_back(DefaultCudaExecutionProvider());
|
||||
|
||||
OpTester test_qorder("QOrderedMatMul", 1, onnxruntime::kMSDomain);
|
||||
test_qorder.AddAttribute("order_B", static_cast<int64_t>(weight_order));
|
||||
test_qorder.AddAttribute("order_A", static_cast<int64_t>(ORDER_ROW));
|
||||
test_qorder.AddAttribute("order_Y", static_cast<int64_t>(ORDER_ROW));
|
||||
test_qorder.AddInput<int8_t>("A", shape_A, vec_A);
|
||||
test_qorder.AddInput<float>("scale_A", {}, {scale_A});
|
||||
test_qorder.AddInput<int8_t>("B", shape_B, vec_B);
|
||||
test_qorder.AddInput<float>("scale_B", {}, {scale_B});
|
||||
test_qorder.AddInput<float>("scale_Y", {}, {scale_Y});
|
||||
if (add_bias) {
|
||||
test_qorder.AddInput<float>("bias", bias_shape, bias);
|
||||
} else {
|
||||
test_qorder.AddOptionalInputEdge<float>();
|
||||
}
|
||||
|
||||
if (scale_C != 0.0f) {
|
||||
test_qorder.AddInput<int8_t>("C", shape_C, vec_C);
|
||||
test_qorder.AddInput<float>("scale_C", {}, {scale_C});
|
||||
}
|
||||
test_qorder.AddOutput<int8_t>("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<int64_t> shape_A = {16, 32};
|
||||
std::vector<int64_t> shape_B = {32, 64};
|
||||
std::vector<int64_t> 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<int64_t> shape_A = {16, 32};
|
||||
std::vector<int64_t> shape_B = {32, 64};
|
||||
std::vector<int64_t> 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<int64_t> shape_A = {16, 32};
|
||||
std::vector<int64_t> shape_B = {32, 64};
|
||||
std::vector<int64_t> 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<int64_t> shape_A = {16, 32};
|
||||
std::vector<int64_t> shape_B = {32, 64};
|
||||
std::vector<int64_t> 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<int64_t> shape_A = {3, 16, 32};
|
||||
std::vector<int64_t> shape_B = {1, 32, 64};
|
||||
std::vector<int64_t> 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<int64_t> shape_A = {2, 16, 32};
|
||||
std::vector<int64_t> shape_B = {1, 32, 64};
|
||||
std::vector<int64_t> 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<int64_t> shape_A = {2, 16, 32};
|
||||
std::vector<int64_t> shape_B = {1, 32, 64};
|
||||
std::vector<int64_t> 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<int64_t> shape_A = {2, 16, 32};
|
||||
std::vector<int64_t> shape_B = {1, 32, 64};
|
||||
std::vector<int64_t> 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<int64_t> shape_A = {2, 16, 32};
|
||||
std::vector<int64_t> shape_B = {1, 32, 64};
|
||||
std::vector<int64_t> 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<int64_t> shape_A = {2, 16, 32};
|
||||
std::vector<int64_t> shape_B = {1, 32, 64};
|
||||
std::vector<int64_t> 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
|
||||
65
onnxruntime/test/contrib_ops/qordered_test_utils.h
Normal file
65
onnxruntime/test/contrib_ops/qordered_test_utils.h
Normal file
|
|
@ -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 <vector>
|
||||
#include <functional>
|
||||
|
||||
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 <typename T>
|
||||
inline std::vector<T> GenData(std::vector<int64_t> const& shape, float scale) {
|
||||
int64_t n = std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies<int64_t>());
|
||||
|
||||
scale = std::is_same<T, int8_t>::value ? 1.0f : scale; // using scale = 1.0f to generate int8_t data,
|
||||
std::vector<T> r(n);
|
||||
RandomValueGenerator random{};
|
||||
std::vector<int> tmp = random.Uniform<int32_t>(shape, -128, 127);
|
||||
for (int64_t i = 0; i < n; i++) {
|
||||
r[i] = static_cast<T>(tmp[i] * scale);
|
||||
}
|
||||
return r;
|
||||
}
|
||||
|
||||
inline void BatchRowColFromShape(std::vector<int64_t> 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<int64_t>());
|
||||
}
|
||||
|
||||
} // namespace test
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
Loading…
Reference in a new issue