saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-26 19:52:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

GatherNDGrad for CPU (#4123)

Browse source 
* GatherNDGrad on CPU

* Remove __CUDA_ARCH__ check in .cc files
This commit is contained in:
pengwa 2020-06-12 02:43:49 +08:00 committed by GitHub
parent 65a682354b
commit e6ccb1ac28
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 210 additions and 45 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

81
onnxruntime/core/providers/cpu/tensor/gather_nd.cc
View file

@ -44,12 +44,8 @@ ONNX_CPU_OPERATOR_KERNEL(
GatherND);
template <typename Tind>
Status GatherNDBase::PrepareForCompute(OpKernelContext* context, Prepare& p, concurrency::ThreadPool* tp) const {
const auto* input_tensor = context->Input<Tensor>(0);
const auto* indices_tensor = context->Input<Tensor>(1);
ORT_ENFORCE(input_tensor != nullptr && indices_tensor != nullptr, "GatherNDBase PrepareForCompute: Input count mismatch");
const auto& input_shape = input_tensor->Shape();
Status GatherNDBase::PrepareForCompute(const TensorShape& input_shape, const Tensor* indices_tensor,
const int64_t bytes_per_value, Prepare& p, concurrency::ThreadPool* tp) const {
const auto& indices_shape = indices_tensor->Shape();
if (indices_shape.NumDimensions() == 0) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "indices tensor must has rank larger than 0");
@ -61,37 +57,18 @@ Status GatherNDBase::PrepareForCompute(OpKernelContext* context, Prepare& p, con
const auto num_batches = input_shape.SizeToDimension(batch_dims_);
const auto input_batch_stride = input_shape.SizeFromDimension(batch_dims_);
const auto num_slices_per_batch = num_slices / num_batches;
int64_t last_indices_dimension = batch_dims_ + num_slice_dims;
if (last_indices_dimension > static_cast<int64_t>(input_shape.NumDimensions())) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"last dimension of indices must not be larger than rank of input tensor");
}
std::vector<int64_t> shape(indices_shape.GetDims().begin(), indices_shape.GetDims().end() - 1);
shape.insert(shape.end(), input_shape.GetDims().begin() + last_indices_dimension, input_shape.GetDims().end());
auto* output_tensor = context->Output(0, TensorShape(std::move(shape)));
std::vector<int64_t> sizes_from_slice_dims(num_slice_dims);
for (int64_t i = 0; i < num_slice_dims; ++i) {
sizes_from_slice_dims[i] = input_shape.SizeFromDimension(batch_dims_ + i + 1);
}
int64_t err_index = 0;
p.element_bytes = input_tensor->DataType()->Size();
p.element_bytes = bytes_per_value;
p.element_count_per_slice = slice_size;
p.bytes_per_slice = p.element_bytes * p.element_count_per_slice;
const auto* indices_data = indices_tensor->Data<Tind>();
p.slice_offsets.assign(num_slices, 0LL);
if (input_tensor->IsDataTypeString()) {
p.input_str_base = static_cast<const std::string*>(input_tensor->DataRaw());
p.output_str_base = static_cast<std::string*>(output_tensor->MutableDataRaw());
} else {
p.input_base = static_cast<const uint8_t*>(input_tensor->DataRaw());
p.output_base = static_cast<uint8_t*>(output_tensor->MutableDataRaw());
}
// Compute the element_offset
auto lambda = [&](int64_t slice_idx) {
const size_t batch_idx = slice_idx / num_slices_per_batch;
@ -125,15 +102,57 @@ Status GatherNDBase::PrepareForCompute(OpKernelContext* context, Prepare& p, con
: ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "invalid index found, index = ", err_index);
}
template Status GatherNDBase::PrepareForCompute<int32_t>(OpKernelContext*, Prepare&, concurrency::ThreadPool*) const;
template Status GatherNDBase::PrepareForCompute<int64_t>(OpKernelContext*, Prepare&, concurrency::ThreadPool*) const;
template Status GatherNDBase::PrepareForCompute<int32_t>(const TensorShape&,
const Tensor*,
const int64_t,
Prepare&,
concurrency::ThreadPool*) const;
template Status GatherNDBase::PrepareForCompute<int64_t>(const TensorShape&,
const Tensor*,
const int64_t,
Prepare&,
concurrency::ThreadPool*) const;
Status GatherND::Compute(OpKernelContext* context) const {
const auto* input_tensor = context->Input<Tensor>(0);
const auto* indices_tensor = context->Input<Tensor>(1);
ORT_ENFORCE(input_tensor != nullptr && indices_tensor != nullptr, "GatherNDBase PrepareForCompute: Input count mismatch");
const auto& input_shape = input_tensor->Shape();
const auto& indices_shape = indices_tensor->Shape();
int64_t last_indices_dimension = batch_dims_ + indices_shape[indices_shape.NumDimensions() - 1];
if (last_indices_dimension > static_cast<int64_t>(input_shape.NumDimensions())) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"last dimension of indices must not be larger than rank of input tensor");
}
std::vector<int64_t> shape(indices_shape.GetDims().begin(), indices_shape.GetDims().end() - 1);
shape.insert(shape.end(), input_shape.GetDims().begin() + last_indices_dimension,
input_shape.GetDims().end());
auto* output_tensor = context->Output(0, TensorShape(std::move(shape)));
Prepare p;
concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
ORT_RETURN_IF_ERROR(context->Input<Tensor>(1)->IsDataType<int32_t>()
? PrepareForCompute<int32_t>(context, p, tp)
: PrepareForCompute<int64_t>(context, p, tp));
if (input_tensor->IsDataTypeString()) {
p.input_str_base = static_cast<const std::string*>(input_tensor->DataRaw());
p.output_str_base = static_cast<std::string*>(output_tensor->MutableDataRaw());
} else {
p.input_base = static_cast<const uint8_t*>(input_tensor->DataRaw());
p.output_base = static_cast<uint8_t*>(output_tensor->MutableDataRaw());
}
auto bytes_per_value = input_tensor->DataType()->Size();
if (indices_tensor->IsDataType<int32_t>()) {
PrepareForCompute<int32_t>(input_shape, indices_tensor, bytes_per_value, p, tp);
} else if (indices_tensor->IsDataType<int64_t>()) {
PrepareForCompute<int64_t>(input_shape, indices_tensor, bytes_per_value, p, tp);
} else {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "indices tensor data type not supported");
}
return nullptr == p.input_str_base ? GatherNumber(p, tp) : GatherString(p, tp);
}

