saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[CUDA] Mod Op Kernel (#12499)

Browse source 
* mod for cuda and rocm

* fix bfloat16 ut

* change bf16 ut number

* fix opset version

* fix op kernel doc
This commit is contained in:
Vincent Wang 2022-08-09 13:05:40 +08:00 committed by GitHub
parent a2dc3e9eac
commit cfa09d16d9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 195 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
docs/OperatorKernels.md
View file

@ -601,6 +601,8 @@ Do not modify directly.*
|Min|*in* data_0:**T**<br> *out* min:**T**|13+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)|
|||12|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)|
|||[6, 11]|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
|Mod|*in* A:**T**<br> *in* B:**T**<br> *out* C:**T**|13+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)|
|||[10, 12]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)|
|Mul|*in* A:**T**<br> *in* B:**T**<br> *out* C:**T**|14+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)|
|||13|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)|
|||[7, 12]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)|

35
onnxruntime/core/providers/cuda/cu_inc/common.cuh
View file

@ -359,6 +359,41 @@ __device__ __inline__ T _Gelu(T a) {
return a * _Normcdf(a);
}
template <typename T>
__device__ __inline__ T _Mod(T a, T b) {
T r = a % b;
T zero = T(0);
if ((r > zero && b < zero) || (r < zero && b > zero)) {
r += b;
}
return r;
}
template <typename T>
__device__ __inline__ T _Fmod(T a, T b) {
return a % b;
}
template <>
__device__ __inline__ float _Fmod(float a, float b) {
return fmodf(a, b);
}
template <>
__device__ __inline__ double _Fmod(double a, double b) {
return fmod(a, b);
}
template <>
__device__ __inline__ half _Fmod(half a, half b) {
return fmodf((float)a, (float)b);
}
template <>
__device__ __inline__ BFloat16 _Fmod(BFloat16 a, BFloat16 b) {
return fmodf((float)a, (float)b);
}
// We would like to use 64-bit integer to support large matrices. However, CUDA seems to support only 32-bit integer
// For now, use int32_t to ensure that both Linux and Windows see this as 32 bit integer type.
#ifndef CUDA_LONG

4
onnxruntime/core/providers/cuda/cuda_execution_provider.cc
View file

@ -833,6 +833,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain,
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, double, ThresholdedRelu);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, MLFloat16, ThresholdedRelu);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, 10, TopK);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, 12, Mod);
// opset 11
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, float, ArgMax);
@ -1217,6 +1218,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain,
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Tanh);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Gemm);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, ReduceSum);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, Mod);
// OpSet 14
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, CumSum);
@ -1709,6 +1711,7 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, uint8_t, QuantizeLinear)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, int8_t, DequantizeLinear)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, uint8_t, DequantizeLinear)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, 12, Mod)>,
// opset 11
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, float, ArgMax)>,
@ -2089,6 +2092,7 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Tanh)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, Gemm)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, ReduceSum)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, Mod)>,
// OpSet 14
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, CumSum)>,

63
onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc
View file

