refactor(//cpp/bin/torchtrtc): Refactor the CLI to make it easier to

extend later

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

cpp/bin/torchtrtc/BUILD

@@ -10,7 +10,14 @@ config_setting(
 cc_binary(
     name = "torchtrtc",
     srcs = [
+        "accuracy.h",
+        "accuracy.cpp",
+        "fileio.h",
+        "fileio.cpp",
+        "luts.h",
         "main.cpp",
+        "parser_util.h",
+        "parser_util.cpp"
     ],
     deps = [
         "//third_party/args",

cpp/bin/torchtrtc/accuracy.cpp

@@ -0,0 +1,27 @@
+#include "accuracy.h"
+
+#include "torch_tensorrt/logging.h"
+#include "torch_tensorrt/torch_tensorrt.h"
+
+namespace torchtrtc {
+namespace accuracy {
+
+bool check_rtol(const at::Tensor& diff, const std::vector<at::Tensor> inputs, float threshold) {
+  double maxValue = 0.0;
+  for (auto& tensor : inputs) {
+    maxValue = fmax(tensor.abs().max().item<float>(), maxValue);
+  }
+  torchtrt::logging::log(
+      torchtrt::logging::Level::kDEBUG,
+      std::string("Max Difference: ") + std::to_string(diff.abs().max().item<float>()));
+  torchtrt::logging::log(
+      torchtrt::logging::Level::kDEBUG, std::string("Acceptable Threshold: ") + std::to_string(threshold));
+  return diff.abs().max().item<float>() <= threshold * maxValue;
+}
+
+bool almost_equal(const at::Tensor& a, const at::Tensor& b, float threshold) {
+  return check_rtol(a - b, {a, b}, threshold);
+}
+
+} // namespace accuracy
+} // namespace torchtrtc

cpp/bin/torchtrtc/accuracy.h

@@ -0,0 +1,18 @@
+#pragma once
+
+#include <stdlib.h>
+#include <iostream>
+#include <sstream>
+#include <vector>
+
+#include "torch/script.h"
+#include "torch/torch.h"
+
+namespace torchtrtc {
+namespace accuracy {
+
+bool check_rtol(const at::Tensor& diff, const std::vector<at::Tensor> inputs, float threshold);
+bool almost_equal(const at::Tensor& a, const at::Tensor& b, float threshold);
+
+} // namespace accuracy
+} // namespace torchtrtc

cpp/bin/torchtrtc/fileio.cpp

@@ -0,0 +1,50 @@
+#include "fileio.h"
+
+namespace torchtrtc {
+namespace fileio {
+
+std::string read_buf(std::string const& path) {
+  std::string buf;
+  std::ifstream stream(path.c_str(), std::ios::binary);
+
+  if (stream) {
+    stream >> std::noskipws;
+    std::copy(std::istream_iterator<char>(stream), std::istream_iterator<char>(), std::back_inserter(buf));
+  }
+
+  return buf;
+}
+
+std::string get_cwd() {
+  char buff[FILENAME_MAX]; // create string buffer to hold path
+  if (getcwd(buff, FILENAME_MAX)) {
+    std::string current_working_dir(buff);
+    return current_working_dir;
+  } else {
+    torchtrt::logging::log(torchtrt::logging::Level::kERROR, "Unable to get current directory");
+    exit(1);
+  }
+}
+
+std::string real_path(std::string path) {
+  auto abs_path = path;
+  char real_path_c[PATH_MAX];
+  char* res = realpath(abs_path.c_str(), real_path_c);
+  if (res) {
+    return std::string(real_path_c);
+  } else {
+    torchtrt::logging::log(torchtrt::logging::Level::kERROR, std::string("Unable to find file ") + abs_path);
+    exit(1);
+  }
+}
+
+std::string resolve_path(std::string path) {
+  auto rpath = path;
+  if (!(rpath.rfind("/", 0) == 0)) {
+    rpath = get_cwd() + '/' + rpath;
+  }
+  return rpath;
+}
+
+} // namespace fileio
+} // namespace torchtrtc

cpp/bin/torchtrtc/fileio.h

@@ -0,0 +1,38 @@
+#pragma once
+#include <stdlib.h>
+#include <iostream>
+#include <sstream>
+
+#ifdef __linux__
+#include <linux/limits.h>
+#else
+#define PATH_MAX 260
+#endif
+
+#if defined(_WIN32)
+#include <direct.h>
+#define getcwd _getcwd
+#define realpath(N, R) _fullpath((R), (N), PATH_MAX)
+#else
+#include <unistd.h>
+#endif
+
+#include "NvInfer.h"
+#include "third_party/args/args.hpp"
+#include "torch/script.h"
+#include "torch/torch.h"
+
+#include "torch_tensorrt/logging.h"
+#include "torch_tensorrt/ptq.h"
+#include "torch_tensorrt/torch_tensorrt.h"
+
+namespace torchtrtc {
+namespace fileio {
+
+std::string read_buf(std::string const& path);
+std::string get_cwd();
+std::string real_path(std::string path);
+std::string resolve_path(std::string path);
+
+} // namespace fileio
+} // namespace torchtrtc

cpp/bin/torchtrtc/luts.h

@@ -0,0 +1,39 @@
+#pragma once
+
+#include "NvInfer.h"
+#include "third_party/args/args.hpp"
+#include "torch/script.h"
+#include "torch/torch.h"
+
+namespace torchtrtc {
+namespace luts {
+
+at::ScalarType to_torch_dtype(torchtrt::DataType dtype) {
+  switch (dtype) {
+    case torchtrt::DataType::kHalf:
+      return at::kHalf;
+    case torchtrt::DataType::kChar:
+      return at::kChar;
+    case torchtrt::DataType::kInt:
+      return at::kInt;
+    case torchtrt::DataType::kBool:
+      return at::kBool;
+    case torchtrt::DataType::kFloat:
+    default:
+      return at::kFloat;
+  }
+}
+
+const std::unordered_map<nvinfer1::DataType, at::ScalarType>& get_trt_at_type_map() {
+  static const std::unordered_map<nvinfer1::DataType, at::ScalarType> trt_at_type_map = {
+      {nvinfer1::DataType::kFLOAT, at::kFloat},
+      {nvinfer1::DataType::kHALF, at::kHalf},
+      {nvinfer1::DataType::kINT32, at::kInt},
+      {nvinfer1::DataType::kINT8, at::kChar},
+      {nvinfer1::DataType::kBOOL, at::kBool},
+  };
+  return trt_at_type_map;
+}
+
+} // namespace luts
+} // namespace torchtrtc