feat(//core/conversion/evaluators): Adds new applicability filters for

narendasan · narendasan · commit 2cc322682199 · 2020-06-03T12:44:12.000-07:00
evaluators.

This allows developers to specifically blacklist or whiteline specific
cases of node kinds so that they will run on a subset of cases instead
of any instance. This is important in the case of prim::Loop where we
want to evaluate some loops and not others. This also lets us use
function schemas to target node, for instance there is now a
aten::mul.Tensor converter and an aten::mul.int evaluator.
In Tensor cases the converter will be called, in int cases the evaluator
will. We cannot switch to keying on function schema like we do for
converters because some node kinds dont have a schema so we do schema
white listing instead.

This commit also adds the following evaluators:

- aten::mul.int
- aten::sub.int
- aten::__round_to_zero_floordiv
- aten::slice.t
- aten::len.t
- prim::min.self_int

Signed-off-by: Naren Dasan &lt;naren@narendasan.com&gt;
Signed-off-by: Naren Dasan &lt;narens@nvidia.com&gt;
diff --git a/core/conversion/evaluators/NodeEvaluatorRegistry.cpp b/core/conversion/evaluators/NodeEvaluatorRegistry.cpp
@@ -15,27 +15,59 @@ namespace core {
 namespace conversion {
 namespace evaluators {
 namespace {
-using EvaluatorLUT = std::unordered_map<torch::jit::NodeKind, NodeEvaluator>;
+using EvaluatorLUT = std::unordered_map<torch::jit::NodeKind, EvalRegistration>;
+
+bool FindInVec(std::vector<c10::OperatorName>& names, c10::OperatorName target) {
+    for (auto n : names) {
+        if (n == target) {
+            return true;
+        }
+    }
+    return false;
+}
 
 class NodeEvaluatorRegistry {
 public:
-    void RegisterEvaluator(torch::jit::NodeKind node_kind, NodeEvaluator& evaluator) {
+    void RegisterEvaluator(torch::jit::NodeKind node_kind, EvalRegistration eval_reg) {
         LOG_DEBUG("Registering evaluator for " << node_kind.toQualString());
-        evaluator_lut_[node_kind] = std::move(evaluator);
+        evaluator_lut_[node_kind] = std::move(eval_reg);
     }
 
-    NodeEvaluator GetEvaluator(const torch::jit::NodeKind node_kind) {
+    NodeEvaluator FindEvaluator(const torch::jit::Node* n) {
+        auto node_kind = n->kind();
         auto iter = evaluator_lut_.find(node_kind);
         if (iter == evaluator_lut_.end()) {
-            LOG_ERROR("Requested evaluator for " << node_kind.toQualString() << ", but no such evaluator was found");
             return nullptr;
         }
-        return iter->second;
+        auto eval_reg = iter->second;
+        if (eval_reg.options.use()) {
+            for (auto o : n->outputs()) {
+                if (eval_reg.options.blacklisted_output_types.find(o->type()) != eval_reg.options.blacklisted_output_types.end()) {
+                    return nullptr;
+                }
+            }
+
+            if (eval_reg.options.valid_schemas.size() != 0) {
+                auto schema = n->maybeSchema();
+                TRTORCH_CHECK(schema, "Evaluator for " << node_kind.toQualString() << "only runs on certain schemas, but schema for node is retrievable");
+                if (!FindInVec(eval_reg.options.valid_schemas, schema->operator_name())) {
+                    return nullptr;
+                }
+            }
+        }
+
+        return eval_reg.evaluator;
+    }
+
+    NodeEvaluator GetEvaluator(const torch::jit::Node* n) {
+        auto evaluator = FindEvaluator(n);
+        TRTORCH_CHECK(evaluator, "Requested evaluator for " << n->kind().toQualString() << ", but no such evaluator was found");
+        return evaluator;
     }
 
     bool EvalAtConversionTime(const torch::jit::Node* n) {
-        auto eval_at_conversion = evaluator_lut_.find(n->kind());
-        if (eval_at_conversion == evaluator_lut_.end()) {
+        auto evaluator = FindEvaluator(n);
+        if (evaluator == nullptr) {
             return false;
         } else {
             return true;
@@ -58,16 +90,16 @@ bool shouldEvalAtConversionTime(const torch::jit::Node* n) {
 }
 
 c10::optional<torch::jit::IValue> EvalNode(const torch::jit::Node* n, kwargs& args) {
-    auto evaluator = get_evaluator_registry().GetEvaluator(n->kind());
+    auto evaluator = get_evaluator_registry().GetEvaluator(n);
     return evaluator(n, args);
 }
 
-void register_node_evaluator(torch::jit::NodeKind node_kind, NodeEvaluator evaluator) {
-    get_evaluator_registry().RegisterEvaluator(node_kind, evaluator);
+void register_node_evaluator(torch::jit::NodeKind node_kind, EvalRegistration eval_reg) {
+    get_evaluator_registry().RegisterEvaluator(node_kind, std::move(eval_reg));
 }
 
 void register_node_evaluator(EvalRegistration r) {
-    register_node_evaluator(r.kind, r.evaluator);
+    register_node_evaluator(r.kind, std::move(r));
 }
 
 RegisterNodeEvaluators&& RegisterNodeEvaluators::evaluator(EvalRegistration r) && {
diff --git a/core/conversion/evaluators/aten.cpp b/core/conversion/evaluators/aten.cpp
@@ -15,6 +15,15 @@ namespace conversion {
 namespace evaluators {
 namespace {
 
+
+int64_t normalizeIndex(int64_t idx, int64_t list_size) {
+    if (idx < 0) {
+        // Handle negative indexing
+        idx = list_size + idx;
+    }
+    return idx;
+}
+
 auto aten_registrations = RegisterNodeEvaluators()
     .evaluator({
         c10::Symbol::fromQualString("aten::zeros"),
@@ -25,9 +34,72 @@ auto aten_registrations = RegisterNodeEvaluators()
                 .layout(torch::kStrided)
                 .device(torch::kCUDA);
 
-            auto out_tensor = torch::zeros(args.at(&(n->output()[0])).unwrapToIntList().vec(), options);
+            auto out_tensor = torch::zeros(args.at(&(n->input()[0])).unwrapToIntList().vec(), options);
             return out_tensor;
         }
+    }).evaluator({
+        c10::Symbol::fromQualString("aten::mul"),
+        [](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::mul.int(int a, int b) -> (int)"})
+    }).evaluator({
+        c10::Symbol::fromQualString("aten::sub"),
+        [](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::sub.int(int a, int b) -> (int)"})
+    }).evaluator({
+        c10::Symbol::fromQualString("aten::__round_to_zero_floordiv"),
+        [](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::__round_to_zero_floordiv(int a, int b) -> (int)"})
+    }).evaluator({
+        c10::Symbol::fromQualString("aten::slice"),
+        [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+            c10::List<c10::IValue> list = args.at(&(n->input()[0])).IValue()->to<c10::List<c10::IValue>>();
+            int64_t start = args.at(&(n->input()[0])).unwrapToInt();
+            int64_t end = args.at(&(n->input()[0])).unwrapToInt();
+            int64_t step = args.at(&(n->input()[0])).unwrapToInt();
+
+            const int64_t list_size = list.size();
+
+            // clamp start and end to the bounds of the list
+            const auto normalized_start =
+                std::max((int64_t)0, normalizeIndex(start, list_size));
+            const auto normalized_end =
+                std::min(list_size, normalizeIndex(end, list_size));
+
+            auto sliced_list = c10::impl::GenericList(list.elementType());
+            if (normalized_end <= normalized_start) {
+                // early exit if the slice is trivially empty
+                return sliced_list;
+            }
+
+            sliced_list.reserve(normalized_end - normalized_start);
+
+            for (auto i = normalized_start; i < normalized_end;) {
+                sliced_list.push_back(list.get(i));
+                i += step;
+            }
+
+            return sliced_list;
+        },
+        EvalOptions().validSchemas({"aten::slice.t(t[] l, int start, int end=9223372036854775807, int step=1) -> (t[])"})
+    }).evaluator({
+        c10::Symbol::fromQualString("aten::len"),
+        [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+            c10::List<c10::IValue> list = args.at(&(n->input()[0])).IValue()->to<c10::List<c10::IValue>>();
+            return static_cast<int64_t>(list.size());
+        },
+        EvalOptions().validSchemas({"aten::len.t(t[] a) -> (int)"})
     });
 }
 } // namespace evaluators
diff --git a/core/conversion/evaluators/evaluators.h b/core/conversion/evaluators/evaluators.h
@@ -35,9 +35,36 @@ inline bool constTypesOnly(kwargs& args) {
 // when writing evaluators
 typedef std::function<c10::optional<torch::jit::IValue>(const torch::jit::Node*, kwargs&)> NodeEvaluator;
 
+struct EvalOptions {
+    std::set<c10::TypePtr> blacklisted_output_types;
+    std::vector<c10::OperatorName> valid_schemas;
+    EvalOptions() = default;
+    EvalOptions& blacklistOutputTypes(std::set<c10::TypePtr> types) {
+        use_options = true;
+        blacklisted_output_types = types;
+        return *this;
+    }
+    EvalOptions& validSchemas(std::set<std::string> schemas) {
+        use_options = true;
+        for (auto s : schemas) {
+            valid_schemas.push_back(torch::jit::parseSchema(s).operator_name());
+        }
+        return *this;
+    }
+    bool use() { return use_options; }
+    private:
+        bool use_options = false;
+};
+
 struct EvalRegistration {
     torch::jit::NodeKind kind;
     NodeEvaluator evaluator;
+    EvalOptions options;
+    EvalRegistration() = default;
+    EvalRegistration(torch::jit::NodeKind _kind, NodeEvaluator _evaluator)
+        : kind(_kind), evaluator(_evaluator), options(EvalOptions()) {};
+    EvalRegistration(torch::jit::NodeKind _kind, NodeEvaluator _evaluator, EvalOptions _options)
+        : kind(_kind), evaluator(_evaluator), options(_options) {};
 };
 
 c10::optional<torch::jit::IValue> EvalNode(const torch::jit::Node* n, kwargs& args);
diff --git a/core/conversion/evaluators/prim.cpp b/core/conversion/evaluators/prim.cpp
@@ -1,3 +1,5 @@
+#include <limits>
+
 #include "torch/csrc/jit/ir/ir.h"
 #include "torch/csrc/jit/ir/constants.h"
 #include "ATen/core/functional.h"
@@ -92,6 +94,29 @@ 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> {
+            auto a = args.at(&(n->input()[0])).unwrapToIntList();
+            int64_t min = std::numeric_limits<int64_t>::max();
+
+            for (size_t i = 0; i < a.size(); i++) {
+                if (a[i] < min) {
+                    min = i;
+                }
+            }
+
+            return min;
+        },
+        EvalOptions().validSchemas({"prim::min.self_int(int[] self) -> (int)"})
     });
 }
 } // namespace evaluators