5
onnxruntime/core/providers/cpu/tensor/gather_nd.h
View file

@ -11,7 +11,7 @@ namespace concurrency {
class ThreadPool;
}
class GatherNDBase {
protected:
public:
struct Prepare {
const uint8_t* input_base;
const std::string* input_str_base;
@ -33,7 +33,8 @@ class GatherNDBase {
}; // struct Prepare
template <typename Tind>
Status PrepareForCompute(OpKernelContext* context, Prepare& p, concurrency::ThreadPool* tp) const;
Status PrepareForCompute(const TensorShape& input_shape, const Tensor* indices_tensor,
const int64_t bytes_per_value, Prepare& p, concurrency::ThreadPool* tp) const;
int64_t batch_dims_;
}; // class GatherNDBase

11
orttraining/orttraining/test/training_ops/cpu/tensor/gather_grad_op_test.cc
View file

@ -2,6 +2,7 @@
// Licensed under the MIT License.
#include "gtest/gtest.h"
#include "test/common/cuda_op_test_utils.h"
#include "test/common/tensor_op_test_utils.h"
#include "test/providers/provider_test_utils.h"
#include "core/util/math.h"
@ -45,9 +46,12 @@ void CalculateOutput(const int64_t stride, const int64_t num_input_before_gather
}
#ifdef USE_CUDA
#if __CUDA_ARCH__ >= 700
//TODO: Currently this cannot pass CI, due to GPU architecture problem
TEST(GatherOpTest, Gather_axis0_indices2d_half) {
if (NeedSkipIfCudaArchLowerThan(700)) {
return;
}
OpTester test("Gather");
test.AddAttribute<int64_t>("axis", 0LL);
test.AddInput<MLFloat16>("data", {3, 3},
@ -64,6 +68,10 @@ TEST(GatherOpTest, Gather_axis0_indices2d_half) {
}
TEST(GatherGradOpTest, GatherGrad_axis0_indices2d_half) {
if (NeedSkipIfCudaArchLowerThan(700)) {
return;
}
OpTester test("GatherGrad", 1, kMSDomain);
test.AddAttribute<int64_t>("axis", 0LL);
test.AddInput<int64_t>("shape", {2},
@ -79,7 +87,6 @@ TEST(GatherGradOpTest, GatherGrad_axis0_indices2d_half) {
test.Run();
}
#endif
#endif
TEST(GatherGradOpTest, GatherGrad_axis0_indices2d_float) {
OpTester test("GatherGrad", 1, kMSDomain);

48
orttraining/orttraining/test/training_ops/cpu/tensor/gather_nd_grad_op_test.cc
View file

@ -9,8 +9,7 @@
namespace onnxruntime {
namespace test {
#ifdef USE_CUDA
TEST(GatherNDGradOpTest, GatherNDGrad_slice_float_int64_t_batch_dims_1) {
TEST(GatherNDGradOpTest, GatherNDGrad_slice_float_int64_t_batch_dims_zero) {
OpTester test("GatherNDGrad", 1, kMSDomain);
test.AddAttribute<int64_t>("batch_dims", 0);
test.AddInput<int64_t>("shape", {3}, {2LL, 2LL, 3LL});
@ -20,8 +19,20 @@ TEST(GatherNDGradOpTest, GatherNDGrad_slice_float_int64_t_batch_dims_1) {
test.Run();
}
TEST(GatherNDGradOpTest, GatherNDGrad_slice_float_int64_t_batch_dims_zero_negative_indices) {
OpTester test("GatherNDGrad", 1, kMSDomain);
test.AddAttribute<int64_t>("batch_dims", 0);
test.AddInput<int64_t>("shape", {3}, {2LL, 2LL, 3LL});
test.AddInput<int64_t>("indices", {2, 2}, {-2LL, -1LL, -1LL, -2LL});
test.AddInput<float>("update", {2, 3}, ValueRange(6, 1.0f));
test.AddOutput<float>("output", {2, 2, 3}, {0, 0, 0, 1, 2, 3, 4, 5, 6, 0, 0, 0});
test.Run();
}
TEST(GatherNDGradOpTest, GatherNDGrad_slice_double_int32_t_batch_dims_3) {
if (!HasCudaEnvironment(600 /*min_cuda_architecture*/)) return;
if (NeedSkipIfCudaArchLowerThan(600)) {
return;
}
OpTester test("GatherNDGrad", 1, kMSDomain);
test.AddAttribute<int64_t>("batch_dims", 1);
@ -32,9 +43,11 @@ TEST(GatherNDGradOpTest, GatherNDGrad_slice_double_int32_t_batch_dims_3) {
test.Run();
}
TEST(GatherNDGradOpTest, GatherNDGrad_slice_half_int32_t_batch_dims_3) {
if (!HasCudaEnvironment(600 /*min_cuda_architecture*/)) return;
if (!HasCudaEnvironment(600)) {
// CPU GatherNDGrad did not support MLFloat16, so we need skip as well.
return;
}
OpTester test("GatherNDGrad", 1, kMSDomain);
test.AddAttribute<int64_t>("batch_dims", 1);
@ -74,7 +87,30 @@ TEST(GatherNDGradOpTest, GatherNDGrad_batch_dims_of_2) {
});
test.Run();
}
#endif
TEST(GatherNDGradOpTest, GatherNDGrad_batch_dims_two_negative_indices) {
OpTester test("GatherNDGrad", 1, kMSDomain);
test.AddAttribute<int64_t>("batch_dims", 2);
test.AddInput<int64_t>("shape", {4}, {2, 2, 2, 3});
test.AddInput<int64_t>(
"indices", {2, 2, 1},
{
-1, // batch 0
-1, // batch 1
-2, // batch 2
1, // batch 3
});
test.AddInput<float>("update", {2, 2, 3}, ValueRange<float>(12));
test.AddOutput<float>(
"output", {2, 2, 2, 3},
{
0, 0, 0, 0, 1, 2, // batch 0
0, 0, 0, 3, 4, 5, // batch 1
6, 7, 8, 0, 0, 0, // batch 2
0, 0, 0, 9, 10, 11, // batch 3
});
test.Run();
}
} // namespace test
} // namespace onnxruntime

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

@ -79,6 +79,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1,
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double, LayerNormalizationGrad);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SliceGrad);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, FastGeluGrad);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, GatherNDGrad);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, RecordEvent);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, WaitEvent);
@ -156,6 +157,7 @@ Status RegisterCpuTrainingKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, GistBinarizeDecoder)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SliceGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, FastGeluGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, GatherNDGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, RecordEvent)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, WaitEvent)>};

