diff --git a/onnxruntime/core/providers/cpu/tensor/gather_nd.cc b/onnxruntime/core/providers/cpu/tensor/gather_nd.cc index 94b8f22539..62778cf28b 100644 --- a/onnxruntime/core/providers/cpu/tensor/gather_nd.cc +++ b/onnxruntime/core/providers/cpu/tensor/gather_nd.cc @@ -44,12 +44,8 @@ ONNX_CPU_OPERATOR_KERNEL( GatherND); template -Status GatherNDBase::PrepareForCompute(OpKernelContext* context, Prepare& p, concurrency::ThreadPool* tp) const { - const auto* input_tensor = context->Input(0); - const auto* indices_tensor = context->Input(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(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 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 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(); p.slice_offsets.assign(num_slices, 0LL); - if (input_tensor->IsDataTypeString()) { - p.input_str_base = static_cast(input_tensor->DataRaw()); - p.output_str_base = static_cast(output_tensor->MutableDataRaw()); - } else { - p.input_base = static_cast(input_tensor->DataRaw()); - p.output_base = static_cast(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(OpKernelContext*, Prepare&, concurrency::ThreadPool*) const; -template Status GatherNDBase::PrepareForCompute(OpKernelContext*, Prepare&, concurrency::ThreadPool*) const; +template Status GatherNDBase::PrepareForCompute(const TensorShape&, + const Tensor*, + const int64_t, + Prepare&, + concurrency::ThreadPool*) const; +template Status GatherNDBase::PrepareForCompute(const TensorShape&, + const Tensor*, + const int64_t, + Prepare&, + concurrency::ThreadPool*) const; Status GatherND::Compute(OpKernelContext* context) const { + const auto* input_tensor = context->Input(0); + const auto* indices_tensor = context->Input(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(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 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(1)->IsDataType() - ? PrepareForCompute(context, p, tp) - : PrepareForCompute(context, p, tp)); + if (input_tensor->IsDataTypeString()) { + p.input_str_base = static_cast(input_tensor->DataRaw()); + p.output_str_base = static_cast(output_tensor->MutableDataRaw()); + } else { + p.input_base = static_cast(input_tensor->DataRaw()); + p.output_base = static_cast(output_tensor->MutableDataRaw()); + } + + auto bytes_per_value = input_tensor->DataType()->Size(); + + if (indices_tensor->IsDataType()) { + PrepareForCompute(input_shape, indices_tensor, bytes_per_value, p, tp); + } else if (indices_tensor->IsDataType()) { + PrepareForCompute(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); } diff --git a/onnxruntime/core/providers/cpu/tensor/gather_nd.h b/onnxruntime/core/providers/cpu/tensor/gather_nd.h index 26c3088b4a..77f15858e2 100644 --- a/onnxruntime/core/providers/cpu/tensor/gather_nd.h +++ b/onnxruntime/core/providers/cpu/tensor/gather_nd.h @@ -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 - 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 diff --git a/orttraining/orttraining/test/training_ops/cpu/tensor/gather_grad_op_test.cc b/orttraining/orttraining/test/training_ops/cpu/tensor/gather_grad_op_test.cc index d34752d26c..67589b5208 100644 --- a/orttraining/orttraining/test/training_ops/cpu/tensor/gather_grad_op_test.cc +++ b/orttraining/orttraining/test/training_ops/cpu/tensor/gather_grad_op_test.cc @@ -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("axis", 0LL); test.AddInput("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("axis", 0LL); test.AddInput("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); diff --git a/orttraining/orttraining/test/training_ops/cpu/tensor/gather_nd_grad_op_test.cc b/orttraining/orttraining/test/training_ops/cpu/tensor/gather_nd_grad_op_test.cc index 02c32a0827..19392478ac 100644 --- a/orttraining/orttraining/test/training_ops/cpu/tensor/gather_nd_grad_op_test.cc +++ b/orttraining/orttraining/test/training_ops/cpu/tensor/gather_nd_grad_op_test.cc @@ -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("batch_dims", 0); test.AddInput("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("batch_dims", 0); + test.AddInput("shape", {3}, {2LL, 2LL, 3LL}); + test.AddInput("indices", {2, 2}, {-2LL, -1LL, -1LL, -2LL}); + test.AddInput("update", {2, 3}, ValueRange(6, 1.0f)); + test.AddOutput("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("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("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("batch_dims", 2); + test.AddInput("shape", {4}, {2, 2, 2, 3}); + test.AddInput( + "indices", {2, 2, 1}, + { + -1, // batch 0 + -1, // batch 1 + -2, // batch 2 + 1, // batch 3 + }); + test.AddInput("update", {2, 2, 3}, ValueRange(12)); + test.AddOutput( + "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 diff --git a/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc b/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc index f24b2b2876..692db2eded 100644 --- a/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc @@ -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, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo}; diff --git a/orttraining/orttraining/training_ops/cpu/tensor/gather_nd_grad.cc b/orttraining/orttraining/training_ops/cpu/tensor/gather_nd_grad.cc new file mode 100644 index 0000000000..1df6dd20e6 --- /dev/null +++ b/orttraining/orttraining/training_ops/cpu/tensor/gather_nd_grad.cc @@ -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(), + DataTypeImpl::GetTensorType()}) + .TypeConstraint("Tind", {DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()}), + GatherNDGrad); + +} // namespace contrib + +#endif + +template +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(p.input_base); + InputT* output_base_casted = reinterpret_cast(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(0); + const auto* indices_tensor = context->Input(1); + const auto* update_tensor = context->Input(2); + + ORT_ENFORCE(shape_tensor != nullptr && indices_tensor != nullptr && update_tensor != nullptr, + "GatherNDGrad::Compute : Input count mismatch"); + + auto shape_data = shape_tensor->Data(); + 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(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(update_tensor->DataRaw()); + p.output_base = static_cast(output_tensor->MutableDataRaw()); + + auto bytes_per_value = update_tensor->DataType()->Size(); + concurrency::ThreadPool* tp = context->GetOperatorThreadPool(); + if (indices_tensor->IsDataType()) { + PrepareForCompute(input_shape, indices_tensor, bytes_per_value, p, tp); + } else if (indices_tensor->IsDataType()) { + PrepareForCompute(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 t_disp(update_tensor->GetElementType()); + t_disp.Invoke(p, update_tensor); + + return Status::OK(); +} + +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cpu/tensor/gather_nd_grad.h b/orttraining/orttraining/training_ops/cpu/tensor/gather_nd_grad.h new file mode 100644 index 0000000000..13ac0b70ce --- /dev/null +++ b/orttraining/orttraining/training_ops/cpu/tensor/gather_nd_grad.h @@ -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(0)); + } + Status Compute(OpKernelContext* context) const override; +}; + +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc b/orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc index 56f0e4e984..1c2a6a662d 100644 --- a/orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc +++ b/orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc @@ -76,13 +76,10 @@ Status DispatchToGatherGradImpl( if (utils::IsPrimitiveDataType(t_data_type)) { return DispatchToGatherGradImplByTin( 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(t_data_type)) { + } else if (utils::IsPrimitiveDataType(t_data_type)) { return DispatchToGatherGradImplByTin( 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); }