mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
Complex multiplication and conjugate contrib ops (#3384)
* adding ComplexMulConj * Adding fp16 support. * adding a util func
This commit is contained in:
Du Li
GitHub
parent
bae1dd7f04
commit
2659f205cc
7 changed files with 421 additions and 0 deletions
63
onnxruntime/contrib_ops/cuda/math/complex_mul.cc
Normal file
63
onnxruntime/contrib_ops/cuda/math/complex_mul.cc
Normal file
|
|
@ -0,0 +1,63 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "complex_mul.h"
|
||||
#include "complex_mul_impl.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace contrib {
|
||||
namespace cuda {
|
||||
#define REGISTER_KERNEL_TYPED(T) \
|
||||
ONNX_OPERATOR_TYPED_KERNEL_EX( \
|
||||
ComplexMul, \
|
||||
kMSDomain, \
|
||||
1, \
|
||||
T, \
|
||||
kCudaExecutionProvider, \
|
||||
KernelDefBuilder() \
|
||||
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
|
||||
ComplexMul<T, false>); \
|
||||
ONNX_OPERATOR_TYPED_KERNEL_EX( \
|
||||
ComplexMulConj, \
|
||||
kMSDomain, \
|
||||
1, \
|
||||
T, \
|
||||
kCudaExecutionProvider, \
|
||||
KernelDefBuilder() \
|
||||
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
|
||||
ComplexMul<T, true>);
|
||||
|
||||
REGISTER_KERNEL_TYPED(float)
|
||||
REGISTER_KERNEL_TYPED(MLFloat16)
|
||||
|
||||
template <typename T, bool is_conj>
|
||||
Status ComplexMul<T, is_conj>::ComputeInternal(OpKernelContext* context) const {
|
||||
for (int index = 0; index < context->InputCount(); ++index) {
|
||||
const Tensor* input = context->Input<Tensor>(index);
|
||||
TensorShape shape = input->Shape();
|
||||
int64_t last_dimension = shape[shape.NumDimensions() - 1];
|
||||
ORT_ENFORCE(last_dimension == 2, "The input_", index, " last dimension is supposed to be 2, but get ", last_dimension);
|
||||
}
|
||||
|
||||
BinaryElementwisePreparation prepare;
|
||||
Prepare(context, &prepare);
|
||||
ComplexMul_Impl<typename ToCudaType<T>::MappedType>(
|
||||
prepare.output_rank_or_simple_broadcast,
|
||||
&prepare.lhs_padded_strides,
|
||||
reinterpret_cast<const typename ToCudaType<T>::MappedType*>(prepare.lhs_tensor->template Data<T>()),
|
||||
&prepare.rhs_padded_strides,
|
||||
reinterpret_cast<const typename ToCudaType<T>::MappedType*>(prepare.rhs_tensor->template Data<T>()),
|
||||
&prepare.fdm_output_strides,
|
||||
prepare.fdm_H,
|
||||
prepare.fdm_C,
|
||||
reinterpret_cast<typename ToCudaType<T>::MappedType*>(prepare.output_tensor->template MutableData<T>()),
|
||||
prepare.output_tensor->Shape().Size(),
|
||||
prepare.lhs_tensor->Shape().Size(),
|
||||
prepare.rhs_tensor->Shape().Size(),
|
||||
is_conj);
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace contrib
|
||||
} // namespace onnxruntime
|
||||
22
onnxruntime/contrib_ops/cuda/math/complex_mul.h
Normal file
22
onnxruntime/contrib_ops/cuda/math/complex_mul.h
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
#include "core/providers/cuda/math/binary_elementwise_ops.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace contrib {
|
||||
namespace cuda {
|
||||
|
||||
using namespace ::onnxruntime::cuda;
|
||||
|
||||
template <typename T, bool is_conj>
|
||||
class ComplexMul : public BinaryElementwise<ShouldBroadcast> {
|
||||
public:
|
||||
ComplexMul(const OpKernelInfo info) : BinaryElementwise{info} {}
|
||||
Status ComputeInternal(OpKernelContext* context) const override;
|
||||
};
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace contrib
|
||||
} // namespace onnxruntime
|
||||
174
onnxruntime/contrib_ops/cuda/math/complex_mul_impl.cu
Normal file
174
onnxruntime/contrib_ops/cuda/math/complex_mul_impl.cu
Normal file
|
|
@ -0,0 +1,174 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "complex_mul.h"
|
||||
#include "complex_mul_impl.h"
|
||||
#include "core/providers/cuda/cu_inc/common.cuh"
|
||||
#include "core/providers/cuda/math/binary_elementwise_ops.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace contrib {
|
||||
namespace cuda {
|
||||
template <typename T>
|
||||
__device__ __inline__ void _ComplexMul(T a0, T a1, T b0, T b1, T* output_data, bool is_conj) {
|
||||
if (is_conj) {
|
||||
T out_real = a0 * b0 + a1 * b1;
|
||||
T out_imag = a1 * b0 - a0 * b1;
|
||||
output_data[0] = out_real;
|
||||
output_data[1] = out_imag;
|
||||
} else {
|
||||
T out_real = a0 * b0 - a1 * b1;
|
||||
T out_imag = a0 * b1 + a1 * b0;
|
||||
output_data[0] = out_real;
|
||||
output_data[1] = out_imag;
|
||||
}
|
||||
};
|
||||
|
||||
// broadcast by computing output coordinate from offset, using fast_divmod
|
||||
template <typename T, bool lhs_need_compute, bool rhs_need_compute, int NumThreadsPerBlock, int NumElementsPerThread>
|
||||
__global__ void _ElementWiseWithStrideTwo(
|
||||
int32_t output_rank,
|
||||
const TArray<int64_t> lhs_padded_strides,
|
||||
const T* lhs_data,
|
||||
const TArray<int64_t> rhs_padded_strides,
|
||||
const T* rhs_data,
|
||||
const TArray<fast_divmod> fdm_output_strides,
|
||||
T* output_data,
|
||||
CUDA_LONG N,
|
||||
int64_t lhs_size,
|
||||
int64_t rhs_size,
|
||||
bool is_conj) {
|
||||
CUDA_LONG start = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x;
|
||||
T a[NumElementsPerThread];
|
||||
T b[NumElementsPerThread];
|
||||
T c[NumElementsPerThread];
|
||||
T d[NumElementsPerThread];
|
||||
|
||||
CUDA_LONG id = start;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < NumElementsPerThread; i++) {
|
||||
if (id < N / 2) {
|
||||
CUDA_LONG lhs_index = (lhs_need_compute ? 0 : id);
|
||||
CUDA_LONG rhs_index = (rhs_need_compute ? 0 : id);
|
||||
// compute indexes with broadcasting rules: https://github.com/onnx/onnx/blob/master/docs/Broadcasting.md
|
||||
CUDA_LONG offset = id;
|
||||
#pragma unroll
|
||||
for (auto dim = 0; dim < fdm_output_strides.GetCapacity(); dim++) {
|
||||
if (dim >= output_rank) {
|
||||
break;
|
||||
}
|
||||
int q, r;
|
||||
fdm_output_strides.data_[dim].divmod(offset, q, r);
|
||||
if (lhs_need_compute) {
|
||||
lhs_index += static_cast<int>(lhs_padded_strides.data_[dim]) * q;
|
||||
}
|
||||
|
||||
if (rhs_need_compute) {
|
||||
rhs_index += static_cast<int>(rhs_padded_strides.data_[dim]) * q;
|
||||
}
|
||||
offset = r;
|
||||
}
|
||||
|
||||
a[i] = lhs_data[(2 * lhs_index) % lhs_size];
|
||||
b[i] = lhs_data[(2 * lhs_index + 1) % lhs_size];
|
||||
c[i] = rhs_data[(2 * rhs_index) % rhs_size];
|
||||
d[i] = rhs_data[(2 * rhs_index + 1) % rhs_size];
|
||||
|
||||
id += NumThreadsPerBlock;
|
||||
}
|
||||
}
|
||||
|
||||
id = start;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < NumElementsPerThread; i++) {
|
||||
if (id < N / 2) {
|
||||
_ComplexMul(a[i], b[i], c[i], d[i], &output_data[2 * id], is_conj);
|
||||
id += NumThreadsPerBlock;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
void ComplexMul_Impl(
|
||||
int32_t output_rank_or_simple_broadcast,
|
||||
const TArray<int64_t>* lhs_padded_strides,
|
||||
const T* lhs_data,
|
||||
const TArray<int64_t>* rhs_padded_strides,
|
||||
const T* rhs_data,
|
||||
const TArray<onnxruntime::cuda::fast_divmod>* fdm_output_strides,
|
||||
const onnxruntime::cuda::fast_divmod& fdm_H,
|
||||
const onnxruntime::cuda::fast_divmod& fdm_C,
|
||||
T* output_data,
|
||||
int64_t count,
|
||||
int64_t lhs_size,
|
||||
int64_t rhs_size,
|
||||
bool is_conj) {
|
||||
if (count == 0) // special case where there's a dim value of 0 in the output shape
|
||||
return;
|
||||
|
||||
int blocksPerGrid = static_cast<int>(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
|
||||
CUDA_LONG N = static_cast<CUDA_LONG>(count);
|
||||
|
||||
if (lhs_padded_strides && rhs_padded_strides && lhs_padded_strides->size_ && rhs_padded_strides->size_)
|
||||
_ElementWiseWithStrideTwo<T, true, true, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
|
||||
output_rank_or_simple_broadcast,
|
||||
*lhs_padded_strides,
|
||||
lhs_data,
|
||||
*rhs_padded_strides,
|
||||
rhs_data,
|
||||
*fdm_output_strides,
|
||||
output_data,
|
||||
N,
|
||||
lhs_size,
|
||||
rhs_size,
|
||||
is_conj);
|
||||
else if (lhs_padded_strides && lhs_padded_strides->size_)
|
||||
_ElementWiseWithStrideTwo<T, true, false, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
|
||||
output_rank_or_simple_broadcast,
|
||||
*lhs_padded_strides,
|
||||
lhs_data,
|
||||
*rhs_padded_strides,
|
||||
rhs_data,
|
||||
*fdm_output_strides,
|
||||
output_data,
|
||||
N,
|
||||
lhs_size,
|
||||
rhs_size,
|
||||
is_conj);
|
||||
else
|
||||
_ElementWiseWithStrideTwo<T, false, true, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
|
||||
output_rank_or_simple_broadcast,
|
||||
*lhs_padded_strides,
|
||||
lhs_data,
|
||||
*rhs_padded_strides,
|
||||
rhs_data,
|
||||
*fdm_output_strides,
|
||||
output_data,
|
||||
N,
|
||||
lhs_size,
|
||||
rhs_size,
|
||||
is_conj);
|
||||
};
|
||||
|
||||
#define SPECIALIZE_STACKEDCOMPLEXMUL_IMPL(T) \
|
||||
template void ComplexMul_Impl<T>( \
|
||||
int32_t output_rank_or_simple_broadcast, \
|
||||
const TArray<int64_t>* lhs_padded_strides, \
|
||||
const T* lhs_data, \
|
||||
const TArray<int64_t>* rhs_padded_strides, \
|
||||
const T* rhs_data, \
|
||||
const TArray<onnxruntime::cuda::fast_divmod>* fdm_output_strides, \
|
||||
const onnxruntime::cuda::fast_divmod& fdm_H, \
|
||||
const onnxruntime::cuda::fast_divmod& fdm_C, \
|
||||
T* output_data, \
|
||||
int64_t count, \
|
||||
int64_t lhs_size, \
|
||||
int64_t rhs_size, \
|
||||
bool is_conj);
|
||||
|
||||
SPECIALIZE_STACKEDCOMPLEXMUL_IMPL(float)
|
||||
SPECIALIZE_STACKEDCOMPLEXMUL_IMPL(half)
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace contrib
|
||||
} // namespace onnxruntime
|
||||
32
onnxruntime/contrib_ops/cuda/math/complex_mul_impl.h
Normal file
32
onnxruntime/contrib_ops/cuda/math/complex_mul_impl.h
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
#include "complex_mul.h"
|
||||
#include "core/providers/cuda/shared_inc/cuda_utils.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace contrib {
|
||||
namespace cuda {
|
||||
|
||||
using namespace ::onnxruntime::cuda;
|
||||
|
||||
template <typename T>
|
||||
void ComplexMul_Impl(
|
||||
int32_t output_rank_or_simple_broadcast,
|
||||
const TArray<int64_t>* lhs_padded_strides,
|
||||
const T* lhs_data,
|
||||
const TArray<int64_t>* rhs_padded_strides,
|
||||
const T* rhs_data,
|
||||
const TArray<onnxruntime::cuda::fast_divmod>* fdm_output_strides,
|
||||
const onnxruntime::cuda::fast_divmod& fdm_H,
|
||||
const onnxruntime::cuda::fast_divmod& fdm_C,
|
||||
T* output_data,
|
||||
int64_t count,
|
||||
int64_t lhs_size,
|
||||
int64_t rhs_size,
|
||||
bool is_conj);
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace contrib
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -23,6 +23,10 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1
|
|||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, Irfft);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, Irfft);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, Irfft);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, ComplexMul);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, ComplexMul);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, ComplexMulConj);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, ComplexMulConj);
|
||||
|
||||
// These ops were experimental ops in onnx domain which have been removed now. We add them here as
|
||||
// contrib ops to maintain backward compatibility
|
||||
|
|
@ -78,6 +82,10 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
|
|||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, Irfft)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, Irfft)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, Irfft)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, ComplexMul)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, ComplexMul)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, ComplexMulConj)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, ComplexMulConj)>,
|
||||
|
||||
// These ops were experimental ops in onnx domain which have been removed now. We add them here as
|
||||
// contrib ops to maintain backward compatibility
|
||||
|
|
|
|||
|
|
@ -994,6 +994,24 @@ Sample echo operator.)DOC");
|
|||
.Output(0, "Y", "output tensor", "T")
|
||||
.TypeConstraint("T", {"tensor(float)", "tensor(double)", "tensor(float16)"}, "Constrain input and output types to float or half tensors.");
|
||||
|
||||
ONNX_CONTRIB_OPERATOR_SCHEMA(ComplexMul)
|
||||
.SetDomain(kMSDomain)
|
||||
.SinceVersion(1)
|
||||
.SetDoc(R"DOC()DOC")
|
||||
.Input(0, "A", "input_0", "T")
|
||||
.Input(1, "B", "input_1", "T")
|
||||
.Output(0, "C", "output tensor", "T")
|
||||
.TypeConstraint("T", {"tensor(float)", "tensor(double)", "tensor(float16)"}, "Constrain input and output types to float or half tensors.");
|
||||
|
||||
ONNX_CONTRIB_OPERATOR_SCHEMA(ComplexMulConj)
|
||||
.SetDomain(kMSDomain)
|
||||
.SinceVersion(1)
|
||||
.SetDoc(R"DOC()DOC")
|
||||
.Input(0, "A", "input_0", "T")
|
||||
.Input(1, "B", "input_1", "T")
|
||||
.Output(0, "C", "output tensor", "T")
|
||||
.TypeConstraint("T", {"tensor(float)", "tensor(double)", "tensor(float16)"}, "Constrain input and output types to float or half tensors.");
|
||||
|
||||
ONNX_CONTRIB_OPERATOR_SCHEMA(ConvTransposeWithDynamicPads)
|
||||
.SetDomain(kMSDomain)
|
||||
.SinceVersion(1)
|
||||
|
|
|
|||
|
|
@ -114,5 +114,109 @@ TEST(BiasGeluTest, Two_One_Dim) {
|
|||
RunBiasGeluTest(input_a_data, input_b_data, {2, 4}, {4});
|
||||
}
|
||||
|
||||
TEST(MathOpTest, ComplexMul) {
|
||||
if (DefaultCudaExecutionProvider() == nullptr) return;
|
||||
|
||||
std::vector<float> input_a_data = {
|
||||
-0.5f, 0.6f};
|
||||
|
||||
std::vector<float> input_b_data = {
|
||||
0.8f, -0.5f, 0.0f, 1.f,
|
||||
0.5f, 0.2f, 0.3f, -0.6f};
|
||||
|
||||
std::vector<float> output_data = {
|
||||
-0.10f, 0.73f,
|
||||
-0.60f, -0.50f,
|
||||
-0.37f, 0.20f,
|
||||
0.21f, 0.48f};
|
||||
|
||||
OpTester tester("ComplexMul", 1, onnxruntime::kMSDomain);
|
||||
tester.AddInput<float>("A", {1, 2}, input_a_data);
|
||||
tester.AddInput<float>("B", {4, 2}, input_b_data);
|
||||
tester.AddOutput<float>("C", {4, 2}, output_data);
|
||||
|
||||
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
|
||||
execution_providers.push_back(DefaultCudaExecutionProvider());
|
||||
tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
|
||||
}
|
||||
|
||||
TEST(MathOpTest, ComplexMulConj) {
|
||||
if (DefaultCudaExecutionProvider() == nullptr) return;
|
||||
|
||||
std::vector<float> input_a_data = {
|
||||
-0.5f, 0.6f};
|
||||
|
||||
std::vector<float> input_b_data = {
|
||||
0.8f, -0.5f, 0.0f, 1.f,
|
||||
0.5f, 0.2f, 0.3f, -0.6f};
|
||||
|
||||
std::vector<float> output_data = {
|
||||
-0.70f, 0.23f,
|
||||
0.60f, 0.50f,
|
||||
-0.13f, 0.40f,
|
||||
-0.51f, -0.12f};
|
||||
|
||||
OpTester tester("ComplexMulConj", 1, onnxruntime::kMSDomain);
|
||||
tester.AddInput<float>("A", {1, 2}, input_a_data);
|
||||
tester.AddInput<float>("B", {4, 2}, input_b_data);
|
||||
tester.AddOutput<float>("C", {4, 2}, output_data);
|
||||
|
||||
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
|
||||
execution_providers.push_back(DefaultCudaExecutionProvider());
|
||||
tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
|
||||
}
|
||||
|
||||
TEST(MathOpTest, ComplexMul_fp16) {
|
||||
if (DefaultCudaExecutionProvider() == nullptr) return;
|
||||
|
||||
std::vector<MLFloat16> input_a_data = {
|
||||
MLFloat16(math::floatToHalf(-0.5f)), MLFloat16(math::floatToHalf(0.6f))};
|
||||
|
||||
std::vector<MLFloat16> input_b_data = {
|
||||
MLFloat16(math::floatToHalf(0.8f)), MLFloat16(math::floatToHalf(-0.5f)), MLFloat16(math::floatToHalf(0.0f)), MLFloat16(math::floatToHalf(1.f)),
|
||||
MLFloat16(math::floatToHalf(0.5f)), MLFloat16(math::floatToHalf(0.2f)), MLFloat16(math::floatToHalf(0.3f)), MLFloat16(math::floatToHalf(-0.6f))};
|
||||
|
||||
std::vector<MLFloat16> output_data = {
|
||||
MLFloat16(math::floatToHalf(-0.10f)), MLFloat16(math::floatToHalf(0.73f)),
|
||||
MLFloat16(math::floatToHalf(-0.60f)), MLFloat16(math::floatToHalf(-0.50f)),
|
||||
MLFloat16(math::floatToHalf(-0.37f)), MLFloat16(math::floatToHalf(0.20f)),
|
||||
MLFloat16(math::floatToHalf(0.21f)), MLFloat16(math::floatToHalf(0.48f))};
|
||||
|
||||
OpTester tester("ComplexMul", 1, onnxruntime::kMSDomain);
|
||||
tester.AddInput<MLFloat16>("A", {1, 2}, input_a_data);
|
||||
tester.AddInput<MLFloat16>("B", {4, 2}, input_b_data);
|
||||
tester.AddOutput<MLFloat16>("C", {4, 2}, output_data);
|
||||
|
||||
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
|
||||
execution_providers.push_back(DefaultCudaExecutionProvider());
|
||||
tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
|
||||
}
|
||||
|
||||
TEST(MathOpTest, ComplexMulConj_fp16) {
|
||||
if (DefaultCudaExecutionProvider() == nullptr) return;
|
||||
|
||||
std::vector<MLFloat16> input_a_data = {
|
||||
MLFloat16(math::floatToHalf(-0.5f)), MLFloat16(math::floatToHalf(0.6f))};
|
||||
|
||||
std::vector<MLFloat16> input_b_data = {
|
||||
MLFloat16(math::floatToHalf(0.8f)), MLFloat16(math::floatToHalf(-0.5f)), MLFloat16(math::floatToHalf(0.0f)), MLFloat16(math::floatToHalf(1.f)),
|
||||
MLFloat16(math::floatToHalf(0.5f)), MLFloat16(math::floatToHalf(0.2f)), MLFloat16(math::floatToHalf(0.3f)), MLFloat16(math::floatToHalf(-0.6f))};
|
||||
|
||||
std::vector<MLFloat16> output_data = {
|
||||
MLFloat16(math::floatToHalf(-0.70f)), MLFloat16(math::floatToHalf(0.23f)),
|
||||
MLFloat16(math::floatToHalf(0.60f)), MLFloat16(math::floatToHalf(0.50f)),
|
||||
MLFloat16(math::floatToHalf(-0.13f)), MLFloat16(math::floatToHalf(0.40f)),
|
||||
MLFloat16(math::floatToHalf(-0.51f)), MLFloat16(math::floatToHalf(-0.12f))};
|
||||
|
||||
OpTester tester("ComplexMulConj", 1, onnxruntime::kMSDomain);
|
||||
tester.AddInput<MLFloat16>("A", {1, 2}, input_a_data);
|
||||
tester.AddInput<MLFloat16>("B", {4, 2}, input_b_data);
|
||||
tester.AddOutput<MLFloat16>("C", {4, 2}, output_data);
|
||||
|
||||
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
|
||||
execution_providers.push_back(DefaultCudaExecutionProvider());
|
||||
tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
|
||||
}
|
||||
|
||||
} // namespace test
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
Loading…
Reference in a new issue