Make sure torch inputs contiguous before passing data pointer to TRT

Signed-off-by: Junjie Bai <[email protected]>

Author: bddppq

core/execution/register_trt_op.cpp

@@ -11,13 +11,16 @@ namespace execution {
 
 std::vector<at::Tensor> RunCudaEngine(nvinfer1::IExecutionContext* ctx, std::pair<uint64_t, uint64_t> io, std::vector<at::Tensor>& inputs) {
     std::vector<void*> gpu_handles;
+
+    std::vector<at::Tensor> contig_inputs{};
+    contig_inputs.reserve(inputs.size());
     for (size_t i = 0; i < inputs.size(); i++) {
-        auto in_gpu = inputs[i].to(at::kCUDA);
-        auto shape = core::util::toDimsPad(in_gpu.sizes(), 4);
-        //in_gpu = inputs[i].reshape(core::util::toVec(shape));
-        //LOG_DEBUG("In shape:" << in_gpu.sizes() );
-        ctx->setBindingDimensions(i, shape);
-        gpu_handles.push_back(in_gpu.data_ptr());
+        auto dims = core::util::toDimsPad(inputs[i].sizes(), 4);
+        auto shape = core::util::toVec(dims);
+        contig_inputs.push_back(inputs[i].to(at::kCUDA).view(shape).contiguous());
+        LOG_DEBUG("In shape:" << shape);
+        ctx->setBindingDimensions(i, dims);
+        gpu_handles.push_back(contig_inputs.back().data_ptr());
     }
 
     TRTORCH_CHECK(ctx->allInputDimensionsSpecified(), "Not enough inputs provided (execution.RunCudaEngine)");

tests/util/run_graph_engine.cpp

@@ -24,13 +24,16 @@ std::vector<at::Tensor> RunEngine(std::string& eng, std::vector<at::Tensor> inpu
     auto ctx = engine->createExecutionContext();
 
     std::vector<void*> gpu_handles;
+
+    std::vector<at::Tensor> contig_inputs{};
+    contig_inputs.reserve(inputs.size());
     for (size_t i = 0; i < inputs.size(); i++) {
-        auto in_gpu = inputs[i].to(at::kCUDA);
-        auto shape = core::util::toDimsPad(in_gpu.sizes(), 1);
-        in_gpu = inputs[i].reshape(core::util::toVec(shape));
-        LOG_DEBUG("In shape:" << in_gpu.sizes() );
-        ctx->setBindingDimensions(i, shape);
-        gpu_handles.push_back(in_gpu.data_ptr());
+        auto dims = core::util::toDimsPad(inputs[i].sizes(), 1);
+        auto shape = core::util::toVec(dims);
+        contig_inputs.push_back(inputs[i].to(at::kCUDA).view(shape).contiguous());
+        LOG_DEBUG("In shape:" << shape);
+        ctx->setBindingDimensions(i, dims);
+        gpu_handles.push_back(contig_inputs.back().data_ptr());
     }
 
     if (!ctx->allInputDimensionsSpecified()) {