83
orttraining/orttraining/training_ops/cpu/tensor/gather_nd_grad.cc Normal file
View file

@ -0,0 +1,83 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "orttraining/training_ops/cpu/tensor/gather_nd_grad.h"
#include "core/providers/cpu/tensor/gather_nd.h"
#include "core/common/common.h"
namespace onnxruntime {
#ifndef DISABLE_CONTRIB_OPS
namespace contrib {
ONNX_OPERATOR_KERNEL_EX(GatherNDGrad, kMSDomain, 1, kCpuExecutionProvider,
KernelDefBuilder()
.TypeConstraint("T", {DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<double>()})
.TypeConstraint("Tind", {DataTypeImpl::GetTensorType<int64_t>(),
DataTypeImpl::GetTensorType<int32_t>()}),
GatherNDGrad);
} // namespace contrib
#endif
template <typename InputT>
struct GatherNDGradComputeImpl {
void operator()(GatherNDBase::Prepare& p, const Tensor* update_tensor) const {
const int64_t grad_size = update_tensor->Shape().Size();
const int64_t slice_size = p.element_count_per_slice;
const InputT* input_base_casted = reinterpret_cast<const InputT*>(p.input_base);
InputT* output_base_casted = reinterpret_cast<InputT*>(p.output_base);
for (int64_t i = 0; i < grad_size; i++) {
uint64_t slice_offset = p.slice_offsets[i / slice_size];
size_t j = i % slice_size;
output_base_casted[slice_offset + j] += input_base_casted[i];
}
}
};
Status GatherNDGrad::Compute(OpKernelContext* context) const {
const auto* shape_tensor = context->Input<Tensor>(0);
const auto* indices_tensor = context->Input<Tensor>(1);
const auto* update_tensor = context->Input<Tensor>(2);
ORT_ENFORCE(shape_tensor != nullptr && indices_tensor != nullptr && update_tensor != nullptr,
"GatherNDGrad::Compute : Input count mismatch");
auto shape_data = shape_tensor->Data<int64_t>();
auto input_shape = TensorShape(shape_data, shape_tensor->SizeInBytes() / sizeof(shape_tensor->DataType()));
const auto& indices_shape = indices_tensor->Shape();
int64_t last_indices_dimension = batch_dims_ + indices_shape[indices_shape.NumDimensions() - 1];
if (last_indices_dimension > static_cast<int64_t>(input_shape.NumDimensions())) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"last dimension of indices must not be larger than rank of input tensor");
}
auto* output_tensor = context->Output(0, input_shape);
memset(output_tensor->MutableDataRaw(), 0, output_tensor->SizeInBytes());
GatherNDBase::Prepare p;
p.input_base = static_cast<const uint8_t*>(update_tensor->DataRaw());
p.output_base = static_cast<uint8_t*>(output_tensor->MutableDataRaw());
auto bytes_per_value = update_tensor->DataType()->Size();
concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
if (indices_tensor->IsDataType<int32_t>()) {
PrepareForCompute<int32_t>(input_shape, indices_tensor, bytes_per_value, p, tp);
} else if (indices_tensor->IsDataType<int64_t>()) {
PrepareForCompute<int64_t>(input_shape, indices_tensor, bytes_per_value, p, tp);
} else {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "indices tensor data type not supported");
}
ORT_RETURN_IF_NOT(nullptr == p.input_str_base);
utils::MLTypeCallDispatcher<GatherNDGradComputeImpl, float, double> t_disp(update_tensor->GetElementType());
t_disp.Invoke(p, update_tensor);
return Status::OK();
}
} // namespace onnxruntime

