feat(//core/conversion/evaluators): adding support for common evaluation

ops

Adds support for:
- prim::shape
- aten::neg
- aten::add
- aten::__getitem__
- aten::append

Fixes:
- prim::min

Removes:
- prim::Loop

Signed-off-by: Naren Dasan <narens@nvidia.com>

Author: narendasan

core/conversion/evaluators/aten.cpp

@@ -37,6 +37,14 @@ auto aten_registrations = RegisterNodeEvaluators()
             auto out_tensor = torch::zeros(args.at(n->input(0)).unwrapToIntList().vec(), options);
             return out_tensor;
         }
+    }).evaluator({
+        c10::Symbol::fromQualString("aten::add"),
+        [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+            auto a = args.at(n->input(0)).unwrapToInt();
+            auto b = args.at(n->input(1)).unwrapToInt();
+            return a + b;
+        },
+        EvalOptions().validSchemas({"aten::add.int(int a, int b) -> (int)"})
     }).evaluator({
         c10::Symbol::fromQualString("aten::mul"),
         [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
@@ -128,6 +136,42 @@ auto aten_registrations = RegisterNodeEvaluators()
             "aten::size(Tensor self) -> (int[])",
             "aten::size.int(Tensor self, int dim) -> (int)"
         })
+    }).evaluator({
+        c10::Symbol::fromQualString("aten::__getitem__"),
+        [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+            auto list = args.at(n->input(0)).unwrapToIntList();
+            auto idx = args.at(n->input(1)).unwrapToInt();
+
+            const int64_t list_size = list.size();
+            const int64_t normalized_idx = normalizeIndex(idx, list_size);
+            TRTORCH_CHECK(normalized_idx >= 0 || normalized_idx < list_size, "List index out of range (aten::__getitem__)");
+            return list.get(normalized_idx);
+        },
+        EvalOptions().validSchemas({
+            "aten::__getitem__.t(t[](a) list, int idx) -> (t(*))",
+        })
+    }).evaluator({
+        c10::Symbol::fromQualString("aten::append"),
+        [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+            auto list = args.at(n->input(0)).unwrapToIntList();
+            auto el = args.at(n->input(1)).unwrapToInt();
+
+            list.push_back(std::move(el));
+            return list;
+        },
+        EvalOptions().validSchemas({
+            "aten::append.t(t[](a!) self, t(c -> *) el) -> (t[](a!))",
+        })
+    }).evaluator({
+        c10::Symbol::fromQualString("aten::neg"),
+        [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+            auto el = args.at(n->input(1)).unwrapToInt();
+
+            return el * -1;
+        },
+        EvalOptions().validSchemas({
+            "aten::neg.int(int a) -> (int)",
+        })
     });
 }
 } // namespace evaluators

core/conversion/evaluators/prim.cpp

@@ -94,14 +94,6 @@ auto prim_registrations = RegisterNodeEvaluators()
                 return c10::optional<torch::jit::IValue>(std::move(torch::jit::IValue(list)));
             }
         }
-    }).evaluator({
-        torch::jit::prim::Loop,
-        [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
-            std::cout << *n << std::endl;
-
-            return {};
-        },
-        EvalOptions().blacklistOutputTypes({c10::TensorType::get()})
     }).evaluator({
          c10::Symbol::fromQualString("prim::min"),
          [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
@@ -110,13 +102,29 @@ auto prim_registrations = RegisterNodeEvaluators()
 
             for (size_t i = 0; i < a.size(); i++) {
                 if (a[i] < min) {
-                    min = i;
+                    min = a[i];
                 }
             }
 
             return min;
         },
         EvalOptions().validSchemas({"prim::min.self_int(int[] self) -> (int)"})
+    }).evaluator({
+        c10::Symbol::fromQualString("prim::shape"),
+        [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+            LOG_WARNING("There may be undefined behavior using dynamic shape and prim::shape");
+            auto tensor_var = args.at(n->input(0));
+            if (tensor_var.isITensor()) {
+                auto tensor = tensor_var.ITensor();
+                return util::toVec(tensor->getDimensions());
+            } else {
+                auto tensor = tensor_var.unwrapToTensor();
+                return tensor.sizes();
+            }
+        },
+        EvalOptions().validSchemas({
+            "prim::shape(Tensor a) -> (int[])"
+        })
     });
 }
 } // namespace evaluators