@ -465,6 +465,69 @@ Status Pow::ComputeInternal(OpKernelContext* context) const {
return s;
}
ONNX_OPERATOR_VERSIONED_KERNEL_EX(
Mod, kOnnxDomain, 10, 12, kCudaExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T",
BuildKernelDefConstraints<int32_t, int64_t, uint32_t, uint64_t, float, double, MLFloat16>()),
Mod);
ONNX_OPERATOR_KERNEL_EX(Mod, kOnnxDomain, 13, kCudaExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", BuildKernelDefConstraints<int32_t, int64_t, uint32_t, uint64_t, float,
double, MLFloat16, BFloat16>()),
Mod);
Status Mod::ComputeInternal(OpKernelContext* context) const {
namespace on = ONNX_NAMESPACE;
BinaryElementwisePreparation prepare;
ORT_RETURN_IF_ERROR(Prepare(context, &prepare));
auto element_type = prepare.lhs_tensor->GetElementType();
ORT_ENFORCE(fmod_ || element_type == on::TensorProto_DataType_INT32 ||
element_type == on::TensorProto_DataType_INT64 || element_type == on::TensorProto_DataType_UINT32 ||
element_type == on::TensorProto_DataType_UINT64,
"Non-fmod can support integer types only.");
#define CASE_MOD_ELEMENT_TYPE(name, onnx_type, data_type) \
case onnx_type: { \
Impl_##name<typename ToCudaType<data_type>::MappedType>( \
Stream(), prepare.output_rank_or_simple_broadcast, &prepare.lhs_padded_strides, \
reinterpret_cast<const typename ToCudaType<data_type>::MappedType*>(prepare.lhs_tensor->Data<data_type>()), \
&prepare.rhs_padded_strides, \
reinterpret_cast<const typename ToCudaType<data_type>::MappedType*>(prepare.rhs_tensor->Data<data_type>()), \
&prepare.fdm_output_strides, prepare.fdm_H, prepare.fdm_C, \
reinterpret_cast<typename ToCudaType<data_type>::MappedType*>( \
prepare.output_tensor->MutableData<data_type>()), \
prepare.output_tensor->Shape().Size()); \
} break
if (fmod_) {
switch (element_type) {
CASE_MOD_ELEMENT_TYPE(Fmod, on::TensorProto_DataType_INT32, int32_t);
CASE_MOD_ELEMENT_TYPE(Fmod, on::TensorProto_DataType_INT64, int64_t);
CASE_MOD_ELEMENT_TYPE(Fmod, on::TensorProto_DataType_UINT32, uint32_t);
CASE_MOD_ELEMENT_TYPE(Fmod, on::TensorProto_DataType_UINT64, uint64_t);
CASE_MOD_ELEMENT_TYPE(Fmod, on::TensorProto_DataType_FLOAT, float);
CASE_MOD_ELEMENT_TYPE(Fmod, on::TensorProto_DataType_DOUBLE, double);
CASE_MOD_ELEMENT_TYPE(Fmod, on::TensorProto_DataType_FLOAT16, MLFloat16);
CASE_MOD_ELEMENT_TYPE(Fmod, on::TensorProto_DataType_BFLOAT16, BFloat16);
default:
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Unsupported element type: ", DataTypeImpl::ToString(prepare.lhs_tensor->DataType()));
}
} else {
switch (element_type) {
CASE_MOD_ELEMENT_TYPE(Mod, on::TensorProto_DataType_INT32, int32_t);
CASE_MOD_ELEMENT_TYPE(Mod, on::TensorProto_DataType_INT64, int64_t);
CASE_MOD_ELEMENT_TYPE(Mod, on::TensorProto_DataType_UINT32, uint32_t);
CASE_MOD_ELEMENT_TYPE(Mod, on::TensorProto_DataType_UINT64, uint64_t);
default:
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Unsupported element type: ", DataTypeImpl::ToString(prepare.lhs_tensor->DataType()));
}
}
#undef CASE_MOD_ELEMENT_TYPE
return Status::OK();
}
//Greater op output tensor type is bool, so it cannot directly fit in the macros
//for other elementwise ops
template <typename T, typename CudaT>

12
onnxruntime/core/providers/cuda/math/binary_elementwise_ops.h
View file

@ -214,6 +214,18 @@ class PRelu final : public BinaryElementwise<ShouldBroadcast> {
Status ComputeInternal(OpKernelContext* context) const override;
};
class Mod final : public BinaryElementwise<ShouldBroadcast> {
public:
Mod(const OpKernelInfo& info) : BinaryElementwise(info) {
int64_t fmod = info.GetAttrOrDefault<int64_t>("fmod", 0LL);
fmod_ = fmod != 0;
}
Status ComputeInternal(OpKernelContext* context) const override;
private:
bool fmod_{false};
};
template <typename T, typename CudaT>
class CompareFunction : public BinaryElementwise<ShouldBroadcast> {
public:

8
onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu
View file

@ -89,6 +89,12 @@ namespace cuda {
SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, double) \
SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, BFloat16)
#define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZIL(x) \
SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, uint32_t) \
SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, uint64_t) \
SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, int32_t) \
SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, int64_t)
#define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T1_ILHFD(x, T) \
SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T1(x, T, int32_t) \
SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T1(x, T, int64_t) \
@ -141,6 +147,8 @@ SPECIALIZED_BINARY_ELEMENTWISE_IMPL(Xor, bool)
SPECIALIZED_BINARY_ELEMENTWISE_IMPL_HFD(PRelu)
SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(Max)
SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(Min)
SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZIL(Mod)
SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(Fmod)
// create declarations for impl for Pow
BINARY_ELEMENTWISE_IMPL_T1(Pow)

4
onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.h
View file