20
orttraining/orttraining/training_ops/cpu/tensor/gather_nd_grad.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"
#include "core/providers/cpu/tensor/gather_nd.h"
namespace onnxruntime {
class GatherNDGrad final : public OpKernel, protected GatherNDBase {
public:
explicit GatherNDGrad(const OpKernelInfo& info) : OpKernel(info) {
info.GetAttrOrDefault("batch_dims", &batch_dims_, static_cast<int64_t>(0));
}
Status Compute(OpKernelContext* context) const override;
};
} // namespace onnxruntime

5
orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc
View file

@ -76,13 +76,10 @@ Status DispatchToGatherGradImpl(
if (utils::IsPrimitiveDataType<float>(t_data_type)) {
return DispatchToGatherGradImplByTin<float>(
tin_data_type, cuda_kernel, num_weights, stride, num_inputs, param_itrs, grad, indices, output);
}
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 700
else if (utils::IsPrimitiveDataType<MLFloat16>(t_data_type)) {
} else if (utils::IsPrimitiveDataType<MLFloat16>(t_data_type)) {
return DispatchToGatherGradImplByTin<MLFloat16>(
tin_data_type, cuda_kernel, num_weights, stride, num_inputs, param_itrs, grad, indices, output);
}
#endif
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "GatherGrad unsupported T type: ", t_data_type);
}