diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index c24c6a56f0..ec0f394ce1 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -711,6 +711,7 @@ Do not modify directly.* |||[1, 9]|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |Transpose|*in* data:**T**
*out* transposed:**T**|13+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||[1, 12]|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| +|Trilu|*in* input:**T**
*in* k:**tensor(int64)**
*out* output:**T**|14+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |Unsqueeze|*in* data:**T**
*in* axes:**tensor(int64)**
*out* expanded:**T**

or

*in* data:**T**
*out* expanded:**T**|13+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||[11, 12]|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||[1, 10]|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| @@ -744,5 +745,6 @@ Do not modify directly.* |Rfft|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| |SkipLayerNormalization|*in* input:**T**
*in* skip:**T**
*in* gamma:**T**
*in* beta:**T**
*in* bias:**T**
*out* output:**T**
*out* mean:**U**
*out* inv_std_var:**U**|1+|**T** = tensor(float), tensor(float16)| |TransposeMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| +|Trilu|*in* X:**T**
*in* k:**tensor(int64)**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| | | | | diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 89fedfab2e..d29ef3e043 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -38,7 +38,6 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, Bias class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BiasDropout); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, NGramRepeatBlock); - // These ops were experimental ops in onnx domain which have been removed now. We add them here as // contrib ops to maintain backward compatibility class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, float, Affine); @@ -77,6 +76,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, double_double, SimplifiedLayerNormalization); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MLFloat16_float, SimplifiedLayerNormalization); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, Inverse); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, Trilu); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, int8_t_MLFloat16, QuantizeLinear); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, uint8_t_MLFloat16, QuantizeLinear); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, int8_t_MLFloat16, DequantizeLinear); @@ -100,7 +100,7 @@ KernelCreateInfo BuildKernelCreateInfo() { Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { static const BuildKernelCreateInfoFn function_table[] = { - BuildKernelCreateInfo, //default entry to avoid the list become empty after ops-reducing + BuildKernelCreateInfo, // default entry to avoid the list become empty after ops-reducing BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -128,46 +128,46 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { 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 - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - 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 + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -175,6 +175,7 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, #if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cuda/tensor/trilu.cc b/onnxruntime/contrib_ops/cuda/tensor/trilu.cc new file mode 100644 index 0000000000..257526e558 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/tensor/trilu.cc @@ -0,0 +1,23 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/cuda/tensor/trilu.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +ONNX_OPERATOR_KERNEL_EX( + Trilu, + kMSDomain, + 1, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .InputMemoryType(OrtMemTypeCPUInput, 1) + .MayInplace(0, 0) + .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()), + onnxruntime::cuda::Trilu); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/tensor/trilu.cc b/onnxruntime/core/providers/cpu/tensor/trilu.cc index dff188b855..7a6864c1cc 100644 --- a/onnxruntime/core/providers/cpu/tensor/trilu.cc +++ b/onnxruntime/core/providers/cpu/tensor/trilu.cc @@ -21,7 +21,8 @@ ONNX_OPERATOR_KERNEL_EX( kCpuExecutionProvider, KernelDefBuilder().MayInplace(0, 0).TypeConstraint("T", BuildKernelDefConstraints()), Trilu); -} // namespace contrib + +} // namespace contrib #endif ONNX_OPERATOR_KERNEL_EX( diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index 14fa9fd737..d2e4fc68dc 100755 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -899,7 +899,7 @@ class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kO class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, 12, Dropout); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, Einsum); -//OpSet 13 +// OpSet 13 class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 14, Pow); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 13, int32_t, Add); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, 13, int64_t, Add); @@ -1182,6 +1182,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, int8_t, ReduceMin); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, uint8_t, ReduceMin); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, int64_t, ReduceMin); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, Trilu); #if defined(CUDA_VERSION) && CUDA_VERSION >= 11000 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, BFloat16, Add); @@ -1191,7 +1192,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, BFloat16, Relu); #endif -//OpSet 15 +// OpSet 15 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 15, Pow); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 15, float, BatchNormalization); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 15, double, BatchNormalization); @@ -1205,7 +1206,7 @@ KernelCreateInfo BuildKernelCreateInfo() { static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { static const BuildKernelCreateInfoFn function_table[] = { - BuildKernelCreateInfo, //default entry to avoid the list become empty after ops-reducing + BuildKernelCreateInfo, // default entry to avoid the list become empty after ops-reducing BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -2033,6 +2034,7 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, #endif + BuildKernelCreateInfo, // OpSet 15 BuildKernelCreateInfo, diff --git a/onnxruntime/core/providers/cuda/tensor/trilu.cc b/onnxruntime/core/providers/cuda/tensor/trilu.cc new file mode 100644 index 0000000000..f9def957e1 --- /dev/null +++ b/onnxruntime/core/providers/cuda/tensor/trilu.cc @@ -0,0 +1,63 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/cuda/tensor/trilu.h" +#include "core/providers/cuda/tensor/trilu_impl.h" +#include "core/providers/cpu/tensor/utils.h" + +using namespace onnxruntime::common; + +namespace onnxruntime { +namespace cuda { + +ONNX_OPERATOR_KERNEL_EX( + Trilu, + kOnnxDomain, + 14, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .InputMemoryType(OrtMemTypeCPUInput, 1) + .MayInplace(0, 0) + .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()), + Trilu); + +Status Trilu::ComputeInternal(OpKernelContext* ctx) const { + const Tensor* input_ptr = ctx->Input(0); + const auto* k = ctx->Input(1); + + int64_t k_val = 0; + if (k) { + ORT_ENFORCE(IsScalarOr1ElementVector(k), "k should be a 1-D or 0-D tensor."); + k_val = *(k->template Data()); + } + if (input_ptr == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "input count mismatch"); + const Tensor& input = *input_ptr; + const auto& shape = input.Shape(); + const auto& input_dims = shape.GetDims(); + int32_t rank = gsl::narrow_cast(input_dims.size()); + if (rank < 2) { + return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Input tensor should have a rank of at least 2"); + } + Tensor* output = ctx->Output(0, shape); + int64_t matrix_size = input_dims[rank - 1] * input_dims[rank - 2]; + if (matrix_size == 0) { + return Status::OK(); + } + const fast_divmod row_col_divmod_indices(gsl::narrow_cast(input_dims[rank - 1])); + const fast_divmod batch_divmod_indices(gsl::narrow_cast(matrix_size)); + + size_t element_size = input.DataType()->Size(); + return TriluImpl( + this->Stream(), + upper_, + element_size, + k_val, + input.DataRaw(), + output->MutableDataRaw(), + gsl::narrow(shape.Size()), + batch_divmod_indices, + row_col_divmod_indices); +} + +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/tensor/trilu.h b/onnxruntime/core/providers/cuda/tensor/trilu.h new file mode 100644 index 0000000000..25271f76a5 --- /dev/null +++ b/onnxruntime/core/providers/cuda/tensor/trilu.h @@ -0,0 +1,21 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/cuda/cuda_kernel.h" + +namespace onnxruntime { +namespace cuda { + +class Trilu final : public CudaKernel { + public: + Trilu(const OpKernelInfo& info) : CudaKernel(info), upper_(info.GetAttrOrDefault("upper", 1) >= 1) { + } + ~Trilu() = default; + Status ComputeInternal(OpKernelContext* context) const override; + + private: + bool upper_; +}; + +} // namespace cuda +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/cuda/tensor/trilu_impl.cu b/onnxruntime/core/providers/cuda/tensor/trilu_impl.cu new file mode 100644 index 0000000000..30b441a2e4 --- /dev/null +++ b/onnxruntime/core/providers/cuda/tensor/trilu_impl.cu @@ -0,0 +1,122 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/cuda/cu_inc/common.cuh" +#include "trilu_impl.h" +#include +namespace onnxruntime { +namespace cuda { + +template +__global__ void TriluKernel( + int64_t k, + const T* input_data, + T* output_data, + const CUDA_LONG N, + const fast_divmod batch_divmod_indices, + const fast_divmod row_col_divmod_indices) { + CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N); + + int row, col; + + row_col_divmod_indices.divmod(batch_divmod_indices.mod(id), row, col); + output_data[id] = upper ? (((row + k) <= col) ? input_data[id] : 0) : (((row + k) >= col) ? input_data[id] : 0); +} + +Status TriluImpl( + cudaStream_t stream, + bool upper, + size_t element_size, + int64_t k, + const void* input_data, + void* output_data, + int N, + const fast_divmod& batch_divmod_indices, + const fast_divmod& row_col_divmod_indices) { + int blocksPerGrid = (int)(ceil(static_cast(N) / GridDim::maxThreadsPerBlock)); + switch (element_size) { + case sizeof(int8_t): + if (upper) { + TriluKernel<<>>( + k, + reinterpret_cast::MappedType*>(input_data), + reinterpret_cast::MappedType*>(output_data), + (CUDA_LONG)N, + batch_divmod_indices, + row_col_divmod_indices); + } else { + TriluKernel<<>>( + k, + reinterpret_cast::MappedType*>(input_data), + reinterpret_cast::MappedType*>(output_data), + (CUDA_LONG)N, + batch_divmod_indices, + row_col_divmod_indices); + } + break; + case sizeof(int16_t): + if (upper) { + TriluKernel<<>>( + k, + reinterpret_cast::MappedType*>(input_data), + reinterpret_cast::MappedType*>(output_data), + (CUDA_LONG)N, + batch_divmod_indices, + row_col_divmod_indices); + } else { + TriluKernel<<>>( + k, + reinterpret_cast::MappedType*>(input_data), + reinterpret_cast::MappedType*>(output_data), + (CUDA_LONG)N, + batch_divmod_indices, + row_col_divmod_indices); + } + break; + case sizeof(int32_t): + if (upper) { + TriluKernel<<>>( + k, + reinterpret_cast::MappedType*>(input_data), + reinterpret_cast::MappedType*>(output_data), + (CUDA_LONG)N, + batch_divmod_indices, + row_col_divmod_indices); + } else { + TriluKernel<<>>( + k, + reinterpret_cast::MappedType*>(input_data), + reinterpret_cast::MappedType*>(output_data), + (CUDA_LONG)N, + batch_divmod_indices, + row_col_divmod_indices); + } + break; + case sizeof(int64_t): + if (upper) { + TriluKernel<<>>( + k, + reinterpret_cast::MappedType*>(input_data), + reinterpret_cast::MappedType*>(output_data), + (CUDA_LONG)N, + batch_divmod_indices, + row_col_divmod_indices); + } else { + TriluKernel<<>>( + k, + reinterpret_cast::MappedType*>(input_data), + reinterpret_cast::MappedType*>(output_data), + (CUDA_LONG)N, + batch_divmod_indices, + row_col_divmod_indices); + } + break; + default: + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Type not supported for transpose on CUDA. Element size was ", + element_size); + } + return Status::OK(); +} + +} // namespace cuda +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/cuda/tensor/trilu_impl.h b/onnxruntime/core/providers/cuda/tensor/trilu_impl.h new file mode 100644 index 0000000000..29309ca20b --- /dev/null +++ b/onnxruntime/core/providers/cuda/tensor/trilu_impl.h @@ -0,0 +1,23 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include "core/providers/cuda/shared_inc/cuda_utils.h" + +namespace onnxruntime { +namespace cuda { + +Status TriluImpl( + cudaStream_t stream, + bool upper, + size_t element_size, + int64_t k, + const void* input_data, + void* output_data, + int N, + const fast_divmod& batch_divmod_indices, + const fast_divmod& row_col_divmod_indices); + +} // namespace cuda +} // namespace onnxruntime diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc index 128312dccb..1fad171dae 100755 --- a/orttraining/orttraining/core/graph/gradient_builder.cc +++ b/orttraining/orttraining/core/graph/gradient_builder.cc @@ -96,6 +96,16 @@ IMPLEMENT_GRADIENT_BUILDER(GetTanhGradient) { {GI(0)})}; } +IMPLEMENT_GRADIENT_BUILDER(GetTriluGradient) { + if (GetSrcNodeInputSize() == 1) { + return std::vector{NodeDef(OpDef{"Trilu", kMSDomain, 1}, {GO(0)}, {GI(0)}, SrcNodeAttributes())}; + } else if (GetSrcNodeInputSize() == 2) { + return std::vector{NodeDef(OpDef{"Trilu", kMSDomain, 1}, {GO(0), I(1)}, {GI(0)}, SrcNodeAttributes())}; + } else { + ORT_ENFORCE(false, "the number of input arguments must be 1 or 2"); + } +} + IMPLEMENT_GRADIENT_BUILDER(GetSqrtGradient) { std::vector result; NodeDef half_constant_node = HalfConstantNode(OElemType(0)); diff --git a/orttraining/orttraining/core/graph/gradient_builder.h b/orttraining/orttraining/core/graph/gradient_builder.h index 2bd8791358..8947f40329 100755 --- a/orttraining/orttraining/core/graph/gradient_builder.h +++ b/orttraining/orttraining/core/graph/gradient_builder.h @@ -77,6 +77,7 @@ DECLARE_GRADIENT_BUILDER(GetPadGradient) DECLARE_GRADIENT_BUILDER(GetIdentityGradient) DECLARE_GRADIENT_BUILDER(GetPythonOpGradient) DECLARE_GRADIENT_BUILDER(GetScatterNDGradient) +DECLARE_GRADIENT_BUILDER(GetTriluGradient) DECLARE_GRADIENT_BUILDER(GetExternalGradient) diff --git a/orttraining/orttraining/core/graph/gradient_builder_registry.cc b/orttraining/orttraining/core/graph/gradient_builder_registry.cc index 97b337f451..6fc1fda644 100755 --- a/orttraining/orttraining/core/graph/gradient_builder_registry.cc +++ b/orttraining/orttraining/core/graph/gradient_builder_registry.cc @@ -108,6 +108,7 @@ void GradientBuilderRegistry::RegisterGradientBuilders() { REGISTER_GRADIENT_BUILDER("Identity", GetIdentityGradient); REGISTER_GRADIENT_BUILDER("PythonOp", GetPythonOpGradient); REGISTER_GRADIENT_BUILDER("ScatterND", GetScatterNDGradient); + REGISTER_GRADIENT_BUILDER("Trilu", GetTriluGradient); REGISTER_GRADIENT_BUILDER("ExternalGradient", GetExternalGradient); }; diff --git a/orttraining/orttraining/test/gradient/gradient_ops_test.cc b/orttraining/orttraining/test/gradient/gradient_ops_test.cc index d32232dd3d..53f885d8eb 100644 --- a/orttraining/orttraining/test/gradient/gradient_ops_test.cc +++ b/orttraining/orttraining/test/gradient/gradient_ops_test.cc @@ -2749,6 +2749,61 @@ TEST(GradientCheckerTest, ScatterNDGrad) { } } +TEST(GradientCheckerTest, TriluGrad) { + float max_error; + GradientChecker gradient_checker; + OpDef op_def{"Trilu", kMSDomain, 1}; + const int M = 3; + const int N = 4; + TensorShape shape = {M, N}; + TensorInfo x_info(shape); + TensorInfo k_info({1}, false, nullptr, DataTypeImpl::GetTensorType()); + TensorInfo y_info(shape); + std::vector x_data = {}; + // Initialize input data + float f = 1.0; + for (int i = 0; i < M; i++) { + for (int j = 0; j < N; j++, f++) { + x_data.push_back(f); + } + } + + // Test without optional input and without attribute + { + ASSERT_STATUS_OK(gradient_checker.ComputeGradientError(op_def, {x_info}, {y_info}, &max_error, {x_data})); + EXPECT_IS_TINY(max_error); + } + { + // Test without optional input and with attribute upper=1 + ASSERT_STATUS_OK(gradient_checker.ComputeGradientError(op_def, {x_info}, {y_info}, &max_error, {x_data}, {MakeAttribute("upper", int64_t(1))})); + EXPECT_IS_TINY(max_error); + } + { + // Test without optional input and with attribute upper=0 + ASSERT_STATUS_OK(gradient_checker.ComputeGradientError(op_def, {x_info}, {y_info}, &max_error, {x_data}, {MakeAttribute("upper", int64_t(0))})); + EXPECT_IS_TINY(max_error); + } + for (int64_t k = -M; k <= M; k++) { + std::vector k_data = {static_cast(k)}; + + // Test with optional input and without attribute + { + ASSERT_STATUS_OK(gradient_checker.ComputeGradientError(op_def, {x_info, k_info}, {y_info}, &max_error, {x_data, k_data})); + EXPECT_IS_TINY(max_error); + } + { + // Test with optional input and with attribute upper=1 + ASSERT_STATUS_OK(gradient_checker.ComputeGradientError(op_def, {x_info, k_info}, {y_info}, &max_error, {x_data, k_data}, {MakeAttribute("upper", int64_t(1))})); + EXPECT_IS_TINY(max_error); + } + { + // Test with optional input and with attribute upper=0 + ASSERT_STATUS_OK(gradient_checker.ComputeGradientError(op_def, {x_info, k_info}, {y_info}, &max_error, {x_data, k_data}, {MakeAttribute("upper", int64_t(0))})); + EXPECT_IS_TINY(max_error); + } + } +} + } // namespace test } // namespace onnxruntime diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index f6c579e353..da7e849eae 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -4373,7 +4373,8 @@ def test_serialize_ortmodule(): pt_out = pt_model.train_step(x_1) ort_out = ort_model.train_step(x_2) _test_helpers.assert_values_are_close(pt_out, ort_out) - _test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model) + _test_helpers.assert_gradients_match_and_reset_gradient( + ort_model, pt_model) pt_out, ort_out = None, None # Serialize ortmodule @@ -4388,4 +4389,42 @@ def test_serialize_ortmodule(): ort_out = ort_model_2.train_step(x_2) assert ort_out is not None _test_helpers.assert_values_are_close(pt_out, ort_out) - _test_helpers.assert_gradients_match_and_reset_gradient(ort_model_2, pt_model) + _test_helpers.assert_gradients_match_and_reset_gradient( + ort_model_2, pt_model) + + +@pytest.mark.parametrize("batch_size, M, N", [(1, 2, 3), (1, 4, 3), (1, 5, 5), (10, 3, 4), (10, 4, 3), (10, 4, 4)]) +@pytest.mark.parametrize("k", [None, -5, -3, -1, 0, 2, 4]) +@pytest.mark.parametrize("has_upper, upper", [(True, 1), (True, 0), (False, 1)]) +def test_trilu_grad(batch_size, M, N, k, has_upper, upper): + class NeuralNetTrilu(torch.nn.Module): + def __init__(self, has_upper, upper): + super(NeuralNetTrilu, self).__init__() + self.upper = upper + self.has_upper = has_upper + + def forward(self, x, k): + if self.has_upper is False or self.upper == 1: + y = torch.triu(x) if k is None else torch.triu(x, k) + else: + y = torch.tril(x) if k is None else torch.tril(x, k) + return y + + def run_step(model, x, k): + prediction = model(x, k) + loss = prediction.sum() + loss.backward() + return prediction, loss + + device = 'cuda' + pt_model = NeuralNetTrilu(has_upper, upper).to(device) + ort_model = ORTModule(copy.deepcopy(pt_model)) + + pt_x = torch.rand((batch_size, M, N), requires_grad=True, device=device) + ort_x = copy.deepcopy(pt_x) + + pt_prediction, pt_loss = run_step(pt_model, pt_x, k) + ort_prediction, ort_loss = run_step(ort_model, ort_x, k) + _test_helpers.assert_values_are_close(pt_prediction, ort_prediction) + _test_helpers.assert_values_are_close(pt_loss, ort_loss) + _test_helpers.assert_values_are_close(pt_x.grad, ort_x.grad)