diff --git a/onnxruntime/contrib_ops/cuda/math/complex_mul.cc b/onnxruntime/contrib_ops/cuda/math/complex_mul.cc new file mode 100644 index 0000000000..1365009619 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/math/complex_mul.cc @@ -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()), \ + ComplexMul); \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + ComplexMulConj, \ + kMSDomain, \ + 1, \ + T, \ + kCudaExecutionProvider, \ + KernelDefBuilder() \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + ComplexMul); + +REGISTER_KERNEL_TYPED(float) +REGISTER_KERNEL_TYPED(MLFloat16) + +template +Status ComplexMul::ComputeInternal(OpKernelContext* context) const { + for (int index = 0; index < context->InputCount(); ++index) { + const Tensor* input = context->Input(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::MappedType>( + prepare.output_rank_or_simple_broadcast, + &prepare.lhs_padded_strides, + reinterpret_cast::MappedType*>(prepare.lhs_tensor->template Data()), + &prepare.rhs_padded_strides, + reinterpret_cast::MappedType*>(prepare.rhs_tensor->template Data()), + &prepare.fdm_output_strides, + prepare.fdm_H, + prepare.fdm_C, + reinterpret_cast::MappedType*>(prepare.output_tensor->template MutableData()), + 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 diff --git a/onnxruntime/contrib_ops/cuda/math/complex_mul.h b/onnxruntime/contrib_ops/cuda/math/complex_mul.h new file mode 100644 index 0000000000..b8cefdd4db --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/math/complex_mul.h @@ -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 +class ComplexMul : public BinaryElementwise { + public: + ComplexMul(const OpKernelInfo info) : BinaryElementwise{info} {} + Status ComputeInternal(OpKernelContext* context) const override; +}; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/math/complex_mul_impl.cu b/onnxruntime/contrib_ops/cuda/math/complex_mul_impl.cu new file mode 100644 index 0000000000..5b75df928e --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/math/complex_mul_impl.cu @@ -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 +__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 +__global__ void _ElementWiseWithStrideTwo( + int32_t output_rank, + const TArray lhs_padded_strides, + const T* lhs_data, + const TArray rhs_padded_strides, + const T* rhs_data, + const TArray 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(lhs_padded_strides.data_[dim]) * q; + } + + if (rhs_need_compute) { + rhs_index += static_cast(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 +void ComplexMul_Impl( + int32_t output_rank_or_simple_broadcast, + const TArray* lhs_padded_strides, + const T* lhs_data, + const TArray* rhs_padded_strides, + const T* rhs_data, + const TArray* 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(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread)); + CUDA_LONG N = static_cast(count); + + if (lhs_padded_strides && rhs_padded_strides && lhs_padded_strides->size_ && rhs_padded_strides->size_) + _ElementWiseWithStrideTwo<<>>( + 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<<>>( + 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<<>>( + 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( \ + int32_t output_rank_or_simple_broadcast, \ + const TArray* lhs_padded_strides, \ + const T* lhs_data, \ + const TArray* rhs_padded_strides, \ + const T* rhs_data, \ + const TArray* 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 diff --git a/onnxruntime/contrib_ops/cuda/math/complex_mul_impl.h b/onnxruntime/contrib_ops/cuda/math/complex_mul_impl.h new file mode 100644 index 0000000000..d48eea9a9f --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/math/complex_mul_impl.h @@ -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 +void ComplexMul_Impl( + int32_t output_rank_or_simple_broadcast, + const TArray* lhs_padded_strides, + const T* lhs_data, + const TArray* rhs_padded_strides, + const T* rhs_data, + const TArray* 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 diff --git a/onnxruntime/contrib_ops/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda_contrib_kernels.cc index eef175e7b8..b8bed20eae 100644 --- a/onnxruntime/contrib_ops/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda_contrib_kernels.cc @@ -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, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, // These ops were experimental ops in onnx domain which have been removed now. We add them here as // contrib ops to maintain backward compatibility diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 598154ac1e..4097f232bb 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -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) diff --git a/onnxruntime/test/contrib_ops/element_wise_ops_test.cc b/onnxruntime/test/contrib_ops/element_wise_ops_test.cc index 10da620c53..a7debe845e 100644 --- a/onnxruntime/test/contrib_ops/element_wise_ops_test.cc +++ b/onnxruntime/test/contrib_ops/element_wise_ops_test.cc @@ -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 input_a_data = { + -0.5f, 0.6f}; + + std::vector input_b_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector 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("A", {1, 2}, input_a_data); + tester.AddInput("B", {4, 2}, input_b_data); + tester.AddOutput("C", {4, 2}, output_data); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST(MathOpTest, ComplexMulConj) { + if (DefaultCudaExecutionProvider() == nullptr) return; + + std::vector input_a_data = { + -0.5f, 0.6f}; + + std::vector input_b_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector 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("A", {1, 2}, input_a_data); + tester.AddInput("B", {4, 2}, input_b_data); + tester.AddOutput("C", {4, 2}, output_data); + + std::vector> 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 input_a_data = { + MLFloat16(math::floatToHalf(-0.5f)), MLFloat16(math::floatToHalf(0.6f))}; + + std::vector 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 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("A", {1, 2}, input_a_data); + tester.AddInput("B", {4, 2}, input_b_data); + tester.AddOutput("C", {4, 2}, output_data); + + std::vector> 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 input_a_data = { + MLFloat16(math::floatToHalf(-0.5f)), MLFloat16(math::floatToHalf(0.6f))}; + + std::vector 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 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("A", {1, 2}, input_a_data); + tester.AddInput("B", {4, 2}, input_b_data); + tester.AddOutput("C", {4, 2}, output_data); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + } // namespace test } // namespace onnxruntime