Address review changes

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

core/conversion/conversion.cpp

@@ -370,19 +370,22 @@ void ConvertBlockToNetDef(
       auto eval = EvaluateNode(ctx, n);
       if (eval) {
         if (n->outputs().size() > 1) { // For ListUnpack scenario
-          if (eval.value().isList()) {
-            LOG_DEBUG(
-                ctx->logger,
-                "Found the evaluated value to be a list" << eval.value() << " for node: " << util::node_info(n));
-            auto eval_list = eval.value().toList();
+          if (eval.value().isTuple()) {
+            auto eval_list = eval.value().toTuple();
             TRTORCH_CHECK(
-                eval_list.size() == n->outputs().size(),
-                "Size of evaluated results: " << eval_list.size() << " and node outputs size: " << n->outputs().size()
+                eval_list->elements().size() == n->outputs().size(),
+                "Size of evaluated results: " << eval_list->elements().size() << " and node outputs size: " << n->outputs().size()
                                               << " must match.");
-            for (int i = 0; i < eval_list.size(); i++) {
-              auto eval_output = eval_list.get(i);
+            for (int i = 0; i < eval_list->elements().size(); i++) {
+              auto eval_output = eval_list.get()->elements()[i];
+              LOG_DEBUG(
+                  ctx->logger,
+                  "Found the evaluated value(s) to be " << eval_output << " for node: " << util::node_info(n));
               ctx->AssociateValueAndIValue(n->output(i), eval_output);
             }
+          } else {
+              TRTORCH_THROW_ERROR(
+                  "Unsupported return type for evaluated node");
           }
         } else if (!eval.value().isTensor()) {
           LOG_DEBUG(ctx->logger, "Found the value to be: " << eval.value());

core/conversion/evaluators/prim.cpp

@@ -36,7 +36,8 @@ auto prim_registrations =
                     [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
                       // Outputs is an IValue which has list of tensors which can be found in ctx->evaluated_value_map
                       const torch::jit::IValue* outputs = args.at(n->input()).IValue();
-                      return *std::move(outputs);
+                      auto outputVec = outputs->toList().vec();
+                      return std::move(c10::ivalue::Tuple::create(outputVec));
                     }})
         .evaluator({torch::jit::prim::ListConstruct,
                     [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {

tests/core/conversion/evaluators/test_prim_evaluators.cpp

@@ -17,4 +17,23 @@ TEST(Evaluators, PrimConstantEvaluatesCorrectly) {
   auto trt_results = trtorch::tests::util::EvaluateGraph(g->block(), {});
 
   ASSERT_TRUE(jit_results[0] == trt_results[0]);
-}
+}
+
+TEST(Evaluators, PrimListUnpackEvaluatesCorrectly) {
+  const auto graph = R"IR(
+      graph():
+        %1 : int = prim::Constant[value=3]()
+        %2 : int = prim::Constant[value=4]()
+        %lc : int[] = prim::ListConstruct(%1, %2)
+        %lu.1 : int, %lu.2 : int = prim::ListUnpack(%lc)
+        return (%lu.1, %lu.2))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto jit_results = trtorch::tests::util::EvaluateGraphJIT(g, {});
+  auto trt_results = trtorch::tests::util::EvaluateGraph(g->block(), {});
+
+  ASSERT_TRUE(jit_results[0] == trt_results[0]);
+  ASSERT_TRUE(jit_results[1] == trt_results[1]);
+}

tests/util/evaluate_graph.cpp

@@ -6,6 +6,7 @@
 #include "core/conversion/evaluators/evaluators.h"
 #include "core/conversion/var/Var.h"
 #include "core/util/prelude.h"
+#include "core/util/jit_util.h"
 
 namespace trtorch {
 namespace tests {
@@ -20,20 +21,30 @@ std::vector<torch::jit::IValue> EvaluateGraph(const torch::jit::Block* b, std::v
   for (size_t i = 0; i < inputs.size(); i++) {
     ctx->AssociateValueAndIValue(b->inputs()[i], inputs[i]);
   }
-
+  LOG_DEBUG("Checking nodes");
   for (const auto n : b->nodes()) {
     TRTORCH_CHECK(
         core::conversion::evaluators::shouldEvalAtConversionTime(n),
         "Test graph contains non evaluatable nodes: " << *n);
     auto eval = core::conversion::EvaluateNode(ctx, n);
     if (eval) {
-      if (!eval.value().isTensor()) {
+      if (eval.value().isTuple()) {
+        auto eval_list = eval.value().toTuple();
+        for (int i = 0; i < eval_list->elements().size(); i++){
+          auto eval_output = eval_list.get()->elements()[i];
+          LOG_DEBUG(
+              ctx->logger,
+              "Found the evaluated value(s) to be " << eval_output << " for node: " << trtorch::core::util::node_info(n));
+          ctx->AssociateValueAndIValue(n->output(i), eval_output);
+        }
+      } else if(!eval.value().isTensor()){
         LOG_DEBUG("Found the value to be: " << eval.value());
-      } else {
+        ctx->AssociateValueAndIValue(n->output(0), eval.value());
+      }else {
         LOG_DEBUG("Found the value to be a tensor (shape " << eval.value().toTensor().sizes() << ')');
+        ctx->AssociateValueAndIValue(n->output(0), eval.value());
       }
     }
-    ctx->AssociateValueAndIValue(n->output(0), eval.value());
   }
 
   std::vector<torch::jit::IValue> outputs;