From 530c4b0a67206e5a57eccdbf8b19808a2e33452f Mon Sep 17 00:00:00 2001 From: andreii Date: Wed, 11 May 2022 08:31:53 -0700 Subject: [PATCH] Converter for Range operator. --- .../tf2tensorrt/convert/convert_nodes.h | 3 + .../tf2tensorrt/convert/convert_nodes_test.cc | 192 ++++++++++++++- .../tf2tensorrt/convert/ops/fill_ops.cc | 218 +++++++++++++++++- .../tf2tensorrt/convert/ops/layer_utils.h | 31 +-- .../tf2tensorrt/convert/ops/like_ops.cc | 25 +- 5 files changed, 426 insertions(+), 43 deletions(-) diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h index 38fa4b7041c90a..b3c597b8a9093c 100644 --- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h @@ -537,6 +537,9 @@ StatusOr ConvertMatMulImpl(OpConverterParams* params, TRT_TensorOrWeights input_b, bool transpose_a, bool transpose_b); +std::string convert_range_error_msg(float start, float limit, float delta); +std::string convert_range_expected_msg(const NodeDef& node_def); + } // namespace convert } // namespace tensorrt } // namespace tensorflow diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc index 4e8da56dd38612..e8095b07222db7 100644 --- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc @@ -28,6 +28,7 @@ limitations under the License. #include #include + #include "absl/algorithm/container.h" #include "absl/base/call_once.h" #include "absl/strings/match.h" @@ -36,8 +37,6 @@ limitations under the License. #include "absl/strings/str_format.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" -#include "third_party/gpus/cuda/include/cuda.h" -#include "third_party/gpus/cuda/include/cuda_runtime_api.h" #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/framework/scope.h" #include "tensorflow/cc/ops/nn_ops_internal.h" @@ -67,6 +66,8 @@ limitations under the License. #include "tensorflow/core/platform/test.h" #include "tensorflow/core/protobuf/config.pb.h" // NOLINT #include "tensorflow/core/public/session.h" +#include "third_party/gpus/cuda/include/cuda.h" +#include "third_party/gpus/cuda/include/cuda_runtime_api.h" #include "third_party/tensorrt/NvInfer.h" namespace tensorflow { @@ -3697,7 +3698,7 @@ TEST_P(OpConverter_FP32_FP16_INT32_Test, ConvertFill) { AddTestWeights("dims", {2}, {2, 2}, DT_INT32); AddTestWeights("value", {1}, {42.0}, tf_type_); RunValidationAndConversion(node_def, error::UNIMPLEMENTED, - "Conversion for Fill is not implemented in" + "Conversion for Fill is not implemented in " "implicit batch mode"); return; } @@ -3737,6 +3738,191 @@ TEST_P(OpConverter_FP32_FP16_INT32_Test, ConvertFill) { } } +TEST_P(OpConverter_FP32_FP16_INT32_Test, ConvertRange) { + auto get_casted_value = [this](const float value, const DataType dtype) { + return dtype == DT_INT32 ? static_cast(value) : value; + }; + + auto set_parameters = [this](const std::array& name, + const std::array, 3>& value, + const std::array& type, + bool all_tensors = false, int shape_idx = -1) { + Reset(); + for (int i = 0; i < 3; i++) { + if (all_tensors) { + std::vector partial_shape_dims = {}; + // The correct partial shape will be provided + // (a) for all parameters, when shape_idx > 3 + // (b) for all parameters, except shape_idx, when shape_idx >= 0 + // (c) for none of the shape_idx < 0 + if (shape_idx > 3 || shape_idx >= 0 && shape_idx != i) { + partial_shape_dims = {1}; + } + AddTestTensor(name[i], {1}, type[i], value[i], partial_shape_dims); + } else { + AddTestWeights(name[i], {1}, value[i], type[i]); + } + } + }; + + const float start = 1.0; + const float limit = 43.0; + const float delta = 2.0; + + const std::array param_name = {"start", "limit", "delta"}; + std::array, 3> param_value; + param_value[0] = {start}; + param_value[1] = {limit}; + param_value[2] = {delta}; + const auto start_type = tf_type_; + std::array param_type = {tf_type_, tf_type_, tf_type_}; + + Scope s = Scope::NewRootScope(); + const auto range = + ops::Range(s.WithOpName("my_range"), + ops::Placeholder(s.WithOpName(param_name[0]), param_type[0]), + ops::Placeholder(s.WithOpName(param_name[1]), param_type[1]), + ops::Placeholder(s.WithOpName(param_name[2]), param_type[2])); + + const NodeDef& node_def = range.operation.node()->def(); + const std::vector param_types{DT_FLOAT, DT_HALF, DT_INT32}; + + // ConverterRange is not implemented for Implicite batch mode. + if (trt_mode_ == TrtTestMode::kImplicitBatch) { + for (bool all_tensors : {false, true}) { + set_parameters(param_name, param_value, param_type, all_tensors); + RunValidationAndConversion(node_def, error::UNIMPLEMENTED, + "Conversion for Range is not implemented in " + "implicit batch mode"); + } + return; + } + + const std::string expected_msg = convert_range_expected_msg(node_def); + { + // We expect that all three (start, limit and delta) are passed as weights + // OR tensors and we reject parameters, if it's not true. + Reset(); + // Passing (start, limit) as weights + for (int i = 0; i < 2; i++) { + AddTestWeights(param_name[i], {1}, param_value[i], param_type[i]); + } + // ... and delta as a tensor + AddTestTensor(param_name[2], {1}, param_type[2], param_value[2]); + + RunValidationAndConversion(node_def, error::INVALID_ARGUMENT, + expected_msg + "passed as weights OR tensors"); + } + + nvinfer1::DataType trt_type; + TF_ASSERT_OK(TfTypeToTrtType(tf_type_, &trt_type)); + const std::string expected = DebugString(trt_type); + + // Reject invalid parameters if delta = 0 (for weights only). + for (auto limit_type : param_types) { + param_type[1] = limit_type; + for (auto delta_type : param_types) { + param_type[2] = delta_type; + param_value[2] = {0}; + + set_parameters(param_name, param_value, param_type); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "The delta parameter of Range operation cannot be equal to 0"); + + // Reject invalid parameters preventing the limit from + // being reached for fixed values of start and delta. + for (int j = 0; j <= 1; j++) { + param_value[j] = {get_casted_value(start, tf_type_)}; + param_value[1 - j] = {get_casted_value(limit, limit_type)}; + param_value[2] = {(2 * j - 1) * get_casted_value(delta, delta_type)}; + set_parameters(param_name, param_value, param_type); + const auto error = convert_range_error_msg( + param_value[0][0], param_value[1][0], param_value[2][0]); + RunValidationAndConversion(node_def, error::INVALID_ARGUMENT, error); + } + + param_value[0] = {start}; + // When passed as tensors, all parameters should be of DT_INT32 type. + if (start_type == DT_INT32 && limit_type == DT_INT32 && + delta_type == DT_INT32) { + if (trt_mode_ == TrtTestMode::kDynamicShape) { + // Wrong dimension for one of parameters. + for (int j = 0; j < 3; j++) { + const string err = + StrCat("Dimension for '", param_name[j], + "' of Range operator should be equal to 1"); + set_parameters(param_name, param_value, param_type, true, j); + RunValidationAndConversion(node_def, error::INVALID_ARGUMENT, err); + } + } + } else { + // When at least one parameter is set as non-integer tensors, + // the following test should fail. + set_parameters(param_name, param_value, param_type, true); + RunValidationAndConversion(node_def, error::UNIMPLEMENTED, + expected_msg + "tensors"); + } + } + } + + // The tests that pass all checks in ConvertRange::Validate(). + const Status status = Status::OK(); + const std::vector int_type{DT_INT32}; + for (bool all_tensors : {false, true}) { + // For now when (start, limit, delta) are passed as tensors + // these tensors should be of DT_INT32 type. + int partial_shape_idx = -1; + if (all_tensors) { + if (start_type != DT_INT32) { + continue; + } + if (trt_mode_ == TrtTestMode::kDynamicShape) { + // The correct partial shape will be provided for all parameters + partial_shape_idx = 3; + } + } + + // For now only parameters of DT_INT32 type could be used when + // they are pased as tensors. + const auto& types = all_tensors ? int_type : param_types; + const auto jEnd = all_tensors ? 0 : 1; + for (auto limit_type : types) { + param_type[1] = limit_type; + for (auto delta_type : types) { + param_type[2] = delta_type; + // Loop for positive and negative deltas. + for (int j = 0; j <= jEnd; j++) { + // Define the expected result which should match the usage + // of DT_INT32 for one of (start, limit, delta). + const int mult = (1 - 2 * j); + param_value[j] = {get_casted_value(start, tf_type_)}; + param_value[1 - j] = {get_casted_value(limit, limit_type)}; + param_value[2] = {mult * get_casted_value(delta, delta_type)}; + + // Create expected output. + std::vector expected_output; + const float limit_curr = param_value[1][0]; + const float delta_curr = param_value[2][0]; + float value = param_value[0][0]; + int num_values = 0; + while (mult * (limit_curr - value) > 0) { + num_values++; + expected_output.push_back(value); + value += delta_curr; + } + + set_parameters(param_name, param_value, param_type, all_tensors, + partial_shape_idx); + const std::vector output_dims = {num_values}; + TestOpConverter("my_range", node_def, output_dims, status, status, + ElementsAreArray(expected_output)); + } + } + } + } +} + TEST_P(OpConverter_FP32_FP16_INT32_Test, ConvertLikeOps) { auto get_node = [&](int value) -> NodeDef { Scope s = Scope::NewRootScope(); diff --git a/tensorflow/compiler/tf2tensorrt/convert/ops/fill_ops.cc b/tensorflow/compiler/tf2tensorrt/convert/ops/fill_ops.cc index 52b54291bc511b..8851c935e337af 100644 --- a/tensorflow/compiler/tf2tensorrt/convert/ops/fill_ops.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/ops/fill_ops.cc @@ -25,15 +25,31 @@ namespace convert { #if IS_TRT_VERSION_GE(8, 2, 0, 0) -class ConvertFill : public OpConverterBase { +template +class ConvertFillBase : public OpConverterBase { public: - explicit ConvertFill(OpConverterParams* params) - : OpConverterBase(params) {} + explicit ConvertFillBase(OpConverterParams* params) + : OpConverterBase(params) {} static constexpr std::array AllowedDataTypes() { return {DataType::DT_FLOAT, DataType::DT_HALF, DataType::DT_INT32}; } + Status ValidateFillBase(const OpConverterParams& params) { + if (params.use_implicit_batch) { + return errors::Unimplemented("Conversion for ", params.node_def.op(), + " is not implemented in" + " implicit batch mode"); + } + return Status::OK(); + } +}; + +class ConvertFill : public ConvertFillBase { + public: + explicit ConvertFill(OpConverterParams* params) + : ConvertFillBase(params) {} + static constexpr std::array InputSpec() { return std::array{ InputArgSpec::Create("dims", TrtInputArg::kBoth), @@ -42,18 +58,13 @@ class ConvertFill : public OpConverterBase { Status Validate() { const auto& params = *this->params_; - - if (params.use_implicit_batch) { - return errors::Unimplemented( - "Conversion for Fill is not implemented in" - "implicit batch mode"); - } + TF_RETURN_IF_ERROR(this->ValidateFillBase(params)); const auto& inputs = params.inputs; const auto& node_def = params.node_def; const TRT_TensorOrWeights& dims_input = inputs.at(0); - nvinfer1::DataType dims_type = dims_input.TrtDType(); + const auto dims_type = dims_input.TrtDType(); if (dims_type != nvinfer1::DataType::kINT32) { return errors::InvalidArgument("The dims parameter of ", node_def.op(), " operation in ", node_def.name(), @@ -62,7 +73,7 @@ class ConvertFill : public OpConverterBase { " type, got ", DebugString(dims_type)); } - int nbDims = dims_input.GetTrtDims().nbDims; + const auto nbDims = dims_input.GetTrtDims().nbDims; if (nbDims < 0) { return errors::InvalidArgument("The shape of parameter ", node_def.op(), " operation in ", node_def.name(), @@ -101,7 +112,192 @@ class ConvertFill : public OpConverterBase { } }; +class ConvertRange : public ConvertFillBase { + public: + explicit ConvertRange(OpConverterParams* params) + : ConvertFillBase(params) {} + + static constexpr std::array InputSpec() { + return std::array{ + InputArgSpec::Create("start", TrtInputArg::kBoth), + InputArgSpec::Create("limit", TrtInputArg::kBoth), + InputArgSpec::Create("delta", TrtInputArg::kBoth)}; + } + + static constexpr const char* NodeDefDataTypeAttributeName() { return ""; } + Status Validate() { + const auto& params = *this->params_; + TF_RETURN_IF_ERROR(this->ValidateFillBase(params)); + + const auto& inputs = params.inputs; + const auto& node_def = params.node_def; + + if (!all_same_types(inputs)) { + return errors::InvalidArgument(convert_range_expected_msg(node_def), + "passed as weights OR tensors"); + } + + if (!all_weights_) { + if (!all_integers(inputs)) { + return errors::Unimplemented(convert_range_expected_msg(node_def), + "tensors"); + } + + for (int i = 0; i < 3; i++) { + const auto& dims = inputs.at(i).GetTrtDims(); + if (dims.nbDims != 1 || dims.d[0] != 1) { + return errors::InvalidArgument("Dimension for '", InputSpec()[i].name, + "' of ", node_def.op(), " operator ", + "should be equal to 1"); + } + } + return Status::OK(); + } + + float param[3]; + for (int i = 0; i < 3; i++) { + const auto& input = inputs.at(i); + switch (input.TrtDType()) { + case nvinfer1::DataType::kFLOAT: + param[i] = get_input_param(input); + break; + case nvinfer1::DataType::kHALF: + param[i] = get_input_param(input); + break; + default: // nvinfer1::DataType::kINT32: + param[i] = get_input_param(input); + } + } + + if ((delta_ = param[2]) == 0) { + return errors::InvalidArgument("The delta parameter of ", node_def.op(), + " operation cannot be equal to 0"); + } + + const auto num_intervals_float = (param[1] - (start_ = param[0])) / delta_; + if (num_intervals_float < 0) { + const auto error = convert_range_error_msg(start_, param[1], delta_); + return errors::InvalidArgument(error); + } + + num_values_ = static_cast(num_intervals_float); + if (start_ + delta_ * num_values_ != param[1]) { + num_values_++; + } + return Status::OK(); + } + + Status Convert() { + const auto& params = *this->params_; + const auto& inputs = params.inputs; + const TRT_TensorOrWeights& input = inputs.at(0); + TRT_TensorOrWeights value_input; + + nvinfer1::Dims trt_dims{1}; + auto builder = TRTNetworkBuilder::Create(params.converter->network(), + params.weight_store); + TRT_ENSURE_OK(builder); + ITensorProxyPtr dims_input_tensor = nullptr; + ITensorProxyPtr beta_tensor = nullptr; + ITensorProxyPtr scalar_tensor = nullptr; + if (!all_weights_) { + StatusOr num = + builder->Sub(/*limit*/ inputs.at(1).tensor()->trt_tensor(), + /*start*/ inputs.at(0).tensor()->trt_tensor()); + + TRT_ENSURE_PTR_OK(num); + beta_tensor = params.inputs.at(2).tensor(); + StatusOr ceil_div = builder->FloorDiv( + (*num)->getOutput(0), beta_tensor->trt_tensor() /*delta*/); + TRT_ENSURE_PTR_OK(ceil_div); + dims_input_tensor = (*ceil_div)->getOutput(0); + dims_input_tensor->setType(nvinfer1::DataType::kINT32); + + nvinfer1::Dims scalar_dims{0}; + TF_RETURN_IF_ERROR(PrepareTensorForShape( + params.converter, params.inputs.at(0), scalar_dims, false, + &scalar_tensor, params.node_def)); + } else { + DimsAdapter value_input_dims(std::vector{1}); + StatusOr value_weights = + params.weight_store->GetTempWeights(input.TrtDType(), + value_input_dims); + + TF_RETURN_IF_ERROR(value_weights.status()); + TF_RETURN_IF_ERROR(value_weights->SetValues(start_)); + value_input = TRT_TensorOrWeights(value_weights.ValueOrDie()); + + trt_dims.d[0] = num_values_; + StatusOr const_layer = + builder->ConstantShape(value_input_dims); + TRT_ENSURE_PTR_OK(const_layer); + dims_input_tensor = (*const_layer)->getOutput(0); + } + + TRT_TensorOrWeights dims_input(dims_input_tensor); + + StatusOr layer = + builder->AddFill(value_input, dims_input, all_weights_, all_weights_, 1, + trt_dims, scalar_tensor, beta_tensor, delta_); + + ITensorProxyPtr output_tensor = (*layer)->getOutput(0); + if (all_integers(inputs)) { + output_tensor->setType(nvinfer1::DataType::kINT32); + } + + this->AddOutput(TRT_TensorOrWeights(output_tensor)); + return Status::OK(); + } + + private: + template + float get_input_param(const TRT_TensorOrWeights& input) { + return static_cast(*input.weights().GetPointer()); + } + + bool all_integers(const std::vector& inputs) const { + for (int i = 0; i < 3; i++) { + if (inputs.at(i).TrtDType() != nvinfer1::DataType::kINT32) { + return false; + } + } + return true; + } + + bool all_same_types(const std::vector& inputs) { + auto i = inputs.size(); + const bool is_weight = inputs.at(--i).is_weights(); + while (i--) { + if (inputs.at(i).is_weights() != is_weight) { + return all_weights_ = false; + } + } + all_weights_ = is_weight; + return true; + } + + float start_; + float delta_; + int num_values_; + bool all_weights_; +}; + +std::string convert_range_error_msg(float start, float limit, float delta) { + const char* format_string = + "For parameters (start, limit) = (%.2f, %.2f) " + "of the Range operation delta cannot be %s, got %.2f"; + return absl::StrFormat(format_string, start, limit, + start < limit ? "negative" : "positive", delta); +} + +std::string convert_range_expected_msg(const NodeDef& node_def) { + return "All parameters (start, limit, delta) of " + node_def.op() + + " operation in " + node_def.name() + " are expected to be "; +} + REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction(), "Fill"); +REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction(), + "Range"); #endif // IS_TRT_VERSION_GE(8, 2, 0, 0) diff --git a/tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h b/tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h index f6329b67d83ab3..37015593e77a77 100644 --- a/tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h +++ b/tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h @@ -483,29 +483,34 @@ class TRTNetworkBuilder { const TRT_TensorOrWeights& dims_input, bool is_value_static, bool is_dims_static, int nbDims, - const nvinfer1::Dims trt_dims) { + const nvinfer1::Dims& trt_dims, + ITensorProxyPtr scalar_tensor = nullptr, + ITensorProxyPtr beta_tensor = nullptr, + const float delta = 0) { // TensorRT IFillLayer requires a rank 0 scalar. - ITensorProxyPtr scalar_tensor; nvinfer1::Dims scalar_dims; scalar_dims.nbDims = 0; - nvinfer1::DataType value_type = value_input.TrtDType(); if (is_value_static) { StatusOr const_layer = WeightsToConstant(value_input.weights().GetTrtWeights(), scalar_dims); if (!const_layer.status().ok()) return const_layer.status(); scalar_tensor = (*const_layer)->getOutput(0); } else { - StatusOr shuffler_layer = - Reshape(value_input.tensor()->trt_tensor(), scalar_dims); - if (!shuffler_layer.status().ok()) return shuffler_layer.status(); - scalar_tensor = (*shuffler_layer)->getOutput(0); + if (scalar_tensor == nullptr) { + StatusOr shuffler_layer = + Reshape(value_input.tensor()->trt_tensor(), scalar_dims); + if (!shuffler_layer.status().ok()) return shuffler_layer.status(); + scalar_tensor = (*shuffler_layer)->getOutput(0); + } } - nvinfer1::Dims beta_shape{1, {nbDims}}; - StatusOr const_layer = - Constant(0, beta_shape, value_type); - TF_RETURN_IF_ERROR(const_layer.status()); - ITensorProxyPtr empty_beta_tensor = (*const_layer)->getOutput(0); + if (beta_tensor == nullptr) { + nvinfer1::Dims beta_shape{1, {nbDims}}; + StatusOr const_layer = + Constant(delta, beta_shape, value_input.TrtDType()); + TF_RETURN_IF_ERROR(const_layer.status()); + beta_tensor = (*const_layer)->getOutput(0); + } nvinfer1::IFillLayer* layer = network_->addFill(trt_dims, nvinfer1::FillOperation::kLINSPACE); @@ -514,7 +519,7 @@ class TRTNetworkBuilder { layer->setInput(0, *dims_input.tensor()->trt_tensor()); } layer->setInput(1, *scalar_tensor->trt_tensor()); - layer->setInput(2, *empty_beta_tensor->trt_tensor()); + layer->setInput(2, *beta_tensor->trt_tensor()); return layer; } diff --git a/tensorflow/compiler/tf2tensorrt/convert/ops/like_ops.cc b/tensorflow/compiler/tf2tensorrt/convert/ops/like_ops.cc index f749ef82f97614..f77afff81bc85a 100644 --- a/tensorflow/compiler/tf2tensorrt/convert/ops/like_ops.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/ops/like_ops.cc @@ -60,38 +60,31 @@ class ConvertLikeOps : public OpConverterBase> { Status Convert() { const auto ¶ms = *this->params_; const auto &inputs = params.inputs; - auto *converter = params.converter; - auto *network = converter->network(); + auto *network = params.converter->network(); const TRT_TensorOrWeights &input = inputs.at(0); + nvinfer1::Dims dims(input.GetTrtDims()); - nvinfer1::Dims dims{0}; - - nvinfer1::DataType value_type = input.TrtDType(); - - std::vector value_input_dims_data = {1}; - DimsAdapter value_input_dims(value_input_dims_data); + const std::vector value_input_dims_data = {1}; + const DimsAdapter value_input_dims(value_input_dims_data); StatusOr value_weights = - params.weight_store->GetTempWeights(value_type, value_input_dims); + params.weight_store->GetTempWeights(input.TrtDType(), value_input_dims); TF_RETURN_IF_ERROR(value_weights.status()); TF_RETURN_IF_ERROR(value_weights->SetValues(V)); TRT_TensorOrWeights value_input(value_weights.ValueOrDie()); - int is_dims_static = true; + const auto is_dims_static = HasStaticShape(dims); + auto builder = TRTNetworkBuilder::Create(network, params.weight_store); ITensorProxyPtr dims_input_tensor; - if (!HasStaticShape(input.GetTrtDims())) { - is_dims_static = false; - auto builder = TRTNetworkBuilder::Create(network, params.weight_store); + if (!is_dims_static) { StatusOr shape_layer = builder->Shape(input.tensor()->trt_tensor()); TF_RETURN_IF_ERROR(shape_layer.status()); dims_input_tensor = (*shape_layer)->getOutput(0); - } else { - dims = input.GetTrtDims(); + dims.nbDims = 0; } TRT_TensorOrWeights dims_input(dims_input_tensor); - auto builder = TRTNetworkBuilder::Create(network, params.weight_store); StatusOr layer = builder->AddFill(value_input, dims_input, true, is_dims_static, input.GetTrtDims().nbDims, dims);