feat(trt_util): from Naren, added unpadDims tool

Signed-off-by: Abhiram Iyer <[email protected]> Signed-off-by: Abhiram Iyer <[email protected]>
pytorch · Jun 26, 2020 · 164a1a6 · 164a1a6
1 parent d33ec82
commit 164a1a6
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 0 deletions.
diff --git a/core/util/trt_util.cpp b/core/util/trt_util.cpp
@@ -82,6 +82,30 @@ nvinfer1::Dims toDimsPad(c10::List<int64_t> l, uint64_t pad_to) {
     return dims;
 }
 
+nvinfer1::Dims unpadDims(const nvinfer1::Dims& d) {
+    nvinfer1::Dims dims;
+
+    int j = 0;
+    bool pad_dims_done = false;
+
+    for (int i = 0; i < d.nbDims; i++) {
+        if (d.d[i] == 1 && !pad_dims_done) {
+            // skip over unecessary dimension
+            continue;
+        } else {
+            dims.d[j] = d.d[i];
+            j++;
+
+            // keep all other dimensions (don't skip over them)
+            pad_dims_done = true;
+        }
+    }
+
+    dims.nbDims = j;
+
+    return dims;
+}
+
 std::vector<int64_t> toVec(nvinfer1::Dims d) {
     std::vector<int64_t> dims;
     for (int i = 0; i < d.nbDims; i++) {

diff --git a/core/util/trt_util.h b/core/util/trt_util.h
@@ -79,6 +79,7 @@ int64_t volume(const nvinfer1::Dims& d);
 
 nvinfer1::Dims toDimsPad(c10::IntArrayRef l, uint64_t pad_to);
 nvinfer1::Dims toDimsPad(c10::List<int64_t> l, uint64_t pad_to);
+nvinfer1::Dims unpadDims(const nvinfer1::Dims& d);
 nvinfer1::Dims toDims(c10::IntArrayRef l);
 nvinfer1::Dims toDims(c10::List<int64_t> l);
 nvinfer1::DimsHW toDimsHW(c10::List<int64_t> l);