@ -26,7 +26,9 @@ namespace cuda {
BINARY_OP_NAME_EXPR(Xor, (a ^ b)) \
BINARY_OP_NAME_EXPR(PRelu, (a > (T)0 ? a : a * b)) \
BINARY_OP_NAME_EXPR(Max, _Max(a, b)) \
BINARY_OP_NAME_EXPR(Min, _Min(a, b))
BINARY_OP_NAME_EXPR(Min, _Min(a, b)) \
BINARY_OP_NAME_EXPR(Mod, _Mod(a, b)) \
BINARY_OP_NAME_EXPR(Fmod, _Fmod(a, b))
// NOTE that cu files are compiled with nvcc and should not refer to any onnxruntime headers
// so struct BinaryElementwisePreparation cannot be used here

35
onnxruntime/core/providers/rocm/cu_inc/common.cuh
View file

@ -280,6 +280,41 @@ __device__ __inline__ T _Gelu(T a) {
return a * _Normcdf(a);
}
template <typename T>
__device__ __inline__ T _Mod(T a, T b) {
T r = a % b;
T zero = T(0);
if ((r > zero && b < zero) || (r < zero && b > zero)) {
r += b;
}
return r;
}
template <typename T>
__device__ __inline__ T _Fmod(T a, T b) {
return a % b;
}
template <>
__device__ __inline__ float _Fmod(float a, float b) {
return fmodf(a, b);
}
template <>
__device__ __inline__ double _Fmod(double a, double b) {
return fmod(a, b);
}
template <>
__device__ __inline__ half _Fmod(half a, half b) {
return fmodf((float)a, (float)b);
}
template <>
__device__ __inline__ BFloat16 _Fmod(BFloat16 a, BFloat16 b) {
return fmodf((float)a, (float)b);
}
// We would like to use 64-bit integer to support large matrices. However, ROCM seems to support only 32-bit integer
// For now, use int32_t to ensure that both Linux and Windows see this as 32 bit integer type.
#ifndef HIP_LONG

4
onnxruntime/core/providers/rocm/rocm_execution_provider.cc
View file

@ -746,6 +746,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain,
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 10, double, ThresholdedRelu);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 10, MLFloat16, ThresholdedRelu);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 10, 10, TopK);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 10, 12, Mod);
// opset 11
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, float, ArgMax);
@ -1130,6 +1131,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain,
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, BFloat16, Tanh);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, BFloat16, Gemm);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, BFloat16, ReduceSum);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, Mod);
// OpSet 14
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 14, CumSum);
@ -1652,6 +1654,7 @@ static Status RegisterRocmKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 10, uint8_t, QuantizeLinear)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 10, int8_t, DequantizeLinear)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 10, uint8_t, DequantizeLinear)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 10, 12, Mod)>,
// opset 11
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 11, float, ArgMax)>,
@ -2046,6 +2049,7 @@ static Status RegisterRocmKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, BFloat16, Tanh)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, BFloat16, Gemm)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, BFloat16, ReduceSum)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 13, Mod)>,
// OpSet 14
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 14, CumSum)>,

31
onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc
View file

@ -33,7 +33,7 @@ void TestFloat16(const char* op_name, const std::vector<int64_t>& lhs_dim,
execution_providers.push_back(DefaultCudaExecutionProvider());
#elif USE_ROCM
execution_providers.push_back(DefaultRocmExecutionProvider());
#endif
#endif
tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
@ -47,7 +47,7 @@ void TestFloat16(const char* op_name, const std::vector<int64_t>& lhs_dim,
execution_providers.push_back(DefaultCudaExecutionProvider());
#elif USE_ROCM
execution_providers.push_back(DefaultRocmExecutionProvider());
#endif
#endif
tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
}
@ -2293,6 +2293,15 @@ TEST(ModOpTest, Fmod_float_mixed_sign) {
test.Run();
}
TEST(ModOpTest, Fmod_double_mixed_sign) {
OpTester test("Mod", ModOp_ver);
test.AddAttribute<int64_t>("fmod", 1);
test.AddInput<double>("X", {6}, {-4.3, 7.2, 5.0, 4.3, -7.2, 8.0});
test.AddInput<double>("Y", {6}, {2.1f, -3.4, 8.0, -2.1, 3.4, 5.0});
test.AddOutput<double>("Z", {6}, {-0.1, 0.4, 5., 0.1, -0.4, 3.});
test.Run();
}
TEST(ModOpTest, Fmod_float16_mixed_sign) {
OpTester test("Mod", ModOp_ver);
test.AddAttribute<int64_t>("fmod", 1);
@ -2305,6 +2314,24 @@ TEST(ModOpTest, Fmod_float16_mixed_sign) {
test.Run();
}
#if defined(USE_CUDA) || defined(USE_ROCM)
TEST(ModOpTest, Fmod_bfloat16_mixed_sign) {
OpTester test("Mod", 13);
test.AddAttribute<int64_t>("fmod", 1);
// Due to BFloat16's precision, if the result is too small, it's not easy get pass for both CUDA and ROCm.
test.AddInput<BFloat16>("X", {4}, MakeBFloat16({8.0f, 5.0f, -8.0f, 8.0f}));
test.AddInput<BFloat16>("Y", {4}, MakeBFloat16({-3.4f, 8.0f, 3.4f, 5.0f}));
test.AddOutput<BFloat16>("Z", {4}, MakeBFloat16({1.2f, 5.f, -1.2f, 3.f}));
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
#ifdef USE_CUDA
execution_providers.push_back(DefaultCudaExecutionProvider());
#elif USE_ROCM
execution_providers.push_back(DefaultRocmExecutionProvider());
#endif
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
#endif
TEST(ModOpTest, Int8_mixed_sign) {
OpTester test("Mod", ModOp_ver);
test.AddInput<int8_t>("X", {6}, {-4, 7, 5, 4, -7, 8});