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

reducealll2 cpu kernel (#5833)

Browse source 
Co-authored-by: Vincent Wang <weicwang@AiFramework2080ti2.corp.microsoft.com>
This commit is contained in:
Vincent Wang 2020-11-19 10:20:05 +08:00 committed by GitHub
parent 7196d4206f
commit 47185b9513
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 77 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
orttraining/orttraining/test/training_ops/cpu/reduction/reduction_ops_test.cc
View file

@ -84,6 +84,7 @@ TEST(AllOpTest, All_1d_large) {
}
}
}
#endif
class ReductionOpTest : public ::testing::TestWithParam<bool> {
protected:
@ -102,6 +103,7 @@ TEST_P(ReductionOpTest, ReduceAllL2) {
test.Run();
}
#ifdef USE_CUDA
TEST_P(ReductionOpTest, ReduceAllL2HalfHalf) {
OpTester test("ReduceAllL2", 1, onnxruntime::kMSDomain, true);
test.SetDeterminism(GetParam());
@ -163,6 +165,7 @@ TEST_P(ReductionOpTest, ReduceAllL2HalfFloat) {
test.AddOutput<float>("reduced", {}, result);
test.Run();
}
#endif
void TestMultiTensorReduce(
const int tensor_count,
@ -226,8 +229,6 @@ TEST_P(ReductionOpTest, ReduceAllL2Many) {
// invoke with and without use_determinism flag for session
INSTANTIATE_TEST_SUITE_P(ReductionOpTestWrapper, ReductionOpTest, ::testing::Bool());
#endif
TEST(ReductionOpTest, ReduceSumTraining_int32) {
OpTester test("ReduceSumTraining", 1, onnxruntime::kMSDomain);
test.AddAttribute("keepdims", (int64_t)1);

4
orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc
View file

@ -107,6 +107,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1,
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_int64_t, Scale);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_int32_t, Scale);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float_float, ReduceAllL2);
#ifdef USE_MPI
// Pipeline communication operators.
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Send);
@ -217,6 +219,8 @@ Status RegisterCpuTrainingKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_int64_t, Scale)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_int32_t, Scale)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float_float, ReduceAllL2)>,
#ifdef USE_MPI
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Send)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Recv)>,

50
orttraining/orttraining/training_ops/cpu/reduction/reduction_all.cc Normal file
View file

@ -0,0 +1,50 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "orttraining/training_ops/cpu/reduction/reduction_all.h"
#include "core/providers/cpu/reduction/reduction_ops.h"
namespace onnxruntime {
namespace contrib {
#define REGISTER_REDUCEALLL2_KERNEL_TYPED(TIn, TOut) \
ONNX_OPERATOR_TYPED_KERNEL_EX(ReduceAllL2, kMSDomain, 1, TIn##_##TOut, kCpuExecutionProvider, \
KernelDefBuilder() \
.TypeConstraint("TIn", DataTypeImpl::GetTensorType<TIn>()) \
.TypeConstraint("TOut", DataTypeImpl::GetTensorType<TOut>()), \
ReduceAllL2<TIn, TOut>);
REGISTER_REDUCEALLL2_KERNEL_TYPED(float, float)
template <typename TIn, typename TOut>
Status ReduceAllL2<TIn, TOut>::Compute(OpKernelContext* ctx) const {
// Get Input tensor count.
const auto total_tensor_count = ctx->InputCount();
std::vector<const TIn*> tensor_pointers(total_tensor_count);
std::vector<int64_t> tensor_sizes(total_tensor_count);
for (int i = 0; i < total_tensor_count; ++i) {
const Tensor* input = ctx->Input<Tensor>(i);
const auto size = input->Shape().Size();
ORT_ENFORCE(size <= std::numeric_limits<int>::max(), "Number of reduced elements (", size,
") exceeds the max allowed value (", std::numeric_limits<int>::max(), ").");
tensor_pointers[i] = input->template Data<TIn>();
tensor_sizes[i] = size;
}
// Allocate output tensor.
Tensor* output = ctx->Output(0, {});
TOut* output_data = output->template MutableData<TOut>();
*output_data = TOut(0.f);
// perform reduction l2norm = sqrt[sum(tensor[i][j]**2)] for i,j over all tensor elements
for (int i = 0; i < total_tensor_count; ++i) {
*output_data +=
ReduceAggregatorSumSquare<TIn, TOut>(tensor_sizes[i], tensor_pointers[i][0]).aggall(tensor_pointers[i]);
}
*output_data = reduce_sqrt<TOut>(*output_data);
return Status::OK();
}
} // namespace contrib
} // namespace onnxruntime

20
orttraining/orttraining/training_ops/cpu/reduction/reduction_all.h Normal file
View file

@ -0,0 +1,20 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
namespace onnxruntime {
namespace contrib {
template <typename TIn, typename TOut>
class ReduceAllL2 final : public OpKernel {
public:
ReduceAllL2(const OpKernelInfo& info) : OpKernel(info) {}
Status Compute(OpKernelContext* context) const override;
};
} // namespace contrib
} // namespace onnxruntime