<commit_before>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h"

#include <string>
#include <vector>

#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_execution_context.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/profiler/lib/traceme.h"

#if GOOGLE_CUDA && GOOGLE_TENSORRT
#include "third_party/tensorrt/NvInfer.h"

namespace tensorflow {
namespace tensorrt {

using absl::StrCat;

ExecutionContext ExecutionContext::Create(nvinfer1::ICudaEngine* cuda_engine) {
  bool has_int32_output = false;
  for (int i = 0; i < cuda_engine->getNbBindings(); i++) {
    if (!cuda_engine->bindingIsInput(i) &&
        cuda_engine->getBindingDataType(i) == nvinfer1::DataType::kINT32) {
      has_int32_output = true;
      break;
    }
  }
  if (!IS_TRT_VERSION_GE(8, 0, 0, 0) && has_int32_output) {
    // TODO(nvbugs/3390469): Remove this workaround when the bug is fixed.
    nvinfer1::IExecutionContext* execution_context =
        cuda_engine->createExecutionContext();
    return ExecutionContext(execution_context, true);
  }

  nvinfer1::IExecutionContext* execution_context =
      cuda_engine->createExecutionContextWithoutDeviceMemory();
  return ExecutionContext(execution_context, false);
}

Status GetTrtBindingShape(const nvinfer1::ICudaEngine* cuda_engine,
                          const nvinfer1::IExecutionContext* execution_context,
                          int binding_index, bool use_implicit_batch,
                          int batch_size, TensorShape& shape) {
  tensorflow::profiler::TraceMe activity(
      "getBindingDimensions", tensorflow::profiler::TraceMeLevel::kInfo);
  nvinfer1::Dims dims =
      use_implicit_batch
          ? cuda_engine->getBindingDimensions(binding_index)
          : execution_context->getBindingDimensions(binding_index);
  if (!use_implicit_batch) {
    if (dims.nbDims == -1) {
      return errors::Internal(
          "Binding index out of range. This can happen if profile is not set, "
          "or the network is invalid for the current profile.");
    }
  }
  TF_RETURN_IF_ERROR(DimsAdapter(dims).TensorShape(
      &shape,
      use_implicit_batch ? std::optional<int>(batch_size) : std::nullopt));
  return Status::OK();
}

Status SetupBindings(nvinfer1::ICudaEngine* cuda_engine, const Tensor& tensor,
                     std::vector<void*>& buffers, int binding_index) {
  tensorflow::profiler::TraceMe activity(
      "SetBindingPointers", tensorflow::profiler::TraceMeLevel::kInfo);
  const auto dtype = cuda_engine->getBindingDataType(binding_index);
  VLOG(2) << "<<<<<<<<< SetupBindings with dtype = " << (int)dtype;
  switch (dtype) {
    case nvinfer1::DataType::kFLOAT:
      buffers[binding_index] = const_cast<float*>(tensor.flat<float>().data());
      break;
    case nvinfer1::DataType::kHALF:
      buffers[binding_index] =
          const_cast<Eigen::half*>(tensor.flat<Eigen::half>().data());
      break;
    case nvinfer1::DataType::kINT8:
      return errors::Internal("INT8 inputs are not supported yet!");
    case nvinfer1::DataType::kINT32:
      buffers[binding_index] = const_cast<int32*>(tensor.flat<int32>().data());
      break;
#if IS_TRT_VERSION_GE(8, 2, 0, 0)
    case nvinfer1::DataType::kBOOL:
      buffers[binding_index] = const_cast<bool*>(tensor.flat<bool>().data());
      break;
#endif
    default:
      return errors::Internal("Unknown TRT data type: ",
                              static_cast<int>(dtype));
  }
  return Status::OK();
}

// Sets up bindings.
Status SetTrtEngineInputs(nvinfer1::ICudaEngine* cuda_engine,
                          nvinfer1::IExecutionContext* execution_context,
                          const int trt_profile_idx,
                          std::vector<void*>& buffers, bool use_implicit_batch,
                          int num_batch,
                          const TrtShapeOptimizationProfile& profiles,
                          OpKernelContext* ctx, const DataVec* input_vec) {
  tensorflow::profiler::TraceMe activity(
      "SetTrtEngineInputs", tensorflow::profiler::TraceMeLevel::kInfo);
  int n_inputs = ctx ? ctx->num_inputs() : (input_vec ? input_vec->size() : 0);
  // Setup engine inputs.
  for (int i = 0; i < n_inputs; i++) {
    const string input_name =
        ctx ? StrCat(IONamePrefixes::kInputPHName, i) : input_vec->at(i).name;
    int binding_index;
    Status status = GetTrtBindingIndex(input_name.c_str(), trt_profile_idx,
                                       cuda_engine, &binding_index);
    if (IS_TRT_VERSION_GE(8, 0, 0, 0)) {
      TF_RETURN_IF_ERROR(status);
    } else if (!status.ok()) {
      // Before TRT 8, an input tensor can be pruned if it is not used by the
      // network (e.g. only its shape is used, but the shape is already defined
      // by the optimization profile by setting min=max). nvbugs/3153064
      VLOG(2) << "Skipping pruned input " << input_name;
      continue;
    }
    const Tensor& input_tensor = ctx ? ctx->input(i) : input_vec->at(i).tensor;
    const TensorShape& input_shape = input_tensor.shape();

    if (use_implicit_batch && ctx) {
      // Ensure all inputs have the same batch size
      if (num_batch != input_shape.dim_size(0)) {
        const string msg =
            StrCat("Input data has inconsistent batch size: ", num_batch,
                   " vs ", input_shape.dim_size(0));
        return errors::NotFound(msg);
      }
    }
    // Set known input dimensions. This is necessary because TRT network
    // could be made with dynamic dimensions.
    if (!use_implicit_batch) {
      TF_RETURN_IF_ERROR(profiles.SetInputShapeBinding(
          i, binding_index, cuda_engine, execution_context));

      if (cuda_engine->isExecutionBinding(binding_index)) {
        tensorflow::profiler::TraceMe activity(
            "SetTrtEngineInputs::setBindingDimensions",
            tensorflow::profiler::TraceMeLevel::kInfo);
        nvinfer1::Dims trt_dims;
        auto adap = DimsAdapter::Create(input_shape);
        TRT_ENSURE_OK(adap);
        VLOG(2) << "Setting binding dimensions for idx " << binding_index;
        bool ret = execution_context->setBindingDimensions(binding_index,
                                                           adap->AsTrtDims());
        if (!ret) {
          VLOG(2) << "Error setting engine input " << binding_index << " "
                  << DebugString(trt_dims);
          return errors::Internal(
              "Binding dimension does not fit selected profile.");
        }
      }
    }
    // Setup input bindings.
    TF_RETURN_IF_ERROR(
        SetupBindings(cuda_engine, input_tensor, buffers, binding_index));
  }

  // Ensure all network dynamic dimensions (if any) are set in execution
  // context.
  if (!execution_context->allInputDimensionsSpecified()) {
    return errors::Internal(<commit_msg>TF-TRT Only set binding dimensions if needed<commit_after>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h"

#include <string>
#include <vector>

#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_execution_context.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/profiler/lib/traceme.h"

#if GOOGLE_CUDA && GOOGLE_TENSORRT
#include "third_party/tensorrt/NvInfer.h"

namespace tensorflow {
namespace tensorrt {

using absl::StrCat;

ExecutionContext ExecutionContext::Create(nvinfer1::ICudaEngine* cuda_engine) {
  bool has_int32_output = false;
  for (int i = 0; i < cuda_engine->getNbBindings(); i++) {
    if (!cuda_engine->bindingIsInput(i) &&
        cuda_engine->getBindingDataType(i) == nvinfer1::DataType::kINT32) {
      has_int32_output = true;
      break;
    }
  }
  if (!IS_TRT_VERSION_GE(8, 0, 0, 0) && has_int32_output) {
    // TODO(nvbugs/3390469): Remove this workaround when the bug is fixed.
    nvinfer1::IExecutionContext* execution_context =
        cuda_engine->createExecutionContext();
    return ExecutionContext(execution_context, true);
  }

  nvinfer1::IExecutionContext* execution_context =
      cuda_engine->createExecutionContextWithoutDeviceMemory();
  return ExecutionContext(execution_context, false);
}

Status GetTrtBindingShape(const nvinfer1::ICudaEngine* cuda_engine,
                          const nvinfer1::IExecutionContext* execution_context,
                          int binding_index, bool use_implicit_batch,
                          int batch_size, TensorShape& shape) {
  tensorflow::profiler::TraceMe activity(
      "getBindingDimensions", tensorflow::profiler::TraceMeLevel::kInfo);
  nvinfer1::Dims dims =
      use_implicit_batch
          ? cuda_engine->getBindingDimensions(binding_index)
          : execution_context->getBindingDimensions(binding_index);
  if (!use_implicit_batch) {
    if (dims.nbDims == -1) {
      return errors::Internal(
          "Binding index out of range. This can happen if profile is not set, "
          "or the network is invalid for the current profile.");
    }
  }
  TF_RETURN_IF_ERROR(DimsAdapter(dims).TensorShape(
      &shape,
      use_implicit_batch ? std::optional<int>(batch_size) : std::nullopt));
  return Status::OK();
}

Status SetupBindings(nvinfer1::ICudaEngine* cuda_engine, const Tensor& tensor,
                     std::vector<void*>& buffers, int binding_index) {
  tensorflow::profiler::TraceMe activity(
      "SetBindingPointers", tensorflow::profiler::TraceMeLevel::kInfo);
  const auto dtype = cuda_engine->getBindingDataType(binding_index);
  VLOG(2) << "<<<<<<<<< SetupBindings with dtype = " << (int)dtype;
  switch (dtype) {
    case nvinfer1::DataType::kFLOAT:
      buffers[binding_index] = const_cast<float*>(tensor.flat<float>().data());
      break;
    case nvinfer1::DataType::kHALF:
      buffers[binding_index] =
          const_cast<Eigen::half*>(tensor.flat<Eigen::half>().data());
      break;
    case nvinfer1::DataType::kINT8:
      return errors::Internal("INT8 inputs are not supported yet!");
    case nvinfer1::DataType::kINT32:
      buffers[binding_index] = const_cast<int32*>(tensor.flat<int32>().data());
      break;
#if IS_TRT_VERSION_GE(8, 2, 0, 0)
    case nvinfer1::DataType::kBOOL:
      buffers[binding_index] = const_cast<bool*>(tensor.flat<bool>().data());
      break;
#endif
    default:
      return errors::Internal("Unknown TRT data type: ",
                              static_cast<int>(dtype));
  }
  return Status::OK();
}

// Sets up bindings.
Status SetTrtEngineInputs(nvinfer1::ICudaEngine* cuda_engine,
                          nvinfer1::IExecutionContext* execution_context,
                          const int trt_profile_idx,
                          std::vector<void*>& buffers, bool use_implicit_batch,
                          int num_batch,
                          const TrtShapeOptimizationProfile& profiles,
                          OpKernelContext* ctx, const DataVec* input_vec) {
  tensorflow::profiler::TraceMe activity(
      "SetTrtEngineInputs", tensorflow::profiler::TraceMeLevel::kInfo);
  int n_inputs = ctx ? ctx->num_inputs() : (input_vec ? input_vec->size() : 0);
  // Setup engine inputs.
  for (int i = 0; i < n_inputs; i++) {
    const string input_name =
        ctx ? StrCat(IONamePrefixes::kInputPHName, i) : input_vec->at(i).name;
    int binding_index;
    Status status = GetTrtBindingIndex(input_name.c_str(), trt_profile_idx,
                                       cuda_engine, &binding_index);
    if (IS_TRT_VERSION_GE(8, 0, 0, 0)) {
      TF_RETURN_IF_ERROR(status);
    } else if (!status.ok()) {
      // Before TRT 8, an input tensor can be pruned if it is not used by the
      // network (e.g. only its shape is used, but the shape is already defined
      // by the optimization profile by setting min=max). nvbugs/3153064
      VLOG(2) << "Skipping pruned input " << input_name;
      continue;
    }
    const Tensor& input_tensor = ctx ? ctx->input(i) : input_vec->at(i).tensor;
    const TensorShape& input_shape = input_tensor.shape();

    if (use_implicit_batch && ctx) {
      // Ensure all inputs have the same batch size
      if (num_batch != input_shape.dim_size(0)) {
        const string msg =
            StrCat("Input data has inconsistent batch size: ", num_batch,
                   " vs ", input_shape.dim_size(0));
        return errors::NotFound(msg);
      }
    }
    // Set known input dimensions. This is necessary because TRT network
    // could be made with dynamic dimensions.
    if (!use_implicit_batch) {
      TF_RETURN_IF_ERROR(profiles.SetInputShapeBinding(
          i, binding_index, cuda_engine, execution_context));

      if (cuda_engine->isExecutionBinding(binding_index)) {
        tensorflow::profiler::TraceMe activity(
            "SetTrtEngineInputs::setBindingDimensions",
            tensorflow::profiler::TraceMeLevel::kInfo);
        auto adap = DimsAdapter::Create(input_shape);
        TRT_ENSURE_OK(adap);
        nvinfer1::Dims trt_dims = adap->AsTrtDims();
        if (execution_context->getBindingDimensions(binding_index) !=
            trt_dims) {
          VLOG(2) << "Setting binding dimensions for idx " << binding_index;
          bool ret =
              execution_context->setBindingDimensions(binding_index, trt_dims);
          if (!ret) {
            VLOG(2) << "Error setting engine input " << binding_index << " "
                    << DebugString(trt_dims);
            return errors::Internal(
                "Binding dimension does not fit selected profile.");
          }
        }
      }
    }
    // Setup input bindings.
    TF_RETURN_IF_ERROR(
        SetupBindings(cuda_engine, input_tensor, buffers, binding_index));
  }

  // Ensure all network dynamic dimensions (if any) are set in execution
  // context.
  if (!execution_context->allInputDimensionsSpecified()) {
    return errors::Internal(
{'fraction_non_alphanumeric': 0.1195810853697239, 'fraction_numerical': 0.010409393843224373, 'mean_word_length': 3.2606814494321257, 'pattern_counts': {'":': 0, '<': 93, '<?xml version=': 0, '>': 66, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3869139', 'n_tokens_mistral': 4755, 'n_tokens_neox': 4703, 'n_words': 1083}