feat(aten::__derive_index): Implement derive index evaluator

Fixes: #834

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

Author: narendasan

core/conversion/evaluators/aten.cpp

@@ -806,7 +806,16 @@ auto aten_registrations TORCHTRT_UNUSED =
                  return 0;
                }
              },
-             EvalOptions().validSchemas({"aten::__range_length(int lo, int hi, int step) -> int"})});
+             EvalOptions().validSchemas({"aten::__range_length(int lo, int hi, int step) -> int"})})
+        .evaluator(
+            {c10::Symbol::fromQualString("aten::__derive_index"),
+             [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+               auto idx = args.at(n->input(0)).unwrapToInt();
+               auto start = args.at(n->input(1)).unwrapToInt();
+               auto step = args.at(n->input(2)).unwrapToInt();
+               return start + idx * step;
+             }});
+
 } // namespace
 } // namespace evaluators
 } // namespace conversion

tests/core/conversion/evaluators/test_aten_evaluators.cpp

@@ -797,3 +797,21 @@ TEST(Evaluators, PowFloatIntEvaluatesCorrectly) {
 
   ASSERT_TRUE(jit_results[0] == trt_results[0]);
 }
+
+TEST(Evaluators, DeriveIndexEvaluatesCorrectly) {
+  const auto graph = R"IR(
+      graph():
+        %1 : int = prim::Constant[value=9]()
+        %2 : int = prim::Constant[value=4]()
+        %3 : int = prim::Constant[value=2]()
+        %4 : int = aten::__derive_index(%1, %2, %3)
+        return (%4))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  auto jit_results = torch_tensorrt::tests::util::EvaluateGraphJIT(g, {});
+  auto trt_results = torch_tensorrt::tests::util::EvaluateGraph(g->block(), {});
+
+  ASSERT_TRUE(jit_results[0] == trt_results[0]);
+}