refactor(//docker): Update docker container for new API

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

Author: narendasan

core/compiler.cpp

@@ -377,7 +377,7 @@ std::string ConvertGraphToTRTEngine(const torch::jit::script::Module& mod, std::
 
   auto engine = conversion::ConvertBlockToEngine(g->block(), cfg.convert_info, static_params);
 
-  return std::move(engine);
+  return engine;
 }
 
 torch::jit::Module CompileGraph(const torch::jit::Module& mod, CompileSpec cfg) {

core/compiler.h

@@ -31,4 +31,4 @@ torch::jit::script::Module EmbedEngineInNewModule(const std::string& engine, run
 void set_device(const int gpu_id);
 
 } // namespace core
-} // namespace trtorch
+} // namespace torch_tensorrt

core/ir/StaticParams.cpp

@@ -20,7 +20,7 @@ StaticParams get_static_params(c10::ArrayRef<torch::jit::Value*> inputs, std::ve
   TRTORCH_CHECK(
       static_params.size() == params.size(),
       "Graph parameter parsing failed, mismatched number of static parameters and IValues")
-  return std::move(static_params);
+  return static_params;
 }
 
 } // namespace ir

core/ir/ir.cpp

@@ -24,7 +24,7 @@ InputSpecMap pair_input_vals_with_specs(std::vector<const torch::jit::Value*> va
     LOG_DEBUG("Paring " << i << ": " << vals[i]->debugName() << " : " << specs[i]);
     a.insert({vals[i], specs[i]});
   }
-  return std::move(a);
+  return a;
 }
 
 std::vector<const torch::jit::Value*> get_tensor_inputs(

core/partitioning/partitioning.cpp

@@ -133,7 +133,7 @@ std::vector<SegmentedBlock> segmentBlocksWithNonTensorInputs(SegmentedBlock& seg
     }
   }
 
-  return std::move(new_seg_blocks);
+  return new_seg_blocks;
 }
 
 void resolveNonTensorInputs(PartitionedGraph& segmented_blocks) { // , std::shared_ptr<torch::jit::Graph> g
@@ -385,7 +385,7 @@ PartitionedGraph segment_graph(torch::jit::Block* block, const PartitionInfo& pa
     finalize_block(segmented_blocks, SegmentedBlock::kTorch, in_prog_pyt_blk_nodes);
   }
 
-  return std::move(segmented_blocks);
+  return segmented_blocks;
 }
 
 PartitionedGraph Partition(

docker/Dockerfile

@@ -1,12 +1,12 @@
-ARG BASE=21.06
+ARG BASE=21.09
 ARG BASE_IMG=nvcr.io/nvidia/pytorch:${BASE}-py3
 FROM ${BASE_IMG} as base
 
-FROM base as trtorch-builder-base
+FROM base as torch-tensorrt-builder-base
 
-# Removing any bazel or trtorch pre-installed from the base image
+# Removing any bazel or torch-tensorrt pre-installed from the base image
 
-RUN rm -rf /opt/pytorch/trtorch /usr/bin/bazel
+RUN rm -rf /opt/torch-tensorrt /usr/bin/bazel
 
 RUN apt-get update && apt-get install --no-install-recommends -y curl gnupg
 RUN curl -fsSL https://bazel.build/bazel-release.pub.gpg | gpg --dearmor > /etc/apt/trusted.gpg.d/bazel.gpg
@@ -20,30 +20,32 @@ RUN cp /usr/lib/x86_64-linux-gnu/libnvinfer.so /usr/lib/x86_64-linux-gnu/libnvin
 
 RUN apt-get update && apt-get install -y --no-install-recommends locales ninja-build && rm -rf /var/lib/apt/lists/* && locale-gen en_US.UTF-8
 
-FROM trtorch-builder-base as trtorch-builder
+FROM torch-tensorrt-builder-base as torch-tensorrt-builder
 
-COPY . /workspace/trtorch/src
-WORKDIR /workspace/trtorch/src
-RUN cp ./docker/WORKSPACE.cu.docker  WORKSPACE
+COPY . /workspace/torch_tensorrt/src
+WORKDIR /workspace/torch_tensorrt/src
+RUN cp ./docker/WORKSPACE.docker WORKSPACE
 
 # This script builds both libtrtorch bin/lib/include tarball and the Pythin wheel, in dist/
 RUN ./docker/dist-build.sh
 
-FROM base as trtorch
+FROM base as torch-tensorrt
 
 # copy source repo
-COPY . /workspace/trtorch
-COPY --from=trtorch-builder  /workspace/trtorch/src/dist/ .
-RUN patch -u /opt/conda/lib/python3.8/site-packages/pytorch_quantization/nn/modules/tensor_quantizer.py -i /workspace/trtorch/docker/qat.patch
+COPY . /workspace/torch_tensorrt
+COPY --from=torch-tensorrt-builder  /workspace/torch_tensorrt/src/dist/ .
+RUN patch -u /opt/conda/lib/python3.8/site-packages/pytorch_quantization/nn/modules/tensor_quantizer.py -i /workspace/torch_tensorrt/docker/qat.patch
 RUN conda init bash
 
 RUN pip3 install ipywidgets --trusted-host pypi.org --trusted-host pypi.python.org --trusted-host=files.pythonhosted.org
 RUN jupyter nbextension enable --py widgetsnbextension
 
-RUN mkdir -p /opt/trtorch && tar xvf libtrtorch.tar.gz --strip-components 2 -C /opt/trtorch --exclude=LICENSE && pip3 install *.whl && rm -fr /workspace/trtorch/dist/*
+RUN pip3 install *.whl && rm -fr /workspace/torch_tensorrt/dist/* *.whl
 
-ENV LD_LIBRARY_PATH /opt/conda/lib/python3.8/site-packages/torch/lib:/opt/trtorch/lib:${LD_LIBRARY_PATH}
-ENV PATH /opt/trtorch/bin:${PATH}
+ENV LD_LIBRARY_PATH /opt/conda/lib/python3.8/site-packages/torch/lib:/opt/conda/lib/python3.8/site-packages/torch_tensorrt/lib:${LD_LIBRARY_PATH}
+ENV PATH /opt/conda/lib/python3.8/site-packages/torch_tensorrt/bin:${PATH}
+
+WORKDIR /workspace
+RUN mv /workspace/torch_tensorrt /opt/pytorch/torch_tensorrt
 
-WORKDIR /workspace/trtorch/
 CMD /bin/bash

docker/WORKSPACE.cu.docker

@@ -2,13 +2,10 @@
 
 mkdir -p dist
 
-bazel build //:libtrtorch --compilation_mode opt
-
 cd py && MAX_JOBS=1 LANG=en_US.UTF-8 LANGUAGE=en_US:en LC_ALL=en_US.UTF-8 python3 setup.py bdist_wheel --use-cxx11-abi
 
 cd ..
 
-cp bazel-bin/libtrtorch.tar.gz dist/
 cp py/dist/* dist/
 
 pip3 install ipywidgets --trusted-host pypi.org --trusted-host pypi.python.org --trusted-host=files.pythonhosted.org
@@ -17,5 +14,5 @@ jupyter nbextension enable --py widgetsnbextension
 pip3 install timm
 
 # test install
-mkdir -p /opt/trtorch && tar xvf dist/libtrtorch.tar.gz --strip-components 2 -C /opt/trtorch --exclude=LICENSE && pip3 uninstall -y trtorch && pip3 install dist/*.whl
+mkdir -p /opt/torch_tensorrt && pip3 uninstall -y torch_tensorrt && pip3 install dist/*.whl